[2/2] Switch snapshot import to StorageZipReader (#40444)

This reads zip snapshots directly from object storage instead of copying to a temporary file first.

GitOrigin-RevId: 80f843f1c27afc8ff4546958953b2a10d0113c33

Author: goffrie

Cargo.lock

@@ -6,7 +6,6 @@ exclude = [ "crates/py_client", "crates/python_client_tests" ]
 [workspace.dependencies]
 application = { path = "crates/application" }
 async_lru = { path = "crates/async_lru" }
-async_zip_reader = { path = "crates/async_zip_reader" }
 authentication = { path = "crates/authentication" }
 aws_s3 = { path = "crates/aws_s3" }
 aws_utils = { path = "crates/aws_utils" }
@@ -55,6 +54,7 @@ shape_inference = { path = "crates/shape_inference" }
 sodium_secretbox = { path = "crates/sodium_secretbox" }
 sqlite = { path = "crates/sqlite" }
 storage = { path = "crates/storage" }
+storage_zip_reader = { path = "crates/storage_zip_reader" }
 sync = { path = "crates/sync" }
 sync_types = { package = "convex_sync_types", path = "crates/convex/sync_types" }
 text_search = { path = "crates/text_search" }

Cargo.toml

@@ -32,6 +32,7 @@ testing = [
     "value/testing",
     "vector/testing",
     "exports/testing",
+    "storage_zip_reader/testing",
 ]
 
 [dependencies]
@@ -41,7 +42,6 @@ async-recursion = { workspace = true }
 async-trait = { workspace = true }
 async_lru = { workspace = true }
 async_zip = { workspace = true }
-async_zip_reader = { workspace = true }
 authentication = { workspace = true }
 aws_s3 = { workspace = true }
 bytes = { workspace = true }
@@ -95,6 +95,7 @@ slugify = "0.1.0"
 smallvec = { workspace = true }
 sourcemap = { workspace = true }
 storage = { workspace = true }
+storage_zip_reader = { workspace = true }
 strum = { workspace = true }
 sync_types = { workspace = true }
 tempfile = { workspace = true }
@@ -115,9 +116,7 @@ vector = { workspace = true }
 libc = { workspace = true }
 
 [dev-dependencies]
-authentication = { workspace = true, features = [
-    "testing",
-] }
+authentication = { workspace = true, features = ["testing"] }
 common = { workspace = true, features = ["testing"] }
 database = { workspace = true, features = ["testing"] }
 errors = { workspace = true, features = ["testing"] }
@@ -140,10 +139,9 @@ runtime = { workspace = true, features = ["testing"] }
 search = { workspace = true, features = ["testing"] }
 shape_inference = { workspace = true, features = ["testing"] }
 storage = { workspace = true, features = ["testing"] }
+storage_zip_reader = { workspace = true, features = ["testing"] }
 udf = { workspace = true, features = ["testing"] }
-usage_tracking = { workspace = true, features = [
-    "testing",
-] }
+usage_tracking = { workspace = true, features = ["testing"] }
 value = { workspace = true, features = ["testing"] }
 vector = { workspace = true, features = ["testing"] }
 

crates/application/Cargo.toml

@@ -3,12 +3,9 @@ use std::{
         BTreeMap,
         BTreeSet,
     },
-    io::Cursor,
     sync::Arc,
 };
 
-use anyhow::Context as _;
-use async_zip_reader::ZipReader;
 use common::components::ComponentId;
 use database::{
     BootstrapComponentsModel,
@@ -27,8 +24,8 @@ use serde_json::json;
 use storage::{
     LocalDirStorage,
     Storage,
-    StorageExt as _,
 };
+use storage_zip_reader::StorageZipArchive;
 use tokio::io::AsyncReadExt as _;
 use value::{
     assert_obj,
@@ -136,22 +133,15 @@ async fn test_export_components(rt: TestRuntime) -> anyhow::Result<()> {
     .await?;
 
     // Check we can get the stored zip.
-    let storage_stream = storage
-        .get(&zip_object_key)
-        .await?
-        .context("object missing from storage")?;
-    let stored_bytes = storage_stream.collect_as_bytes().await?;
-    let mut zip_reader = ZipReader::new(Cursor::new(stored_bytes)).await?;
+    let zip_reader = StorageZipArchive::open(storage.clone(), &zip_object_key).await?;
     let mut zip_entries = BTreeMap::new();
-    let filenames: Vec<_> = zip_reader.file_names().await?;
-    for (i, filename) in filenames.into_iter().enumerate() {
-        let entry_reader = zip_reader.by_index(i).await?;
+    for entry in zip_reader.entries() {
         let mut entry_contents = String::new();
-        entry_reader
-            .read()
+        zip_reader
+            .read_entry(entry.clone())
             .read_to_string(&mut entry_contents)
             .await?;
-        zip_entries.insert(filename, entry_contents);
+        zip_entries.insert(entry.name.clone(), entry_contents);
     }
     assert_eq!(zip_entries, expected_export_entries);
 
@@ -198,22 +188,15 @@ async fn test_export_unmounted_components(rt: TestRuntime) -> anyhow::Result<()>
     .await?;
 
     // Check we can get the stored zip.
-    let storage_stream = storage
-        .get(&zip_object_key)
-        .await?
-        .context("object missing from storage")?;
-    let stored_bytes = storage_stream.collect_as_bytes().await?;
-    let mut zip_reader = ZipReader::new(Cursor::new(stored_bytes)).await?;
+    let zip_reader = StorageZipArchive::open(storage.clone(), &zip_object_key).await?;
     let mut zip_entries = BTreeSet::new();
-    let filenames: Vec<_> = zip_reader.file_names().await?;
-    for (i, filename) in filenames.into_iter().enumerate() {
-        let entry_reader = zip_reader.by_index(i).await?;
+    for entry in zip_reader.entries() {
         let mut entry_contents = String::new();
-        entry_reader
-            .read()
+        zip_reader
+            .read_entry(entry.clone())
             .read_to_string(&mut entry_contents)
             .await?;
-        zip_entries.insert(filename);
+        zip_entries.insert(entry.name.clone());
     }
     assert_eq!(zip_entries, expected_export_entries);
 

crates/application/src/exports/tests.rs

@@ -92,10 +92,7 @@ use shape_inference::{
     export_context::GeneratedSchema,
     ProdConfigWithOptionalFields,
 };
-use storage::{
-    Storage,
-    StorageExt,
-};
+use storage::Storage;
 use sync_types::{
     backoff::Backoff,
     Timestamp,
@@ -412,17 +409,17 @@ impl<RT: Runtime> SnapshotImportExecutor<RT> {
                 snapshot_import.component_path.clone(),
             )
         };
-        let body_stream = move || {
-            let object_key = object_key.clone();
-            async move {
-                let reader = match object_key.clone() {
-                    Ok(key) => self.snapshot_imports_storage.get_fq_object(&key).await?,
-                    Err(key) => self.snapshot_imports_storage.get(&key).await?,
-                };
-                reader.with_context(|| format!("Missing import object {object_key:?}"))
-            }
+        let fq_key = match &object_key {
+            Ok(key) => key.clone(),
+            Err(key) => self.snapshot_imports_storage.fully_qualified_key(key),
         };
-        let objects = parse_objects(format.clone(), component_path.clone(), body_stream).boxed();
+        let objects = parse_objects(
+            format.clone(),
+            component_path.clone(),
+            self.snapshot_imports_storage.clone(),
+            fq_key,
+        )
+        .boxed();
 
         let component_id = prepare_component_for_import(&self.database, &component_path).await?;
         // Remapping could be more extensive here, it's just relatively simple to handle