feat(admin): add data export tab (Axum streaming, SQL/CSV)

新增 GET /api/database/export Axum 流式路由(镜像 upload.rs 的 cookie→admin
鉴权),支持按表/按查询导出。CSV 用 COPY ... TO STDOUT WITH CSV 流式
(大表不 OOM),SQL 逐行拼 INSERT(含类型化转义)。表名白名单 + 查询
AST 只读强制。前端 ExportTab 触发 window.open 下载。
This commit is contained in:
xfy 2026-06-29 19:01:16 +08:00
parent a49f47c8a6
commit 5c4b2a53d1
4 changed files with 389 additions and 3 deletions

234
src/api/database/export.rs Normal file
View File

@ -0,0 +1,234 @@
//! 数据导出Axum 流式路由(大文件不走 JSON 序列化)。
//!
//! 鉴权镜像 [`crate::api::upload`]:从 cookie 取 session校验管理员。
//! 两种来源:按表导出(表名白名单)/ 按查询导出强制只读AST 校验)。
//! 两种格式CSVCOPY TO STDOUT 流式)/ SQL逐行拼 INSERT
#![allow(clippy::unused_unit)]
use axum::body::Body;
use axum::extract::Query;
use axum::http::{header, HeaderMap, StatusCode};
use axum::response::IntoResponse;
use futures::StreamExt;
use serde::Deserialize;
use crate::auth::session::parse_session_token;
#[derive(Deserialize)]
pub struct ExportParams {
/// `table:<name>` 或 `query:<SELECT ...>`。
pub source: String,
/// `sql` 或 `csv`。
pub format: String,
/// 是否带表头CSV/ 列名SQL INSERT。默认 true。
pub include_columns: Option<bool>,
}
/// POST/GET /api/database/export —— 流式导出。
pub async fn export_data(
headers: HeaderMap,
Query(params): Query<ExportParams>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
// 1. 鉴权cookie → session → admin
let cookie_header = headers
.get("cookie")
.and_then(|h| h.to_str().ok())
.unwrap_or("");
let token = parse_session_token(cookie_header).map(str::to_string);
let token = match token {
Some(t) => t,
None => return Err((StatusCode::UNAUTHORIZED, "未登录".to_string())),
};
let user = match crate::api::auth::get_user_by_token(&token).await {
Ok(Some(u)) => u,
_ => return Err((StatusCode::UNAUTHORIZED, "会话已过期".to_string())),
};
if user.role != crate::models::user::UserRole::Admin {
return Err((StatusCode::FORBIDDEN, "权限不足".to_string()));
}
// 2. 解析来源 + 白名单/只读校验
let include_columns = params.include_columns.unwrap_or(true);
let source_sql = parse_source(&params.source)?;
match params.format.as_str() {
"csv" => export_csv(source_sql, include_columns).await,
"sql" => export_sql(source_sql, include_columns).await,
_ => Err((StatusCode::BAD_REQUEST, "不支持的格式".to_string())),
}
.map(|(body, content_type, filename)| {
let disposition = format!("attachment; filename=\"{}\"", filename);
let disposition_value = axum::http::HeaderValue::from_str(&disposition)
.unwrap_or_else(|_| axum::http::HeaderValue::from_static("attachment"));
let content_type_value = axum::http::HeaderValue::from_str(content_type)
.unwrap_or_else(|_| axum::http::HeaderValue::from_static("application/octet-stream"));
(
StatusCode::OK,
[
(header::CONTENT_TYPE, content_type_value),
(header::CONTENT_DISPOSITION, disposition_value),
],
body,
)
})
}
/// 解析导出来源,返回可直接执行的内部 SQL。
/// - `table:posts` → 校验表名合法后,`SELECT * FROM "posts"`(只读)。
/// - `query:SELECT ...` → 校验为只读语句后原样返回。
fn parse_source(source: &str) -> Result<String, (StatusCode, String)> {
if let Some(table) = source.strip_prefix("table:") {
// 表名白名单:仅允许标识符字符,防注入
let t = table.trim();
if t.is_empty() || !is_simple_ident(t) {
return Err((StatusCode::BAD_REQUEST, "无效的表名".to_string()));
}
Ok(format!("SELECT * FROM \"{}\"", t))
} else if let Some(query) = source.strip_prefix("query:") {
// 只读校验sqlparser 解析后所有语句均为 Query/Explain
let dialect = sqlparser::dialect::PostgreSqlDialect {};
let parsed = sqlparser::parser::Parser::parse_sql(&dialect, query)
.map_err(|_| (StatusCode::BAD_REQUEST, "SQL 解析失败".to_string()))?;
if parsed.iter().any(|s| !is_read_only_ast(s)) {
return Err((
StatusCode::BAD_REQUEST,
"导出查询必须是只读SELECT/EXPLAIN".to_string(),
));
}
Ok(query.to_string())
} else {
Err((
StatusCode::BAD_REQUEST,
"source 必须是 table:<name> 或 query:<sql>".to_string(),
))
}
}
/// 简单标识符校验(字母数字下划线,防 SQL 注入与路径穿越)。
fn is_simple_ident(s: &str) -> bool {
!s.is_empty()
&& s.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_')
}
/// 判断 AST 是否只读SELECT/EXPLAIN
fn is_read_only_ast(stmt: &sqlparser::ast::Statement) -> bool {
use sqlparser::ast::Statement;
matches!(stmt, Statement::Query(_) | Statement::Explain { .. })
}
/// CSV 导出:用 COPY ... TO STDOUT WITH CSV 流式(大表不 OOM
async fn export_csv(
source_sql: String,
include_columns: bool,
) -> Result<(Body, &'static str, String), (StatusCode, String)> {
let client = crate::db::pool::get_conn()
.await
.map_err(|_| (StatusCode::SERVICE_UNAVAILABLE, "数据库不可用".to_string()))?;
let header_clause = if include_columns { "HEADER" } else { "" };
let copy_stmt = format!("COPY ({}) TO STDOUT WITH CSV {}", source_sql, header_clause);
let stream = client
.copy_out(&copy_stmt)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("COPY 失败:{e}")))?;
// tokio_postgres 的 copy_out 流产出 Bytes直接转 axum Body。
let mapped = stream.map(|res| res.map_err(std::io::Error::other));
let body = Body::from_stream(mapped);
Ok((body, "text/csv; charset=utf-8", "export.csv".to_string()))
}
/// SQL 导出:逐行拼 INSERT纯数据不含 DDL
async fn export_sql(
source_sql: String,
include_columns: bool,
) -> Result<(Body, &'static str, String), (StatusCode, String)> {
let client = crate::db::pool::get_conn()
.await
.map_err(|_| (StatusCode::SERVICE_UNAVAILABLE, "数据库不可用".to_string()))?;
let rows = client
.query(&source_sql, &[])
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询失败:{e}")))?;
// 表名(从 SQL 简单解析,或用 "export"
let table_name = "export";
let columns: Vec<String> = rows
.first()
.map(|r| r.columns().iter().map(|c| c.name().to_string()).collect())
.unwrap_or_default();
let mut out = String::new();
out.push_str("-- Yggdrasil 数据导出(仅数据,不含 schema\n");
let col_clause = if include_columns && !columns.is_empty() {
format!("({})", columns.join(", "))
} else {
String::new()
};
for r in &rows {
let vals: Vec<String> = (0..r.len())
.map(|i| sql_quote_cell(r, i))
.collect();
out.push_str(&format!(
"INSERT INTO {} {} VALUES ({});\n",
table_name,
col_clause,
vals.join(", ")
));
}
let body = Body::from(out);
Ok((
body,
"application/sql; charset=utf-8",
"export.sql".to_string(),
))
}
/// 把一个单元格值转成 SQL 字面量字符串单引号转义数字原样NULL
fn sql_quote_cell(row: &tokio_postgres::Row, idx: usize) -> String {
let ty = row.columns().get(idx).map(|c| c.type_().name()).unwrap_or("");
match ty {
"int2" => row
.try_get::<_, Option<i16>>(idx)
.ok()
.flatten()
.map(|v| v.to_string())
.unwrap_or_else(|| "NULL".into()),
"int4" => row
.try_get::<_, Option<i32>>(idx)
.ok()
.flatten()
.map(|v| v.to_string())
.unwrap_or_else(|| "NULL".into()),
"int8" => row
.try_get::<_, Option<i64>>(idx)
.ok()
.flatten()
.map(|v| v.to_string())
.unwrap_or_else(|| "NULL".into()),
"float4" | "float8" => row
.try_get::<_, Option<f64>>(idx)
.ok()
.flatten()
.map(|v| v.to_string())
.unwrap_or_else(|| "NULL".into()),
"bool" => row
.try_get::<_, Option<bool>>(idx)
.ok()
.flatten()
.map(|v| if v { "TRUE" } else { "FALSE" }.to_string())
.unwrap_or_else(|| "NULL".into()),
_ => {
// 字符串/时间戳等:单引号转义
match row.try_get::<_, Option<String>>(idx) {
Ok(Some(s)) => format!("'{}'", s.replace('\'', "''")),
_ => "NULL".into(),
}
}
}
}

View File

@ -14,3 +14,5 @@ pub mod system_status;
pub mod sql_console;
/// SQL 补全用 schema 拉取。
pub mod schema;
/// 数据导出 Axum 流式处理器。
pub mod export;

View File

@ -388,6 +388,19 @@ fn main() {
))
.layer(axum::middleware::from_fn(crate::api::csrf::csrf_middleware));
// 数据导出:流式响应,走 GET + query参数较短
// 鉴权在 handler 内部从 cookie 校验 adminCSRF 最外层拦截非法来源。
let export_route = axum::Router::new()
.route(
"/api/database/export",
axum::routing::get(crate::api::database::export::export_data),
)
.layer(TimeoutLayer::with_status_code(
StatusCode::REQUEST_TIMEOUT,
Duration::from_secs(120),
))
.layer(axum::middleware::from_fn(crate::api::csrf::csrf_middleware));
// Dioxus 应用路由:自动挂载所有 server function 并渲染前端组件
let dioxus_app =
axum::Router::new().serve_dioxus_application(config, router::AppRouter);
@ -425,8 +438,8 @@ fn main() {
axum::routing::get(|| async { StatusCode::NOT_FOUND }),
);
// 合并upload 路由保持自己独立的 300s 超时;app routes 加可选压缩/30sstatic routes 无任何中间件
let router = upload_route.merge(app_routes).merge(static_routes);
// 合并upload 路由保持自己独立的 300s 超时;export 路由 120sapp routes 加可选压缩/30sstatic routes 无任何中间件
let router = upload_route.merge(export_route).merge(app_routes).merge(static_routes);
Ok(router)
});

View File

@ -62,7 +62,7 @@ pub fn System() -> Element {
SystemTab::DbStatus => rsx! { DbStatusTab {} },
SystemTab::ServerStatus => rsx! { ServerStatusTab {} },
SystemTab::SqlConsole => rsx! { SqlConsoleTab {} },
SystemTab::Export => rsx! { div { class: "text-paper-secondary py-8", "数据导出(待实现)" } },
SystemTab::Export => rsx! { ExportTab {} },
SystemTab::Backup => rsx! { div { class: "text-paper-secondary py-8", "备份恢复(待实现)" } },
}
}
@ -900,3 +900,140 @@ fn SqlConsoleTab() -> Element {
}
}
}
/// 数据导出 tab按表/按查询导出 SQL/CSV走 Axum 流式下载。
#[allow(non_snake_case)]
fn ExportTab() -> Element {
use crate::components::ui::ADMIN_CARD_CLASS;
// 导出模式:"table" / "query"
let mut mode = use_signal(|| "table".to_string());
let mut table_name = use_signal(String::new);
let mut query = use_signal(String::new);
let mut format = use_signal(|| "csv".to_string());
let mut include_columns = use_signal(|| true);
// 触发下载:构造 GET /api/database/export?... URL 并打开
let do_export = move || {
#[cfg(target_arch = "wasm32")]
{
let source = if mode().as_str() == "table" {
format!("table:{}", table_name.read().trim())
} else {
format!("query:{}", query.read())
};
let url = format!(
"/api/database/export?source={}&format={}&include_columns={}",
urlencode(&source),
format(),
include_columns(),
);
if let Some(window) = web_sys::window() {
let _ = window.open_with_url(&url);
}
}
};
rsx! {
div { class: "space-y-4 max-w-2xl",
div { class: "{ADMIN_CARD_CLASS} p-4 space-y-4",
// 模式选择
div { class: "flex items-center gap-4",
label { class: "flex items-center gap-2 text-sm text-paper-primary",
input {
r#type: "radio",
name: "export-mode",
checked: mode() == "table",
onchange: move |_| mode.set("table".to_string()),
}
"按表导出"
}
label { class: "flex items-center gap-2 text-sm text-paper-primary",
input {
r#type: "radio",
name: "export-mode",
checked: mode() == "query",
onchange: move |_| mode.set("query".to_string()),
}
"按查询导出"
}
}
// 表名输入
if mode().as_str() == "table" {
div {
label { class: "block text-sm text-paper-secondary mb-1", "表名" }
input {
r#type: "text",
class: "w-full px-3 py-2 text-sm border border-paper-border rounded bg-paper-theme text-paper-primary font-mono",
placeholder: "如 posts",
value: "{table_name}",
oninput: move |e| table_name.set(e.value()),
}
p { class: "text-xs text-paper-secondary mt-1", "仅支持 public schema 下的用户表,表名需为合法标识符" }
}
} else {
// 查询输入
div {
label { class: "block text-sm text-paper-secondary mb-1", "SELECT 查询(只读)" }
textarea {
class: "w-full px-3 py-2 text-sm border border-paper-border rounded bg-paper-theme text-paper-primary font-mono",
rows: "4",
placeholder: "SELECT id, title FROM posts WHERE published = true",
value: "{query}",
oninput: move |e| query.set(e.value()),
}
}
}
// 格式 + 选项
div { class: "flex flex-wrap items-center gap-4",
div { class: "flex items-center gap-2",
span { class: "text-sm text-paper-secondary", "格式" }
select {
class: "text-sm border border-paper-border rounded px-2 py-1 bg-paper-theme text-paper-primary",
value: "{format}",
onchange: move |e| format.set(e.value()),
option { value: "csv", "CSV" }
option { value: "sql", "SQL (INSERT)" }
}
}
label { class: "flex items-center gap-1 text-sm text-paper-secondary",
input {
r#type: "checkbox",
class: "mr-1",
checked: include_columns(),
onchange: move |e| include_columns.set(e.checked()),
}
"包含列名CSV 表头 / INSERT 列清单)"
}
}
button {
class: "px-4 py-1.5 text-sm bg-paper-accent text-paper-theme rounded hover:brightness-110 transition",
onclick: move |_| do_export(),
"导出并下载"
}
}
p { class: "text-xs text-paper-secondary",
"导出走流式响应大表不会占满内存。SQL 格式仅含 INSERT 语句(不含 DDL/schema"
}
}
}
}
/// 简易 URL 编码(避免引入新依赖;仅编码导出参数里的特殊字符)。
/// 仅在 WASM 前端的导出按钮里用。
#[cfg(target_arch = "wasm32")]
fn urlencode(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
out.push(b as char);
}
b' ' => out.push('+'),
_ => out.push_str(&format!("%{:02X}", b)),
}
}
out
}