linkrs_server/http/handlers/
import.rs

1//! HTTP handler for data import operations
2//!
3//! Parses an uploaded file (CSV with header, JSON array, or JSONL) into DML
4//! statements and loads them through the group-commit batch window:
5//! each `batch_size` group of statements shares one write timestamp, one WAL
6//! fsync, and one commit point, so a large import amortizes fsync across
7//! statements instead of syncing per row.
8
9use axum::{
10    extract::{Extension, Multipart, State},
11    response::Json as JsonResponse,
12};
13use serde::{Deserialize, Serialize};
14
15use crate::http::{error::HttpError, state::AppState};
16use crate::storage::{
17    StorageClient, StorageOperationContextOps, StorageSchemaContextOps, StorageSyncContextOps,
18};
19
20#[derive(Debug, Deserialize, utoipa::ToSchema)]
21pub struct ImportForm {
22    pub space: String,
23    pub format: String,
24    pub target_type: String,
25    pub target_name: String,
26    pub batch_size: Option<usize>,
27}
28
29#[derive(Debug, Serialize, utoipa::ToSchema)]
30pub struct ImportResponse {
31    pub success: bool,
32    pub message: String,
33    pub rows_imported: usize,
34    pub rows_failed: usize,
35}
36
37#[derive(Debug, Serialize, utoipa::ToSchema)]
38pub struct ImportStatusResponse {
39    pub job_id: String,
40    pub status: String,
41    pub rows_imported: usize,
42    pub rows_failed: usize,
43    pub started_at: Option<String>,
44    pub completed_at: Option<String>,
45}
46
47/// Build one `INSERT VERTEX <tag> (fields) VALUES "<vid>":(values)` statement
48/// from a CSV record.
49fn vertex_statement_from_csv(
50    target_name: &str,
51    headers: &csv::StringRecord,
52    record: &csv::StringRecord,
53) -> Result<String, HttpError> {
54    let fields: Vec<&str> = headers.iter().collect();
55    let values: Vec<String> = record.iter().map(format_value).collect();
56    let vid = record
57        .get(0)
58        .ok_or_else(|| HttpError::BadRequest("CSV vertex rows must start with the VID".into()))?
59        .to_string();
60    Ok(format!(
61        "INSERT VERTEX {} ({}) VALUES \"{}\":({})",
62        target_name,
63        fields.join(", "),
64        vid,
65        values.join(", ")
66    ))
67}
68
69/// Build one `INSERT EDGE <type> (fields) VALUES "<src>"->"<dst>":(values)`
70/// statement from a CSV record. Columns 0/1 are the source/destination VIDs,
71/// remaining columns are properties.
72fn edge_statement_from_csv(
73    target_name: &str,
74    headers: &csv::StringRecord,
75    record: &csv::StringRecord,
76) -> Result<String, HttpError> {
77    let src = record.get(0).ok_or_else(|| {
78        HttpError::BadRequest("CSV edge rows must start with the source VID".into())
79    })?;
80    let dst = record
81        .get(1)
82        .ok_or_else(|| HttpError::BadRequest("CSV edge rows must have a destination VID".into()))?;
83    let fields: Vec<&str> = headers.iter().skip(2).collect();
84    let values: Vec<String> = record.iter().skip(2).map(format_value).collect();
85    Ok(format!(
86        "INSERT EDGE {} ({}) VALUES \"{}\"->\"{}\":({})",
87        target_name,
88        fields.join(", "),
89        src,
90        dst,
91        values.join(", ")
92    ))
93}
94
95/// Build one INSERT statement from a JSON object. Vertex objects use `_id`
96/// as the VID; edge objects use `_src`/`_dst`.
97fn statement_from_json_object(
98    target_type: &str,
99    target_name: &str,
100    obj: &serde_json::Map<String, serde_json::Value>,
101) -> Result<String, HttpError> {
102    match target_type {
103        "tag" | "vertex" => {
104            let vid = obj
105                .get("_id")
106                .and_then(|v| v.as_str())
107                .ok_or_else(|| HttpError::BadRequest("JSON vertex objects must have _id".into()))?;
108            let mut fields = Vec::new();
109            let mut values = Vec::new();
110            for (key, value) in obj {
111                if key == "_id" {
112                    continue;
113                }
114                fields.push(key.as_str());
115                values.push(json_value_to_gql(value));
116            }
117            Ok(format!(
118                "INSERT VERTEX {} ({}) VALUES \"{}\":({})",
119                target_name,
120                fields.join(", "),
121                vid,
122                values.join(", ")
123            ))
124        }
125        "edge" => {
126            let src = obj
127                .get("_src")
128                .and_then(|v| v.as_str())
129                .ok_or_else(|| HttpError::BadRequest("JSON edge objects must have _src".into()))?;
130            let dst = obj
131                .get("_dst")
132                .and_then(|v| v.as_str())
133                .ok_or_else(|| HttpError::BadRequest("JSON edge objects must have _dst".into()))?;
134            let mut fields = Vec::new();
135            let mut values = Vec::new();
136            for (key, value) in obj {
137                if key == "_src" || key == "_dst" {
138                    continue;
139                }
140                fields.push(key.as_str());
141                values.push(json_value_to_gql(value));
142            }
143            Ok(format!(
144                "INSERT EDGE {} ({}) VALUES \"{}\"->\"{}\":({})",
145                target_name,
146                fields.join(", "),
147                src,
148                dst,
149                values.join(", ")
150            ))
151        }
152        other => Err(HttpError::BadRequest(format!(
153            "Unsupported target_type: {other} (expected 'tag' or 'edge')"
154        ))),
155    }
156}
157
158fn format_value(value: &str) -> String {
159    if value.is_empty() {
160        return "NULL".to_string();
161    }
162    format!("\"{}\"", value.replace('"', "\\\""))
163}
164
165fn json_value_to_gql(value: &serde_json::Value) -> String {
166    match value {
167        serde_json::Value::Null => "NULL".to_string(),
168        serde_json::Value::Bool(b) => b.to_string(),
169        serde_json::Value::Number(n) => n.to_string(),
170        serde_json::Value::String(s) => format!("\"{}\"", s.replace('"', "\\\"")),
171        serde_json::Value::Array(arr) => {
172            let items: Vec<String> = arr.iter().map(json_value_to_gql).collect();
173            format!("[{}]", items.join(", "))
174        }
175        serde_json::Value::Object(_) => {
176            format!("\"{}\"", serde_json::to_string(value).unwrap_or_default())
177        }
178    }
179}
180
181/// Parse the uploaded file bytes into DML statements.
182fn parse_statements(
183    format: &str,
184    target_type: &str,
185    target_name: &str,
186    data: &[u8],
187) -> Result<Vec<String>, HttpError> {
188    match format {
189        "csv" => {
190            let mut reader = csv::ReaderBuilder::new()
191                .has_headers(true)
192                .from_reader(data);
193            let headers = reader
194                .headers()
195                .map_err(|e| HttpError::BadRequest(format!("CSV header parse failed: {e}")))?
196                .clone();
197            let mut statements = Vec::new();
198            for record in reader.records() {
199                let record = record
200                    .map_err(|e| HttpError::BadRequest(format!("CSV record parse failed: {e}")))?;
201                let statement = match target_type {
202                    "edge" => edge_statement_from_csv(target_name, &headers, &record)?,
203                    _ => vertex_statement_from_csv(target_name, &headers, &record)?,
204                };
205                statements.push(statement);
206            }
207            Ok(statements)
208        }
209        "json" | "jsonl" => {
210            let text = std::str::from_utf8(data)
211                .map_err(|e| HttpError::BadRequest(format!("File is not valid UTF-8: {e}")))?;
212            let mut statements = Vec::new();
213            if format == "jsonl" {
214                for (idx, line) in text.lines().enumerate() {
215                    let line = line.trim();
216                    if line.is_empty() {
217                        continue;
218                    }
219                    let obj: serde_json::Map<String, serde_json::Value> =
220                        serde_json::from_str(line).map_err(|e| {
221                            HttpError::BadRequest(format!(
222                                "JSONL line {} parse failed: {e}",
223                                idx + 1
224                            ))
225                        })?;
226                    statements.push(statement_from_json_object(target_type, target_name, &obj)?);
227                }
228            } else {
229                let value: serde_json::Value = serde_json::from_str(text)
230                    .map_err(|e| HttpError::BadRequest(format!("JSON parse failed: {e}")))?;
231                let objects: Vec<&serde_json::Map<String, serde_json::Value>> = match &value {
232                    serde_json::Value::Array(items) => items
233                        .iter()
234                        .map(|item| {
235                            item.as_object().ok_or_else(|| {
236                                HttpError::BadRequest(
237                                    "JSON array items must be objects".to_string(),
238                                )
239                            })
240                        })
241                        .collect::<Result<_, _>>()?,
242                    serde_json::Value::Object(_) => {
243                        vec![value.as_object().expect("checked above")]
244                    }
245                    _ => {
246                        return Err(HttpError::BadRequest(
247                            "JSON import expects an object or an array of objects".to_string(),
248                        ))
249                    }
250                };
251                for obj in objects {
252                    statements.push(statement_from_json_object(target_type, target_name, obj)?);
253                }
254            }
255            Ok(statements)
256        }
257        other => Err(HttpError::BadRequest(format!(
258            "Unsupported format: {other} (expected 'csv', 'json' or 'jsonl')"
259        ))),
260    }
261}
262
263#[utoipa::path(
264    post,
265    path = "/v1/import",
266    tag = "Import",
267    responses(
268        (status = 200, body = ImportResponse, description = "File imported"),
269        (status = 500, description = "Internal error")
270    )
271)]
272pub async fn import_file<
273    S: StorageClient
274        + StorageSchemaContextOps
275        + StorageSyncContextOps
276        + StorageOperationContextOps
277        + crate::storage::AutoCommitBatchOps
278        + crate::storage::AutoCommitGroupOps
279        + Clone
280        + Send
281        + Sync
282        + 'static,
283>(
284    State(state): State<AppState<S>>,
285    Extension(session_id): Extension<i64>,
286    mut multipart: Multipart,
287) -> Result<JsonResponse<ImportResponse>, HttpError> {
288    let mut space = String::new();
289    let mut format = "csv".to_string();
290    let mut target_type = "tag".to_string();
291    let mut target_name = String::new();
292    let mut batch_size = Some(1000usize);
293    let mut file_data: Option<Vec<u8>> = None;
294
295    while let Some(field) = multipart
296        .next_field()
297        .await
298        .map_err(|e| HttpError::BadRequest(format!("Multipart error: {e}")))?
299    {
300        let name = field.name().unwrap_or("").to_string();
301
302        match name.as_str() {
303            "space" => {
304                space = field
305                    .text()
306                    .await
307                    .map_err(|e| HttpError::BadRequest(format!("Invalid space field: {e}")))?;
308            }
309            "format" => {
310                format = field
311                    .text()
312                    .await
313                    .map_err(|e| HttpError::BadRequest(format!("Invalid format field: {e}")))?;
314            }
315            "target_type" => {
316                target_type = field.text().await.map_err(|e| {
317                    HttpError::BadRequest(format!("Invalid target_type field: {e}"))
318                })?;
319            }
320            "target_name" => {
321                target_name = field.text().await.map_err(|e| {
322                    HttpError::BadRequest(format!("Invalid target_name field: {e}"))
323                })?;
324            }
325            "batch_size" => {
326                let bs = field
327                    .text()
328                    .await
329                    .map_err(|e| HttpError::BadRequest(format!("Invalid batch_size field: {e}")))?;
330                batch_size = bs.parse().ok();
331            }
332            "file" => {
333                file_data = Some(
334                    field
335                        .bytes()
336                        .await
337                        .map_err(|e| HttpError::BadRequest(format!("Failed to read file: {e}")))?
338                        .to_vec(),
339                );
340            }
341            _ => {}
342        }
343    }
344
345    if space.is_empty() {
346        return Err(HttpError::BadRequest("Missing 'space' field".to_string()));
347    }
348    if target_name.is_empty() {
349        return Err(HttpError::BadRequest(
350            "Missing 'target_name' field".to_string(),
351        ));
352    }
353    let file_data =
354        file_data.ok_or_else(|| HttpError::BadRequest("Missing 'file' field".to_string()))?;
355
356    let statements = parse_statements(&format, &target_type, &target_name, &file_data)?;
357    if statements.is_empty() {
358        return Ok(JsonResponse(ImportResponse {
359            success: true,
360            message: "File contained no importable rows".to_string(),
361            rows_imported: 0,
362            rows_failed: 0,
363        }));
364    }
365
366    let group_size = batch_size.unwrap_or(1000).max(1);
367    let outcomes = state
368        .server
369        .get_graph_service()
370        .execute_batch_grouped(session_id, &statements, group_size)
371        .await;
372
373    let mut rows_imported = 0;
374    let mut rows_failed = 0;
375    for outcome in &outcomes {
376        match outcome {
377            Ok(_) => rows_imported += 1,
378            Err(_) => rows_failed += 1,
379        }
380    }
381
382    let success = rows_failed == 0;
383    Ok(JsonResponse(ImportResponse {
384        success,
385        message: if success {
386            format!("Imported {rows_imported} rows into {target_type} '{target_name}'")
387        } else {
388            format!("Imported {rows_imported} rows, {rows_failed} failed (group size {group_size})")
389        },
390        rows_imported,
391        rows_failed,
392    }))
393}
394
395#[utoipa::path(
396    get,
397    path = "/v1/import/{id}",
398    tag = "Import",
399    params(("id" = String, Path, description = "Import job id")),
400    responses(
401        (status = 200, body = ImportStatusResponse, description = "Import job status"),
402        (status = 500, description = "Internal error")
403    )
404)]
405pub async fn import_status<
406    S: StorageClient
407        + StorageSchemaContextOps
408        + StorageSyncContextOps
409        + StorageOperationContextOps
410        + Clone
411        + Send
412        + Sync
413        + 'static,
414>(
415    State(_state): State<AppState<S>>,
416) -> Result<JsonResponse<ImportStatusResponse>, HttpError> {
417    // Imports are executed synchronously; the status endpoint reports the
418    // completion contract kept for API compatibility.
419    Ok(JsonResponse(ImportStatusResponse {
420        job_id: String::new(),
421        status: "unknown".to_string(),
422        rows_imported: 0,
423        rows_failed: 0,
424        started_at: None,
425        completed_at: None,
426    }))
427}
428
429#[cfg(test)]
430mod tests {
431    use super::*;
432
433    #[test]
434    fn parse_csv_vertex_statements() {
435        let csv = "vid,name,age\n1,Alice,30\n2,Bob,25\n";
436        let statements = parse_statements("csv", "tag", "Person", csv.as_bytes()).unwrap();
437        assert_eq!(statements.len(), 2);
438        assert_eq!(
439            statements[0],
440            "INSERT VERTEX Person (vid, name, age) VALUES \"1\":(\"1\", \"Alice\", \"30\")"
441        );
442        assert_eq!(
443            statements[1],
444            "INSERT VERTEX Person (vid, name, age) VALUES \"2\":(\"2\", \"Bob\", \"25\")"
445        );
446    }
447
448    #[test]
449    fn parse_csv_edge_statements() {
450        let csv = "src,dst,weight\n1,2,0.5\n3,4,1.5\n";
451        let statements = parse_statements("csv", "edge", "Follows", csv.as_bytes()).unwrap();
452        assert_eq!(statements.len(), 2);
453        assert_eq!(
454            statements[0],
455            "INSERT EDGE Follows (weight) VALUES \"1\"->\"2\":(\"0.5\")"
456        );
457    }
458
459    #[test]
460    fn parse_json_array_vertex_statements() {
461        let json = r#"[{"_id":"v1","name":"Alice","age":30},{"_id":"v2","name":"Bob","age":25}]"#;
462        let statements = parse_statements("json", "tag", "Person", json.as_bytes()).unwrap();
463        assert_eq!(statements.len(), 2);
464        // serde_json::Map iterates keys in sorted order.
465        assert_eq!(
466            statements[0],
467            "INSERT VERTEX Person (age, name) VALUES \"v1\":(30, \"Alice\")"
468        );
469    }
470
471    #[test]
472    fn parse_jsonl_edge_statements() {
473        let jsonl = "{\"_src\":\"1\",\"_dst\":\"2\",\"weight\":0.5}\n{\"_src\":\"3\",\"_dst\":\"4\",\"weight\":1.5}\n";
474        let statements = parse_statements("jsonl", "edge", "Follows", jsonl.as_bytes()).unwrap();
475        assert_eq!(statements.len(), 2);
476        assert_eq!(
477            statements[0],
478            "INSERT EDGE Follows (weight) VALUES \"1\"->\"2\":(0.5)"
479        );
480    }
481
482    #[test]
483    fn parse_escapes_quotes_in_values() {
484        let csv = "vid,name\n1,\"Al\"\"ice\"\n";
485        let statements = parse_statements("csv", "tag", "Person", csv.as_bytes()).unwrap();
486        assert_eq!(statements.len(), 1);
487        // The embedded quote is escaped for the INSERT statement literal.
488        assert!(statements[0].contains("Al\\\"ice"), "{}", statements[0]);
489    }
490
491    #[test]
492    fn parse_rejects_unsupported_format() {
493        assert!(parse_statements("yaml", "tag", "Person", b"a: 1").is_err());
494    }
495
496    #[test]
497    fn parse_rejects_json_edge_without_src() {
498        let json = r#"[{"_dst":"2"}]"#;
499        assert!(parse_statements("json", "edge", "Follows", json.as_bytes()).is_err());
500    }
501
502    #[test]
503    fn parse_empty_csv_yields_no_statements() {
504        let statements = parse_statements("csv", "tag", "Person", b"vid,name\n").unwrap();
505        assert!(statements.is_empty());
506    }
507}
linkrs_server/http/handlers/import.rs

linkrs_server/http/handlers/
import.rs
