feat(snapshots): Download via archive endpoint with build-and-poll

CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ### Fixes
 
+- (snapshots) Update `snapshots download` to use the new archive endpoint, triggering the archive build and polling until it is ready before downloading ([#3344](https://github.com/getsentry/sentry-cli/pull/3344))
 - (snapshots) Show a clear "project was renamed" error instead of a cryptic JSON parse failure when uploading to a renamed project slug ([#3341](https://github.com/getsentry/sentry-cli/pull/3341))
 
 ## 3.5.1

src/api/mod.rs

@@ -1067,14 +1067,37 @@ impl AuthenticatedApi<'_> {
         }
     }
 
+    pub fn get_snapshot_archive_status(
+        &self,
+        org: &str,
+        snapshot_id: &str,
+    ) -> ApiResult<SnapshotArchiveStatus> {
+        let path = format!(
+            "/organizations/{}/preprodartifacts/snapshots/{}/archive/",
+            PathArg(org),
+            PathArg(snapshot_id),
+        );
+        self.get(&path)?.convert()
+    }
+
+    pub fn trigger_snapshot_archive_build(&self, org: &str, snapshot_id: &str) -> ApiResult<()> {
+        let path = format!(
+            "/organizations/{}/preprodartifacts/snapshots/{}/archive/",
+            PathArg(org),
+            PathArg(snapshot_id),
+        );
+        self.request(Method::Post, &path)?.send()?.into_result()?;
+        Ok(())
+    }
+
     pub fn download_snapshot_zip(
         &self,
         org: &str,
         snapshot_id: &str,
         dst: &mut std::fs::File,
     ) -> ApiResult<ApiResponse> {
         let path = format!(
-            "/organizations/{}/preprodartifacts/snapshots/{}/download/",
+            "/organizations/{}/preprodartifacts/snapshots/{}/archive/?download",
             PathArg(org),
             PathArg(snapshot_id),
         );
@@ -2123,6 +2146,11 @@ pub struct LatestBaseSnapshotResponse {
     pub image_count: u64,
 }
 
+#[derive(Deserialize)]
+pub struct SnapshotArchiveStatus {
+    pub ready: bool,
+}
+
 /// Upload options returned by the snapshots upload-options endpoint.
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]

src/commands/snapshots/download.rs

@@ -1,15 +1,20 @@
 use std::fs;
 use std::io;
 use std::path::PathBuf;
+use std::thread::sleep;
+use std::time::{Duration, Instant};
 
 use anyhow::{bail, Result};
 use clap::{Arg, ArgMatches, Command};
 
-use crate::api::Api;
+use crate::api::{Api, AuthenticatedApi};
 use crate::config::Config;
 use crate::utils::args::ArgExt as _;
 use crate::utils::fs::{path_as_url, TempFile};
 
+const POLL_INTERVAL: Duration = Duration::from_secs(2);
+const POLL_TIMEOUT: Duration = Duration::from_secs(300);
+
 pub fn make_command(command: Command) -> Command {
     command
         .about("Download baseline snapshot images from Sentry.")
@@ -89,6 +94,8 @@ pub fn execute(matches: &ArgMatches) -> Result<()> {
         _ => bail!("Exactly one of --app-id or --snapshot-id must be provided"),
     };
 
+    wait_for_archive(&api, &org, &snapshot_id)?;
+
     eprintln!("Downloading snapshot {snapshot_id}...");
     let tmp = TempFile::create()?;
     let mut tmp_file = tmp.open()?;
@@ -130,3 +137,25 @@ pub fn execute(matches: &ArgMatches) -> Result<()> {
 
     Ok(())
 }
+
+fn wait_for_archive(api: &AuthenticatedApi, org: &str, snapshot_id: &str) -> Result<()> {
+    if api.get_snapshot_archive_status(org, snapshot_id)?.ready {
+        return Ok(());
+    }
+
+    api.trigger_snapshot_archive_build(org, snapshot_id)?;
+    eprintln!("Building snapshot archive...");
+
+    let start = Instant::now();
+    while !api.get_snapshot_archive_status(org, snapshot_id)?.ready {
+        if start.elapsed() >= POLL_TIMEOUT {
+            bail!(
+                "Snapshot archive was not ready after {}s. The build may still be running; try again shortly.",
+                POLL_TIMEOUT.as_secs()
+            );
+        }
+        sleep(POLL_INTERVAL);
+    }
+
+    Ok(())
+}

tests/integration/snapshots.rs

@@ -1,4 +1,16 @@
-use crate::integration::{MockEndpointBuilder, TestManager};
+use std::io::{Cursor, Write as _};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+use crate::integration::{AssertCommand, MockEndpointBuilder, TestManager};
+
+fn snapshot_zip_bytes() -> Vec<u8> {
+    let mut zip = zip::ZipWriter::new(Cursor::new(Vec::new()));
+    zip.start_file("snapshot.png", zip::write::SimpleFileOptions::default())
+        .unwrap();
+    zip.write_all(b"fake png bytes").unwrap();
+    zip.finish().unwrap().into_inner()
+}
 
 #[test]
 fn command_snapshots_diff_help() {
@@ -20,6 +32,86 @@ fn command_snapshots_upload_help() {
     TestManager::new().register_trycmd_test("snapshots/snapshots-upload-help.trycmd");
 }
 
+#[test]
+fn command_snapshots_download_ready() {
+    let output = tempfile::tempdir().unwrap();
+    TestManager::new()
+        .mock_endpoint(
+            MockEndpointBuilder::new(
+                "GET",
+                "/api/0/organizations/wat-org/preprodartifacts/snapshots/123/archive/",
+            )
+            .with_response_body(r#"{"ready":true}"#),
+        )
+        .mock_endpoint(
+            MockEndpointBuilder::new(
+                "GET",
+                "/api/0/organizations/wat-org/preprodartifacts/snapshots/123/archive/?download",
+            )
+            .with_response_body(snapshot_zip_bytes()),
+        )
+        .assert_cmd(vec![
+            "snapshots",
+            "download",
+            "--org",
+            "wat-org",
+            "--snapshot-id",
+            "123",
+            "--output",
+            output.path().to_str().unwrap(),
+        ])
+        .with_default_token()
+        .run_and_assert(AssertCommand::Success);
+}
+
+#[test]
+fn command_snapshots_download_builds_then_downloads() {
+    let output = tempfile::tempdir().unwrap();
+    let probe_count = Arc::new(AtomicUsize::new(0));
+    TestManager::new()
+        .mock_endpoint(
+            MockEndpointBuilder::new(
+                "GET",
+                "/api/0/organizations/wat-org/preprodartifacts/snapshots/123/archive/",
+            )
+            .expect(2)
+            .with_response_fn(move |_| {
+                if probe_count.fetch_add(1, Ordering::SeqCst) == 0 {
+                    br#"{"ready":false}"#.to_vec()
+                } else {
+                    br#"{"ready":true}"#.to_vec()
+                }
+            }),
+        )
+        .mock_endpoint(
+            MockEndpointBuilder::new(
+                "POST",
+                "/api/0/organizations/wat-org/preprodartifacts/snapshots/123/archive/",
+            )
+            .with_status(202)
+            .with_response_body(r#"{"detail":"Building your snapshot archive."}"#),
+        )
+        .mock_endpoint(
+            MockEndpointBuilder::new(
+                "GET",
+                "/api/0/organizations/wat-org/preprodartifacts/snapshots/123/archive/?download",
+            )
+            .with_response_body(snapshot_zip_bytes()),
+        )
+        .assert_cmd(vec![
+            "snapshots",
+            "download",
+            "--org",
+            "wat-org",
+            "--snapshot-id",
+            "123",
+            "--output",
+            output.path().to_str().unwrap(),
+        ])
+        .with_default_token()
+        .run_and_assert(AssertCommand::Success);
+}
+
 #[test]
 fn command_snapshots_upload_renamed_project() {
     TestManager::new()