[storage] change github api to restore epoch archive (#394)

Co-authored-by: Sandra di Bittern <[email protected]>
Co-authored-by: Isa LeHind <[email protected]>

Author: 3 people

tools/storage/src/download_bundle.rs

@@ -10,14 +10,35 @@ use std::fs;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
+// Update GitHubContent structure for contents API
 #[derive(Deserialize, Debug)]
+#[allow(dead_code)]
 struct GitHubContent {
     download_url: Option<String>,
     #[serde(rename = "type")]
     content_type: String,
     name: String,
 }
 
+// Add TreeResponse structures for the tree API
+#[derive(Deserialize, Debug)]
+#[allow(dead_code)]
+struct TreeItem {
+    path: String,
+    mode: String,
+    #[serde(rename = "type")]
+    item_type: String,
+    sha: String,
+    url: Option<String>,
+}
+
+#[derive(Deserialize, Debug)]
+struct TreeResponse {
+    tree: Vec<TreeItem>,
+    truncated: bool,
+}
+
+// Rest of the original data structures
 #[derive(Debug)]
 pub struct EpochFolders {
     pub epoch_ending: String,
@@ -63,13 +84,13 @@ fn find_closest_transaction_folder(
         }
     }
 
-    // println!("For target version {}, found candidates:", target_version);
-    // if let Some((v, name)) = version_below {
-    //     println!("  Below target: {} ({})", name, v);
-    // }
-    // if let Some((v, name)) = version_above {
-    //     println!("  Above target: {} ({})", name, v);
-    // }
+    println!("For target version {}, found candidates:", target_version);
+    if let Some((v, name)) = version_below {
+        println!("  Below target: {} ({})", name, v);
+    }
+    if let Some((v, name)) = version_above {
+        println!("  Above target: {} ({})", name, v);
+    }
 
     // Choose the version below target
     version_below
@@ -84,37 +105,51 @@ pub async fn find_closest_epoch_folder(
     branch: &str,
     target_epoch: u64,
 ) -> Result<EpochFolders> {
+    // Update to use the Git Tree API instead of Contents API
     let api_url = format!(
-        "https://api.github.com/repos/{}/{}/contents/snapshots?ref={}",
+        "https://api.github.com/repos/{}/{}/git/trees/{}:snapshots",
         owner, repo, branch
     );
 
-    let contents: Vec<GitHubContent> = client
+    let response = client
         .get(&api_url)
         .header("User-Agent", "libra-framework-downloader")
         .send()
         .await
-        .context("Failed to list snapshots directory")?
+        .context("Failed to list snapshots directory")?;
+
+    let tree_response: TreeResponse = response
         .json()
         .await
         .context("Failed to parse snapshots directory contents")?;
 
+    if tree_response.truncated {
+        info!("Warning: GitHub Tree API response is truncated. Some folders might be missing.");
+    }
+
     // Separate folders by type
     let mut epoch_ending_folders: Vec<(u64, String)> = Vec::new();
     let mut state_epoch_folders: Vec<(u64, String)> = Vec::new();
     let mut transaction_folders: Vec<(u64, String)> = Vec::new();
 
-    for item in contents {
-        if item.content_type != "dir" {
+    for item in tree_response.tree {
+        // Only consider tree items (directories)
+        if item.item_type != "tree" {
             continue;
         }
 
-        if let Some(epoch) = parse_epoch_ending_number(&item.name) {
-            epoch_ending_folders.push((epoch, item.name));
-        } else if let Some((epoch, _version)) = parse_state_epoch_info(&item.name) {
-            state_epoch_folders.push((epoch, item.name));
-        } else if let Some(version) = parse_transaction_number(&item.name) {
-            transaction_folders.push((version, item.name));
+        // Extract just the folder name from the path
+        let folder_name = match Path::new(&item.path).file_name() {
+            Some(name) => name.to_string_lossy().to_string(),
+            None => continue,
+        };
+
+        if let Some(epoch) = parse_epoch_ending_number(&folder_name) {
+            epoch_ending_folders.push((epoch, folder_name));
+        } else if let Some((epoch, _version)) = parse_state_epoch_info(&folder_name) {
+            state_epoch_folders.push((epoch, folder_name));
+        } else if let Some(version) = parse_transaction_number(&folder_name) {
+            transaction_folders.push((version, folder_name));
         }
     }
 
@@ -189,62 +224,77 @@ pub async fn download_github_folder(
     output_dir: &str,
 ) -> Result<()> {
     let client = reqwest::Client::new();
-    let mut pending_dirs = vec![(path.to_string(), output_dir.to_string())];
 
     // Create the root output directory first
     fs::create_dir_all(output_dir)?;
 
-    while let Some((current_path, current_dir)) = pending_dirs.pop() {
-        let api_url = format!(
-            "https://api.github.com/repos/{}/{}/contents/{}?ref={}",
-            owner, repo, current_path, branch
-        );
+    // Use the Git Tree API with recursive flag to get all contents at once
+    let api_url = format!(
+        "https://api.github.com/repos/{}/{}/git/trees/{}:{}?recursive=1",
+        owner, repo, branch, path
+    );
+
+    info!("Downloading tree from: {}", api_url);
+
+    let response = client
+        .get(&api_url)
+        .header("User-Agent", "libra-framework-downloader")
+        .send()
+        .await
+        .context("Failed to send tree request")?;
+
+    let tree_response: TreeResponse = response
+        .json()
+        .await
+        .context("Failed to parse JSON tree response")?;
+
+    if tree_response.truncated {
+        info!("Warning: Response was truncated, not all files will be downloaded");
+    }
+
+    // Get the base path to properly handle nested directories
+    let base_path = Path::new(path)
+        .file_name()
+        .map_or(String::new(), |name| name.to_string_lossy().to_string());
+    let base_dir = Path::new(output_dir).join(base_path);
+    fs::create_dir_all(&base_dir)?;
+
+    // Process files from the tree
+    for item in tree_response.tree {
+        // Skip if not a blob (file)
+        if item.item_type != "blob" {
+            continue;
+        }
+
+        // The path in tree response is relative to the requested path
+        let relative_path = item.path;
 
-        info!("Downloading from: {}", api_url);
+        // Compute where to save the file
+        let output_path = base_dir.join(&relative_path);
+
+        // Create parent directory if it doesn't exist
+        if let Some(parent) = output_path.parent() {
+            fs::create_dir_all(parent)?;
+        }
 
-        let contents: Vec<GitHubContent> = client
-            .get(&api_url)
+        // Download the file content using raw GitHub URL
+        let content_url = format!(
+            "https://raw.githubusercontent.com/{}/{}/{}/{}/{}",
+            owner, repo, branch, path, relative_path
+        );
+
+        let content = client
+            .get(&content_url)
             .header("User-Agent", "libra-framework-downloader")
             .send()
             .await
-            .context("Failed to send request")?
-            .json()
+            .with_context(|| format!("Failed to download file: {}", relative_path))?
+            .bytes()
             .await
-            .context("Failed to parse JSON response")?;
-
-        // Extract the last component of the path to create the current directory
-        let current_folder = Path::new(&current_path)
-            .file_name()
-            .map(|s| s.to_string_lossy().to_string())
-            .unwrap_or_default();
-
-        // Create full directory path including the current folder
-        let full_dir_path = Path::new(&current_dir).join(&current_folder);
-        fs::create_dir_all(&full_dir_path)?;
-
-        for item in contents {
-            let output_path = full_dir_path.join(&item.name);
-
-            if item.content_type == "file" {
-                if let Some(download_url) = item.download_url {
-                    // println!("Downloading file: {}", item.name);
-                    let content = client
-                        .get(&download_url)
-                        .header("User-Agent", "libra-framework-downloader")
-                        .send()
-                        .await?
-                        .bytes()
-                        .await?;
-
-                    fs::write(&output_path, content)
-                        .with_context(|| format!("Failed to write file: {}", item.name))?;
-                }
-            } else if item.content_type == "dir" {
-                println!("Processing directory: {}", item.name);
-                let new_path = format!("{}/{}", current_path, item.name);
-                pending_dirs.push((new_path, full_dir_path.to_str().unwrap().to_string()));
-            }
-        }
+            .with_context(|| format!("Failed to read bytes from: {}", relative_path))?;
+
+        fs::write(&output_path, content)
+            .with_context(|| format!("Failed to write file: {}", relative_path))?;
     }
 
     Ok(())

tools/storage/src/restore_bundle.rs

@@ -161,12 +161,19 @@ impl RestoreBundle {
         for entry in file_list.flatten() {
             let content = fs::read_to_string(&entry)?;
             let updated_content = Self::update_manifest_paths(&content);
-            fs::write(&entry, &updated_content)?; // Add & here
-            verify_valid_transaction_list(&entry, self.version)?;
+            fs::write(&entry, &updated_content)?;
 
-            self.transaction_manifest = entry;
+            if verify_valid_transaction_list(&entry, self.version) {
+                self.transaction_manifest = entry;
+                return Ok(());
+            }
         }
-        Ok(())
+
+        // If we get here, no valid transaction manifest was found
+        bail!(
+            "No valid transaction manifest found for version {}",
+            self.version
+        );
     }
 
     fn update_manifest_paths(manifest_content: &str) -> String {
@@ -204,21 +211,31 @@ impl RestoreBundle {
     }
 }
 
-pub fn verify_valid_transaction_list(
-    transaction_manifest: &Path,
-    version: u64,
-) -> anyhow::Result<()> {
-    let s = fs::read_to_string(transaction_manifest)?;
-    let tm: TransactionBackup = serde_json::from_str(&s)?;
-
-    if tm.last_version < version {
-        bail!("the transaction you are looking for is newer than the last version in this bundle. Get a newer transaction backup");
-    };
-
-    if tm.first_version > version {
-        bail!("the transaction you are looking for is older than the last version in this bundle. Get an older transaction backup.");
+pub fn verify_valid_transaction_list(transaction_manifest: &Path, version: u64) -> bool {
+    match fs::read_to_string(transaction_manifest) {
+        Ok(s) => match serde_json::from_str::<TransactionBackup>(&s) {
+            Ok(tm) => {
+                if version > tm.last_version {
+                    info!("The transaction you are looking for is newer than the last version in this bundle. Get a newer transaction backup");
+                    return false;
+                };
+
+                if version < tm.first_version {
+                    info!("The transaction you are looking for is older than the last version in this bundle. Get an older transaction backup.");
+                    return false;
+                }
+                true
+            }
+            Err(e) => {
+                info!("Failed to parse transaction manifest: {}", e);
+                false
+            }
+        },
+        Err(e) => {
+            info!("Failed to read transaction manifest file: {}", e);
+            false
+        }
     }
-    Ok(())
 }
 
 #[test]