From 73bea35d4c2ecb43085fe2cef9a969ef7dbbe95c Mon Sep 17 00:00:00 2001
From: Stuart Axelbrooke <stuart@axelbrooke.com>
Date: Mon, 4 Aug 2025 11:33:57 -0700
Subject: [PATCH] Add mermaid endpoint

---
 databuild/service/handlers.rs | 73 +++++++++++++++++++++++++++++++++++
 databuild/service/mod.rs      | 13 +------
 2 files changed, 74 insertions(+), 12 deletions(-)

diff --git a/databuild/service/handlers.rs b/databuild/service/handlers.rs
index 1c44f7e..680e4c6 100644
--- a/databuild/service/handlers.rs
+++ b/databuild/service/handlers.rs
@@ -1679,4 +1679,77 @@ pub async fn get_job_run_metrics(
             ))
         }
     }
+}
+
+/// Request for build mermaid diagram endpoint
+#[derive(Deserialize, JsonSchema)]
+pub struct BuildMermaidRequest {
+    pub build_request_id: String,
+}
+
+/// Response for build mermaid diagram endpoint
+#[derive(serde::Serialize, JsonSchema)]
+pub struct BuildMermaidResponse {
+    pub mermaid: String,
+}
+
+/// Get Mermaid diagram for a specific build request ID
+pub async fn get_build_mermaid_diagram(
+    State(service): State<ServiceState>,
+    Path(BuildMermaidRequest { build_request_id }): Path<BuildMermaidRequest>,
+) -> Result<Json<BuildMermaidResponse>, (StatusCode, Json<ErrorResponse>)> {
+    info!("Generating mermaid diagram for build request {}", build_request_id);
+    
+    // Get build events for this build request
+    let events = match service.event_log.get_build_request_events(&build_request_id, None).await {
+        Ok(events) => events,
+        Err(e) => {
+            error!("Failed to get build events for {}: {}", build_request_id, e);
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                Json(ErrorResponse {
+                    error: format!("Failed to get build events: {}", e),
+                }),
+            ));
+        }
+    };
+    
+    if events.is_empty() {
+        return Err((
+            StatusCode::NOT_FOUND,
+            Json(ErrorResponse {
+                error: "Build request not found".to_string(),
+            }),
+        ));
+    }
+    
+    // Find job graph event to get the graph structure
+    let job_graph = events.iter()
+        .find_map(|event| {
+            match &event.event_type {
+                Some(crate::build_event::EventType::JobGraphEvent(graph_event)) => {
+                    graph_event.job_graph.as_ref()
+                }
+                _ => None,
+            }
+        });
+    
+    match job_graph {
+        Some(graph) => {
+            // Generate mermaid diagram with current status
+            let mermaid_diagram = mermaid_utils::generate_mermaid_with_status(graph, &events);
+            
+            Ok(Json(BuildMermaidResponse {
+                mermaid: mermaid_diagram,
+            }))
+        }
+        None => {
+            Err((
+                StatusCode::NOT_FOUND,
+                Json(ErrorResponse {
+                    error: "No job graph found for this build request".to_string(),
+                }),
+            ))
+        }
+    }
 }
\ No newline at end of file
diff --git a/databuild/service/mod.rs b/databuild/service/mod.rs
index 4e77391..d6158b2 100644
--- a/databuild/service/mod.rs
+++ b/databuild/service/mod.rs
@@ -47,18 +47,6 @@ pub struct BuildRequestResponse {
     pub build_request_id: String,
 }
 
-#[derive(Debug, Serialize, Deserialize, JsonSchema)]
-pub struct BuildStatusResponse {
-    pub build_request_id: String,
-    pub status: String,
-    pub requested_partitions: Vec<String>,
-    pub created_at: i64,
-    pub updated_at: i64,
-    pub events: Vec<BuildEventSummary>,
-    pub job_graph: Option<serde_json::Value>,
-    pub mermaid_diagram: Option<String>,
-}
-
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
 pub struct BuildEventSummary {
     pub event_id: String,
@@ -272,6 +260,7 @@ impl BuildGraphService {
             .api_route("/api/v1/builds", post(handlers::submit_build_request))
             .api_route("/api/v1/builds", get(handlers::list_builds_repository))
             .api_route("/api/v1/builds/:build_request_id", get(handlers::get_build_detail))
+            .api_route("/api/v1/builds/:build_request_id/mermaid", get(handlers::get_build_mermaid_diagram))
             .api_route("/api/v1/builds/:build_request_id", delete(handlers::cancel_build_repository))
             .api_route("/api/v1/partitions", get(handlers::list_partitions_repository))
             .api_route("/api/v1/partitions/:partition_ref", get(handlers::get_partition_detail))