yggdrasil/api/database/
export.rs

1//! 数据导出：Axum 流式路由（大文件不走 JSON 序列化）。
2//!
3//! 鉴权走 [`crate::api::auth::require_admin_from_headers`]：从 cookie 取 session，校验管理员。
4//! 两种来源：按表导出（表名白名单）/ 按查询导出（强制只读，AST 校验）。
5//! 两种格式：CSV（COPY TO STDOUT 流式）/ SQL（逐行拼 INSERT）。
6
7#![allow(clippy::unused_unit)]
8
9use axum::body::Body;
10use axum::extract::Query;
11use axum::http::{header, HeaderMap, StatusCode};
12use axum::response::IntoResponse;
13use futures::StreamExt;
14use serde::Deserialize;
15use tokio_postgres::types::{FromSql, Type};
16
17use super::read_only::is_read_only_statement;
18
19#[derive(Deserialize)]
20pub struct ExportParams {
21    /// `table:<name>` 或 `query:<SELECT ...>`。
22    pub source: String,
23    /// `sql` 或 `csv`。
24    pub format: String,
25    /// 是否带表头（CSV）/ 列名（SQL INSERT）。默认 true。
26    pub include_columns: Option<bool>,
27}
28
29/// GET /api/database/export —— 流式导出。
30///
31/// 前端用 `window.open` 触发下载，因此保持 GET；路由挂
32/// [`crate::api::csrf::same_origin_read_middleware`] 拒绝跨站触发。
33pub async fn export_data(
34    headers: HeaderMap,
35    Query(params): Query<ExportParams>,
36) -> Result<impl IntoResponse, (StatusCode, String)> {
37    // 1. 鉴权：cookie → session → admin
38    crate::api::auth::require_admin_from_headers(&headers)
39        .await
40        .map_err(|(status, msg)| (status, msg.to_string()))?;
41
42    // 2. 解析来源 + 白名单/只读校验
43    let include_columns = params.include_columns.unwrap_or(true);
44    let source = parse_source(&params.source)?;
45
46    match params.format.as_str() {
47        "csv" => export_csv(source.sql, include_columns).await,
48        "sql" => export_sql(source, include_columns).await,
49        _ => Err((StatusCode::BAD_REQUEST, "不支持的格式".to_string())),
50    }
51    .map(|(body, content_type, filename)| {
52        let disposition = format!("attachment; filename=\"{}\"", filename);
53        let disposition_value = axum::http::HeaderValue::from_str(&disposition)
54            .unwrap_or_else(|_| axum::http::HeaderValue::from_static("attachment"));
55        let content_type_value = axum::http::HeaderValue::from_str(content_type)
56            .unwrap_or_else(|_| axum::http::HeaderValue::from_static("application/octet-stream"));
57        (
58            StatusCode::OK,
59            [
60                (header::CONTENT_TYPE, content_type_value),
61                (header::CONTENT_DISPOSITION, disposition_value),
62            ],
63            body,
64        )
65    })
66}
67
68/// 解析后的导出来源。
69#[derive(Debug)]
70struct ExportSource {
71    /// 内部 SQL。
72    sql: String,
73    /// SQL 导出的 INSERT 目标表名。
74    table_name: String,
75    /// 是否全是查询语句：查询能包进子查询逐列转 text，EXPLAIN 不能。
76    query_only: bool,
77}
78
79/// 解析导出来源。
80/// - `table:posts` → 校验表名合法后，`SELECT * FROM "posts"`（只读）+ 表名 "posts"。
81/// - `query:SELECT ...` → 校验为只读语句后原样返回，表名为 "export"。
82fn parse_source(source: &str) -> Result<ExportSource, (StatusCode, String)> {
83    if let Some(table) = source.strip_prefix("table:") {
84        // 表名白名单：仅允许标识符字符，防注入
85        let t = table.trim();
86        if t.is_empty() || !is_simple_ident(t) {
87            return Err((StatusCode::BAD_REQUEST, "无效的表名".to_string()));
88        }
89        Ok(ExportSource {
90            sql: format!("SELECT * FROM \"{}\"", t),
91            table_name: t.to_string(),
92            query_only: true,
93        })
94    } else if let Some(query) = source.strip_prefix("query:") {
95        // 只读校验：sqlparser 解析后所有语句都必须是无副作用的查询或非 ANALYZE 的 EXPLAIN
96        let dialect = sqlparser::dialect::PostgreSqlDialect {};
97        let parsed = sqlparser::parser::Parser::parse_sql(&dialect, query)
98            .map_err(|_| (StatusCode::BAD_REQUEST, "SQL 解析失败".to_string()))?;
99        if parsed.iter().any(|s| !is_read_only_statement(s)) {
100            return Err((
101                StatusCode::BAD_REQUEST,
102                "导出查询必须是只读（SELECT/EXPLAIN）".to_string(),
103            ));
104        }
105        Ok(ExportSource {
106            sql: query.to_string(),
107            table_name: "export".to_string(),
108            query_only: parsed
109                .iter()
110                .all(|s| matches!(s, sqlparser::ast::Statement::Query(_))),
111        })
112    } else {
113        Err((
114            StatusCode::BAD_REQUEST,
115            "source 必须是 table:<name> 或 query:<sql>".to_string(),
116        ))
117    }
118}
119
120/// 简单标识符校验（字母数字下划线，防 SQL 注入与路径穿越）。
121fn is_simple_ident(s: &str) -> bool {
122    !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
123}
124
125/// CSV 导出：用 COPY ... TO STDOUT WITH CSV 流式（大表不 OOM）。
126///
127/// 流式响应在函数返回后才消费完，无法在同一连接上以事务包住整个 COPY，
128/// 这里只依赖 [`parse_source`] 的 AST 只读校验。
129async fn export_csv(
130    source_sql: String,
131    include_columns: bool,
132) -> Result<(Body, &'static str, String), (StatusCode, String)> {
133    let client = crate::db::pool::get_conn()
134        .await
135        .map_err(|_| (StatusCode::SERVICE_UNAVAILABLE, "数据库不可用".to_string()))?;
136
137    let header_clause = if include_columns { "HEADER" } else { "" };
138    let copy_stmt = format!("COPY ({}) TO STDOUT WITH CSV {}", source_sql, header_clause);
139    let stream = client.copy_out(&copy_stmt).await.map_err(|e| {
140        tracing::error!("Export COPY failed: {}", e);
141        (StatusCode::INTERNAL_SERVER_ERROR, "COPY 失败".to_string())
142    })?;
143
144    // tokio_postgres 的 copy_out 流产出 Bytes，直接转 axum Body。
145    let mapped = stream.map(|res| res.map_err(std::io::Error::other));
146    let body = Body::from_stream(mapped);
147    Ok((body, "text/csv; charset=utf-8", "export.csv".to_string()))
148}
149
150/// 单个响应块的字节上限；缓冲达到后即推给客户端。
151const SQL_EXPORT_CHUNK_BYTES: usize = 64 * 1024;
152
153/// SQL 导出：逐行拼 INSERT（纯数据，不含 DDL）。
154///
155/// 流式读取：行在只读事务中逐批产出、按块推送，大表不再整表进内存。与 CSV 相同，
156/// 慢消费者可能触发 `statement_timeout` 导致下载被截断；客户端中途断开时任务结束，
157/// 事务随之丢弃并回滚。
158///
159/// 单元格一律取数据库的文本表示（[`text_cast_query`]），timestamptz、uuid、数组、
160/// jsonb、bytea 等类型原样导出，回灌时由目标列解析。
161async fn export_sql(
162    source: ExportSource,
163    include_columns: bool,
164) -> Result<(Body, &'static str, String), (StatusCode, String)> {
165    let mut client = crate::db::pool::get_conn()
166        .await
167        .map_err(|_| (StatusCode::SERVICE_UNAVAILABLE, "数据库不可用".to_string()))?;
168
169    // 事务与行流都借用连接，必须在同一个任务里；任务起步结果经 oneshot 回报，
170    // 这样查询失败仍能在产生响应前返回错误。
171    let (ready_tx, ready_rx) =
172        tokio::sync::oneshot::channel::<Result<(), (StatusCode, &'static str)>>();
173    let (tx_chunks, rx_chunks) =
174        tokio::sync::mpsc::channel::<Result<axum::body::Bytes, std::io::Error>>(8);
175    const QUERY_FAILED: (StatusCode, &str) = (StatusCode::INTERNAL_SERVER_ERROR, "查询失败");
176
177    tokio::spawn(async move {
178        let no_params = || std::iter::empty::<&(dyn tokio_postgres::types::ToSql + Sync)>();
179        // 纵深防御：AST 校验之外再放进只读事务，查询里若混入写操作由数据库拒绝。
180        // 事务未提交即被丢弃时 tokio_postgres 会发送 ROLLBACK，连接归还池时不残留事务。
181        let tx = match client.build_transaction().read_only(true).start().await {
182            Ok(tx) => tx,
183            Err(e) => {
184                tracing::error!("Export transaction failed: {}", e);
185                let _ = ready_tx.send(Err(QUERY_FAILED));
186                return;
187            }
188        };
189        // 先只做预备拿到原始列名与类型：类型决定字面量写法，列名进 INSERT 列清单。
190        let stmt = match tx.prepare(&source.sql).await {
191            Ok(stmt) => stmt,
192            Err(e) => {
193                tracing::error!("Export query failed: {}", e);
194                let _ = ready_tx.send(Err(QUERY_FAILED));
195                return;
196            }
197        };
198        let types: Vec<Type> = stmt.columns().iter().map(|c| c.type_().clone()).collect();
199        let col_clause = if include_columns && !types.is_empty() {
200            let names: Vec<&str> = stmt.columns().iter().map(|c| c.name()).collect();
201            format!("({})", names.join(", "))
202        } else {
203            String::new()
204        };
205        let rows = if source.query_only && !types.is_empty() {
206            tx.query_raw(
207                text_cast_query(&source.sql, types.len()).as_str(),
208                no_params(),
209            )
210            .await
211        } else if types.iter().all(<&str as FromSql>::accepts) {
212            // EXPLAIN 不能做子查询；它的 text 列可直接按文本读出。
213            tx.query_raw(&stmt, no_params()).await
214        } else {
215            let _ = ready_tx.send(Err((
216                StatusCode::BAD_REQUEST,
217                "该语句的结果列无法导出为 SQL",
218            )));
219            return;
220        };
221        let rows = match rows {
222            Ok(rows) => rows,
223            Err(e) => {
224                tracing::error!("Export query failed: {}", e);
225                let _ = ready_tx.send(Err(QUERY_FAILED));
226                return;
227            }
228        };
229        let _ = ready_tx.send(Ok(()));
230
231        let mut rows = std::pin::pin!(rows);
232        let mut buf = String::from("-- Yggdrasil 数据导出（仅数据，不含 schema）\n");
233        while let Some(row) = rows.next().await {
234            let vals = match row.and_then(|r| row_literals(&r, &types)) {
235                Ok(vals) => vals,
236                Err(e) => {
237                    // 与 CSV 中途出错一样截断下载。
238                    tracing::error!("Export row failed: {e}");
239                    let _ = tx_chunks.send(Err(std::io::Error::other("查询失败"))).await;
240                    return;
241                }
242            };
243            buf.push_str(&format!(
244                "INSERT INTO {} {} VALUES ({});\n",
245                source.table_name,
246                col_clause,
247                vals.join(", ")
248            ));
249            if buf.len() >= SQL_EXPORT_CHUNK_BYTES {
250                let chunk = axum::body::Bytes::from(std::mem::take(&mut buf));
251                // 客户端断开：直接结束，事务随 drop 回滚。
252                if tx_chunks.send(Ok(chunk)).await.is_err() {
253                    return;
254                }
255            }
256        }
257        if !buf.is_empty()
258            && tx_chunks
259                .send(Ok(axum::body::Bytes::from(buf)))
260                .await
261                .is_err()
262        {
263            return;
264        }
265        if let Err(e) = tx.rollback().await {
266            tracing::error!("Export transaction rollback failed: {e}");
267        }
268    });
269
270    match ready_rx.await {
271        Ok(Ok(())) => {}
272        Ok(Err((status, msg))) => return Err((status, msg.to_string())),
273        Err(_) => return Err((QUERY_FAILED.0, QUERY_FAILED.1.to_string())),
274    }
275
276    let body = Body::from_stream(tokio_stream::wrappers::ReceiverStream::new(rx_chunks));
277    Ok((
278        body,
279        "application/sql; charset=utf-8",
280        "export.sql".to_string(),
281    ))
282}
283
284/// 把查询包进子查询，按位置给列起别名后逐列 `::text`，由数据库输出任意类型的文本表示；
285/// 按位置取列，重名列也不冲突。
286fn text_cast_query(sql: &str, columns: usize) -> String {
287    let aliases: Vec<String> = (1..=columns).map(|i| format!("c{i}")).collect();
288    let casts: Vec<String> = aliases.iter().map(|a| format!("q.{a}::text")).collect();
289    format!(
290        "SELECT {} FROM ({}) AS q({})",
291        casts.join(", "),
292        trim_statement_tail(sql),
293        aliases.join(", ")
294    )
295}
296
297/// 去掉语句末尾的分号、空白与注释：嵌进子查询后，末尾的 `;` 或行注释会破坏外层语法。
298fn trim_statement_tail(sql: &str) -> &str {
299    use sqlparser::tokenizer::{Token, Tokenizer};
300    let dialect = sqlparser::dialect::PostgreSqlDialect {};
301    let Ok(tokens) = Tokenizer::new(&dialect, sql).tokenize_with_location() else {
302        return sql;
303    };
304    let Some(last) = tokens
305        .iter()
306        .rev()
307        .find(|t| !matches!(t.token, Token::Whitespace(_) | Token::SemiColon))
308    else {
309        return "";
310    };
311    // 分词器的位置按字符计列、遇 '\n' 换行，结束位置不含该字符。
312    let end = (last.span.end.line, last.span.end.column);
313    let (mut line, mut column) = (1, 1);
314    for (i, c) in sql.char_indices() {
315        if (line, column) == end {
316            return &sql[..i];
317        }
318        if c == '\n' {
319            line += 1;
320            column = 1;
321        } else {
322            column += 1;
323        }
324    }
325    sql
326}
327
328/// 一行文本单元格转成 SQL 字面量。
329fn row_literals(
330    row: &tokio_postgres::Row,
331    types: &[Type],
332) -> Result<Vec<String>, tokio_postgres::Error> {
333    types
334        .iter()
335        .enumerate()
336        .map(|(i, ty)| Ok(sql_literal(ty, row.try_get::<_, Option<&str>>(i)?)))
337        .collect()
338}
339
340/// 把单元格的文本表示转成 SQL 字面量：NULL 原样；布尔写 TRUE/FALSE；有限数值不加引号；
341/// 其余（含 NaN/Infinity）单引号转义，插入时由目标列类型解析。
342fn sql_literal(ty: &Type, text: Option<&str>) -> String {
343    let Some(s) = text else {
344        return "NULL".into();
345    };
346    match ty.name() {
347        // `bool::text` 输出 true/false，布尔输出函数为 t/f。
348        "bool" => if s.starts_with('t') { "TRUE" } else { "FALSE" }.into(),
349        "int2" | "int4" | "int8" | "oid" | "float4" | "float8" | "numeric"
350            if s.parse::<f64>().is_ok_and(f64::is_finite) =>
351        {
352            s.into()
353        }
354        _ => format!("'{}'", s.replace('\'', "''")),
355    }
356}
357
358#[cfg(all(test, feature = "server"))]
359mod tests {
360    use super::*;
361
362    // ---- is_simple_ident：表名白名单（防 SQL 注入与路径穿越）----
363
364    #[test]
365    fn simple_ident_accepts_legal_names() {
366        assert!(is_simple_ident("posts"));
367        assert!(is_simple_ident("user_posts"));
368        assert!(is_simple_ident("table1"));
369        assert!(is_simple_ident("a"));
370        assert!(is_simple_ident("_private"));
371        assert!(is_simple_ident("ABC123xyz"));
372    }
373
374    #[test]
375    fn simple_ident_rejects_empty_and_whitespace() {
376        assert!(!is_simple_ident(""));
377        assert!(!is_simple_ident("   "));
378    }
379
380    #[test]
381    fn simple_ident_rejects_injection_vectors() {
382        // SQL 注入：引号、分号、注释、空格分隔
383        assert!(!is_simple_ident("posts; DROP TABLE users"));
384        assert!(!is_simple_ident("posts' OR '1'='1"));
385        assert!(!is_simple_ident("posts--"));
386        assert!(!is_simple_ident("public.posts"));
387        assert!(!is_simple_ident("posts where 1=1"));
388        // 路径穿越
389        assert!(!is_simple_ident("../etc/passwd"));
390        assert!(!is_simple_ident("..\\windows"));
391        assert!(!is_simple_ident("/etc/passwd"));
392        // 反引号 / 双引号标识符语法
393        assert!(!is_simple_ident("`posts`"));
394        assert!(!is_simple_ident("\"posts\""));
395        // 连字符（常见合法表名片段，但本白名单禁止，防 `a-b` 被 SQL 解析为减法）
396        assert!(!is_simple_ident("my-posts"));
397    }
398
399    #[test]
400    fn simple_ident_rejects_unicode_and_non_ascii() {
401        assert!(!is_simple_ident("文章"));
402        assert!(!is_simple_ident("posts²"));
403        assert!(!is_simple_ident("café"));
404    }
405
406    // ---- parse_source：来源解析 + 只读校验（安全入口）----
407
408    #[test]
409    fn parse_table_source_builds_select_star() {
410        let source = parse_source("table:posts").unwrap();
411        assert_eq!(source.sql, "SELECT * FROM \"posts\"");
412        assert_eq!(source.table_name, "posts");
413        assert!(source.query_only);
414    }
415
416    #[test]
417    fn parse_table_source_trims_whitespace() {
418        let source = parse_source("table:  posts  ").unwrap();
419        assert_eq!(source.sql, "SELECT * FROM \"posts\"");
420        assert_eq!(source.table_name, "posts");
421    }
422
423    #[test]
424    fn parse_table_source_rejects_bad_name() {
425        // 非法表名（注入/穿越/空）必须在拼进 SQL 前被拒，防 SQL 注入。
426        for bad in [
427            "table:",
428            "table:   ",
429            "table:posts; DROP TABLE users",
430            "table:../secret",
431            "table:a b",
432            "table:\"x\"",
433        ] {
434            let (code, msg) = parse_source(bad).expect_err(bad);
435            assert_eq!(code, StatusCode::BAD_REQUEST);
436            assert!(msg.contains("表名"), "{bad}: {msg}");
437        }
438    }
439
440    #[test]
441    fn parse_query_source_accepts_select_and_explain() {
442        for query in [
443            "query:SELECT * FROM posts",
444            "query:SELECT id, title FROM posts WHERE id > 0",
445            "query:SELECT 1",
446            // 复杂但只读的查询
447            "query:WITH t AS (SELECT 1) SELECT * FROM t",
448        ] {
449            assert!(parse_source(query).unwrap().query_only, "{query}");
450        }
451        // EXPLAIN 放行，但不能包进子查询做逐列文本转换。
452        assert!(
453            !parse_source("query:EXPLAIN SELECT * FROM posts")
454                .unwrap()
455                .query_only
456        );
457    }
458
459    #[test]
460    fn parse_query_source_rejects_all_write_operations() {
461        // 每一种写/破坏操作都必须被 AST 校验拦截——这是数据导出的核心安全不变量。
462        for evil in [
463            "query:INSERT INTO posts VALUES (1)",
464            "query:UPDATE posts SET title='x'",
465            "query:DELETE FROM posts",
466            "query:DROP TABLE posts",
467            "query:TRUNCATE posts",
468            "query:CREATE TABLE evil (id int)",
469            "query:ALTER TABLE posts ADD COLUMN x int",
470            "query:GRANT SELECT ON posts TO public",
471            // 即便嵌在 SELECT 里，子查询的写操作也要拒
472            "query:INSERT INTO logs VALUES (1) RETURNING 1",
473        ] {
474            let (code, msg) = parse_source(evil).expect_err(evil);
475            assert_eq!(code, StatusCode::BAD_REQUEST, "{evil}");
476            // 写操作要么被 AST 校验拒（"只读"），要么 sqlparser 解析阶段就失败
477            assert!(
478                msg.contains("只读") || msg.contains("解析失败"),
479                "{evil}: {msg}"
480            );
481        }
482    }
483
484    #[test]
485    fn parse_query_source_rejects_explain_analyze_and_hidden_writes() {
486        // EXPLAIN ANALYZE 会真正执行内层语句；data-modifying CTE、SELECT INTO、
487        // 行锁子句外形都是 SELECT，但都有副作用。
488        for evil in [
489            "query:EXPLAIN ANALYZE DELETE FROM posts",
490            "query:EXPLAIN ANALYZE DELETE FROM posts WHERE id=1",
491            "query:EXPLAIN (ANALYZE) DELETE FROM posts WHERE id = 1",
492            "query:EXPLAIN (ANALYZE true) SELECT 1",
493            "query:EXPLAIN DELETE FROM posts WHERE id = 1",
494            "query:WITH d AS (DELETE FROM posts RETURNING *) SELECT * FROM d",
495            "query:SELECT * INTO t FROM posts",
496            "query:SELECT * FROM posts FOR UPDATE",
497        ] {
498            let (code, msg) = parse_source(evil).expect_err(evil);
499            assert_eq!(code, StatusCode::BAD_REQUEST, "{evil}");
500            assert!(msg.contains("只读"), "{evil}: {msg}");
501        }
502    }
503
504    #[test]
505    fn parse_query_source_accepts_plain_explain_options_and_ctes() {
506        for ok in [
507            "query:SELECT 1",
508            "query:EXPLAIN SELECT * FROM posts",
509            "query:EXPLAIN (FORMAT JSON) SELECT * FROM posts",
510            "query:WITH t AS (SELECT id FROM posts) SELECT * FROM t",
511        ] {
512            assert!(parse_source(ok).is_ok(), "{ok}");
513        }
514    }
515
516    #[test]
517    fn parse_query_source_rejects_unparseable_sql() {
518        for bad in [
519            "query:这不是SQL",
520            "query:SELECT FROM",
521            "query:@#$%",
522            "query:1+1",
523        ] {
524            assert!(parse_source(bad).is_err(), "{bad} 应解析失败或被拒");
525        }
526    }
527
528    #[test]
529    fn parse_query_source_allows_empty_result_set() {
530        // sqlparser 容忍纯分号（解析为零语句），`any()` 在空集上为 false → 放行。
531        // 锁定该契约：这只产生空导出，无数据泄露，属可接受行为；
532        // 若未来想收紧为"空查询拒绝"，此测试会提醒你更新。
533        assert!(parse_source("query:;;;").is_ok());
534    }
535
536    #[test]
537    fn parse_source_rejects_unknown_prefix() {
538        let (code, msg) = parse_source("unknown:posts").expect_err("未知前缀");
539        assert_eq!(code, StatusCode::BAD_REQUEST);
540        assert!(msg.contains("source"));
541        // 完全无前缀的裸字符串也应被拒
542        assert!(parse_source("posts").is_err());
543        assert!(parse_source("").is_err());
544    }
545
546    // ---- SQL 导出：单元格字面量 ----
547
548    #[test]
549    fn sql_literal_keeps_values_parseable() {
550        let lit = |ty: Type, s: Option<&str>| sql_literal(&ty, s);
551        assert_eq!(lit(Type::TIMESTAMPTZ, None), "NULL");
552        assert_eq!(lit(Type::BOOL, Some("true")), "TRUE");
553        assert_eq!(lit(Type::BOOL, Some("f")), "FALSE");
554        assert_eq!(lit(Type::INT8, Some("-42")), "-42");
555        assert_eq!(lit(Type::FLOAT8, Some("1e+20")), "1e+20");
556        assert_eq!(lit(Type::NUMERIC, Some("1.50")), "1.50");
557        // 非有限数值不是合法的裸字面量，必须加引号。
558        assert_eq!(lit(Type::FLOAT8, Some("NaN")), "'NaN'");
559        assert_eq!(lit(Type::FLOAT4, Some("-Infinity")), "'-Infinity'");
560        assert_eq!(lit(Type::NUMERIC, Some("Infinity")), "'Infinity'");
561        assert_eq!(lit(Type::TEXT, Some("it's")), "'it''s'");
562        assert_eq!(lit(Type::TEXT, Some("42")), "'42'");
563        assert_eq!(
564            lit(Type::TIMESTAMPTZ, Some("2026-11-03 04:34:56.789+00")),
565            "'2026-11-03 04:34:56.789+00'"
566        );
567    }
568
569    #[test]
570    fn trim_statement_tail_drops_only_trailing_terminators() {
571        for (sql, want) in [
572            ("SELECT 1", "SELECT 1"),
573            ("SELECT 1;", "SELECT 1"),
574            ("SELECT 1 ; ;\n", "SELECT 1"),
575            ("SELECT 1; -- 注释\n", "SELECT 1"),
576            ("SELECT 1 /* 尾注 */ ;", "SELECT 1"),
577            ("SELECT '中文;'\n  FROM t -- x", "SELECT '中文;'\n  FROM t"),
578            // 字符串里的分号、注释符号不算语句结尾。
579            ("SELECT ';' AS s;", "SELECT ';' AS s"),
580            ("SELECT 'a--b'", "SELECT 'a--b'"),
581            (";;", ""),
582        ] {
583            assert_eq!(trim_statement_tail(sql), want, "{sql:?}");
584        }
585    }
586
587    #[test]
588    #[ignore = "requires DATABASE_URL; uses a connection-local temporary table"]
589    fn export_sql_database_round_trips_typed_values() {
590        crate::db::TEST_DATABASE_RUNTIME.block_on(export_sql_round_trip_impl());
591    }
592
593    async fn export_sql_body(
594        source: &str,
595        include_columns: bool,
596    ) -> Result<String, (StatusCode, String)> {
597        let (body, _, _) = export_sql(parse_source(source).unwrap(), include_columns).await?;
598        let bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();
599        Ok(String::from_utf8(bytes.to_vec()).unwrap())
600    }
601
602    /// 导出的 INSERT 回灌到同结构的表后逐列不变；非标量列不能被写成 NULL。
603    async fn export_sql_round_trip_impl() {
604        let _guard = crate::db::TEST_DATABASE_LOCK.lock().await;
605        const ROW: &str = "SELECT TIMESTAMPTZ '2026-11-03 12:34:56.789+08' AS created_at, \
606             '0f8fad5b-d9cb-469f-a165-70867728950e'::uuid AS id, \
607             ARRAY['a', 'b''c', NULL] AS tags, '{\"k\": [1, 2]}'::jsonb AS meta, \
608             '\\x00ff'::bytea AS raw, DATE '2026-11-03' AS day, \
609             INTERVAL '1 day 2 hours' AS span, 'NaN'::float8 AS nan, 0.1::float4 AS f4, \
610             12345678901234567890.5 AS big, true AS ok, 'it''s' AS title, NULL::int AS missing";
611        let client = crate::db::pool::get_conn().await.unwrap();
612        client
613            .batch_execute(&format!("CREATE TEMP TABLE export AS {ROW} WITH NO DATA"))
614            .await
615            .unwrap();
616        // 末尾分号与行注释不影响子查询包装。
617        let dump = export_sql_body(&format!("query:{ROW}; -- 注释\n"), true)
618            .await
619            .unwrap();
620        client.batch_execute(&dump).await.unwrap();
621        let exported: String = client
622            .query_one("SELECT e::text FROM export e", &[])
623            .await
624            .unwrap()
625            .get(0);
626        let expected: String = client
627            .query_one(&format!("SELECT o::text FROM ({ROW}) o"), &[])
628            .await
629            .unwrap()
630            .get(0);
631        assert_eq!(exported, expected, "{dump}");
632        client.batch_execute("DROP TABLE export").await.unwrap();
633
634        // 重名列按位置取值。
635        let dump = export_sql_body("query:SELECT 1 AS a, 'x' AS a", false)
636            .await
637            .unwrap();
638        assert!(dump.contains(" VALUES (1, 'x');\n"), "{dump}");
639
640        // EXPLAIN 不能包成子查询：文本格式直接导出，非文本格式明确拒绝。
641        let plan = export_sql_body("query:EXPLAIN SELECT 1", true)
642            .await
643            .unwrap();
644        assert!(plan.contains(" VALUES ('Result"), "{plan}");
645        let (code, _) = export_sql_body("query:EXPLAIN (FORMAT JSON) SELECT 1", true)
646            .await
647            .unwrap_err();
648        assert_eq!(code, StatusCode::BAD_REQUEST);
649    }
650}
yggdrasil/api/database/export.rs

yggdrasil/api/database/
export.rs
