[TRUNK-12756] Update bundle upload status (round 2) (#119)

* update bundle upload status round 2

* trunk check -y

api/src/lib.rs

@@ -19,6 +19,25 @@ pub struct CreateBundleUploadResponse {
     pub key: String,
 }
 
+#[derive(Debug, Serialize, Clone, Deserialize, PartialEq, Eq)]
+pub enum BundleUploadStatus {
+    #[serde(rename = "PENDING")]
+    Pending,
+    #[serde(rename = "UPLOAD_COMPLETE")]
+    UploadComplete,
+    #[serde(rename = "UPLOAD_FAILED")]
+    UploadFailed,
+    #[serde(rename = "DRY_RUN")]
+    DryRun,
+}
+
+#[derive(Debug, Serialize, Clone, Deserialize, PartialEq, Eq)]
+pub struct UpdateBundleUploadRequest {
+    pub id: String,
+    #[serde(rename = "uploadStatus")]
+    pub upload_status: BundleUploadStatus,
+}
+
 #[derive(Debug, Serialize, Clone, Deserialize, PartialEq, Eq)]
 pub struct CreateRepoRequest {
     pub repo: RepoUrlParts,

cli-tests/src/upload.rs

@@ -4,7 +4,10 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use api::{CreateRepoRequest, GetQuarantineBulkTestStatusRequest};
+use api::{
+    BundleUploadStatus, CreateRepoRequest, GetQuarantineBulkTestStatusRequest,
+    UpdateBundleUploadRequest,
+};
 use assert_cmd::Command;
 use assert_matches::assert_matches;
 use context::repo::RepoUrlParts as Repo;
@@ -81,7 +84,7 @@ async fn upload_bundle() {
         .failure();
 
     let requests = state.requests.lock().unwrap().clone();
-    assert_eq!(requests.len(), 4);
+    assert_eq!(requests.len(), 5);
     let mut requests_iter = requests.into_iter();
 
     assert_eq!(
@@ -176,6 +179,14 @@ async fn upload_bundle() {
     assert_eq!(bundled_file.owners, ["@user"]);
     assert_eq!(bundled_file.team, None);
 
+    assert_eq!(
+        requests_iter.next().unwrap(),
+        RequestPayload::UpdateBundleUpload(UpdateBundleUploadRequest {
+            id: "test-bundle-upload-id".to_string(),
+            upload_status: BundleUploadStatus::UploadComplete
+        }),
+    );
+
     assert_eq!(
         requests_iter.next().unwrap(),
         RequestPayload::CreateRepo(CreateRepoRequest {

cli/src/clients.rs

@@ -2,8 +2,8 @@ use std::{format, path::PathBuf};
 
 use anyhow::Context;
 use api::{
-    CreateBundleUploadRequest, CreateBundleUploadResponse, CreateRepoRequest,
-    GetQuarantineBulkTestStatusRequest, QuarantineConfig,
+    BundleUploadStatus, CreateBundleUploadRequest, CreateBundleUploadResponse, CreateRepoRequest,
+    GetQuarantineBulkTestStatusRequest, QuarantineConfig, UpdateBundleUploadRequest,
 };
 use context::repo::RepoUrlParts as Repo;
 
@@ -48,6 +48,36 @@ pub async fn create_trunk_repo(
     Ok(())
 }
 
+pub async fn update_bundle_upload_status(
+    origin: &str,
+    api_token: &str,
+    id: &str,
+    upload_status: &BundleUploadStatus,
+) -> anyhow::Result<()> {
+    let client = reqwest::Client::new();
+    let resp = client
+        .patch(format!("{}/v1/metrics/updateBundleUpload", origin))
+        .timeout(TRUNK_API_TIMEOUT)
+        .header(reqwest::header::CONTENT_TYPE, "application/json")
+        .header(TRUNK_API_TOKEN_HEADER, api_token)
+        .json(&UpdateBundleUploadRequest {
+            id: id.to_owned(),
+            upload_status: upload_status.to_owned(),
+        })
+        .send()
+        .await
+        .map_err(|e| anyhow::anyhow!(e).context("Failed to update bundle upload status"))?;
+
+    if resp.status() != reqwest::StatusCode::OK {
+        return Err(
+            anyhow::anyhow!("{}: {}", resp.status(), status_code_help(resp.status()))
+                .context("Failed to update bundle upload status"),
+        );
+    }
+
+    Ok(())
+}
+
 pub async fn create_bundle_upload_intent(
     origin: &str,
     api_token: &str,

cli/src/main.rs

@@ -4,13 +4,14 @@ use std::time::{SystemTime, UNIX_EPOCH};
 #[cfg(target_os = "macos")]
 use xcresult::XCResult;
 
+use api::BundleUploadStatus;
 use clap::{Args, Parser, Subcommand};
 use context::repo::BundleRepo;
 use tokio_retry::strategy::ExponentialBackoff;
 use tokio_retry::Retry;
 use trunk_analytics_cli::bundler::BundlerUtil;
 use trunk_analytics_cli::clients::{
-    create_bundle_upload_intent, create_trunk_repo, put_bundle_to_s3,
+    create_bundle_upload_intent, create_trunk_repo, put_bundle_to_s3, update_bundle_upload_status,
 };
 use trunk_analytics_cli::codeowners::CodeOwners;
 use trunk_analytics_cli::constants::{
@@ -318,7 +319,7 @@ async fn run_upload(
         org: org_url_slug.clone(),
         repo: repo.clone(),
         cli_version,
-        bundle_upload_id: upload.id,
+        bundle_upload_id: upload.id.clone(),
         tags,
         file_sets,
         envs,
@@ -357,14 +358,42 @@ async fn run_upload(
     log::info!("Flushed temporary tarball to {:?}", bundle_time_file);
 
     if dry_run {
+        if let Err(e) = update_bundle_upload_status(
+            &api_address,
+            &token,
+            &upload.id,
+            &BundleUploadStatus::DryRun,
+        )
+        .await
+        {
+            log::warn!("Failed to update bundle upload status: {}", e);
+        } else {
+            log::debug!("Updated bundle upload status to DRY_RUN");
+        }
         log::info!("Dry run, skipping upload.");
         return Ok(exit_code);
     }
 
-    Retry::spawn(default_delay(), || {
+    let upload_status = Retry::spawn(default_delay(), || {
         put_bundle_to_s3(&upload.url, &bundle_time_file)
     })
-    .await?;
+    .await
+    .map(|_| BundleUploadStatus::UploadComplete)
+    .unwrap_or_else(|e| {
+        log::error!("Failed to upload bundle to S3 after retries: {}", e);
+        BundleUploadStatus::UploadFailed
+    });
+    if let Err(e) =
+        update_bundle_upload_status(&api_address, &token, &upload.id, &upload_status).await
+    {
+        log::warn!(
+            "Failed to update bundle upload status to {:#?}: {}",
+            upload_status,
+            e
+        )
+    } else {
+        log::debug!("Updated bundle upload status to {:#?}", upload_status)
+    }
 
     let remote_urls = vec![repo.repo_url.clone()];
     Retry::spawn(default_delay(), || {

test_utils/src/mock_server.rs

@@ -8,14 +8,14 @@ use std::{
 
 use api::{
     CreateBundleUploadRequest, CreateBundleUploadResponse, CreateRepoRequest,
-    GetQuarantineBulkTestStatusRequest, QuarantineConfig,
+    GetQuarantineBulkTestStatusRequest, QuarantineConfig, UpdateBundleUploadRequest,
 };
 use axum::{
     body::Bytes,
     extract::State,
     http::StatusCode,
     response::Response,
-    routing::{any, post, put},
+    routing::{any, patch, post, put},
     {Json, Router},
 };
 use tempfile::tempdir;
@@ -25,6 +25,7 @@ use tokio::{net::TcpListener, spawn};
 pub enum RequestPayload {
     CreateRepo(CreateRepoRequest),
     CreateBundleUpload(CreateBundleUploadRequest),
+    UpdateBundleUpload(UpdateBundleUploadRequest),
     GetQuarantineBulkTestStatus(GetQuarantineBulkTestStatusRequest),
     S3Upload(PathBuf),
 }
@@ -56,6 +57,10 @@ pub async fn spawn_mock_server() -> SharedMockServerState {
             "/v1/metrics/createBundleUpload",
             post(create_bundle_handler),
         )
+        .route(
+            "/v1/metrics/updateBundleUpload",
+            patch(update_bundle_handler),
+        )
         .route(
             "/v1/metrics/getQuarantineConfig",
             post(get_quarantining_config_handler),
@@ -116,6 +121,21 @@ async fn create_bundle_handler(
     })
 }
 
+#[axum::debug_handler]
+async fn update_bundle_handler(
+    State(state): State<SharedMockServerState>,
+    Json(update_bundle_upload_request): Json<UpdateBundleUploadRequest>,
+) -> Response<String> {
+    state
+        .requests
+        .lock()
+        .unwrap()
+        .push(RequestPayload::UpdateBundleUpload(
+            update_bundle_upload_request,
+        ));
+    Response::new(String::from("OK"))
+}
+
 #[axum::debug_handler]
 async fn get_quarantining_config_handler(
     State(state): State<SharedMockServerState>,