linkrs_server/http/handlers/
export.rs

1//! HTTP handler for data export operations
2
3use axum::body::{Body, Bytes};
4use axum::http::{header, StatusCode};
5use axum::{
6    extract::{Extension, Query, State},
7    response::Response,
8};
9use serde::Deserialize;
10use tokio_stream::wrappers::ReceiverStream;
11
12use super::query::result_size_limit;
13use crate::http::{error::HttpError, state::AppState};
14use crate::storage::{
15    StorageClient, StorageOperationContextOps, StorageSchemaContextOps, StorageSyncContextOps,
16};
17use crate::value::to_json as value_to_json;
18
19#[derive(Debug, Deserialize, utoipa::ToSchema)]
20pub struct ExportQuery {
21    pub space: Option<String>,
22    pub format: Option<String>,
23    pub query: Option<String>,
24    pub all: Option<String>,
25    pub session_id: Option<i64>,
26}
27
28#[utoipa::path(
29    get,
30    path = "/v1/export",
31    tag = "Export",
32    params(
33        ("space" = Option<String>, Query, description = "Space to export (reserved, session space applies)"),
34        ("format" = Option<String>, Query, description = "Export format: csv or jsonl"),
35        ("query" = Option<String>, Query, description = "Query selecting rows to export"),
36        ("all" = Option<String>, Query, description = "Reserved"),
37        ("session_id" = Option<i64>, Query, description = "Session id executing the export query")
38    ),
39    responses(
40        (status = 200, description = "Export file download"),
41        (status = 400, description = "Invalid export request"),
42        (status = 500, description = "Internal error")
43    )
44)]
45/// Re-execute a single query and stream the result out as a download.
46///
47/// Rows are pulled chunk-at-a-time and encoded straight into the response
48/// body, so peak memory stays flat regardless of result size. This is a
49/// fresh execution, not a snapshot of what the browser buffered: data may
50/// have changed between display and export. Only `csv` and `jsonl` are
51/// supported; `jsonl` carries one object per line because a JSON array
52/// cannot be appended to without backtracking.
53///
54/// Transport note: this route sits behind the session-header auth
55/// middleware, so browsers fetch it with script (which can attach the
56/// session header) and buffer the body before saving. A native anchor
57/// download is not used because it cannot send the header, and passing the
58/// session through the URL would leak it via history and logs.
59pub async fn export_data<
60    S: StorageClient
61        + StorageSchemaContextOps
62        + StorageSyncContextOps
63        + StorageOperationContextOps
64        + Clone
65        + Send
66        + Sync
67        + 'static,
68>(
69    State(state): State<AppState<S>>,
70    Extension(caller_session_id): Extension<i64>,
71    Query(params): Query<ExportQuery>,
72) -> Result<Response, HttpError> {
73    let format = params.format.as_deref().unwrap_or("csv");
74    let (content_type, file_ext) = match format {
75        "csv" => ("text/csv", "csv"),
76        "jsonl" => ("application/x-ndjson", "jsonl"),
77        other => {
78            return Err(HttpError::bad_request(format!(
79                "Unsupported export format '{other}'; use csv or jsonl"
80            )));
81        }
82    };
83    let query = params
84        .query
85        .as_deref()
86        .map(str::trim)
87        .filter(|q| !q.is_empty())
88        .ok_or_else(|| HttpError::bad_request("Export requires a non-empty `query` parameter"))?;
89    let session_id = params
90        .session_id
91        .ok_or_else(|| HttpError::bad_request("Export requires a `session_id` parameter"))?;
92    super::authz::require_session_owner_or_admin(&state, caller_session_id, session_id)?;
93
94    let server = state.server.clone();
95    let row_limit = result_size_limit(&state);
96
97    // Set up execution before responding so setup failures surface as HTTP
98    // errors instead of a truncated download.
99    let stream_result = server
100        .get_graph_service()
101        .execute_stream(session_id, query)
102        .await
103        .map_err(|e| {
104            if e.message().contains("Invalid session ID") {
105                HttpError::bad_request(e.message().to_string())
106            } else {
107                HttpError::InternalError(e.message().to_string())
108            }
109        })?;
110
111    let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(100);
112    let as_csv = format == "csv";
113    tokio::task::spawn_blocking(move || {
114        let mut columns: Option<Vec<String>> = stream_result.column_names();
115        let mut header_sent = false;
116        let mut exported_rows: usize = 0;
117
118        // Fail the download when the configured row ceiling is hit: the
119        // client discards the partial file and reports the limit.
120        let send_failure = |tx: &tokio::sync::mpsc::Sender<Result<Bytes, std::io::Error>>,
121                            message: String| {
122            stream_result.cancel();
123            let _ = tx.blocking_send(Err(std::io::Error::other(message)));
124        };
125
126        loop {
127            match stream_result.next_chunk() {
128                Ok(Some(chunk)) => {
129                    let names = chunk.col_names();
130                    if columns.is_none() {
131                        columns = Some(names.clone());
132                    }
133                    if as_csv && !header_sent {
134                        let header = csv_header_line(columns.as_ref());
135                        if tx
136                            .blocking_send(Ok(Bytes::from(format!("{header}\n"))))
137                            .is_err()
138                        {
139                            stream_result.cancel();
140                            return;
141                        }
142                        header_sent = true;
143                    }
144                    for row in chunk.rows {
145                        if let Some(limit) = row_limit {
146                            if exported_rows >= limit {
147                                send_failure(
148                                    &tx,
149                                    format!("Result size limit of {limit} rows exceeded"),
150                                );
151                                return;
152                            }
153                        }
154                        let values: Vec<serde_json::Value> =
155                            row.into_iter().map(value_to_json).collect();
156                        let line = if as_csv {
157                            encode_csv_line(&names, &values)
158                        } else {
159                            encode_jsonl_line(&names, &values)
160                        };
161                        exported_rows += 1;
162                        if tx.blocking_send(Ok(Bytes::from(line))).is_err() {
163                            stream_result.cancel();
164                            return;
165                        }
166                    }
167                }
168                Ok(None) => {
169                    // Empty result still yields a header-only CSV.
170                    if as_csv && !header_sent {
171                        let header = csv_header_line(columns.as_ref());
172                        let _ = tx.blocking_send(Ok(Bytes::from(format!("{header}\n"))));
173                    }
174                    return;
175                }
176                Err(e) => {
177                    send_failure(&tx, format!("Export query failed: {e}"));
178                    return;
179                }
180            }
181        }
182    });
183
184    let response = Response::builder()
185        .status(StatusCode::OK)
186        .header(header::CONTENT_TYPE, content_type)
187        .header(
188            header::CONTENT_DISPOSITION,
189            format!("attachment; filename=\"export.{file_ext}\""),
190        )
191        .body(Body::from_stream(ReceiverStream::new(rx)))
192        .map_err(|e| HttpError::InternalError(format!("Failed to build response: {}", e)))?;
193
194    Ok(response)
195}
196
197/// Render the CSV header from the column list.
198fn csv_header_line(columns: Option<&Vec<String>>) -> String {
199    columns
200        .map(|cols| {
201            cols.iter()
202                .map(|c| csv_escape(c))
203                .collect::<Vec<_>>()
204                .join(",")
205        })
206        .unwrap_or_default()
207}
208
209/// Render one CSV line: cells in column order, objects as compact JSON.
210fn encode_csv_line(columns: &[String], values: &[serde_json::Value]) -> String {
211    let mut cells = Vec::with_capacity(columns.len());
212    for (i, _col) in columns.iter().enumerate() {
213        let cell = values.get(i).map(cell_text).unwrap_or_default();
214        cells.push(csv_escape(&cell));
215    }
216    cells.join(",") + "\n"
217}
218
219/// Render one JSONL line: a column-name to value object.
220fn encode_jsonl_line(columns: &[String], values: &[serde_json::Value]) -> String {
221    let obj: serde_json::Map<String, serde_json::Value> = columns
222        .iter()
223        .enumerate()
224        .map(|(i, col)| {
225            (
226                col.clone(),
227                values.get(i).cloned().unwrap_or(serde_json::Value::Null),
228            )
229        })
230        .collect();
231    serde_json::to_string(&serde_json::Value::Object(obj)).unwrap_or_else(|_| "{}".to_string())
232        + "\n"
233}
234
235/// Plain-text rendering of a cell: nulls vanish, objects become JSON.
236fn cell_text(value: &serde_json::Value) -> String {
237    match value {
238        serde_json::Value::Null => String::new(),
239        serde_json::Value::String(s) => s.clone(),
240        serde_json::Value::Bool(b) => b.to_string(),
241        serde_json::Value::Number(n) => n.to_string(),
242        serde_json::Value::Array(_) | serde_json::Value::Object(_) => value.to_string(),
243    }
244}
245
246/// Quote a field when it carries a comma, quote, or line break.
247fn csv_escape(field: &str) -> String {
248    if field.contains(',') || field.contains('"') || field.contains('\n') || field.contains('\r') {
249        format!("\"{}\"", field.replace('"', "\"\""))
250    } else {
251        field.to_string()
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258    use serde_json::json;
259
260    #[test]
261    fn csv_escapes_delimiters_quotes_and_line_breaks() {
262        assert_eq!(csv_escape("plain"), "plain");
263        assert_eq!(csv_escape("a,b"), "\"a,b\"");
264        assert_eq!(csv_escape("say \"hi\""), "\"say \"\"hi\"\"\"");
265        assert_eq!(csv_escape("line1\nline2"), "\"line1\nline2\"");
266        assert_eq!(csv_escape("line1\r\nline2"), "\"line1\r\nline2\"");
267        assert_eq!(csv_escape(""), "");
268    }
269
270    #[test]
271    fn csv_header_renders_once_for_empty_results() {
272        assert_eq!(
273            csv_header_line(Some(&vec!["a".to_string(), "b,c".to_string()])),
274            "a,\"b,c\""
275        );
276        assert_eq!(csv_header_line(None), "");
277    }
278
279    #[test]
280    fn csv_line_orders_cells_and_stringifies_objects() {
281        let columns = vec![
282            "name".to_string(),
283            "meta".to_string(),
284            "missing".to_string(),
285        ];
286        let values = vec![json!("a,b"), json!({"k": 1})];
287        assert_eq!(
288            encode_csv_line(&columns, &values),
289            "\"a,b\",\"{\"\"k\"\":1}\",\n"
290        );
291    }
292
293    #[test]
294    fn jsonl_line_maps_columns_to_values() {
295        let columns = vec!["a".to_string(), "b".to_string()];
296        let values = vec![json!(1), json!("x")];
297        assert_eq!(
298            encode_jsonl_line(&columns, &values),
299            "{\"a\":1,\"b\":\"x\"}\n"
300        );
301    }
302}
linkrs_server/http/handlers/export.rs

linkrs_server/http/handlers/
export.rs
