linkrs_server/http/handlers/
migration_progress.rs

1use std::collections::HashMap;
2use std::convert::Infallible;
3use std::sync::{Arc, RwLock};
4
5use axum::extract::{Path, Query, State};
6use axum::response::sse::{Event, Sse};
7use tokio_stream::wrappers::BroadcastStream;
8use tokio_stream::StreamExt;
9
10use crate::http::state::AppState;
11use crate::storage::{
12    StorageClient, StorageOperationContextOps, StorageSchemaContextOps, StorageSyncContextOps,
13};
14use linkrs_migration::MigrationEvent;
15
16type Sender = tokio::sync::broadcast::Sender<MigrationEvent>;
17type Receiver = tokio::sync::broadcast::Receiver<MigrationEvent>;
18
19static HUB: std::sync::OnceLock<Arc<RwLock<HashMap<String, Sender>>>> = std::sync::OnceLock::new();
20
21fn hub() -> Arc<RwLock<HashMap<String, Sender>>> {
22    HUB.get_or_init(|| Arc::new(RwLock::new(HashMap::new())))
23        .clone()
24}
25
26fn hub_key(space: &str, label: &str, is_edge: bool) -> String {
27    format!("{}/{}/{}", space, label, is_edge)
28}
29
30pub fn get_or_create_sender(space: &str, label: &str, is_edge: bool) -> Sender {
31    let key = hub_key(space, label, is_edge);
32    let hub_arc = hub();
33    {
34        let read = hub_arc.read().unwrap();
35        if let Some(s) = read.get(&key) {
36            return s.clone();
37        }
38    }
39    let mut write = hub_arc.write().unwrap();
40    if let Some(s) = write.get(&key) {
41        return s.clone();
42    }
43    let (tx, _rx) = tokio::sync::broadcast::channel(100);
44    write.insert(key, tx.clone());
45    tx
46}
47
48pub fn subscribe(space: &str, label: &str, is_edge: bool) -> Receiver {
49    get_or_create_sender(space, label, is_edge).subscribe()
50}
51
52/// Adapt a broadcast sender into an `EventSubscriptions` callback so a
53/// running migration fans its lifecycle events out to SSE subscribers.
54pub fn event_bridge(sender: Sender) -> Arc<dyn Fn(&MigrationEvent) + Send + Sync> {
55    Arc::new(move |event: &MigrationEvent| {
56        let _ = sender.send(event.clone());
57    })
58}
59
60fn migration_event_to_sse(event: MigrationEvent) -> Event {
61    match event {
62        MigrationEvent::Started { plan } => {
63            let data = serde_json::json!({
64                "type": "started",
65                "plan_hash": plan.plan_hash,
66                "space": plan.target.space,
67                "label": plan.target.label,
68                "is_edge": plan.target.is_edge,
69            });
70            Event::default().event("started").data(data.to_string())
71        }
72        MigrationEvent::StepStarted { step_idx } => {
73            let data = serde_json::json!({
74                "type": "step_started",
75                "step_idx": step_idx,
76            });
77            Event::default()
78                .event("step_started")
79                .data(data.to_string())
80        }
81        MigrationEvent::StepCompleted { step_idx, rows } => {
82            let data = serde_json::json!({
83                "type": "step_completed",
84                "step_idx": step_idx,
85                "rows": rows,
86            });
87            Event::default()
88                .event("step_completed")
89                .data(data.to_string())
90        }
91        MigrationEvent::Completed { report } => {
92            let data = serde_json::json!({
93                "type": "completed",
94                "success": report.success,
95                "steps_completed": report.steps_completed,
96                "rows_migrated": report.rows_migrated,
97                "errors": report.errors,
98            });
99            Event::default().event("completed").data(data.to_string())
100        }
101        MigrationEvent::Failed { error } => {
102            let data = serde_json::json!({
103                "type": "failed",
104                "error": error,
105            });
106            Event::default().event("failed").data(data.to_string())
107        }
108        MigrationEvent::RolledBack { report } => {
109            let data = serde_json::json!({
110                "type": "rolled_back",
111                "success": report.success,
112                "steps_completed": report.steps_completed,
113                "rows_migrated": report.rows_migrated,
114            });
115            Event::default().event("rolled_back").data(data.to_string())
116        }
117    }
118}
119
120#[derive(serde::Deserialize, utoipa::ToSchema)]
121pub struct MigrationProgressQuery {
122    pub is_edge: Option<bool>,
123}
124
125#[utoipa::path(
126    get,
127    path = "/v1/migration/stream/{space}/{label}",
128    tag = "Migration",
129    params(
130        ("space" = String, Path, description = "Space name"),
131        ("label" = String, Path, description = "Tag or edge type name"),
132        ("is_edge" = Option<bool>, Query, description = "Whether the label is an edge type")
133    ),
134    responses(
135        (status = 200, description = "Migration progress event stream", content_type = "text/event-stream")
136    )
137)]
138/// SSE endpoint to stream migration progress for a given space/label.
139///
140/// Example: GET /v1/migration/stream/{space}/{label}?is_edge=false
141pub async fn migration_progress_stream<
142    S: StorageClient
143        + StorageSchemaContextOps
144        + StorageSyncContextOps
145        + StorageOperationContextOps
146        + Clone
147        + Send
148        + Sync
149        + 'static,
150>(
151    State(_state): State<AppState<S>>,
152    Path((space, label)): Path<(String, String)>,
153    Query(query): Query<MigrationProgressQuery>,
154) -> Sse<impl tokio_stream::Stream<Item = Result<Event, Infallible>> + Send + 'static> {
155    let is_edge = query.is_edge.unwrap_or(false);
156    let rx = subscribe(&space, &label, is_edge);
157
158    let stream = BroadcastStream::new(rx).filter_map(|res| match res {
159        Ok(ev) => Some(Ok(migration_event_to_sse(ev))),
160        Err(tokio_stream::wrappers::errors::BroadcastStreamRecvError::Lagged(_)) => None,
161    });
162
163    Sse::new(stream).keep_alive(
164        axum::response::sse::KeepAlive::new()
165            .interval(std::time::Duration::from_secs(10))
166            .text("keepalive"),
167    )
168}
linkrs_server/http/handlers/migration_progress.rs

linkrs_server/http/handlers/
migration_progress.rs
