Accept either ID or Snapshot TS for the get_zip_export endpoint (#31213)

We will want to prefer ID eventually, but accept both for now.

GitOrigin-RevId: 9b1b4c70a04cb09ea2d021d6e99613bf2bd090cd

Author: nipunn1313

Cargo.lock

@@ -39,6 +39,7 @@ deno_core_icudata = "0.73.0"
 derive_more = "0.99"
 divan = "0.1.14"
 dotenvy = "0.15.7"
+either = "1.9.0"
 elliptic-curve = { version = "0.12.1", features = [ "std", "pem" ] }
 encoding_rs = "0.8.32"
 p256 = { version = "0.11.1", features = [ "ecdh" ] }

Cargo.toml

@@ -21,6 +21,7 @@ common = { path = "../common" }
 convex_macro = { path = "../convex_macro" }
 csv-async = { workspace = true }
 database = { path = "../database" }
+either = { workspace = true }
 errors = { path = "../errors" }
 events = { path = "../events" }
 file_storage = { path = "../file_storage" }

crates/application/Cargo.toml

@@ -53,7 +53,6 @@ use common::{
     },
     document::{
         DocumentUpdate,
-        ParsedDocument,
         CREATION_TIME_FIELD_PATH,
     },
     errors::{
@@ -122,6 +121,7 @@ use database::{
     Transaction,
     WriteSource,
 };
+use either::Either;
 use errors::{
     ErrorMetadata,
     ErrorMetadataAnyhowExt,
@@ -1373,39 +1373,32 @@ impl<RT: Runtime> Application<RT> {
     pub async fn get_zip_export(
         &self,
         identity: Identity,
-        snapshot_ts: Timestamp,
+        id: Either<DeveloperDocumentId, Timestamp>,
     ) -> anyhow::Result<(StorageGetStream, String)> {
-        let stream = self.get_export_inner(identity, snapshot_ts).await?;
-        let filename = format!(
-            // This should match the format in SnapshotExport.tsx.
-            "snapshot_{}_{snapshot_ts}.zip",
-            self.instance_name
-        );
-        Ok((stream, filename))
-    }
-
-    async fn get_export_inner(
-        &self,
-        identity: Identity,
-        snapshot_ts: Timestamp,
-    ) -> anyhow::Result<StorageGetStream> {
-        let object_key = {
+        let (object_key, snapshot_ts) = {
             let mut tx = self.begin(identity).await?;
-            let export_doc = ExportsModel::new(&mut tx)
-                .completed_export_at_ts(snapshot_ts)
-                .await?;
-            let export: ParsedDocument<Export> = export_doc
-                .context(ErrorMetadata::not_found(
-                    "ExportNotFound",
-                    format!("The requested export {snapshot_ts} was not found"),
-                ))?
-                .try_into()?;
+            let export = match id {
+                Either::Left(id) => ExportsModel::new(&mut tx).get(id).await?,
+                Either::Right(ts) => {
+                    ExportsModel::new(&mut tx)
+                        .completed_export_at_ts(ts)
+                        .await?
+                },
+            }
+            .context(ErrorMetadata::not_found(
+                "ExportNotFound",
+                format!("The requested export {id} was not found"),
+            ))?;
             match export.into_value() {
-                Export::Completed { zip_object_key, .. } => zip_object_key,
+                Export::Completed {
+                    zip_object_key,
+                    start_ts,
+                    ..
+                } => (zip_object_key, start_ts),
                 Export::Failed { .. } | Export::InProgress { .. } | Export::Requested { .. } => {
                     anyhow::bail!(ErrorMetadata::bad_request(
                         "ExportNotComplete",
-                        format!("The requested export {snapshot_ts} has not completed"),
+                        format!("The requested export {id} has not completed"),
                     ))
                 },
             }
@@ -1418,7 +1411,13 @@ impl<RT: Runtime> Application<RT> {
                     "ExportNotFound",
                     format!("The requested export {snapshot_ts}/{object_key:?} was not found"),
                 ))?;
-        Ok(storage_get_stream)
+
+        let filename = format!(
+            // This should match the format in SnapshotExport.tsx.
+            "snapshot_{}_{snapshot_ts}.zip",
+            self.instance_name
+        );
+        Ok((storage_get_stream, filename))
     }
 
     /// Returns the cloud export key - fully qualified to the instance.

crates/application/src/lib.rs

@@ -21,6 +21,7 @@ clap = { workspace = true }
 cmd_util = { path = "../../crates/cmd_util" }
 common = { path = "../common" }
 database = { path = "../database" }
+either = { workspace = true }
 errors = { path = "../errors" }
 events = { path = "../events" }
 file_storage = { path = "../file_storage" }

crates/local_backend/Cargo.toml

@@ -192,7 +192,7 @@ pub fn router(st: LocalAppState) -> Router {
 
     let snapshot_export_routes = Router::new()
         .route("/request/zip", post(request_zip_export))
-        .route("/zip/:snapshot_ts", get(get_zip_export));
+        .route("/zip/:id", get(get_zip_export));
 
     let api_routes = Router::new()
         .merge(cli_routes)

crates/local_backend/src/router.rs

@@ -24,6 +24,7 @@ use common::{
         HttpResponseError,
     },
 };
+use either::Either;
 use errors::ErrorMetadata;
 use http::StatusCode;
 use model::exports::types::{
@@ -33,6 +34,7 @@ use model::exports::types::{
 use serde::Deserialize;
 use storage::StorageGetStream;
 use sync_types::Timestamp;
+use value::DeveloperDocumentId;
 
 use crate::{
     admin::must_be_admin_with_write_access,
@@ -77,28 +79,31 @@ pub async fn request_zip_export(
 
 #[derive(Deserialize)]
 pub struct ZipExportRequest {
-    // Timestamp the snapshot export started at
-    snapshot_ts: String,
+    // The ID of the snapshot
+    id: String,
 }
 
 #[debug_handler]
 pub async fn get_zip_export(
     State(st): State<LocalAppState>,
     ExtractIdentity(identity): ExtractIdentity,
-    Path(ZipExportRequest { snapshot_ts }): Path<ZipExportRequest>,
+    Path(ZipExportRequest { id }): Path<ZipExportRequest>,
 ) -> Result<impl IntoResponse, HttpResponseError> {
     must_be_admin_with_write_access(&identity)?;
-    let ts: Timestamp = snapshot_ts.parse().context(ErrorMetadata::bad_request(
-        "BadSnapshotTimestamp",
-        "Snapshot timestamp did not parse to a timestamp.",
-    ))?;
+    let id: Either<DeveloperDocumentId, Timestamp> = match id.parse() {
+        Ok(id) => Either::Left(id),
+        Err(_) => Either::Right(id.parse().context(ErrorMetadata::bad_request(
+            "BadSnapshotId",
+            "Snapshot Id did not parse to an ID.",
+        ))?),
+    };
     let (
         StorageGetStream {
             content_length,
             stream,
         },
         filename,
-    ) = st.application.get_zip_export(identity, ts).await?;
+    ) = st.application.get_zip_export(identity, id).await?;
     let content_length = ContentLength(content_length as u64);
     Ok((
         TypedHeader(content_length),

crates/local_backend/src/snapshot_export.rs

@@ -215,7 +215,7 @@ impl<'a, RT: Runtime> ExportsModel<'a, RT> {
     pub async fn completed_export_at_ts(
         &mut self,
         snapshot_ts: Timestamp,
-    ) -> anyhow::Result<Option<ResolvedDocument>> {
+    ) -> anyhow::Result<Option<ParsedDocument<Export>>> {
         let index_range = IndexRange {
             index_name: EXPORTS_BY_STATE_AND_TS_INDEX.clone(),
             range: vec![
@@ -229,7 +229,11 @@ impl<'a, RT: Runtime> ExportsModel<'a, RT> {
         };
         let query = Query::index_range(index_range);
         let mut query_stream = ResolvedQuery::new(self.tx, TableNamespace::Global, query)?;
-        query_stream.expect_at_most_one(self.tx).await
+        query_stream
+            .expect_at_most_one(self.tx)
+            .await?
+            .map(|doc| doc.try_into())
+            .transpose()
     }
 
     pub async fn get(