Update BundleUpload status for observability

cli/src/clients.rs

@@ -5,7 +5,7 @@ use anyhow::Context;
 
 use crate::types::{
     BundleUploadLocation, CreateBundleUploadRequest, CreateRepoRequest,
-    GetQuarantineBulkTestStatusRequest, QuarantineConfig, Repo,
+    GetQuarantineBulkTestStatusRequest, QuarantineConfig, UpdateBundleUploadRequest, Repo,
 };
 use crate::utils::status_code_help;
 
@@ -47,6 +47,46 @@ pub async fn create_trunk_repo(
 
     Ok(())
 }
+pub async fn update_bundle_upload_status(
+    origin: &str,
+    api_token: &str,
+    id: &str,
+    upload_status: &str,
+) -> anyhow::Result<()> {
+    let client = reqwest::Client::new();
+    let resp = match client
+        .patch(format!("{}/v1/metrics/updateBundleUpload/{}", origin, id))
+        .timeout(TRUNK_API_TIMEOUT)
+        .header(reqwest::header::CONTENT_TYPE, "application/json")
+        .header(TRUNK_API_TOKEN_HEADER, api_token)
+        .json(&UpdateBundleUploadRequest{
+            upload_status: upload_status.to_owned()
+        })
+        .send()
+        .await
+    {
+        Ok(resp) => resp,
+        Err(e) => return Err(anyhow::anyhow!(e).context("Failed to update bundle upload status")),
+    };
+
+    if resp.status().is_success() {
+        return Ok(());
+    }
+
+    if resp.status().is_client_error() {
+        return Err(anyhow::anyhow!(
+            "Failed to update bundle upload status. Client error: {}",
+            resp.status()
+        ));
+    }
+
+    log::warn!(
+        "Failed to update bundle upload status. {}: {}",
+        resp.status(),
+        status_code_help(resp.status())
+    );
+    Ok(())
+}
 
 pub async fn get_bundle_upload_location(
     origin: &str,

cli/src/main.rs

@@ -7,7 +7,7 @@ use tokio_retry::strategy::ExponentialBackoff;
 use tokio_retry::Retry;
 use trunk_analytics_cli::bundler::BundlerUtil;
 use trunk_analytics_cli::clients::{
-    create_trunk_repo, get_bundle_upload_location, put_bundle_to_s3,
+    create_trunk_repo, get_bundle_upload_location, put_bundle_to_s3, update_bundle_upload_status
 };
 use trunk_analytics_cli::codeowners::CodeOwners;
 use trunk_analytics_cli::constants::{
@@ -285,15 +285,30 @@ async fn run_upload(
     .await?;
 
     if dry_run {
+        if let Some(upload) = upload_op {
+            if let Err(e) = update_bundle_upload_status(&api_address, &token, &upload.id, "DRY_RUN").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);
     }
 
     if let Some(upload) = upload_op {
-        Retry::spawn(default_delay(), || {
-            put_bundle_to_s3(&upload.url, &bundle_time_file)
-        })
-        .await?;
+        let upload_status = match Retry::spawn(default_delay(), || put_bundle_to_s3(&upload.url, &bundle_time_file)).await {
+            Ok(_) => "UPLOAD_COMPLETE",
+            Err(e) => {
+                log::error!("Failed to upload bundle to S3 after retries: {}", e);
+                "UPLOAD_FAILURE"
+            }
+        };
+
+        match update_bundle_upload_status(&api_address, &token, &upload.id, upload_status).await {
+            Ok(_) => log::debug!("Updated bundle upload status to {}", upload_status),
+            Err(e) => log::warn!("Failed to update bundle upload status to {}: {}", upload_status, e),
+        }
     }
 
     let remote_urls = vec![repo.repo_url.clone()];

cli/src/types.rs

@@ -33,6 +33,12 @@ pub struct CreateBundleUploadRequest {
     pub org_url_slug: String,
 }
 
+#[derive(Debug, Serialize, Clone, Deserialize, PartialEq, Eq)]
+pub struct UpdateBundleUploadRequest {
+    #[serde(rename = "uploadStatus")]
+    pub upload_status: String,
+}
+
 #[derive(Debug, Serialize, Clone, Deserialize, PartialEq, Eq)]
 pub struct GetQuarantineBulkTestStatusRequest {
     pub repo: Repo,
@@ -103,6 +109,7 @@ pub struct QuarantineConfig {
 pub struct BundleUploadLocation {
     pub url: String,
     pub key: String,
+    pub id: String,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]

cli/tests/test_utils/mock_server.rs

@@ -8,20 +8,21 @@ use axum::body::Bytes;
 use axum::extract::State;
 use axum::http::StatusCode;
 use axum::response::Response;
-use axum::routing::{any, post, put};
+use axum::routing::{any, patch, post, put};
 use axum::{Json, Router};
 use tempfile::tempdir;
 use tokio::net::TcpListener;
 use tokio::spawn;
 use trunk_analytics_cli::types::{
     BundleUploadLocation, CreateBundleUploadRequest, CreateRepoRequest,
-    GetQuarantineBulkTestStatusRequest, QuarantineConfig,
+    GetQuarantineBulkTestStatusRequest, QuarantineConfig, UpdateBundleUploadRequest
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum RequestPayload {
     CreateRepo(CreateRepoRequest),
     CreateBundleUpload(CreateBundleUploadRequest),
+    UpdateBundleUpload(UpdateBundleUploadRequest),
     GetQuarantineBulkTestStatus(GetQuarantineBulkTestStatusRequest),
     S3Upload(PathBuf),
 }
@@ -54,6 +55,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),
@@ -112,9 +117,27 @@ async fn create_bundle_handler(
     Json(BundleUploadLocation {
         url: format!("{host}/s3upload"),
         key: String::from("unused"),
+        id: String::from("some-arbitrary-test-id"),
     })
 }
 
+#[allow(dead_code)] // TODO: move this to its own crate to get rid of the need for this
+#[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"))
+}
+
+
 #[allow(dead_code)] // TODO: move this to its own crate to get rid of the need for this
 #[axum::debug_handler]
 async fn get_quarantining_config_handler(