From 453e471c5e5319df99242f9ac12d02ee976d8b28 Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 17:10:24 +0800 Subject: [PATCH 01/10] fix(core): classify vector REST query risk --- crates/dbx-core/src/query_execution_sql.rs | 100 +++++++++++++++++++++ crates/dbx-core/src/sql_risk.rs | 29 +++++- 2 files changed, 126 insertions(+), 3 deletions(-) diff --git a/crates/dbx-core/src/query_execution_sql.rs b/crates/dbx-core/src/query_execution_sql.rs index 2fb79dc545..8ddb47dde9 100644 --- a/crates/dbx-core/src/query_execution_sql.rs +++ b/crates/dbx-core/src/query_execution_sql.rs @@ -255,6 +255,9 @@ pub fn is_write_sql_for_database(sql: &str, database_type: DatabaseType) -> bool if let Some(risk) = classify_search_engine_query_risk(sql, database_type) { return risk != SearchEngineQueryRisk::ReadOnly; } + if let Some(risk) = classify_vector_query_risk(sql, database_type) { + return risk != SearchEngineQueryRisk::ReadOnly; + } is_write_sql_with_database_type(sql, Some(database_type)) } @@ -327,6 +330,84 @@ pub(crate) fn classify_search_engine_query_risk( } } +/// 对向量数据库的 REST 风格查询进行风险分类。 +/// +/// `POST` 在 Milvus、Qdrant 和 Chroma 中既可能是检索，也可能是写入，不能仅按 HTTP +/// 方法判断。未知端点一律视为高风险，防止 MCP 通用查询绕过专用向量工具的权限边界。 +pub(crate) fn classify_vector_query_risk(source: &str, database_type: DatabaseType) -> Option { + if !matches!( + database_type, + DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb + ) { + return None; + } + let source = strip_leading_search_engine_comments(source); + let request_line = source.lines().next()?.trim(); + let mut parts = request_line.split_whitespace(); + let method = parts.next()?.to_ascii_uppercase(); + let path = parts.next()?.split('?').next().unwrap_or("").trim_end_matches('/').to_ascii_lowercase(); + + if matches!(method.as_str(), "GET" | "HEAD" | "OPTIONS") { + return Some(SearchEngineQueryRisk::ReadOnly); + } + + let risk = match database_type { + DatabaseType::Milvus => match (method.as_str(), path.as_str()) { + ( + "POST", + "/v2/vectordb/entities/search" + | "/v2/vectordb/entities/query" + | "/v2/vectordb/entities/get" + | "/v2/vectordb/collections/list" + | "/v2/vectordb/collections/describe" + | "/v2/vectordb/databases/list" + | "/v2/vectordb/indexes/list" + | "/v2/vectordb/indexes/describe", + ) => SearchEngineQueryRisk::ReadOnly, + ("POST", "/v2/vectordb/entities/insert" | "/v2/vectordb/entities/upsert") => SearchEngineQueryRisk::Write, + ("POST", "/v2/vectordb/entities/delete") => SearchEngineQueryRisk::Dangerous, + ("POST" | "PUT" | "PATCH" | "DELETE", _) => SearchEngineQueryRisk::Dangerous, + _ => return None, + }, + DatabaseType::Qdrant => { + let read_post = ["/scroll", "/search", "/query", "/recommend", "/discover", "/count"] + .iter() + .any(|suffix| path.ends_with(suffix)); + if method == "POST" && read_post { + SearchEngineQueryRisk::ReadOnly + } else if method == "PUT" && path.contains("/points") { + SearchEngineQueryRisk::Write + } else if matches!(method.as_str(), "POST" | "PUT" | "PATCH" | "DELETE") { + SearchEngineQueryRisk::Dangerous + } else { + return None; + } + } + DatabaseType::ChromaDb => { + let read_post = ["/get", "/query", "/count"].iter().any(|suffix| path.ends_with(suffix)); + let safe_write_post = ["/add", "/update", "/upsert"].iter().any(|suffix| path.ends_with(suffix)); + if method == "POST" && read_post { + SearchEngineQueryRisk::ReadOnly + } else if method == "POST" && safe_write_post { + SearchEngineQueryRisk::Write + } else if matches!(method.as_str(), "POST" | "PUT" | "PATCH" | "DELETE") { + SearchEngineQueryRisk::Dangerous + } else { + return None; + } + } + DatabaseType::Weaviate => { + if matches!(method.as_str(), "POST" | "PUT" | "PATCH" | "DELETE") { + SearchEngineQueryRisk::Dangerous + } else { + return None; + } + } + _ => return None, + }; + Some(risk) +} + fn strip_leading_search_engine_comments(input: &str) -> &str { let mut rest = input; loop { @@ -1832,6 +1913,25 @@ mod tests { )); } + #[test] + fn classifies_vector_rest_queries_without_treating_every_post_as_read_only() { + assert!(!is_write_sql_for_database( + "POST /v2/vectordb/entities/search\n{\"collectionName\":\"semantic_cards\"}", + DatabaseType::Milvus, + )); + assert!(is_write_sql_for_database( + "POST /v2/vectordb/entities/upsert\n{\"collectionName\":\"semantic_cards\"}", + DatabaseType::Milvus, + )); + assert!(is_write_sql_for_database( + "POST /v2/vectordb/entities/delete\n{\"filter\":\"semantic_batch_id == 'x'\"}", + DatabaseType::Milvus, + )); + assert!(is_write_sql_for_database("POST /v2/vectordb/collections/drop\n{}", DatabaseType::Milvus,)); + assert!(!is_write_sql_for_database("POST /collections/cards/points/search\n{}", DatabaseType::Qdrant,)); + assert!(is_write_sql_for_database("PUT /collections/cards/points\n{}", DatabaseType::Qdrant,)); + } + #[test] fn excludes_victoriametrics_from_sql_query_paths() { assert!(!supports_sql_query(DatabaseType::VictoriaMetrics)); diff --git a/crates/dbx-core/src/sql_risk.rs b/crates/dbx-core/src/sql_risk.rs index 5bd7fa5be6..b2589717fe 100644 --- a/crates/dbx-core/src/sql_risk.rs +++ b/crates/dbx-core/src/sql_risk.rs @@ -684,7 +684,9 @@ pub fn classify_sql_risk(sql: &str, dialect: &str) -> Result { /// Classify SQL risk using both the parser dialect and the concrete database /// type so dialect-specific write forms cannot be mistaken for read queries. pub fn classify_sql_risk_for_database(sql: &str, database_type: DatabaseType) -> Result { - if let Some(risk) = crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type) { + if let Some(risk) = crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type) + .or_else(|| crate::query_execution_sql::classify_vector_query_risk(sql, database_type)) + { return Ok(match risk { crate::query_execution_sql::SearchEngineQueryRisk::ReadOnly => SqlRisk::ReadOnly, crate::query_execution_sql::SearchEngineQueryRisk::Write => SqlRisk::Write, @@ -701,7 +703,9 @@ pub fn classify_sql_risk_for_database(sql: &str, database_type: DatabaseType) -> /// and single-table UPDATE/DELETE statements with an effective predicate; /// broader or opaque mutations require central high-risk permission. pub fn is_dangerous_sql_for_database(sql: &str, database_type: DatabaseType) -> bool { - if let Some(risk) = crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type) { + if let Some(risk) = crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type) + .or_else(|| crate::query_execution_sql::classify_vector_query_risk(sql, database_type)) + { return risk == crate::query_execution_sql::SearchEngineQueryRisk::Dangerous; } let database_type_name = format!("{database_type:?}"); @@ -730,7 +734,9 @@ pub fn is_dangerous_sql_for_database(sql: &str, database_type: DatabaseType) -> /// A `USE` statement mutates pooled/session state and could redirect later SQL, /// so it is forbidden independently of read/write and high-risk permissions. pub fn mcp_sql_has_forbidden_database_switch(sql: &str, database_type: DatabaseType) -> bool { - if crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type).is_some() { + if crate::query_execution_sql::classify_search_engine_query_risk(sql, database_type).is_some() + || crate::query_execution_sql::classify_vector_query_risk(sql, database_type).is_some() + { return false; } let database_type_name = format!("{database_type:?}"); @@ -1389,4 +1395,21 @@ mod tests { assert!(is_dangerous_sql_for_database("DELETE /products", database_type)); } } + + #[test] + fn classifies_vector_rest_risk_by_method_and_path() { + assert_eq!( + classify_sql_risk_for_database("POST /v2/vectordb/entities/search\n{}", DatabaseType::Milvus).unwrap(), + SqlRisk::ReadOnly, + ); + assert_eq!( + classify_sql_risk_for_database("POST /v2/vectordb/entities/upsert\n{}", DatabaseType::Milvus).unwrap(), + SqlRisk::Write, + ); + assert_eq!( + classify_sql_risk_for_database("POST /v2/vectordb/entities/delete\n{}", DatabaseType::Milvus).unwrap(), + SqlRisk::Ddl, + ); + assert!(is_dangerous_sql_for_database("POST /v2/vectordb/collections/drop\n{}", DatabaseType::Milvus,)); + } } From a1421a782d2c2d4023853b5faa1c9a3c0b1c075b Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 17:10:35 +0800 Subject: [PATCH 02/10] feat(mcp): add governed imports and Milvus semantics --- Cargo.lock | 1 + crates/dbx-mcp/Cargo.toml | 3 +- crates/dbx-mcp/src/backend.rs | 81 +- crates/dbx-mcp/src/enterprise_tools.rs | 1225 ++++++++++++++++++++++++ crates/dbx-mcp/src/lib.rs | 1 + crates/dbx-mcp/src/server.rs | 694 +++++++++++++- crates/dbx-mcp/tests/protocol.rs | 35 +- docs/content/docs/mcp.cn.mdx | 46 +- docs/content/docs/mcp.mdx | 46 +- packages/mcp-server/README.md | 44 +- 10 files changed, 2163 insertions(+), 13 deletions(-) create mode 100644 crates/dbx-mcp/src/enterprise_tools.rs diff --git a/Cargo.lock b/Cargo.lock index bf7875d5ab..32c7fa7708 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2164,6 +2164,7 @@ dependencies = [ "schemars 1.2.2", "serde", "serde_json", + "sha2 0.10.9", "tempfile", "tokio", "tokio-rustls 0.26.4", diff --git a/crates/dbx-mcp/Cargo.toml b/crates/dbx-mcp/Cargo.toml index 152af9dfbe..152fb3159c 100644 --- a/crates/dbx-mcp/Cargo.toml +++ b/crates/dbx-mcp/Cargo.toml @@ -22,7 +22,8 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "rus schemars = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" -tokio = { version = "1", features = ["fs", "macros", "rt-multi-thread"] } +sha2 = "0.10" +tokio = { version = "1", features = ["fs", "io-util", "macros", "rt-multi-thread", "sync", "time"] } uuid = { version = "1", features = ["v4"] } url = "2" diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 837036a3b8..1a9336c8f7 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -1,7 +1,10 @@ use std::{ collections::HashMap, path::{Path, PathBuf}, - sync::Arc, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, }; use async_trait::async_trait; @@ -12,6 +15,7 @@ use dbx_core::{ db::{redis_driver::RedisCommandResult, ColumnInfo, IndexInfo, TableInfo}, models::connection::{ConnectionConfig, DatabaseType}, storage::{DesktopSettings, McpGlobalPolicy, McpGlobalPolicyState, Storage}, + table_import::{TableImportProgress, TableImportRequest, TableImportSummary}, }; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -254,6 +258,18 @@ pub trait DbxBackend: Send + Sync { let _ = (connection, database, options); Err("Documentation snapshots are not supported by this backend.".to_string()) } + + /// 执行已经由 MCP prepare 固化并复验的本地文件导入。 + async fn import_table_file_for_mcp( + &self, + connection: &ConnectionConfig, + request: TableImportRequest, + cancelled: Arc, + progress: Arc, + ) -> Result { + let _ = (connection, request, cancelled, progress); + Err("IMPORT_UNSUPPORTED_IN_WEB_MODE_V1: v1 文件导入仅支持本地 DBX Desktop/MCP 模式。".to_string()) + } } pub struct LocalBackend { @@ -672,6 +688,69 @@ impl DbxBackend for LocalBackend { .await } + async fn import_table_file_for_mcp( + &self, + connection: &ConnectionConfig, + mut request: TableImportRequest, + cancelled: Arc, + progress: Arc, + ) -> Result { + // prepare/start 复验后再复制到任务私有快照，避免源文件在后台读取期间被替换。 + let expected_sha256 = request + .source_ref + .as_deref() + .filter(|value| value.len() == 64) + .ok_or_else(|| "IMPORT_SOURCE_HASH_REQUIRED: 导入计划缺少源文件 SHA-256。".to_string())? + .to_string(); + let snapshot_root = self.data_dir.join("tmp").join("mcp_import"); + tokio::fs::create_dir_all(&snapshot_root).await.map_err(|error| format!("创建导入快照目录失败：{error}"))?; + let snapshot_dir = snapshot_root.join(&request.import_id); + tokio::fs::create_dir(&snapshot_dir).await.map_err(|error| format!("创建任务快照目录失败：{error}"))?; + let extension = Path::new(&request.file_path).extension().and_then(|value| value.to_str()).unwrap_or("data"); + let snapshot_path = snapshot_dir.join(format!("source.{extension}")); + let snapshot_result = async { + tokio::fs::copy(&request.file_path, &snapshot_path) + .await + .map_err(|error| format!("复制导入源快照失败：{error}"))?; + let identity = crate::enterprise_tools::file_identity(snapshot_path.clone()) + .await + .map_err(|error| error.to_string())?; + if identity.sha256 != expected_sha256 { + return Err("IMPORT_FILE_CHANGED: 源文件在启动导入时发生变化；未写入数据库。".to_string()); + } + request.file_path = snapshot_path.to_string_lossy().to_string(); + let database = (!request.database.trim().is_empty()).then_some(request.database.as_str()); + let client_session_id = dbx_core::table_import::table_import_client_session_id(&request.import_id); + let pool_key = self + .state + .get_or_create_pool_for_session(&request.connection_id, database, Some(&client_session_id)) + .await?; + let cancellation = cancelled.clone(); + let result = dbx_core::table_import::import_table_file_core( + &self.state, + &request, + &connection.db_type, + &pool_key, + move |_| { + let cancellation = cancellation.clone(); + Box::pin(async move { cancellation.load(Ordering::Acquire) }) + }, + move |update| progress(update), + ) + .await; + let cleanup = + self.state.detach_client_session_pool(&request.connection_id, database, &client_session_id).await; + if let Err(error) = cleanup { + return Err(format!("导入连接清理失败：{error}")); + } + result + } + .await; + let _ = tokio::fs::remove_file(&snapshot_path).await; + let _ = tokio::fs::remove_dir(&snapshot_dir).await; + snapshot_result + } + async fn execute_redis_command( &self, connection: &ConnectionConfig, diff --git a/crates/dbx-mcp/src/enterprise_tools.rs b/crates/dbx-mcp/src/enterprise_tools.rs new file mode 100644 index 0000000000..d4b96520b8 --- /dev/null +++ b/crates/dbx-mcp/src/enterprise_tools.rs @@ -0,0 +1,1225 @@ +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + ffi::OsStr, + io::Read, + path::{Path, PathBuf}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use dbx_core::table_import::{ + TableImportColumnMapping, TableImportMode, TableImportParseOptions, TableImportPhase, TableImportPreview, + TableImportProgress, TableImportRequest, TableImportSourceFormat, TableImportStatus, TableImportSummary, + TableImportTextEncoding, +}; +use rmcp::schemars; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; +use tokio::sync::RwLock; +use uuid::Uuid; + +pub const FORMAT_VERSION: u8 = 1; +pub const DEFAULT_PLAN_TTL: Duration = Duration::from_secs(30 * 60); +const DEFAULT_IMPORT_FILE_SIZE_BYTES: u64 = 512 * 1024 * 1024; +const DEFAULT_SEMANTIC_FILE_SIZE_BYTES: u64 = 64 * 1024 * 1024; +const DEFAULT_PREVIEW_ROWS: usize = 20; +pub const STRUCTURE_PROFILE_ROWS: usize = 100; +const MAX_PREVIEW_ROWS: usize = STRUCTURE_PROFILE_ROWS; +const DEFAULT_CELL_CHAR_LIMIT: usize = 1_000; +const MAX_CELL_CHAR_LIMIT: usize = 4_000; +const DEFAULT_VECTOR_TOP_K: usize = 12; +const HARD_VECTOR_TOP_K: usize = 50; +const DEFAULT_VECTOR_DIMENSION: usize = 1_024; +const DEFAULT_VECTOR_UPSERT_BATCH: usize = 200; +const MAX_VECTOR_UPSERT_BATCH: usize = 500; +const MAX_JSONL_RECORD_BYTES: usize = 512 * 1024; + +const IMPORT_EXTENSIONS: &[&str] = &["xlsx", "xlsm", "xls", "csv", "tsv", "json"]; +const SEMANTIC_EXTENSIONS: &[&str] = &["jsonl"]; +const DEFAULT_VECTOR_FILTER_FIELDS: &[&str] = + &["business_domain", "dataset_id", "template_version", "card_type", "semantic_batch_id"]; +const DEFAULT_VECTOR_OUTPUT_FIELDS: &[&str] = &[ + "card_id", + "card_type", + "business_domain", + "dataset_id", + "template_version", + "title", + "content", + "aliases", + "approval_status", + "effective_from", + "effective_to", + "source_uri", + "source_checksum", + "semantic_version", + "embedding_model", + "semantic_batch_id", +]; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EnterpriseToolError { + pub code: &'static str, + pub message: String, +} + +impl EnterpriseToolError { + pub fn new(code: &'static str, message: impl Into) -> Self { + Self { code, message: message.into() } + } +} + +impl std::fmt::Display for EnterpriseToolError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "{}: {}", self.code, self.message) + } +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum McpImportSourceFormat { + Csv, + Tsv, + Json, + Excel, +} + +impl From for TableImportSourceFormat { + fn from(value: McpImportSourceFormat) -> Self { + match value { + McpImportSourceFormat::Csv => Self::Csv, + McpImportSourceFormat::Tsv => Self::Tsv, + McpImportSourceFormat::Json => Self::Json, + McpImportSourceFormat::Excel => Self::Excel, + } + } +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize, schemars::JsonSchema)] +pub struct McpImportParseOptions { + #[schemars(extend("type" = "string"))] + pub delimiter: Option, + #[schemars(extend("type" = "string"))] + pub encoding: Option, + #[schemars(extend("type" = "boolean"))] + pub has_header: Option, + #[schemars(extend("type" = "integer"))] + pub title_row: Option, + #[schemars(extend("type" = "integer"))] + pub data_start_row: Option, + #[schemars(extend("type" = "integer"))] + pub last_data_row: Option, + #[schemars(extend("type" = "boolean"))] + pub trim_values: Option, + #[schemars(extend("type" = "boolean"))] + pub empty_string_as_null: Option, + #[schemars(extend("type" = "string"))] + pub sheet_name: Option, + #[schemars(extend("type" = "integer"))] + pub sheet_index: Option, +} + +impl McpImportParseOptions { + pub fn into_core(self) -> Result { + let encoding = match self.encoding.as_deref().map(str::trim).filter(|value| !value.is_empty()) { + None | Some("auto") => Some(TableImportTextEncoding::Auto), + Some(value) if value.eq_ignore_ascii_case("utf-8") || value.eq_ignore_ascii_case("utf8") => { + Some(TableImportTextEncoding::Utf8) + } + Some(value) if value.eq_ignore_ascii_case("gbk") || value.eq_ignore_ascii_case("gb18030") => { + Some(TableImportTextEncoding::Gbk) + } + Some(value) if value.eq_ignore_ascii_case("utf-16le") => Some(TableImportTextEncoding::Utf16Le), + Some(value) if value.eq_ignore_ascii_case("utf-16be") => Some(TableImportTextEncoding::Utf16Be), + Some(value) => { + return Err(EnterpriseToolError::new( + "IMPORT_ENCODING_UNSUPPORTED", + format!("不支持编码 {value}；允许 auto、utf-8、gbk、utf-16le、utf-16be。"), + )) + } + }; + if self.title_row == Some(0) || self.data_start_row == Some(0) || self.last_data_row == Some(0) { + return Err(EnterpriseToolError::new("IMPORT_ROW_OUT_OF_RANGE", "行号从 1 开始，不能为 0。")); + } + Ok(TableImportParseOptions { + delimiter: self.delimiter, + encoding, + has_header: self.has_header, + title_row: self.title_row, + data_start_row: self.data_start_row, + last_data_row: self.last_data_row, + trim_values: self.trim_values, + empty_string_as_null: self.empty_string_as_null, + sheet_name: self.sheet_name, + sheet_index: self.sheet_index, + ..Default::default() + }) + } +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct McpImportColumnMapping { + pub source_column: String, + pub target_column: String, + #[schemars(extend("type" = "string"))] + pub target_data_type: Option, +} + +impl From for TableImportColumnMapping { + fn from(value: McpImportColumnMapping) -> Self { + Self { + source_column: value.source_column, + target_column: value.target_column, + target_data_type: value.target_data_type, + } + } +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct PreviewImportFileRequest { + pub file_path: String, + #[schemars(extend("type" = "string"))] + pub source_format: Option, + #[serde(default)] + pub parse_options: McpImportParseOptions, + #[schemars(extend("type" = "integer"))] + pub preview_rows: Option, + #[schemars(extend("type" = "integer"))] + pub cell_char_limit: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct PrepareTableImportRequest { + #[serde(flatten)] + pub selector: crate::server::ConnectionSelector, + #[schemars(extend("type" = "string"))] + pub database: Option, + pub schema: String, + pub table: String, + pub template_version: String, + pub file_path: String, + #[schemars(extend("type" = "string"))] + pub source_format: Option, + #[serde(default)] + pub parse_options: McpImportParseOptions, + pub mappings: Vec, + #[serde(default)] + pub create_table: bool, + #[schemars(extend("type" = "integer"))] + pub batch_size: Option, + #[schemars(extend("type" = "string"))] + pub date_time_format: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct StartTableImportRequest { + pub plan_id: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct ImportStatusRequest { + pub import_id: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct VectorSearchRequest { + #[serde(flatten)] + pub selector: crate::server::ConnectionSelector, + #[schemars(extend("type" = "string"))] + pub database: Option, + pub collection: String, + pub active_at: String, + #[schemars(extend("type" = "string"))] + pub semantic_version: Option, + pub embedding: Vec, + #[schemars(extend("type" = "integer"))] + pub top_k: Option, + #[serde(default)] + pub filters: BTreeMap, + #[schemars(extend("type" = "array"))] + pub output_fields: Option>, +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct VectorUpsertFileRequest { + #[serde(flatten)] + pub selector: crate::server::ConnectionSelector, + #[schemars(extend("type" = "string"))] + pub database: Option, + pub collection: String, + pub semantic_batch_id: String, + pub file_path: String, + #[schemars(extend("type" = "integer"))] + pub batch_size: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +pub struct VectorDeleteByBatchRequest { + #[serde(flatten)] + pub selector: crate::server::ConnectionSelector, + #[schemars(extend("type" = "string"))] + pub database: Option, + pub collection: String, + pub semantic_batch_id: String, +} + +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct FileIdentity { + pub canonical_path: String, + pub size_bytes: u64, + pub modified_nanos: u128, + pub sha256: String, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct PreparedImportPlan { + pub plan_id: String, + pub plan_digest: String, + pub created_at_ms: u128, + pub expires_at_ms: u128, + pub connection_id: String, + pub connection_name: String, + pub database: String, + pub schema: String, + pub table: String, + pub template_version: String, + pub file: FileIdentity, + pub structure_fingerprint: String, + pub source_format: Option, + pub parse_options: TableImportParseOptions, + pub mappings: Vec, + pub create_table: bool, + pub batch_size: usize, + pub date_time_format: Option, + #[serde(skip)] + consumed: bool, +} + +impl PreparedImportPlan { + pub fn to_import_request(&self, import_id: String) -> TableImportRequest { + TableImportRequest { + import_id, + connection_id: self.connection_id.clone(), + database: self.database.clone(), + schema: self.schema.clone(), + table: self.table.clone(), + file_path: self.file.canonical_path.clone(), + source_ref: Some(self.file.sha256.clone()), + source_format: self.source_format, + parse_options: self.parse_options.clone(), + mappings: self.mappings.clone(), + mode: TableImportMode::Append, + create_table: self.create_table, + batch_size: self.batch_size, + date_time_format: self.date_time_format.clone(), + prepared_source: None, + retain_source: true, + } + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ImportJobSnapshot { + pub import_id: String, + pub plan_id: String, + pub status: TableImportStatus, + pub phase: TableImportPhase, + pub rows_imported: usize, + pub total_rows: usize, + pub total_rows_exact: bool, + pub bytes_read: u64, + pub total_bytes: u64, + pub elapsed_ms: u128, + pub error: Option, + pub summary: Option, +} + +impl ImportJobSnapshot { + fn initial(import_id: String, plan_id: String, total_bytes: u64) -> Self { + Self { + import_id, + plan_id, + status: TableImportStatus::Running, + phase: TableImportPhase::Preparing, + rows_imported: 0, + total_rows: 0, + total_rows_exact: false, + bytes_read: 0, + total_bytes, + elapsed_ms: 0, + error: None, + summary: None, + } + } + + pub fn apply_progress(&mut self, progress: TableImportProgress) { + self.status = progress.status; + self.phase = progress.phase; + self.rows_imported = progress.rows_imported; + self.total_rows = progress.total_rows; + self.total_rows_exact = progress.total_rows_exact; + self.bytes_read = progress.bytes_read; + self.total_bytes = progress.total_bytes; + self.elapsed_ms = progress.elapsed_ms; + self.error = progress.error; + } +} + +pub struct ImportJob { + pub snapshot: Mutex, + pub cancelled: Arc, +} + +#[derive(Default)] +pub struct EnterpriseRuntime { + plans: RwLock>, + jobs: RwLock>>, +} + +impl EnterpriseRuntime { + pub async fn insert_plan(&self, plan: PreparedImportPlan) { + let now = unix_epoch_millis(); + let mut plans = self.plans.write().await; + plans.retain(|_, plan| !plan.consumed && plan.expires_at_ms > now); + plans.insert(plan.plan_id.clone(), plan); + } + + pub async fn consume_plan(&self, plan_id: &str) -> Result { + let mut plans = self.plans.write().await; + let plan = plans + .get_mut(plan_id) + .ok_or_else(|| EnterpriseToolError::new("IMPORT_PLAN_NOT_FOUND", "导入计划不存在，或已经被清理。"))?; + if plan.expires_at_ms <= unix_epoch_millis() { + return Err(EnterpriseToolError::new("IMPORT_PLAN_EXPIRED", "导入计划已超过 30 分钟有效期。")); + } + if plan.consumed { + return Err(EnterpriseToolError::new("IMPORT_PLAN_ALREADY_USED", "导入计划只能启动一次。")); + } + plan.consumed = true; + Ok(plan.clone()) + } + + pub async fn create_job(&self, plan: &PreparedImportPlan) -> Arc { + let import_id = format!("mcp-import-{}", Uuid::new_v4()); + let job = Arc::new(ImportJob { + snapshot: Mutex::new(ImportJobSnapshot::initial( + import_id.clone(), + plan.plan_id.clone(), + plan.file.size_bytes, + )), + cancelled: Arc::new(AtomicBool::new(false)), + }); + self.jobs.write().await.insert(import_id, job.clone()); + job + } + + pub async fn job(&self, import_id: &str) -> Option> { + self.jobs.read().await.get(import_id).cloned() + } + + pub async fn cancel_job(&self, import_id: &str) -> Result { + let job = self + .job(import_id) + .await + .ok_or_else(|| EnterpriseToolError::new("IMPORT_JOB_NOT_FOUND", "没有找到指定导入任务。"))?; + let snapshot = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).clone(); + if snapshot.status != TableImportStatus::Running { + return Err(EnterpriseToolError::new("IMPORT_JOB_ALREADY_TERMINAL", "导入任务已经结束，不能再取消。")); + } + job.cancelled.store(true, Ordering::Release); + Ok(snapshot) + } +} + +pub fn configured_import_roots() -> Result, EnterpriseToolError> { + let value = std::env::var_os("DBX_MCP_IMPORT_ROOTS").ok_or_else(|| { + EnterpriseToolError::new( + "IMPORT_ROOTS_NOT_CONFIGURED", + "必须先配置 DBX_MCP_IMPORT_ROOTS，MCP 才能读取本地导入文件。", + ) + })?; + let roots = std::env::split_paths(&value) + .filter_map(|path| std::fs::canonicalize(path).ok()) + .filter(|path| path.parent().is_some()) + .collect::>(); + if roots.is_empty() { + return Err(EnterpriseToolError::new("IMPORT_ROOTS_NOT_CONFIGURED", "DBX_MCP_IMPORT_ROOTS 中没有可用目录。")); + } + Ok(roots) +} + +pub fn validate_import_file(file_path: &str, semantic_jsonl: bool) -> Result { + let roots = configured_import_roots()?; + let max_bytes = if semantic_jsonl { + env_u64("DBX_MCP_SEMANTIC_FILE_MAX_BYTES", DEFAULT_SEMANTIC_FILE_SIZE_BYTES) + } else { + env_u64("DBX_MCP_IMPORT_FILE_MAX_BYTES", DEFAULT_IMPORT_FILE_SIZE_BYTES) + }; + validate_import_file_with_roots(file_path, &roots, semantic_jsonl, max_bytes) +} + +pub fn validate_import_file_with_roots( + file_path: &str, + roots: &[PathBuf], + semantic_jsonl: bool, + max_bytes: u64, +) -> Result { + let path = Path::new(file_path); + if !path.is_absolute() { + return Err(EnterpriseToolError::new("IMPORT_PATH_NOT_ABSOLUTE", "文件路径必须是绝对路径。")); + } + let link_metadata = std::fs::symlink_metadata(path) + .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法读取文件元数据：{error}")))?; + if link_metadata.file_type().is_symlink() { + return Err(EnterpriseToolError::new("IMPORT_SYMLINK_REJECTED", "不允许通过符号链接导入文件。")); + } + let canonical = std::fs::canonicalize(path) + .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法解析文件路径：{error}")))?; + if !roots.iter().filter_map(|root| std::fs::canonicalize(root).ok()).any(|root| canonical.starts_with(root)) { + return Err(EnterpriseToolError::new( + "IMPORT_PATH_OUTSIDE_ROOTS", + "文件不在 DBX_MCP_IMPORT_ROOTS 允许目录中。", + )); + } + let metadata = std::fs::metadata(&canonical) + .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法读取文件元数据：{error}")))?; + if !metadata.file_type().is_file() { + return Err(EnterpriseToolError::new("IMPORT_NOT_REGULAR_FILE", "导入源必须是普通文件。")); + } + if metadata.len() > max_bytes { + return Err(EnterpriseToolError::new( + "IMPORT_FILE_TOO_LARGE", + format!("文件大小 {} 字节，超过限制 {} 字节。", metadata.len(), max_bytes), + )); + } + let extension = canonical.extension().and_then(OsStr::to_str).unwrap_or("").to_ascii_lowercase(); + let allowed = if semantic_jsonl { SEMANTIC_EXTENSIONS } else { IMPORT_EXTENSIONS }; + if !allowed.contains(&extension.as_str()) { + return Err(EnterpriseToolError::new("IMPORT_FILE_TYPE_UNSUPPORTED", format!("不支持 .{extension} 文件。"))); + } + Ok(canonical) +} + +pub async fn file_identity(path: PathBuf) -> Result { + tokio::task::spawn_blocking(move || file_identity_sync(&path)) + .await + .map_err(|error| EnterpriseToolError::new("IMPORT_HASH_FAILED", error.to_string()))? +} + +fn file_identity_sync(path: &Path) -> Result { + let metadata = std::fs::metadata(path) + .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法读取文件元数据：{error}")))?; + let modified_nanos = metadata + .modified() + .ok() + .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok()) + .map(|duration| duration.as_nanos()) + .unwrap_or_default(); + let mut file = std::fs::File::open(path) + .map_err(|error| EnterpriseToolError::new("IMPORT_HASH_FAILED", format!("无法打开文件：{error}")))?; + let mut hasher = Sha256::new(); + let mut buffer = vec![0_u8; 1024 * 1024]; + loop { + let read = file + .read(&mut buffer) + .map_err(|error| EnterpriseToolError::new("IMPORT_HASH_FAILED", format!("读取文件失败：{error}")))?; + if read == 0 { + break; + } + hasher.update(&buffer[..read]); + } + Ok(FileIdentity { + canonical_path: path.to_string_lossy().to_string(), + size_bytes: metadata.len(), + modified_nanos, + sha256: format!("{:x}", hasher.finalize()), + }) +} + +pub fn preview_limit(value: Option) -> Result { + let value = value.unwrap_or(DEFAULT_PREVIEW_ROWS); + if value == 0 || value > MAX_PREVIEW_ROWS { + return Err(EnterpriseToolError::new( + "IMPORT_PREVIEW_LIMIT_INVALID", + format!("preview_rows 必须在 1 到 {MAX_PREVIEW_ROWS} 之间。"), + )); + } + Ok(value) +} + +pub fn cell_char_limit(value: Option) -> Result { + let value = value.unwrap_or(DEFAULT_CELL_CHAR_LIMIT); + if value == 0 || value > MAX_CELL_CHAR_LIMIT { + return Err(EnterpriseToolError::new( + "IMPORT_CELL_LIMIT_INVALID", + format!("cell_char_limit 必须在 1 到 {MAX_CELL_CHAR_LIMIT} 之间。"), + )); + } + Ok(value) +} + +pub fn sanitize_preview(mut preview: TableImportPreview, char_limit: usize) -> TableImportPreview { + for row in &mut preview.rows { + for cell in row { + if let Value::String(value) = cell { + let mut chars = value.chars(); + let truncated = chars.by_ref().take(char_limit).collect::(); + if chars.next().is_some() { + *value = format!("{truncated}…"); + } + } + } + } + preview +} + +pub fn structure_fingerprint(preview: &TableImportPreview, parse_options: &TableImportParseOptions) -> String { + let value = json!({ + "fileType": preview.file_type, + "columns": preview.columns, + "columnTypeHints": column_type_hints(preview), + "sheets": preview.sheets, + "parseOptions": parse_options, + }); + sha256_bytes(value.to_string().as_bytes()) +} + +pub fn column_type_hints(preview: &TableImportPreview) -> Vec { + (0..preview.columns.len()) + .map(|index| { + let kinds = preview + .rows + .iter() + .filter_map(|row| row.get(index)) + .filter_map(|value| match value { + Value::Null => None, + Value::Bool(_) => Some("boolean"), + Value::Number(number) if number.is_i64() || number.is_u64() => Some("integer"), + Value::Number(_) => Some("decimal"), + Value::String(_) => Some("text"), + Value::Array(_) | Value::Object(_) => Some("json"), + }) + .collect::>(); + if kinds.is_empty() { + "empty".to_string() + } else { + kinds.into_iter().collect::>().join("|") + } + }) + .collect() +} + +pub fn validate_staging_target(schema: &str, table: &str) -> Result<(), EnterpriseToolError> { + validate_identifier(schema, "schema")?; + validate_identifier(table, "table")?; + let allowed = comma_list_env("DBX_MCP_IMPORT_STAGING_SCHEMAS", &["staging"]); + if !allowed.iter().any(|candidate| candidate == schema) { + return Err(EnterpriseToolError::new( + "IMPORT_TARGET_NOT_STAGING", + format!("MCP 导入只能写入隔离 staging schema；当前允许：{}。", allowed.join(", ")), + )); + } + Ok(()) +} + +pub fn validate_mappings( + mappings: &[TableImportColumnMapping], + source_columns: &[String], +) -> Result<(), EnterpriseToolError> { + if mappings.is_empty() { + return Err(EnterpriseToolError::new("IMPORT_MAPPING_REQUIRED", "至少需要一个字段映射。")); + } + let source = source_columns.iter().collect::>(); + let mut targets = HashSet::new(); + for mapping in mappings { + if !source.contains(&mapping.source_column) { + return Err(EnterpriseToolError::new( + "IMPORT_SOURCE_COLUMN_NOT_FOUND", + format!("源字段 {} 不存在。", mapping.source_column), + )); + } + validate_identifier(&mapping.target_column, "target_column")?; + if let Some(data_type) = mapping.target_data_type.as_deref() { + validate_mcp_postgres_type(data_type)?; + } + if !targets.insert(mapping.target_column.as_str()) { + return Err(EnterpriseToolError::new( + "IMPORT_TARGET_COLUMN_DUPLICATED", + format!("目标字段 {} 被重复映射。", mapping.target_column), + )); + } + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub fn build_plan( + connection_id: String, + connection_name: String, + database: String, + schema: String, + table: String, + template_version: String, + file: FileIdentity, + structure_fingerprint: String, + source_format: Option, + parse_options: TableImportParseOptions, + mappings: Vec, + create_table: bool, + batch_size: usize, + date_time_format: Option, +) -> Result { + if template_version.trim().is_empty() { + return Err(EnterpriseToolError::new("IMPORT_TEMPLATE_VERSION_REQUIRED", "必须绑定已批准的模板版本。")); + } + let created_at_ms = unix_epoch_millis(); + let expires_at_ms = created_at_ms + DEFAULT_PLAN_TTL.as_millis(); + let plan_id = format!("plan-{}", Uuid::new_v4()); + let digest_source = json!({ + "planId": plan_id, + "connectionId": connection_id, + "database": database, + "schema": schema, + "table": table, + "templateVersion": template_version, + "file": file, + "structureFingerprint": structure_fingerprint, + "sourceFormat": source_format, + "parseOptions": parse_options, + "mappings": mappings, + "createTable": create_table, + "batchSize": batch_size, + "dateTimeFormat": date_time_format, + "expiresAtMs": expires_at_ms, + }); + let plan_digest = sha256_bytes(digest_source.to_string().as_bytes()); + Ok(PreparedImportPlan { + plan_id, + plan_digest, + created_at_ms, + expires_at_ms, + connection_id, + connection_name, + database, + schema, + table, + template_version, + file, + structure_fingerprint, + source_format, + parse_options, + mappings, + create_table, + batch_size, + date_time_format, + consumed: false, + }) +} + +pub async fn revalidate_plan_file(plan: &PreparedImportPlan) -> Result<(), EnterpriseToolError> { + let canonical = validate_import_file(&plan.file.canonical_path, false)?; + let current = file_identity(canonical).await?; + if current != plan.file { + return Err(EnterpriseToolError::new( + "IMPORT_FILE_CHANGED", + "文件在 prepare 与 start 之间发生变化；请重新预览并准备导入计划。", + )); + } + Ok(()) +} + +pub fn validate_vector_collection(collection: &str) -> Result<(), EnterpriseToolError> { + validate_identifier(collection, "collection")?; + let allowed = comma_list_env("DBX_MCP_VECTOR_COLLECTIONS", &["semantic_cards"]); + if !allowed.iter().any(|candidate| candidate == collection) { + return Err(EnterpriseToolError::new( + "VECTOR_COLLECTION_NOT_ALLOWED", + format!("集合 {collection} 不在允许列表中。"), + )); + } + Ok(()) +} + +pub fn vector_top_k(top_k: Option) -> Result { + let configured = env_usize("DBX_MCP_VECTOR_TOP_K_MAX", 20).clamp(1, HARD_VECTOR_TOP_K); + let top_k = top_k.unwrap_or(DEFAULT_VECTOR_TOP_K); + if top_k == 0 || top_k > configured { + return Err(EnterpriseToolError::new("VECTOR_TOP_K_INVALID", format!("top_k 必须在 1 到 {configured} 之间。"))); + } + Ok(top_k) +} + +pub fn validate_embedding(embedding: &[f32]) -> Result<(), EnterpriseToolError> { + let dimension = env_usize("DBX_MCP_VECTOR_DIMENSION", DEFAULT_VECTOR_DIMENSION).clamp(1, 65_536); + if embedding.len() != dimension { + return Err(EnterpriseToolError::new( + "VECTOR_DIMENSION_MISMATCH", + format!("向量维度为 {}，预期 {dimension}。", embedding.len()), + )); + } + if embedding.iter().any(|value| !value.is_finite()) { + return Err(EnterpriseToolError::new("VECTOR_VALUE_INVALID", "向量包含 NaN 或无穷值。")); + } + Ok(()) +} + +pub fn vector_output_fields(requested: Option>) -> Result, EnterpriseToolError> { + let allowed = comma_list_env("DBX_MCP_VECTOR_OUTPUT_FIELDS", DEFAULT_VECTOR_OUTPUT_FIELDS); + let requested = + requested.unwrap_or_else(|| DEFAULT_VECTOR_OUTPUT_FIELDS.iter().map(|value| value.to_string()).collect()); + if requested.is_empty() { + return Err(EnterpriseToolError::new("VECTOR_OUTPUT_FIELDS_REQUIRED", "至少需要一个输出字段。")); + } + let mut unique = HashSet::new(); + for field in &requested { + validate_identifier(field, "output_field")?; + if field == "embedding" || !allowed.iter().any(|allowed| allowed == field) { + return Err(EnterpriseToolError::new( + "VECTOR_OUTPUT_FIELD_NOT_ALLOWED", + format!("不允许返回字段 {field}。"), + )); + } + if !unique.insert(field) { + return Err(EnterpriseToolError::new("VECTOR_OUTPUT_FIELD_DUPLICATED", format!("输出字段 {field} 重复。"))); + } + } + Ok(requested) +} + +pub fn build_milvus_filter( + active_at: &str, + semantic_version: Option<&str>, + filters: &BTreeMap, +) -> Result { + if !valid_effective_timestamp(active_at) { + return Err(EnterpriseToolError::new( + "VECTOR_ACTIVE_AT_INVALID", + "active_at 必须是 YYYY-MM-DD 或 RFC 3339 时间。", + )); + } + let allowed = comma_list_env("DBX_MCP_VECTOR_FILTER_FIELDS", DEFAULT_VECTOR_FILTER_FIELDS); + let mut clauses = vec![ + format!("approval_status == {}", json_string("approved")), + format!("effective_from <= {}", json_string(active_at)), + format!("(effective_to == \"\" or effective_to >= {})", json_string(active_at)), + ]; + if let Some(semantic_version) = semantic_version.map(str::trim).filter(|value| !value.is_empty()) { + clauses.push(format!("semantic_version == {}", json_string(semantic_version))); + } + for (field, value) in filters { + validate_identifier(field, "filter_field")?; + if field == "approval_status" || field == "semantic_version" || !allowed.iter().any(|allowed| allowed == field) + { + return Err(EnterpriseToolError::new( + "VECTOR_FILTER_FIELD_NOT_ALLOWED", + format!("不允许过滤字段 {field}。"), + )); + } + clauses.push(filter_clause(field, value)?); + } + Ok(clauses.join(" and ")) +} + +pub fn read_semantic_jsonl(path: &Path, semantic_batch_id: &str) -> Result, EnterpriseToolError> { + if semantic_batch_id.trim().is_empty() { + return Err(EnterpriseToolError::new("SEMANTIC_BATCH_ID_REQUIRED", "semantic_batch_id 不能为空。")); + } + let source = std::fs::read_to_string(path) + .map_err(|error| EnterpriseToolError::new("VECTOR_JSONL_READ_FAILED", format!("读取 JSONL 失败：{error}")))?; + let expected_dimension = env_usize("DBX_MCP_VECTOR_DIMENSION", DEFAULT_VECTOR_DIMENSION); + let mut records = Vec::new(); + for (index, line) in source.lines().enumerate() { + let line_number = index + 1; + let line = line.trim(); + if line.is_empty() { + continue; + } + if line.len() > MAX_JSONL_RECORD_BYTES { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_RECORD_TOO_LARGE", + format!("JSONL 第 {line_number} 行超过 {} 字节。", MAX_JSONL_RECORD_BYTES), + )); + } + let mut record = serde_json::from_str::(line).map_err(|error| { + EnterpriseToolError::new("VECTOR_JSONL_INVALID", format!("JSONL 第 {line_number} 行无效：{error}")) + })?; + let object = record.as_object_mut().ok_or_else(|| { + EnterpriseToolError::new("VECTOR_JSONL_INVALID", format!("JSONL 第 {line_number} 行必须是对象。")) + })?; + for required in ["card_id", "approval_status", "semantic_version", "effective_from", "embedding"] { + if !object.contains_key(required) { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_MISSING", + format!("JSONL 第 {line_number} 行缺少 {required}。"), + )); + } + } + if object.get("approval_status").and_then(Value::as_str) != Some("approved") { + return Err(EnterpriseToolError::new( + "VECTOR_CARD_NOT_APPROVED", + format!("JSONL 第 {line_number} 行不是 approved 语义卡。"), + )); + } + for field in ["card_id", "semantic_version"] { + if object.get(field).and_then(Value::as_str).is_none_or(|value| value.trim().is_empty()) { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_INVALID", + format!("JSONL 第 {line_number} 行的 {field} 必须是非空字符串。"), + )); + } + } + if object.get("effective_from").and_then(Value::as_str).is_none_or(|value| !valid_effective_timestamp(value)) { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_INVALID", + format!("JSONL 第 {line_number} 行的 effective_from 不是有效日期或 RFC 3339 时间。"), + )); + } + match object.get("effective_to") { + None | Some(Value::Null) => { + object.insert("effective_to".to_string(), Value::String(String::new())); + } + Some(Value::String(value)) if value.is_empty() || valid_effective_timestamp(value) => {} + Some(_) => { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_INVALID", + format!("JSONL 第 {line_number} 行的 effective_to 必须为 null、空字符串或有效日期/RFC 3339 时间。"), + )); + } + } + if let Some(existing) = object.get("semantic_batch_id").and_then(Value::as_str) { + if existing != semantic_batch_id { + return Err(EnterpriseToolError::new( + "SEMANTIC_BATCH_ID_MISMATCH", + format!("JSONL 第 {line_number} 行的 semantic_batch_id 不一致。"), + )); + } + } + object.insert("semantic_batch_id".to_string(), Value::String(semantic_batch_id.to_string())); + object.entry("effective_to".to_string()).or_insert_with(|| Value::String(String::new())); + let embedding = object.get("embedding").and_then(Value::as_array).ok_or_else(|| { + EnterpriseToolError::new( + "VECTOR_JSONL_EMBEDDING_INVALID", + format!("JSONL 第 {line_number} 行 embedding 必须是数组。"), + ) + })?; + if embedding.len() != expected_dimension + || embedding.iter().any(|value| value.as_f64().is_none_or(|v| !v.is_finite())) + { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_EMBEDDING_INVALID", + format!("JSONL 第 {line_number} 行 embedding 必须是 {expected_dimension} 维有限数值数组。"), + )); + } + records.push(record); + } + if records.is_empty() { + return Err(EnterpriseToolError::new("VECTOR_JSONL_EMPTY", "JSONL 中没有可 upsert 的语义卡。")); + } + Ok(records) +} + +pub fn vector_upsert_batch_size(value: Option) -> Result { + let value = value.unwrap_or(DEFAULT_VECTOR_UPSERT_BATCH); + if value == 0 || value > MAX_VECTOR_UPSERT_BATCH { + return Err(EnterpriseToolError::new( + "VECTOR_UPSERT_BATCH_INVALID", + format!("batch_size 必须在 1 到 {MAX_VECTOR_UPSERT_BATCH} 之间。"), + )); + } + Ok(value) +} + +pub fn milvus_search_query( + database: &str, + collection: &str, + embedding: &[f32], + top_k: usize, + filter: &str, + output_fields: &[String], +) -> String { + format!( + "POST /v2/vectordb/entities/search\n{}", + json!({ + "dbName": if database.is_empty() { "default" } else { database }, + "collectionName": collection, + "data": [embedding], + "annsField": "embedding", + "limit": top_k, + "filter": filter, + "outputFields": output_fields, + }) + ) +} + +pub fn milvus_upsert_query(database: &str, collection: &str, records: &[Value]) -> String { + format!( + "POST /v2/vectordb/entities/upsert\n{}", + json!({ + "dbName": if database.is_empty() { "default" } else { database }, + "collectionName": collection, + "data": records, + }) + ) +} + +pub fn milvus_delete_batch_query(database: &str, collection: &str, semantic_batch_id: &str) -> String { + format!( + "POST /v2/vectordb/entities/delete\n{}", + json!({ + "dbName": if database.is_empty() { "default" } else { database }, + "collectionName": collection, + "filter": format!("semantic_batch_id == {}", json_string(semantic_batch_id)), + }) + ) +} + +pub fn query_result_rows(result: dbx_core::db::QueryResult) -> Vec { + result + .rows + .into_iter() + .map(|row| Value::Object(result.columns.iter().cloned().zip(row).collect::>())) + .collect() +} + +fn filter_clause(field: &str, value: &Value) -> Result { + match value { + Value::String(value) => Ok(format!("{field} == {}", json_string(value))), + Value::Number(value) => Ok(format!("{field} == {value}")), + Value::Bool(value) => Ok(format!("{field} == {value}")), + Value::Array(values) if !values.is_empty() && values.len() <= 50 => { + let values = values + .iter() + .map(|value| match value { + Value::String(value) => Ok(json_string(value)), + Value::Number(value) => Ok(value.to_string()), + Value::Bool(value) => Ok(value.to_string()), + _ => Err(EnterpriseToolError::new( + "VECTOR_FILTER_VALUE_INVALID", + "过滤数组只允许字符串、数字或布尔值。", + )), + }) + .collect::, _>>()?; + Ok(format!("{field} in [{}]", values.join(", "))) + } + _ => Err(EnterpriseToolError::new( + "VECTOR_FILTER_VALUE_INVALID", + "过滤值只允许标量，或最多 50 个标量的非空数组。", + )), + } +} + +fn validate_identifier(value: &str, field: &str) -> Result<(), EnterpriseToolError> { + let valid = !value.is_empty() + && value.len() <= 63 + && value + .chars() + .enumerate() + .all(|(index, ch)| ch == '_' || ch.is_ascii_alphanumeric() && (index > 0 || !ch.is_ascii_digit())); + if !valid { + return Err(EnterpriseToolError::new( + "IDENTIFIER_INVALID", + format!("{field} 只能使用 1～63 位英文字母、数字或下划线，且不能以数字开头。"), + )); + } + Ok(()) +} + +fn validate_mcp_postgres_type(data_type: &str) -> Result<(), EnterpriseToolError> { + let normalized = data_type.trim().to_ascii_uppercase().replace(' ', ""); + let simple = matches!(normalized.as_str(), "TEXT" | "TIMESTAMPTZ" | "DATE" | "BOOLEAN" | "JSONB"); + let numeric = normalized + .strip_prefix("NUMERIC(") + .and_then(|value| value.strip_suffix(')')) + .and_then(|value| value.split_once(',')) + .and_then(|(precision, scale)| Some((precision.parse::().ok()?, scale.parse::().ok()?))) + .is_some_and(|(precision, scale)| (1..=100).contains(&precision) && scale <= precision); + if !simple && !numeric { + return Err(EnterpriseToolError::new( + "IMPORT_TARGET_TYPE_NOT_ALLOWED", + "v1 MCP 导入类型只允许 TEXT、NUMERIC(p,s)、TIMESTAMPTZ、DATE、BOOLEAN 或 JSONB。", + )); + } + Ok(()) +} + +fn json_string(value: &str) -> String { + serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string()) +} + +fn valid_effective_timestamp(value: &str) -> bool { + let value = value.trim(); + (10..=40).contains(&value.len()) + && value.as_bytes().get(4) == Some(&b'-') + && value.as_bytes().get(7) == Some(&b'-') + && value.chars().all(|ch| ch.is_ascii_digit() || matches!(ch, '-' | ':' | 'T' | 'Z' | '+' | '.')) +} + +fn sha256_bytes(value: &[u8]) -> String { + let mut hasher = Sha256::new(); + hasher.update(value); + format!("{:x}", hasher.finalize()) +} + +fn unix_epoch_millis() -> u128 { + SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_millis() +} + +fn env_u64(name: &str, default: u64) -> u64 { + std::env::var(name).ok().and_then(|value| value.parse().ok()).filter(|value| *value > 0).unwrap_or(default) +} + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name).ok().and_then(|value| value.parse().ok()).filter(|value| *value > 0).unwrap_or(default) +} + +fn comma_list_env(name: &str, defaults: &[&str]) -> Vec { + let values = std::env::var(name) + .ok() + .map(|value| { + value.split(',').map(str::trim).filter(|value| !value.is_empty()).map(str::to_string).collect::>() + }) + .filter(|values| !values.is_empty()); + values.unwrap_or_else(|| defaults.iter().map(|value| value.to_string()).collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn file_policy_rejects_outside_root_and_symlink() { + let root = tempfile::tempdir().unwrap(); + let outside = tempfile::tempdir().unwrap(); + let source = root.path().join("input.csv"); + std::fs::write(&source, "id,name\n1,Ada\n").unwrap(); + assert_eq!( + validate_import_file_with_roots(source.to_str().unwrap(), &[root.path().to_path_buf()], false, 1024) + .unwrap(), + std::fs::canonicalize(source).unwrap() + ); + + let outside_file = outside.path().join("outside.csv"); + std::fs::write(&outside_file, "id\n1\n").unwrap(); + assert_eq!( + validate_import_file_with_roots(outside_file.to_str().unwrap(), &[root.path().to_path_buf()], false, 1024,) + .unwrap_err() + .code, + "IMPORT_PATH_OUTSIDE_ROOTS" + ); + + #[cfg(unix)] + { + let link = root.path().join("link.csv"); + std::os::unix::fs::symlink(&outside_file, &link).unwrap(); + assert_eq!( + validate_import_file_with_roots(link.to_str().unwrap(), &[root.path().to_path_buf()], false, 1024,) + .unwrap_err() + .code, + "IMPORT_SYMLINK_REJECTED" + ); + } + } + + #[test] + fn milvus_filter_forces_approval_and_semantic_version() { + let filters = BTreeMap::from([("business_domain".to_string(), json!("交易"))]); + let filter = build_milvus_filter("2026-08-25T10:00:00+08:00", Some("semantic-v3"), &filters).unwrap(); + assert!(filter.contains("approval_status == \"approved\"")); + assert!(filter.contains("effective_from <= \"2026-08-25T10:00:00+08:00\"")); + assert!(filter.contains("semantic_version == \"semantic-v3\"")); + assert!(filter.contains("business_domain == \"交易\"")); + + let forbidden = BTreeMap::from([("approval_status".to_string(), json!("draft"))]); + assert_eq!( + build_milvus_filter("2026-08-25", None, &forbidden).unwrap_err().code, + "VECTOR_FILTER_FIELD_NOT_ALLOWED" + ); + } + + #[test] + fn mappings_reject_missing_and_duplicate_targets() { + let source = vec!["订单号".to_string(), "金额".to_string()]; + let duplicated = vec![ + TableImportColumnMapping { + source_column: "订单号".to_string(), + target_column: "order_id".to_string(), + target_data_type: None, + }, + TableImportColumnMapping { + source_column: "金额".to_string(), + target_column: "order_id".to_string(), + target_data_type: None, + }, + ]; + assert_eq!(validate_mappings(&duplicated, &source).unwrap_err().code, "IMPORT_TARGET_COLUMN_DUPLICATED"); + + let unsafe_type = vec![TableImportColumnMapping { + source_column: "订单号".to_string(), + target_column: "order_id".to_string(), + target_data_type: Some("TEXT DEFAULT current_user".to_string()), + }]; + assert_eq!(validate_mappings(&unsafe_type, &source).unwrap_err().code, "IMPORT_TARGET_TYPE_NOT_ALLOWED"); + } + + #[test] + fn semantic_jsonl_normalizes_null_effective_to() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("cards.jsonl"); + let record = json!({ + "card_id": "metric-order-gmv", + "approval_status": "approved", + "semantic_version": "semantic-v1", + "effective_from": "2026-08-25T00:00:00+08:00", + "effective_to": null, + "embedding": vec![0.0_f32; DEFAULT_VECTOR_DIMENSION], + }); + std::fs::write(&path, format!("{}\n", record)).unwrap(); + + let records = read_semantic_jsonl(&path, "semantic-batch-1").unwrap(); + assert_eq!(records[0].get("effective_to"), Some(&json!(""))); + assert_eq!(records[0].get("semantic_batch_id"), Some(&json!("semantic-batch-1"))); + } + + #[tokio::test] + async fn prepared_plan_is_single_use() { + let plan = build_plan( + "connection-1".to_string(), + "运营组数据写入".to_string(), + "enterprise".to_string(), + "staging".to_string(), + "orders_batch_1".to_string(), + "orders-v1".to_string(), + FileIdentity { + canonical_path: "/allowed/orders.csv".to_string(), + size_bytes: 42, + modified_nanos: 1, + sha256: "a".repeat(64), + }, + "structure-v1".to_string(), + Some(TableImportSourceFormat::Csv), + TableImportParseOptions::default(), + vec![TableImportColumnMapping { + source_column: "订单号".to_string(), + target_column: "order_id".to_string(), + target_data_type: Some("TEXT".to_string()), + }], + true, + 1_000, + None, + ) + .unwrap(); + let plan_id = plan.plan_id.clone(); + let runtime = EnterpriseRuntime::default(); + runtime.insert_plan(plan).await; + + assert_eq!(runtime.consume_plan(&plan_id).await.unwrap().plan_id, plan_id); + assert_eq!(runtime.consume_plan(&plan_id).await.unwrap_err().code, "IMPORT_PLAN_ALREADY_USED"); + } +} diff --git a/crates/dbx-mcp/src/lib.rs b/crates/dbx-mcp/src/lib.rs index 35c30452f8..51e6da52a3 100644 --- a/crates/dbx-mcp/src/lib.rs +++ b/crates/dbx-mcp/src/lib.rs @@ -1,4 +1,5 @@ pub mod backend; +mod enterprise_tools; pub mod paths; pub mod server; pub mod session; diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 1b07498ca3..bc8a3f197f 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -5,11 +5,20 @@ use rmcp::{ model::{CallToolResult, ContentBlock, Implementation, ServerCapabilities, ServerInfo}, schemars, tool, tool_handler, tool_router, ServerHandler, }; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use serde_json::json; use uuid::Uuid; use crate::backend::{format_query_result, new_connection_config, parse_database_type, ConnectionSummary, DbxBackend}; +use crate::enterprise_tools::{ + build_milvus_filter, build_plan, cell_char_limit, column_type_hints, file_identity, milvus_delete_batch_query, + milvus_search_query, milvus_upsert_query, preview_limit, query_result_rows, read_semantic_jsonl, + revalidate_plan_file, sanitize_preview, structure_fingerprint, validate_embedding, validate_import_file, + validate_mappings, validate_staging_target, validate_vector_collection, vector_output_fields, vector_top_k, + vector_upsert_batch_size, EnterpriseRuntime, EnterpriseToolError, ImportStatusRequest, PrepareTableImportRequest, + PreviewImportFileRequest, StartTableImportRequest, VectorDeleteByBatchRequest, VectorSearchRequest, + VectorUpsertFileRequest, FORMAT_VERSION, STRUCTURE_PROFILE_ROWS, +}; use crate::mongo::{self, MongoCommand, MongoSafetyError}; use crate::session::{McpSession, McpSessionStore}; use dbx_core::{ @@ -30,7 +39,7 @@ use dbx_core::{ #[derive(Debug, Deserialize, schemars::JsonSchema)] pub struct ListConnectionsRequest {} -#[derive(Debug, Deserialize, schemars::JsonSchema)] +#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] pub struct ConnectionSelector { #[schemars(description = "Unique ID of the DBX connection")] #[schemars(extend("type" = "string"))] @@ -193,6 +202,8 @@ pub struct DbxMcpServer { backend: Arc, scope: McpScope, sessions: Arc, + enterprise: Arc, + web_mode: bool, tool_router: ToolRouter, } @@ -258,7 +269,14 @@ impl DbxMcpServer { tool_router.disable_route("dbx_open_table"); tool_router.disable_route("dbx_execute_and_show"); } - Self { backend, scope, sessions: McpSessionStore::new(), tool_router } + Self { + backend, + scope, + sessions: McpSessionStore::new(), + enterprise: Arc::new(EnterpriseRuntime::default()), + web_mode, + tool_router, + } } async fn close_backend_sessions_best_effort(&self, sessions: Vec) { @@ -354,6 +372,510 @@ impl DbxMcpServer { } } + #[tool( + name = "dbx_preview_import_file", + description = "Preview an allowed local Excel/CSV/TSV/JSON file without writing data. Returns bounded rows, sheets, hashes, columns, row counts and a structure fingerprint. Local DBX Desktop/MCP mode only." + )] + async fn preview_import_file(&self, Parameters(request): Parameters) -> CallToolResult { + if self.web_mode { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_UNSUPPORTED_IN_WEB_MODE_V1", + "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", + )); + } + let path = match validate_import_file(&request.file_path, false) { + Ok(path) => path, + Err(error) => return enterprise_error(error), + }; + let identity = match file_identity(path.clone()).await { + Ok(identity) => identity, + Err(error) => return enterprise_error(error), + }; + let parse_options = match request.parse_options.into_core() { + Ok(options) => options, + Err(error) => return enterprise_error(error), + }; + let preview_rows = match preview_limit(request.preview_rows) { + Ok(limit) => limit, + Err(error) => return enterprise_error(error), + }; + let char_limit = match cell_char_limit(request.cell_char_limit) { + Ok(limit) => limit, + Err(error) => return enterprise_error(error), + }; + let preview = dbx_core::table_import::preview_table_import_file_with_request( + dbx_core::table_import::TableImportPreviewRequest { + file_path: path.to_string_lossy().to_string(), + source_ref: Some(identity.sha256.clone()), + source_format: request.source_format.map(Into::into), + parse_options: parse_options.clone(), + preview_limit: Some(STRUCTURE_PROFILE_ROWS), + }, + ) + .await; + let mut preview = match preview { + Ok(preview) => sanitize_preview(preview, char_limit), + Err(error) => { + return enterprise_error(EnterpriseToolError::new("IMPORT_PREVIEW_FAILED", error)); + } + }; + let fingerprint = structure_fingerprint(&preview, &parse_options); + let type_hints = column_type_hints(&preview); + preview.rows.truncate(preview_rows); + let title_row = parse_options.title_row.or_else(|| parse_options.has_header.unwrap_or(true).then_some(1)); + let data_start_row = parse_options.data_start_row.unwrap_or_else(|| title_row.map_or(1, |row| row + 1)); + let last_row = data_start_row.saturating_add(preview.total_rows.saturating_sub(1)); + let summary = format!( + "已只读剖析 {}：{} 列，{} 行，返回 {} 行预览。", + preview.file_name, + preview.columns.len(), + preview.total_rows, + preview.rows.len() + ); + structured_success( + summary, + json!({ + "file": identity, + "structureFingerprint": fingerprint, + "columnTypeHints": type_hints, + "usedRange": { + "firstRow": title_row.unwrap_or(data_start_row), + "lastRow": last_row, + "firstColumn": if preview.columns.is_empty() { 0 } else { 1 }, + "lastColumn": preview.columns.len(), + }, + "preview": preview, + }), + ) + } + + #[tool( + name = "dbx_prepare_table_import", + description = "Validate an approved-template mapping and immutable local source, then create a single-use import plan valid for 30 minutes. This tool never writes data." + )] + async fn prepare_table_import(&self, Parameters(request): Parameters) -> CallToolResult { + if self.web_mode { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_UNSUPPORTED_IN_WEB_MODE_V1", + "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", + )); + } + if let Err(error) = validate_staging_target(&request.schema, &request.table) { + return enterprise_error(error); + } + let resolved = match self.resolve_connection(&request.selector).await { + Ok(resolved) => resolved, + Err(error) => return error, + }; + let database = match self.resolve_database(request.database.clone(), &resolved.connection) { + Ok(database) => database, + Err(error) => return error, + }; + if resolved.connection.db_type != DatabaseType::Postgres { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_REQUIRES_POSTGRES_V1", + "v1 MCP 表导入只允许 PostgreSQL staging 连接。", + )); + } + if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &database) { + return error; + } + let path = match validate_import_file(&request.file_path, false) { + Ok(path) => path, + Err(error) => return enterprise_error(error), + }; + let identity = match file_identity(path.clone()).await { + Ok(identity) => identity, + Err(error) => return enterprise_error(error), + }; + let parse_options = match request.parse_options.into_core() { + Ok(options) => options, + Err(error) => return enterprise_error(error), + }; + let source_format = request.source_format.map(Into::into); + let preview = dbx_core::table_import::preview_table_import_file_with_request( + dbx_core::table_import::TableImportPreviewRequest { + file_path: path.to_string_lossy().to_string(), + source_ref: Some(identity.sha256.clone()), + source_format, + parse_options: parse_options.clone(), + preview_limit: Some(STRUCTURE_PROFILE_ROWS), + }, + ) + .await; + let preview = match preview { + Ok(preview) => preview, + Err(error) => return enterprise_error(EnterpriseToolError::new("IMPORT_PREVIEW_FAILED", error)), + }; + let mappings = request.mappings.into_iter().map(Into::into).collect::>(); + if let Err(error) = validate_mappings(&mappings, &preview.columns) { + return enterprise_error(error); + } + if !request.create_table { + let columns = match self + .backend + .get_columns(&resolved.connection, &database, &request.schema, &request.table) + .await + { + Ok(columns) => columns, + Err(error) => return enterprise_error(EnterpriseToolError::new("IMPORT_TARGET_CHECK_FAILED", error)), + }; + let target_columns = + columns.iter().map(|column| column.name.as_str()).collect::>(); + if let Some(missing) = + mappings.iter().find(|mapping| !target_columns.contains(mapping.target_column.as_str())) + { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_TARGET_COLUMN_NOT_FOUND", + format!("目标表缺少字段 {}。", missing.target_column), + )); + } + } + let batch_size = request.batch_size.unwrap_or(1_000); + if !(1..=50_000).contains(&batch_size) { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_BATCH_SIZE_INVALID", + "batch_size 必须在 1 到 50000 之间。", + )); + } + let fingerprint = structure_fingerprint(&preview, &parse_options); + let plan = match build_plan( + resolved.connection.id.clone(), + resolved.connection.name.clone(), + database, + request.schema, + request.table, + request.template_version, + identity, + fingerprint, + source_format, + parse_options, + mappings, + request.create_table, + batch_size, + request.date_time_format, + ) { + Ok(plan) => plan, + Err(error) => return enterprise_error(error), + }; + self.enterprise.insert_plan(plan.clone()).await; + structured_success( + format!("导入计划 {} 已准备；绑定模板 {}，30 分钟内可启动一次。", plan.plan_id, plan.template_version), + json!({ "plan": plan }), + ) + } + + #[tool( + name = "dbx_start_table_import", + description = "Start a previously prepared single-use PostgreSQL staging import. The file, target, parsing options and mappings are revalidated and cannot be replaced." + )] + async fn start_table_import(&self, Parameters(request): Parameters) -> CallToolResult { + if self.web_mode { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_UNSUPPORTED_IN_WEB_MODE_V1", + "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", + )); + } + let plan = match self.enterprise.consume_plan(request.plan_id.trim()).await { + Ok(plan) => plan, + Err(error) => return enterprise_error(error), + }; + if let Err(error) = revalidate_plan_file(&plan).await { + return enterprise_error(error); + } + let selector = ConnectionSelector { connection_id: Some(plan.connection_id.clone()), connection_name: None }; + let resolved = match self.resolve_connection(&selector).await { + Ok(resolved) => resolved, + Err(error) => return error, + }; + if resolved.connection.db_type != DatabaseType::Postgres { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_REQUIRES_POSTGRES_V1", + "计划绑定的连接已不再是 PostgreSQL。", + )); + } + if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &plan.database) { + return error; + } + let job = self.enterprise.create_job(&plan).await; + let import_id = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).import_id.clone(); + let import_request = plan.to_import_request(import_id.clone()); + let backend = self.backend.clone(); + let connection = resolved.connection.clone(); + let task_job = job.clone(); + tokio::spawn(async move { + let progress_job = task_job.clone(); + let progress = Arc::new(move |mut update: dbx_core::table_import::TableImportProgress| { + if update.status == dbx_core::table_import::TableImportStatus::Done { + update.status = dbx_core::table_import::TableImportStatus::Running; + update.phase = dbx_core::table_import::TableImportPhase::Finalizing; + } + progress_job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).apply_progress(update); + }); + let result = backend + .import_table_file_for_mcp(&connection, import_request, task_job.cancelled.clone(), progress) + .await; + let mut snapshot = task_job.snapshot.lock().unwrap_or_else(|error| error.into_inner()); + match result { + Ok(summary) => { + snapshot.status = dbx_core::table_import::TableImportStatus::Done; + snapshot.phase = dbx_core::table_import::TableImportPhase::Done; + snapshot.rows_imported = summary.rows_imported; + snapshot.total_rows = summary.total_rows; + snapshot.elapsed_ms = summary.elapsed_ms; + snapshot.summary = Some(summary); + snapshot.error = None; + } + Err(error) => { + snapshot.status = if task_job.cancelled.load(std::sync::atomic::Ordering::Acquire) { + dbx_core::table_import::TableImportStatus::Cancelled + } else { + dbx_core::table_import::TableImportStatus::Error + }; + snapshot.error = Some(error); + } + } + }); + structured_success( + format!("导入任务 {import_id} 已启动；只会写入隔离 staging。"), + json!({ "importId": import_id, "planId": plan.plan_id, "status": "running" }), + ) + } + + #[tool( + name = "dbx_get_import_status", + description = "Get machine-readable progress, row counts, bytes, elapsed time, errors and final summary for an MCP table import." + )] + async fn get_import_status(&self, Parameters(request): Parameters) -> CallToolResult { + if self.web_mode { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_UNSUPPORTED_IN_WEB_MODE_V1", + "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", + )); + } + let Some(job) = self.enterprise.job(request.import_id.trim()).await else { + return enterprise_error(EnterpriseToolError::new("IMPORT_JOB_NOT_FOUND", "没有找到指定导入任务。")); + }; + let snapshot = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).clone(); + structured_success( + format!("导入任务 {} 当前状态：{:?} / {:?}。", snapshot.import_id, snapshot.status, snapshot.phase), + json!({ "job": snapshot }), + ) + } + + #[tool( + name = "dbx_cancel_import", + description = "Request cancellation of an MCP staging import. Cancellation stops subsequent batches and never publishes staging data to core or mart." + )] + async fn cancel_import(&self, Parameters(request): Parameters) -> CallToolResult { + if self.web_mode { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_UNSUPPORTED_IN_WEB_MODE_V1", + "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", + )); + } + let snapshot = match self.enterprise.cancel_job(request.import_id.trim()).await { + Ok(snapshot) => snapshot, + Err(error) => return enterprise_error(error), + }; + structured_success( + format!("已向导入任务 {} 发送取消请求。", snapshot.import_id), + json!({ "importId": snapshot.import_id, "cancellationRequested": true, "lastStatus": snapshot.status }), + ) + } + + #[tool( + name = "dbx_vector_search", + description = "Search an allowed Milvus semantic collection with a bounded vector and safe equality filters. Approval status and semantic version are always enforced." + )] + async fn vector_search(&self, Parameters(request): Parameters) -> CallToolResult { + let resolved = match self.resolve_connection(&request.selector).await { + Ok(resolved) => resolved, + Err(error) => return error, + }; + if resolved.connection.db_type != DatabaseType::Milvus { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_REQUIRES_MILVUS", + "dbx_vector_search 只接受 Milvus 连接。", + )); + } + if let Err(error) = validate_vector_collection(&request.collection) { + return enterprise_error(error); + } + if let Err(error) = validate_embedding(&request.embedding) { + return enterprise_error(error); + } + let top_k = match vector_top_k(request.top_k) { + Ok(value) => value, + Err(error) => return enterprise_error(error), + }; + let output_fields = match vector_output_fields(request.output_fields) { + Ok(fields) => fields, + Err(error) => return enterprise_error(error), + }; + let filter = + match build_milvus_filter(&request.active_at, request.semantic_version.as_deref(), &request.filters) { + Ok(filter) => filter, + Err(error) => return enterprise_error(error), + }; + let database = match self.resolve_database(request.database, &resolved.connection) { + Ok(database) => database, + Err(error) => return error, + }; + let query = + milvus_search_query(&database, &request.collection, &request.embedding, top_k, &filter, &output_fields); + match self.backend.execute_query(&resolved.connection, &database, &query, Some(top_k), Some(30)).await { + Ok(result) => { + let rows = query_result_rows(result); + structured_success( + format!("Milvus 语义检索返回 {} 条已批准结果。", rows.len()), + json!({ + "collection": request.collection, + "activeAt": request.active_at, + "semanticVersion": request.semantic_version, + "topK": top_k, + "filter": filter, + "rows": rows, + }), + ) + } + Err(error) => enterprise_error(EnterpriseToolError::new("VECTOR_SEARCH_FAILED", error)), + } + } + + #[tool( + name = "dbx_vector_upsert_file", + description = "Upsert approved semantic cards from one allowed JSONL file into an allowed Milvus collection. Requires MCP safe-write permission." + )] + async fn vector_upsert_file(&self, Parameters(request): Parameters) -> CallToolResult { + let resolved = match self.resolve_connection(&request.selector).await { + Ok(resolved) => resolved, + Err(error) => return error, + }; + if resolved.connection.db_type != DatabaseType::Milvus { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_REQUIRES_MILVUS", + "dbx_vector_upsert_file 只接受 Milvus 连接。", + )); + } + let database = match self.resolve_database(request.database, &resolved.connection) { + Ok(database) => database, + Err(error) => return error, + }; + if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &database) { + return error; + } + if let Err(error) = validate_vector_collection(&request.collection) { + return enterprise_error(error); + } + let path = match validate_import_file(&request.file_path, true) { + Ok(path) => path, + Err(error) => return enterprise_error(error), + }; + let identity = match file_identity(path.clone()).await { + Ok(identity) => identity, + Err(error) => return enterprise_error(error), + }; + let semantic_batch_id = request.semantic_batch_id.clone(); + let records = match tokio::task::spawn_blocking(move || read_semantic_jsonl(&path, &semantic_batch_id)).await { + Ok(Ok(records)) => records, + Ok(Err(error)) => return enterprise_error(error), + Err(error) => { + return enterprise_error(EnterpriseToolError::new("VECTOR_JSONL_READ_FAILED", error.to_string())) + } + }; + let current_identity = match file_identity(std::path::PathBuf::from(&identity.canonical_path)).await { + Ok(identity) => identity, + Err(error) => return enterprise_error(error), + }; + if current_identity != identity { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_JSONL_CHANGED", + "语义 JSONL 在校验期间发生变化；未向 Milvus 写入。", + )); + } + let batch_size = match vector_upsert_batch_size(request.batch_size) { + Ok(batch_size) => batch_size, + Err(error) => return enterprise_error(error), + }; + let mut upserted = 0usize; + for chunk in records.chunks(batch_size) { + let query = milvus_upsert_query(&database, &request.collection, chunk); + if let Err(error) = + self.backend.execute_query(&resolved.connection, &database, &query, Some(1), Some(60)).await + { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_UPSERT_FAILED", + format!("已提交 {upserted} 条后失败；语义版本尚未激活，可按 semantic_batch_id 撤销：{error}"), + )); + } + upserted += chunk.len(); + } + structured_success( + format!("已向 Milvus upsert {upserted} 张 approved 语义卡。"), + json!({ + "collection": request.collection, + "semanticBatchId": request.semantic_batch_id, + "upserted": upserted, + "source": identity, + }), + ) + } + + #[tool( + name = "dbx_vector_delete_by_batch", + description = "Delete one unpublished Milvus semantic batch by exact semantic_batch_id. Requires high-risk MCP permission; arbitrary filter expressions are never accepted." + )] + async fn vector_delete_by_batch( + &self, + Parameters(request): Parameters, + ) -> CallToolResult { + let resolved = match self.resolve_connection(&request.selector).await { + Ok(resolved) => resolved, + Err(error) => return error, + }; + if resolved.connection.db_type != DatabaseType::Milvus { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_REQUIRES_MILVUS", + "dbx_vector_delete_by_batch 只接受 Milvus 连接。", + )); + } + let database = match self.resolve_database(request.database, &resolved.connection) { + Ok(database) => database, + Err(error) => return error, + }; + if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &database) { + return error; + } + if !resolved.policy.allow_dangerous_sql { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_HIGH_RISK_WRITE_BLOCKED", + "按批删除语义向量需要开启 DBX MCP 高风险写入权限。", + )); + } + if let Err(error) = validate_vector_collection(&request.collection) { + return enterprise_error(error); + } + if request.semantic_batch_id.trim().is_empty() || request.semantic_batch_id.len() > 200 { + return enterprise_error(EnterpriseToolError::new( + "SEMANTIC_BATCH_ID_INVALID", + "semantic_batch_id 必须是 1～200 个字符。", + )); + } + let query = milvus_delete_batch_query(&database, &request.collection, &request.semantic_batch_id); + match self.backend.execute_query(&resolved.connection, &database, &query, Some(1), Some(30)).await { + Ok(_) => structured_success( + format!("已删除未发布语义批次 {}。", request.semantic_batch_id), + json!({ + "collection": request.collection, + "semanticBatchId": request.semantic_batch_id, + "deleted": true, + }), + ), + Err(error) => enterprise_error(EnterpriseToolError::new("VECTOR_DELETE_FAILED", error)), + } + } + #[tool( name = "dbx_execute_query", description = "Execute a SQL query on a database connection (max 100 rows returned)" @@ -1033,10 +1555,64 @@ fn text(value: impl Into) -> CallToolResult { CallToolResult::success(vec![ContentBlock::text(value)]) } +fn structured_success(summary: impl Into, data: serde_json::Value) -> CallToolResult { + let summary = summary.into(); + let mut result = CallToolResult::structured(json!({ + "formatVersion": FORMAT_VERSION, + "ok": true, + "summary": summary, + "data": data, + })); + result.content = vec![ContentBlock::text(summary)]; + result +} + +fn enterprise_error(error: EnterpriseToolError) -> CallToolResult { + let summary = format!("错误 [{}]：{}", error.code, error.message); + let mut result = CallToolResult::structured_error(json!({ + "formatVersion": FORMAT_VERSION, + "ok": false, + "error": { + "code": error.code, + "message": error.message, + }, + })); + result.content = vec![ContentBlock::text(summary)]; + result +} + fn tool_error(code: &str, message: impl Into) -> CallToolResult { CallToolResult::error(vec![ContentBlock::text(format!("Error [{code}]: {}", message.into()))]) } +// CallToolResult 是 MCP 原生错误载荷；保持不装箱可让所有工具直接返回一致错误。 +#[allow(clippy::result_large_err)] +fn validate_safe_write_connection( + connection: &dbx_core::models::connection::ConnectionConfig, + policy: &McpGlobalPolicy, + database: &str, +) -> Result<(), CallToolResult> { + if policy.read_only { + return Err(enterprise_error(EnterpriseToolError::new( + "MCP_READ_ONLY", + "DBX 全局 MCP 只读模式已开启，写入被阻断。", + ))); + } + if connection.read_only { + return Err(enterprise_error(EnterpriseToolError::new( + "CONNECTION_READ_ONLY", + format!("连接 {} 已启用只读保护。", connection.name), + ))); + } + if is_production_database(connection, database) { + return Err(enterprise_error(EnterpriseToolError::new( + "PRODUCTION_WRITE_BLOCKED", + "MCP 不允许对标记为生产的数据库执行文件或向量写入。", + ))); + } + Ok(()) +} + fn backend_tool_error(default_code: &str, error: impl Into) -> CallToolResult { let error = error.into(); for code in [ @@ -1327,6 +1903,8 @@ mod tests { use dbx_core::models::connection::ConnectionConfig; use std::collections::HashSet; + static IMPORT_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + struct FakeBackend { connections: Vec, recorded_arguments: std::sync::Mutex>, @@ -1471,7 +2049,7 @@ mod tests { let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); let tools = server.tool_router.list_all(); let names = tools.iter().map(|tool| tool.name.as_ref()).collect::>(); - assert_eq!(tools.len(), 13); + assert_eq!(tools.len(), 21); assert!(names.contains(&"dbx_list_connections")); assert!(names.contains(&"dbx_list_tables")); assert!(names.contains(&"dbx_describe_table")); @@ -1485,6 +2063,91 @@ mod tests { assert!(names.contains(&"dbx_execute_and_show")); assert!(names.contains(&"dbx_open_session")); assert!(names.contains(&"dbx_close_session")); + assert!(names.contains(&"dbx_preview_import_file")); + assert!(names.contains(&"dbx_prepare_table_import")); + assert!(names.contains(&"dbx_start_table_import")); + assert!(names.contains(&"dbx_get_import_status")); + assert!(names.contains(&"dbx_cancel_import")); + assert!(names.contains(&"dbx_vector_search")); + assert!(names.contains(&"dbx_vector_upsert_file")); + assert!(names.contains(&"dbx_vector_delete_by_batch")); + } + + #[tokio::test] + async fn import_plan_rejects_source_changed_after_prepare() { + let _environment_guard = IMPORT_ENV_LOCK.lock().unwrap(); + let previous_roots = std::env::var_os("DBX_MCP_IMPORT_ROOTS"); + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("orders.csv"); + std::fs::write(&source, "order_id,amount\nA001,10.00\n").unwrap(); + std::env::set_var("DBX_MCP_IMPORT_ROOTS", directory.path()); + + let postgres = connection("write-pg", "运营组数据写入", "postgres", "enterprise"); + let server = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend { connections: vec![postgres], ..Default::default() }), + McpScope::default(), + false, + ); + let preview = server + .preview_import_file(Parameters(PreviewImportFileRequest { + file_path: source.to_string_lossy().to_string(), + source_format: None, + parse_options: Default::default(), + preview_rows: Some(5), + cell_char_limit: Some(100), + })) + .await; + assert_eq!(preview.is_error, Some(false)); + assert_eq!(preview.structured_content.as_ref().and_then(|value| value.get("formatVersion")), Some(&json!(1))); + + let prepared = server + .prepare_table_import(Parameters(PrepareTableImportRequest { + selector: ConnectionSelector { connection_id: Some("write-pg".to_string()), connection_name: None }, + database: Some("enterprise".to_string()), + schema: "staging".to_string(), + table: "orders_batch_1".to_string(), + template_version: "orders-v1".to_string(), + file_path: source.to_string_lossy().to_string(), + source_format: None, + parse_options: Default::default(), + mappings: vec![ + crate::enterprise_tools::McpImportColumnMapping { + source_column: "order_id".to_string(), + target_column: "order_id".to_string(), + target_data_type: Some("TEXT".to_string()), + }, + crate::enterprise_tools::McpImportColumnMapping { + source_column: "amount".to_string(), + target_column: "amount".to_string(), + target_data_type: Some("NUMERIC(20,2)".to_string()), + }, + ], + create_table: true, + batch_size: Some(1_000), + date_time_format: None, + })) + .await; + assert_eq!(prepared.is_error, Some(false)); + let plan_id = prepared + .structured_content + .as_ref() + .and_then(|value| value.pointer("/data/plan/planId")) + .and_then(serde_json::Value::as_str) + .unwrap() + .to_string(); + + std::fs::write(&source, "order_id,amount\nA001,10.00\nA002,20.00\n").unwrap(); + let started = server.start_table_import(Parameters(StartTableImportRequest { plan_id })).await; + assert_eq!(started.is_error, Some(true)); + assert_eq!( + started.structured_content.as_ref().and_then(|value| value.pointer("/error/code")), + Some(&json!("IMPORT_FILE_CHANGED")) + ); + + match previous_roots { + Some(value) => std::env::set_var("DBX_MCP_IMPORT_ROOTS", value), + None => std::env::remove_var("DBX_MCP_IMPORT_ROOTS"), + } } #[test] @@ -1517,7 +2180,15 @@ mod tests { let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); let tools = server.tool_router.list_all(); - for tool_name in ["dbx_execute_query", "dbx_list_tables", "dbx_open_session"] { + for tool_name in [ + "dbx_execute_query", + "dbx_list_tables", + "dbx_open_session", + "dbx_prepare_table_import", + "dbx_vector_search", + "dbx_vector_upsert_file", + "dbx_vector_delete_by_batch", + ] { let tool = tools.iter().find(|tool| tool.name == tool_name).expect("selector tool should be registered"); let properties = tool .input_schema @@ -1563,6 +2234,17 @@ mod tests { ("dbx_remove_connection", &["connection_id"]), ("dbx_execute_redis_command", &["db"]), ("dbx_get_schema_context", &["database", "schema", "max_tables"]), + ("dbx_preview_import_file", &["source_format", "preview_rows", "cell_char_limit"]), + ( + "dbx_prepare_table_import", + &["connection_id", "connection_name", "database", "source_format", "batch_size", "date_time_format"], + ), + ( + "dbx_vector_search", + &["connection_id", "connection_name", "database", "semantic_version", "top_k", "output_fields"], + ), + ("dbx_vector_upsert_file", &["connection_id", "connection_name", "database", "batch_size"]), + ("dbx_vector_delete_by_batch", &["connection_id", "connection_name", "database"]), ]; for (tool_name, fields) in checks { @@ -1640,7 +2322,7 @@ mod tests { false, ); let names = server.tool_router.list_all().into_iter().map(|tool| tool.name).collect::>(); - assert_eq!(names.len(), 8); + assert_eq!(names.len(), 16); assert!(!names.iter().any(|name| name == "dbx_add_connection")); assert!(!names.iter().any(|name| name == "dbx_duplicate_connection")); assert!(!names.iter().any(|name| name == "dbx_remove_connection")); diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index 547474de61..d364d6f9bb 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -165,13 +165,21 @@ async fn initializes_lists_tools_and_calls_a_tool() { let tools = client.peer().list_tools(None).await.expect("list tools"); let names = tools.tools.iter().map(|tool| tool.name.as_ref()).collect::>(); - assert_eq!(names.len(), 13); + assert_eq!(names.len(), 21); assert!(names.contains(&"dbx_list_connections")); assert!(names.contains(&"dbx_duplicate_connection")); assert!(names.contains(&"dbx_execute_redis_command")); assert!(names.contains(&"dbx_execute_and_show")); assert!(names.contains(&"dbx_open_session")); assert!(names.contains(&"dbx_close_session")); + assert!(names.contains(&"dbx_preview_import_file")); + assert!(names.contains(&"dbx_prepare_table_import")); + assert!(names.contains(&"dbx_start_table_import")); + assert!(names.contains(&"dbx_get_import_status")); + assert!(names.contains(&"dbx_cancel_import")); + assert!(names.contains(&"dbx_vector_search")); + assert!(names.contains(&"dbx_vector_upsert_file")); + assert!(names.contains(&"dbx_vector_delete_by_batch")); let result = client.peer().call_tool(CallToolRequestParams::new("dbx_list_connections")).await.expect("call tool"); let response = result.content[0].as_text().expect("text response"); @@ -181,6 +189,31 @@ async fn initializes_lists_tools_and_calls_a_tool() { server_task.abort(); } +#[tokio::test] +async fn web_import_tool_returns_versioned_structured_error() { + let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); + let server = DbxMcpServer::with_runtime_options(Arc::new(EmptyBackend), McpScope::default(), true); + let server_task = tokio::spawn(async move { server.serve(server_transport).await }); + let client = ().serve(client_transport).await.expect("initialize MCP client"); + + let result = client + .peer() + .call_tool( + CallToolRequestParams::new("dbx_preview_import_file") + .with_arguments(Map::from_iter([("file_path".to_string(), json!("/tmp/input.xlsx"))])), + ) + .await + .expect("call import preview tool"); + + assert_eq!(result.is_error, Some(true)); + let structured = result.structured_content.expect("structured error"); + assert_eq!(structured.get("formatVersion"), Some(&json!(1))); + assert_eq!(structured.pointer("/error/code"), Some(&json!("IMPORT_UNSUPPORTED_IN_WEB_MODE_V1"))); + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} + #[tokio::test] async fn enforces_global_connection_scope_and_read_only_policy() { let backend = PolicyBackend { diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index bb2161988d..fde17e844c 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -87,12 +87,13 @@ DeepSeek Harness 通过 Cordis 插件条目加载 MCP Server，不读取 `mcpSer ## 工具列表 -DBX MCP 当前提供 12 个工具： +DBX MCP 当前提供 21 个工具： | 工具 | 说明 | | --- | --- | | `dbx_list_connections` | 列出当前 MCP 会话可见的连接 | | `dbx_add_connection` | 添加连接到 DBX 存储 | +| `dbx_duplicate_connection` | 复制已有连接且不暴露其密钥 | | `dbx_remove_connection` | 从 DBX 存储删除连接 | | `dbx_list_tables` | 列出表、视图或集合 | | `dbx_describe_table` | 返回列定义和表元数据 | @@ -103,9 +104,43 @@ DBX MCP 当前提供 12 个工具： | `dbx_execute_redis_command` | 执行 Redis 命令 | | `dbx_open_table` | 在运行中的 DBX 桌面端打开表 | | `dbx_execute_and_show` | 执行查询并在 DBX 中展示结果 | +| `dbx_preview_import_file` | 只读预览并计算允许目录内 Excel/CSV/TSV/JSON 的指纹 | +| `dbx_prepare_table_import` | 创建有效 30 分钟、不可变且只能使用一次的 PostgreSQL staging 导入计划 | +| `dbx_start_table_import` | 启动已经准备并复验的 staging 导入 | +| `dbx_get_import_status` | 查看导入进度、行数、错误和最终摘要 | +| `dbx_cancel_import` | 请求在后续 staging 批次写入前取消任务 | +| `dbx_vector_search` | 在允许的 Milvus 集合中按批准状态和生效时间检索语义卡 | +| `dbx_vector_upsert_file` | 从允许目录中的 JSONL upsert 已批准语义卡 | +| `dbx_vector_delete_by_batch` | 在完全访问权限下删除一个明确的未发布语义批次 | 启用连接作用域后，修改连接和桌面 UI 工具会被隐藏。 +## 本地文件导入与 Milvus 工具 + +导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个允许目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。 + +`dbx_prepare_table_import` 会把连接、已批准模板版本、规范路径、大小、修改时间、SHA-256、解析参数、字段映射和 staging 目标固化为不可变摘要。计划 30 分钟后过期，且只能启动一次。`dbx_start_table_import` 在后台启动前会重新核对完整文件身份。导入目标仅限 `DBX_MCP_IMPORT_STAGING_SCHEMAS` 中的 PostgreSQL Schema（默认为 `staging`）；发布到业务 Schema 仍由调用方通过独立治理流程完成。 + +Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索请求必须提供 `active_at`；服务端始终追加 `approval_status == "approved"` 以及 `effective_from`/`effective_to` 生效时间过滤，并可选追加精确 `semantic_version`。JSONL upsert 需要“数据读写”权限；删除只能按精确 `semantic_batch_id` 执行，并要求“完全访问”权限。通过 `dbx_execute_query` 提交的通用向量 REST 请求也会按端点分类，因此写入或破坏性 `POST` 不能绕过 MCP 策略。 + +向量检索的核心参数形状如下： + +```jsonc +{ + "connection_name": "运营组语义库", + "database": "default", + "collection": "semantic_cards", + "active_at": "2026-08-25T10:00:00+08:00", + "semantic_version": "semantic-v3", + "embedding": [/* 必须恰好包含 1024 个有限浮点数 */], + "top_k": 12, + "filters": { "business_domain": "交易" }, + "output_fields": ["card_id", "card_type", "title", "content", "source_uri"] +} +``` + +仅依赖生效日期选择活动语义卡时可以省略 `semantic_version`，但 `active_at` 仍为必填。`embedding` 的实际长度必须等于 `DBX_MCP_VECTOR_DIMENSION`；上面数组内的注释只是文档缩写，不是可直接提交的工具参数。 + ## 有状态查询会话 普通 `dbx_execute_query` 调用彼此独立。需要保持数据库 Session 状态时，先调用 `dbx_open_session`，再把返回的 `sessionId` 传给后续 `dbx_execute_query`： @@ -184,6 +219,15 @@ DBX 在 **设置 → MCP** 中保存一份权威策略，并在每次请求时 | `DBX_MCP_SCOPE_CONNECTION_IDS` | 兼容旧配置：限制为多个连接 ID | | `DBX_MCP_SCOPE_CONNECTION_NAME` | 限制为一个连接名称 | | `DBX_MCP_SCOPE_DATABASE` | 限制为一个数据库 | +| `DBX_MCP_IMPORT_ROOTS` | 导入和语义 JSONL 工具可读取的本地目录 allowlist，使用平台路径分隔符 | +| `DBX_MCP_IMPORT_STAGING_SCHEMAS` | 允许导入的 PostgreSQL staging Schema，逗号分隔，默认 `staging` | +| `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Excel/CSV/TSV/JSON 最大字节数，默认 512 MiB | +| `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | 语义 JSONL 最大字节数，默认 64 MiB | +| `DBX_MCP_VECTOR_COLLECTIONS` | Milvus 集合 allowlist，逗号分隔，默认 `semantic_cards` | +| `DBX_MCP_VECTOR_DIMENSION` | 固定向量维度，默认 `1024` | +| `DBX_MCP_VECTOR_TOP_K_MAX` | Milvus Top K 上限，默认 `20`，硬上限 50 | +| `DBX_MCP_VECTOR_FILTER_FIELDS` | 允许用于安全等值过滤的附加字段，逗号分隔 | +| `DBX_MCP_VECTOR_OUTPUT_FIELDS` | 向量检索允许返回的标量字段，逗号分隔 | | `DBX_MCP_DEBUG_SQL` | 临时诊断时输出 SQL | ## 常见问题 diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index 36bf36700b..e3a95cacf3 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -87,12 +87,13 @@ Run `dsh web --dump-config` to verify the composed configuration, then restart D ## Tools -DBX MCP currently provides 12 tools: +DBX MCP currently provides 21 tools: | Tool | Description | | --- | --- | | `dbx_list_connections` | List connections visible to the MCP session | | `dbx_add_connection` | Add a connection to DBX storage | +| `dbx_duplicate_connection` | Duplicate an existing connection without exposing its secret | | `dbx_remove_connection` | Remove a connection from DBX storage | | `dbx_list_tables` | List tables, views, or collections | | `dbx_describe_table` | Return columns and table metadata | @@ -103,9 +104,43 @@ DBX MCP currently provides 12 tools: | `dbx_execute_redis_command` | Execute a Redis command | | `dbx_open_table` | Open a table in the running DBX desktop app | | `dbx_execute_and_show` | Execute a query and display the result in DBX | +| `dbx_preview_import_file` | Read-only preview and fingerprint of an allowed Excel/CSV/TSV/JSON file | +| `dbx_prepare_table_import` | Create a 30-minute, immutable, single-use PostgreSQL staging import plan | +| `dbx_start_table_import` | Start an already prepared staging import | +| `dbx_get_import_status` | Read import progress, row counts, errors, and final summary | +| `dbx_cancel_import` | Request cancellation before subsequent staging batches are written | +| `dbx_vector_search` | Search an allowed Milvus collection with bounded output and enforced approval/effective-time filters | +| `dbx_vector_upsert_file` | Upsert approved semantic cards from one allowed JSONL file | +| `dbx_vector_delete_by_batch` | Delete one exact unpublished semantic batch with full-access permission | Connection-scoped sessions hide connection-mutating and desktop UI tools. +## Local file imports and Milvus tools + +The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. + +`dbx_prepare_table_import` binds the connection, approved template version, canonical path, size, modification time, SHA-256, parsing options, mappings, and staging target into an immutable digest. The returned plan expires after 30 minutes and can be started once. `dbx_start_table_import` rechecks the complete file identity before launching a background import. Import targets are limited to PostgreSQL schemas listed in `DBX_MCP_IMPORT_STAGING_SCHEMAS` (default `staging`); publication to business schemas remains the caller's separately governed workflow. + +Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires `active_at`, always adds `approval_status == "approved"` plus `effective_from`/`effective_to` activity filters, and optionally applies an exact `semantic_version`. Upsert reads bounded `.jsonl` files and requires Data read/write permission. Delete accepts only an exact `semantic_batch_id` and requires Full access. Generic vector REST passed to `dbx_execute_query` is also classified by endpoint, so write or destructive POST requests cannot bypass MCP policy. + +The core vector-search arguments are: + +```jsonc +{ + "connection_name": "Operations semantics", + "database": "default", + "collection": "semantic_cards", + "active_at": "2026-08-25T10:00:00+08:00", + "semantic_version": "semantic-v3", + "embedding": [/* exactly 1024 finite floating-point values */], + "top_k": 12, + "filters": { "business_domain": "transactions" }, + "output_fields": ["card_id", "card_type", "title", "content", "source_uri"] +} +``` + +Omit `semantic_version` when effective dates alone select the active cards. `active_at` remains required. `embedding` must contain exactly `DBX_MCP_VECTOR_DIMENSION` finite values; the abbreviated comment above is documentation only, not valid tool input. + ## Stateful Query Sessions Regular `dbx_execute_query` calls are independent. When a workflow must preserve database session state, call `dbx_open_session` first and pass its returned `sessionId` to later `dbx_execute_query` calls. Typical uses include: @@ -184,6 +219,15 @@ Updated servers do not let `DBX_MCP_ALLOW_WRITES` or `DBX_MCP_ALLOW_DANGEROUS_SQ | `DBX_MCP_SCOPE_CONNECTION_IDS` | Compatibility scope for multiple connection IDs | | `DBX_MCP_SCOPE_CONNECTION_NAME` | Restrict the session to one connection name | | `DBX_MCP_SCOPE_DATABASE` | Restrict the session to one database | +| `DBX_MCP_IMPORT_ROOTS` | Platform-separated allowlist of local directories readable by import and semantic JSONL tools | +| `DBX_MCP_IMPORT_STAGING_SCHEMAS` | Comma-separated PostgreSQL staging schema allowlist (default `staging`) | +| `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Maximum Excel/CSV/TSV/JSON source size (default 512 MiB) | +| `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | Maximum semantic JSONL size (default 64 MiB) | +| `DBX_MCP_VECTOR_COLLECTIONS` | Comma-separated Milvus collection allowlist (default `semantic_cards`) | +| `DBX_MCP_VECTOR_DIMENSION` | Required embedding dimension (default `1024`) | +| `DBX_MCP_VECTOR_TOP_K_MAX` | Maximum Milvus Top K, hard-capped at 50 (default `20`) | +| `DBX_MCP_VECTOR_FILTER_FIELDS` | Comma-separated additional safe equality-filter fields | +| `DBX_MCP_VECTOR_OUTPUT_FIELDS` | Comma-separated scalar fields that vector search may return | | `DBX_MCP_DEBUG_SQL` | Include SQL in temporary diagnostics | ## Troubleshooting diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index 6d88d81889..7b7bd23de5 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -17,7 +17,7 @@ The MCP protocol, connection loading, SQL safety, schema access, Redis support, ## Features -- **10 MCP tools** for connections, schemas, SQL, Redis, and DBX UI integration +- **21 MCP tools** for connections, schemas, SQL, sessions, bounded local imports, Milvus semantics, Redis, and DBX UI integration - **Precompiled native binaries** with no local Rust, Cargo, Python, or C/C++ build requirement - **No `better-sqlite3` runtime dependency** and no Node native-addon ABI coupling - **Local, Web, and Docker modes** using the same tool interface @@ -143,14 +143,25 @@ Ask the MCP client to: | --- | --- | | `dbx_list_connections` | List connections visible to the MCP session | | `dbx_add_connection` | Add a connection to DBX storage | +| `dbx_duplicate_connection` | Duplicate a connection without exposing its secret | | `dbx_remove_connection` | Remove a connection from DBX storage | | `dbx_list_tables` | List tables, views, collections, or message queue topics | | `dbx_describe_table` | Return columns and table metadata | | `dbx_get_schema_context` | Return compact schema context suitable for an AI model | | `dbx_execute_query` | Execute SQL or a supported MongoDB shell command, returning at most 100 rows | +| `dbx_open_session` | Open a stateful SQL session pinned to one backend connection | +| `dbx_close_session` | Close a stateful session and release its connection | | `dbx_execute_redis_command` | Execute a Redis command | | `dbx_open_table` | Open a table in the running DBX desktop application | | `dbx_execute_and_show` | Execute a query and display the result in the DBX desktop application | +| `dbx_preview_import_file` | Preview and fingerprint an allowed local Excel/CSV/TSV/JSON file | +| `dbx_prepare_table_import` | Create an immutable, 30-minute, single-use PostgreSQL staging import plan | +| `dbx_start_table_import` | Start an already prepared staging import | +| `dbx_get_import_status` | Read background import progress and final summary | +| `dbx_cancel_import` | Request cancellation of a staging import | +| `dbx_vector_search` | Search approved, active semantic cards in an allowed Milvus collection | +| `dbx_vector_upsert_file` | Upsert approved semantic cards from an allowed JSONL file | +| `dbx_vector_delete_by_batch` | Delete one exact unpublished semantic batch with Full access | When connection scoping is enabled, mutating connection tools and desktop UI tools are hidden. @@ -172,6 +183,8 @@ DBX connection storage defaults to: Override the directory with `DBX_DATA_DIR`. +Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Paths must be absolute regular files inside one configured root. Import plans bind SHA-256, parsing, mapping, template version, connection, and staging target; start revalidates the source and imports from a private task snapshot. Import tools support local mode only and return `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` in Web mode. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. + ### Agent/JDBC databases Databases such as Dameng, KingbaseES, Oracle, DB2, Hive, Trino, Snowflake, SAP HANA, and other DBX Agent profiles use DBX's Java agent infrastructure rather than a Node.js database driver. @@ -276,6 +289,13 @@ SQL text is not included in normal MCP errors or logged by default. Enable tempo | `DBX_MCP_SCOPE_CONNECTION_IDS` | Compatibility scope for multiple connection IDs | | `DBX_MCP_SCOPE_CONNECTION_NAME` | Restrict tools to one connection name | | `DBX_MCP_SCOPE_DATABASE` | Restrict tools to one database | +| `DBX_MCP_IMPORT_ROOTS` | Platform-separated allowlist of local import/semantic-file directories | +| `DBX_MCP_IMPORT_STAGING_SCHEMAS` | Comma-separated PostgreSQL staging schema allowlist (default `staging`) | +| `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Maximum tabular import source size (default 512 MiB) | +| `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | Maximum semantic JSONL size (default 64 MiB) | +| `DBX_MCP_VECTOR_COLLECTIONS` | Comma-separated Milvus collection allowlist (default `semantic_cards`) | +| `DBX_MCP_VECTOR_DIMENSION` | Required embedding dimension (default `1024`) | +| `DBX_MCP_VECTOR_TOP_K_MAX` | Maximum Milvus Top K (default `20`, hard cap 50) | | `DBX_MCP_DEBUG_SQL` | Include SQL in temporary diagnostics | | `DBX_MCP_BINARY` | Override the native binary used by the npm launcher | @@ -378,7 +398,7 @@ MCP 协议、连接读取、SQL 安全检查、Schema、Redis、MongoDB、Web ### 主要能力 -- 10 个 MCP 工具 +- 21 个 MCP 工具，覆盖连接、Schema、SQL、会话、本地受控导入、Milvus 语义、Redis 和 DBX 桌面操作 - 不依赖 `better-sqlite3`，没有 Node 原生模块 ABI 问题 - 支持本地 DBX、DBX Web 和 Docker - 支持预编译原生二进制和离线运行 @@ -461,14 +481,25 @@ MCP 配置： | --- | --- | | `dbx_list_connections` | 列出当前 MCP 会话可见的连接 | | `dbx_add_connection` | 添加 DBX 连接配置 | +| `dbx_duplicate_connection` | 复制已有连接且不暴露其密钥 | | `dbx_remove_connection` | 删除 DBX 连接配置 | | `dbx_list_tables` | 列出表、视图、集合或消息队列 Topic | | `dbx_describe_table` | 获取字段和表结构 | | `dbx_get_schema_context` | 获取适合 AI 使用的紧凑 Schema 上下文 | | `dbx_execute_query` | 执行 SQL 或支持的 MongoDB Shell 命令，最多返回 100 行 | +| `dbx_open_session` | 打开固定到一个后端连接的有状态 SQL 会话 | +| `dbx_close_session` | 关闭有状态会话并释放连接 | | `dbx_execute_redis_command` | 执行 Redis 命令 | | `dbx_open_table` | 在 DBX 桌面端打开表 | | `dbx_execute_and_show` | 执行查询并在 DBX 桌面端展示结果 | +| `dbx_preview_import_file` | 预览并计算允许目录内 Excel/CSV/TSV/JSON 的指纹 | +| `dbx_prepare_table_import` | 创建有效 30 分钟、不可变且只能使用一次的 PostgreSQL staging 导入计划 | +| `dbx_start_table_import` | 启动已经准备并复验的 staging 导入 | +| `dbx_get_import_status` | 查看后台导入进度和最终摘要 | +| `dbx_cancel_import` | 请求取消 staging 导入 | +| `dbx_vector_search` | 在允许的 Milvus 集合中检索已批准且生效的语义卡 | +| `dbx_vector_upsert_file` | 从允许目录中的 JSONL upsert 已批准语义卡 | +| `dbx_vector_delete_by_batch` | 在完全访问权限下删除一个明确的未发布语义批次 | ### 本地数据目录 @@ -478,6 +509,8 @@ MCP 配置： 通过 `DBX_DATA_DIR` 覆盖默认目录。Windows 便携版应指向 `DBX.exe` 同级、包含 `dbx.db` 的 `data` 文件夹。 +本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。文件必须是允许目录内的绝对路径普通文件。prepare 会固化 SHA-256、解析参数、映射、模板版本、连接和 staging 目标；start 会重新复验，并从任务私有快照执行导入。导入工具仅支持本地模式，Web 模式稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 + ### DBX Web / Docker ```json @@ -559,6 +592,13 @@ MongoDB 更新和删除在未启用完全访问时必须提供可验证有效的 | `DBX_MCP_SCOPE_CONNECTION_IDS` | 兼容旧配置：限制到多个连接 ID | | `DBX_MCP_SCOPE_CONNECTION_NAME` | 限制到指定连接名称 | | `DBX_MCP_SCOPE_DATABASE` | 限制到指定数据库 | +| `DBX_MCP_IMPORT_ROOTS` | 本地导入/语义文件目录 allowlist，使用平台路径分隔符 | +| `DBX_MCP_IMPORT_STAGING_SCHEMAS` | PostgreSQL staging Schema allowlist，逗号分隔，默认 `staging` | +| `DBX_MCP_IMPORT_FILE_MAX_BYTES` | 表格导入源最大字节数，默认 512 MiB | +| `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | 语义 JSONL 最大字节数，默认 64 MiB | +| `DBX_MCP_VECTOR_COLLECTIONS` | Milvus 集合 allowlist，逗号分隔，默认 `semantic_cards` | +| `DBX_MCP_VECTOR_DIMENSION` | 固定向量维度，默认 `1024` | +| `DBX_MCP_VECTOR_TOP_K_MAX` | Milvus Top K 上限，默认 `20`，硬上限 50 | | `DBX_MCP_DEBUG_SQL` | 临时输出 SQL 诊断信息 | | `DBX_MCP_BINARY` | 覆盖 npm 启动器使用的原生文件 | From 505924ccfe8fa53f2cbc837d320105e8794a0305 Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 18:05:48 +0800 Subject: [PATCH 03/10] fix(mcp): harden governed imports and semantic writes --- Cargo.lock | 2 + crates/dbx-core/src/table_import.rs | 59 +- crates/dbx-mcp/Cargo.toml | 2 + crates/dbx-mcp/src/backend.rs | 14 +- crates/dbx-mcp/src/enterprise_tools.rs | 974 +++++++++++++++++++++---- crates/dbx-mcp/src/server.rs | 343 +++++++-- docs/content/docs/mcp.cn.mdx | 15 +- docs/content/docs/mcp.mdx | 15 +- packages/mcp-server/README.md | 6 +- 9 files changed, 1200 insertions(+), 230 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 32c7fa7708..9f7af1f1da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2155,6 +2155,7 @@ name = "dbx-mcp" version = "0.4.74" dependencies = [ "async-trait", + "csv", "dbx-core", "dirs", "reqwest 0.12.28", @@ -2166,6 +2167,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "tempfile", + "time", "tokio", "tokio-rustls 0.26.4", "url", diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index 6a5b3f9934..aa59fe0626 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -65,6 +65,8 @@ pub fn table_import_client_session_id(import_id: &str) -> String { pub struct ParsedImportFile { pub columns: Vec, pub rows: Vec>, + /// 每个已解析数据行在源文件中的 1 基行号；JSON 使用 1 基记录序号。 + pub source_row_numbers: Vec, pub total_rows: usize, pub effective_encoding: Option, } @@ -282,6 +284,7 @@ pub struct TableImportPreview { pub size_bytes: u64, pub columns: Vec, pub rows: Vec>, + pub source_row_numbers: Vec, pub total_rows: usize, pub total_rows_exact: bool, pub source_fingerprint: String, @@ -926,6 +929,7 @@ fn parse_csv_reader_inner( count_all_rows: bool, ) -> Result { let mut rows = Vec::new(); + let mut source_row_numbers = Vec::new(); let mut total_rows = 0; let mut columns = Vec::new(); let mut record = csv::StringRecord::new(); @@ -954,6 +958,7 @@ fn parse_csv_reader_inner( total_rows += 1; if rows.len() < preview_limit { rows.push(delimited_record_to_row(&record, columns.len(), config)); + source_row_numbers.push(row_number); } if !count_all_rows && rows.len() >= preview_limit { break; @@ -965,7 +970,7 @@ fn parse_csv_reader_inner( if total_rows == 0 { return Err("Import file has no data rows in the selected row range".to_string()); } - Ok(ParsedImportFile { columns, rows, total_rows, effective_encoding: Some(effective_encoding) }) + Ok(ParsedImportFile { columns, rows, source_row_numbers, total_rows, effective_encoding: Some(effective_encoding) }) } fn parse_delimited_preview_file_with_options( @@ -1073,7 +1078,14 @@ pub fn parse_json_bytes_with_options( .collect::>() }) .collect::>(); - return Ok(ParsedImportFile { columns, rows, total_rows: items.len(), effective_encoding: None }); + let source_row_numbers = (1..=rows.len()).collect(); + return Ok(ParsedImportFile { + columns, + rows, + source_row_numbers, + total_rows: items.len(), + effective_encoding: None, + }); } if all_arrays { @@ -1092,7 +1104,14 @@ pub fn parse_json_bytes_with_options( .collect::>() }) .collect::>(); - return Ok(ParsedImportFile { columns, rows, total_rows: items.len(), effective_encoding: None }); + let source_row_numbers = (1..=rows.len()).collect(); + return Ok(ParsedImportFile { + columns, + rows, + source_row_numbers, + total_rows: items.len(), + effective_encoding: None, + }); } Err("JSON rows must all be objects or all be arrays; mixed row shapes are not supported".to_string()) @@ -2608,6 +2627,9 @@ fn parse_xlsx_preview_file_with_options( if last_preview_row < row_range.data_start_row { return Err("Import file has no data rows in the selected row range".to_string()); } + let source_row_numbers = (row_range.data_start_row..=last_preview_row) + .map(|relative_row| start_row + relative_row - 1) + .collect::>(); let rows = (row_range.data_start_row..=last_preview_row) .map(|relative_row| { let absolute_row = start_row + relative_row - 1; @@ -2626,7 +2648,10 @@ fn parse_xlsx_preview_file_with_options( if rows.is_empty() { return Err("Import file has no data rows in the selected row range".to_string()); } - Ok((ParsedImportFile { columns, total_rows: rows.len(), rows, effective_encoding: None }, sheets)) + Ok(( + ParsedImportFile { columns, total_rows: rows.len(), rows, source_row_numbers, effective_encoding: None }, + sheets, + )) } fn xlsx_cell_styles( @@ -3335,6 +3360,7 @@ where let empty_string_as_null = options.empty_string_as_null.unwrap_or(true); let mut columns = Vec::new(); let mut rows = Vec::new(); + let mut source_row_numbers = Vec::new(); let mut total_rows = 0; for (index, source_row) in range.rows().enumerate() { let row_number = index + 1; @@ -3386,6 +3412,7 @@ where row.push(value); } rows.push(row); + source_row_numbers.push(range_start_row + row_number); } if columns.is_empty() { return Err("Import file has no columns in the selected row range".to_string()); @@ -3393,7 +3420,7 @@ where if total_rows == 0 { return Err("Import file has no data rows in the selected row range".to_string()); } - Ok(ParsedImportFile { columns, rows, total_rows, effective_encoding: None }) + Ok(ParsedImportFile { columns, rows, source_row_numbers, total_rows, effective_encoding: None }) } pub fn parse_xlsx_file(path: &str, preview_limit: usize) -> Result { @@ -5278,6 +5305,8 @@ fn validated_prepared_import_source( Some(ParsedImportFile { columns: prepared.columns.clone(), rows: prepared.rows.clone(), + // 已准备预览不保存绝对源行位置；宁可明确缺失，也不伪造行号。治理导入不走此路径。 + source_row_numbers: Vec::new(), total_rows: prepared.total_rows, effective_encoding: prepared.effective_encoding, }) @@ -5309,6 +5338,7 @@ pub async fn preview_table_import_file_with_request( size_bytes: metadata.len(), columns: parsed.columns, rows: parsed.rows, + source_row_numbers: parsed.source_row_numbers, total_rows: parsed.total_rows, total_rows_exact, source_fingerprint, @@ -7980,6 +8010,7 @@ mod tests { assert_eq!(preview.total_rows, 2); assert!(!preview.total_rows_exact); assert_eq!(preview.rows[0], vec![serde_json::json!("1"), serde_json::json!("北京")]); + assert_eq!(preview.source_row_numbers, vec![2, 3]); } #[test] @@ -7996,6 +8027,7 @@ mod tests { assert_eq!(preview.columns, vec!["id", "name"]); assert_eq!(preview.rows, vec![vec![serde_json::json!("1"), serde_json::json!("Ada")]]); assert_eq!(preview.total_rows, 1); + assert_eq!(preview.source_row_numbers, vec![2]); } #[test] @@ -9292,6 +9324,7 @@ mod tests { let data = ParsedImportFile { columns: target_types.iter().map(|(column, _)| column.to_string()).collect(), rows, + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -9701,6 +9734,8 @@ mod tests { assert_eq!(parsed.total_rows, 2); assert_eq!(parsed.rows[0], vec![serde_json::json!(1), serde_json::json!("Ada")]); assert_eq!(parsed.rows[1], vec![serde_json::json!(2), serde_json::json!("Grace")]); + assert_eq!(parsed.source_row_numbers, vec![3, 4]); + assert_eq!(preview.source_row_numbers, vec![3, 4]); assert_eq!(preview.rows, parsed.rows); let _ = std::fs::remove_file(path); } @@ -9734,6 +9769,7 @@ mod tests { serde_json::json!({ "source": "json" }), ], ], + source_row_numbers: Vec::new(), total_rows: 2, effective_encoding: None, }; @@ -9772,6 +9808,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["id".to_string()], rows: vec![vec![serde_json::json!(1)]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -9791,6 +9828,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["notes".to_string()], rows: vec![vec![serde_json::json!("long text")]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -9822,6 +9860,7 @@ mod tests { serde_json::json!("2026-07-07 08:15:00"), serde_json::json!("invoice"), ]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -9853,6 +9892,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["code".to_string(), "amount".to_string()], rows: vec![vec![serde_json::json!("1001"), serde_json::json!("12.5")]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -9918,6 +9958,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["name".to_string()], rows: vec![vec![serde_json::json!("Ada")]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -9953,6 +9994,7 @@ mod tests { vec![serde_json::json!(2), serde_json::json!("O'Hara"), serde_json::json!("y")], vec![serde_json::json!(3), serde_json::Value::Null, serde_json::json!("z")], ], + source_row_numbers: Vec::new(), total_rows: 3, effective_encoding: None, }; @@ -9982,6 +10024,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["id".to_string()], rows: vec![vec![serde_json::json!(1)], vec![serde_json::json!(2)]], + source_row_numbers: Vec::new(), total_rows: 2, effective_encoding: None, }; @@ -10957,6 +11000,7 @@ mod tests { vec![serde_json::json!(2), serde_json::json!("Grace")], vec![serde_json::json!(3), serde_json::Value::Null], ], + source_row_numbers: Vec::new(), total_rows: 3, effective_encoding: None, }; @@ -10980,6 +11024,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["payload".to_string()], rows: (0..4).map(|index| vec![serde_json::json!(format!("{index}{}", "x".repeat(180 * 1024)))]).collect(), + source_row_numbers: Vec::new(), total_rows: 4, effective_encoding: None, }; @@ -11013,6 +11058,7 @@ mod tests { serde_json::json!("2026-05-12T00:00:00+00:00"), serde_json::json!("2026-05-12T00:00:00+00:00"), ]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -11047,6 +11093,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["created_at".to_string()], rows: vec![vec![serde_json::json!("2024/2/25 13:02:15")]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -11082,6 +11129,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["created_at".to_string()], rows: vec![vec![imported_value]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; @@ -11176,6 +11224,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["name".to_string()], rows: vec![vec![serde_json::json!("Tiếng Việt")]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, }; diff --git a/crates/dbx-mcp/Cargo.toml b/crates/dbx-mcp/Cargo.toml index 152fb3159c..fd7814b8e2 100644 --- a/crates/dbx-mcp/Cargo.toml +++ b/crates/dbx-mcp/Cargo.toml @@ -15,6 +15,7 @@ path = "src/main.rs" [dependencies] async-trait = "0.1" +csv = "1" dbx-core = { path = "../dbx-core", default-features = false, features = ["sqlite-bundled"] } dirs = "6" rmcp = { version = "2.2.0", features = ["client", "transport-io"] } @@ -24,6 +25,7 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" sha2 = "0.10" tokio = { version = "1", features = ["fs", "io-util", "macros", "rt-multi-thread", "sync", "time"] } +time = { version = "0.3", features = ["formatting"] } uuid = { version = "1", features = ["v4"] } url = "2" diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 1a9336c8f7..ab085eeb84 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -264,10 +264,11 @@ pub trait DbxBackend: Send + Sync { &self, connection: &ConnectionConfig, request: TableImportRequest, + plan_id: String, cancelled: Arc, progress: Arc, ) -> Result { - let _ = (connection, request, cancelled, progress); + let _ = (connection, request, plan_id, cancelled, progress); Err("IMPORT_UNSUPPORTED_IN_WEB_MODE_V1: v1 文件导入仅支持本地 DBX Desktop/MCP 模式。".to_string()) } } @@ -692,6 +693,7 @@ impl DbxBackend for LocalBackend { &self, connection: &ConnectionConfig, mut request: TableImportRequest, + plan_id: String, cancelled: Arc, progress: Arc, ) -> Result { @@ -708,6 +710,7 @@ impl DbxBackend for LocalBackend { tokio::fs::create_dir(&snapshot_dir).await.map_err(|error| format!("创建任务快照目录失败：{error}"))?; let extension = Path::new(&request.file_path).extension().and_then(|value| value.to_str()).unwrap_or("data"); let snapshot_path = snapshot_dir.join(format!("source.{extension}")); + let normalized_path = snapshot_dir.join("normalized.csv"); let snapshot_result = async { tokio::fs::copy(&request.file_path, &snapshot_path) .await @@ -719,6 +722,14 @@ impl DbxBackend for LocalBackend { return Err("IMPORT_FILE_CHANGED: 源文件在启动导入时发生变化；未写入数据库。".to_string()); } request.file_path = snapshot_path.to_string_lossy().to_string(); + request = crate::enterprise_tools::build_governed_import_snapshot( + request, + &plan_id, + &expected_sha256, + &normalized_path, + ) + .await + .map_err(|error| error.to_string())?; let database = (!request.database.trim().is_empty()).then_some(request.database.as_str()); let client_session_id = dbx_core::table_import::table_import_client_session_id(&request.import_id); let pool_key = self @@ -746,6 +757,7 @@ impl DbxBackend for LocalBackend { result } .await; + let _ = tokio::fs::remove_file(&normalized_path).await; let _ = tokio::fs::remove_file(&snapshot_path).await; let _ = tokio::fs::remove_dir(&snapshot_dir).await; snapshot_result diff --git a/crates/dbx-mcp/src/enterprise_tools.rs b/crates/dbx-mcp/src/enterprise_tools.rs index d4b96520b8..dffde397f4 100644 --- a/crates/dbx-mcp/src/enterprise_tools.rs +++ b/crates/dbx-mcp/src/enterprise_tools.rs @@ -19,7 +19,8 @@ use rmcp::schemars; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use sha2::{Digest, Sha256}; -use tokio::sync::RwLock; +use time::{format_description::well_known::Rfc3339, OffsetDateTime, UtcOffset}; +use tokio::sync::{OwnedSemaphorePermit, RwLock, Semaphore}; use uuid::Uuid; pub const FORMAT_VERSION: u8 = 1; @@ -27,8 +28,7 @@ pub const DEFAULT_PLAN_TTL: Duration = Duration::from_secs(30 * 60); const DEFAULT_IMPORT_FILE_SIZE_BYTES: u64 = 512 * 1024 * 1024; const DEFAULT_SEMANTIC_FILE_SIZE_BYTES: u64 = 64 * 1024 * 1024; const DEFAULT_PREVIEW_ROWS: usize = 20; -pub const STRUCTURE_PROFILE_ROWS: usize = 100; -const MAX_PREVIEW_ROWS: usize = STRUCTURE_PROFILE_ROWS; +const MAX_PREVIEW_ROWS: usize = 100; const DEFAULT_CELL_CHAR_LIMIT: usize = 1_000; const MAX_CELL_CHAR_LIMIT: usize = 4_000; const DEFAULT_VECTOR_TOP_K: usize = 12; @@ -37,6 +37,10 @@ const DEFAULT_VECTOR_DIMENSION: usize = 1_024; const DEFAULT_VECTOR_UPSERT_BATCH: usize = 200; const MAX_VECTOR_UPSERT_BATCH: usize = 500; const MAX_JSONL_RECORD_BYTES: usize = 512 * 1024; +const DEFAULT_PLAN_CAPACITY: usize = 512; +const DEFAULT_JOB_CAPACITY: usize = 256; +const DEFAULT_TERMINAL_JOB_LIMIT: usize = 128; +const JOB_TTL_MS: u128 = 24 * 60 * 60 * 1_000; const IMPORT_EXTENSIONS: &[&str] = &["xlsx", "xlsm", "xls", "csv", "tsv", "json"]; const SEMANTIC_EXTENSIONS: &[&str] = &["jsonl"]; @@ -56,10 +60,14 @@ const DEFAULT_VECTOR_OUTPUT_FIELDS: &[&str] = &[ "effective_to", "source_uri", "source_checksum", + "content_checksum", "semantic_version", "embedding_model", + "embedding_revision", "semantic_batch_id", ]; +const RESERVED_STAGING_COLUMNS: &[&str] = + &["import_id", "plan_id", "source_sha", "source_row_number", "source_row_hash", "loaded_at"]; #[derive(Debug, Clone, PartialEq, Eq)] pub struct EnterpriseToolError { @@ -163,20 +171,10 @@ impl McpImportParseOptions { #[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] pub struct McpImportColumnMapping { - pub source_column: String, + #[schemars(description = "1-based source column position returned by dbx_preview_import_file")] + pub source_position: usize, + pub source_name: String, pub target_column: String, - #[schemars(extend("type" = "string"))] - pub target_data_type: Option, -} - -impl From for TableImportColumnMapping { - fn from(value: McpImportColumnMapping) -> Self { - Self { - source_column: value.source_column, - target_column: value.target_column, - target_data_type: value.target_data_type, - } - } } #[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] @@ -198,8 +196,6 @@ pub struct PrepareTableImportRequest { pub selector: crate::server::ConnectionSelector, #[schemars(extend("type" = "string"))] pub database: Option, - pub schema: String, - pub table: String, pub template_version: String, pub file_path: String, #[schemars(extend("type" = "string"))] @@ -207,8 +203,6 @@ pub struct PrepareTableImportRequest { #[serde(default)] pub parse_options: McpImportParseOptions, pub mappings: Vec, - #[serde(default)] - pub create_table: bool, #[schemars(extend("type" = "integer"))] pub batch_size: Option, #[schemars(extend("type" = "string"))] @@ -265,6 +259,8 @@ pub struct VectorDeleteByBatchRequest { pub database: Option, pub collection: String, pub semantic_batch_id: String, + #[schemars(description = "Must be explicitly false; published semantic batches cannot be deleted by this tool")] + pub published: bool, } #[derive(Debug, Clone, Serialize, PartialEq, Eq)] @@ -329,6 +325,10 @@ impl PreparedImportPlan { pub struct ImportJobSnapshot { pub import_id: String, pub plan_id: String, + pub relation: String, + pub created_at_ms: u128, + pub finished_at_ms: Option, + pub cancellation_requested: bool, pub status: TableImportStatus, pub phase: TableImportPhase, pub rows_imported: usize, @@ -342,17 +342,21 @@ pub struct ImportJobSnapshot { } impl ImportJobSnapshot { - fn initial(import_id: String, plan_id: String, total_bytes: u64) -> Self { + fn initial(import_id: String, plan: &PreparedImportPlan) -> Self { Self { import_id, - plan_id, + plan_id: plan.plan_id.clone(), + relation: format!("{}.{}", plan.schema, plan.table), + created_at_ms: unix_epoch_millis(), + finished_at_ms: None, + cancellation_requested: false, status: TableImportStatus::Running, phase: TableImportPhase::Preparing, rows_imported: 0, total_rows: 0, total_rows_exact: false, bytes_read: 0, - total_bytes, + total_bytes: plan.file.size_bytes, elapsed_ms: 0, error: None, summary: None, @@ -377,18 +381,43 @@ pub struct ImportJob { pub cancelled: Arc, } -#[derive(Default)] pub struct EnterpriseRuntime { plans: RwLock>, jobs: RwLock>>, + import_slots: Arc, + semantic_write_slots: Arc, + plan_capacity: usize, + job_capacity: usize, + terminal_job_limit: usize, +} + +impl Default for EnterpriseRuntime { + fn default() -> Self { + Self { + plans: RwLock::new(HashMap::new()), + jobs: RwLock::new(HashMap::new()), + import_slots: Arc::new(Semaphore::new(env_usize("DBX_MCP_IMPORT_CONCURRENCY", 2).clamp(1, 16))), + semantic_write_slots: Arc::new(Semaphore::new(1)), + plan_capacity: env_usize("DBX_MCP_IMPORT_PLAN_CAPACITY", DEFAULT_PLAN_CAPACITY).clamp(1, 4_096), + job_capacity: env_usize("DBX_MCP_IMPORT_JOB_CAPACITY", DEFAULT_JOB_CAPACITY).clamp(1, 2_048), + terminal_job_limit: env_usize("DBX_MCP_IMPORT_TERMINAL_LIMIT", DEFAULT_TERMINAL_JOB_LIMIT).clamp(1, 1_024), + } + } } impl EnterpriseRuntime { - pub async fn insert_plan(&self, plan: PreparedImportPlan) { + pub async fn insert_plan(&self, plan: PreparedImportPlan) -> Result<(), EnterpriseToolError> { let now = unix_epoch_millis(); let mut plans = self.plans.write().await; plans.retain(|_, plan| !plan.consumed && plan.expires_at_ms > now); + if plans.len() >= self.plan_capacity { + return Err(EnterpriseToolError::new( + "IMPORT_PLAN_CAPACITY_REACHED", + "待启动导入计划已达到容量上限；请等待计划过期或启动已有计划。", + )); + } plans.insert(plan.plan_id.clone(), plan); + Ok(()) } pub async fn consume_plan(&self, plan_id: &str) -> Result { @@ -397,6 +426,7 @@ impl EnterpriseRuntime { .get_mut(plan_id) .ok_or_else(|| EnterpriseToolError::new("IMPORT_PLAN_NOT_FOUND", "导入计划不存在，或已经被清理。"))?; if plan.expires_at_ms <= unix_epoch_millis() { + plans.remove(plan_id); return Err(EnterpriseToolError::new("IMPORT_PLAN_EXPIRED", "导入计划已超过 30 分钟有效期。")); } if plan.consumed { @@ -406,35 +436,79 @@ impl EnterpriseRuntime { Ok(plan.clone()) } - pub async fn create_job(&self, plan: &PreparedImportPlan) -> Arc { + pub async fn try_import_permit(&self) -> Result { + self.import_slots + .clone() + .try_acquire_owned() + .map_err(|_| EnterpriseToolError::new("IMPORT_CONCURRENCY_LIMIT", "并发导入已达到上限，请稍后重试。")) + } + + pub async fn semantic_write_permit(&self) -> Result { + self.semantic_write_slots + .clone() + .acquire_owned() + .await + .map_err(|_| EnterpriseToolError::new("VECTOR_WRITE_COORDINATOR_CLOSED", "语义写入协调器不可用。")) + } + + pub async fn create_job(&self, plan: &PreparedImportPlan) -> Result, EnterpriseToolError> { let import_id = format!("mcp-import-{}", Uuid::new_v4()); let job = Arc::new(ImportJob { - snapshot: Mutex::new(ImportJobSnapshot::initial( - import_id.clone(), - plan.plan_id.clone(), - plan.file.size_bytes, - )), + snapshot: Mutex::new(ImportJobSnapshot::initial(import_id.clone(), plan)), cancelled: Arc::new(AtomicBool::new(false)), }); - self.jobs.write().await.insert(import_id, job.clone()); - job + let mut jobs = self.jobs.write().await; + cleanup_jobs(&mut jobs, self.terminal_job_limit); + if jobs.len() >= self.job_capacity { + return Err(EnterpriseToolError::new( + "IMPORT_JOB_CAPACITY_REACHED", + "导入任务记录已达到容量上限；请等待终态任务过期。", + )); + } + jobs.insert(import_id, job.clone()); + Ok(job) } pub async fn job(&self, import_id: &str) -> Option> { - self.jobs.read().await.get(import_id).cloned() + let mut jobs = self.jobs.write().await; + cleanup_jobs(&mut jobs, self.terminal_job_limit); + jobs.get(import_id).cloned() } - pub async fn cancel_job(&self, import_id: &str) -> Result { + pub async fn cancel_job(&self, import_id: &str) -> Result<(ImportJobSnapshot, bool), EnterpriseToolError> { let job = self .job(import_id) .await .ok_or_else(|| EnterpriseToolError::new("IMPORT_JOB_NOT_FOUND", "没有找到指定导入任务。"))?; - let snapshot = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).clone(); + let mut snapshot = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()); if snapshot.status != TableImportStatus::Running { return Err(EnterpriseToolError::new("IMPORT_JOB_ALREADY_TERMINAL", "导入任务已经结束，不能再取消。")); } - job.cancelled.store(true, Ordering::Release); - Ok(snapshot) + let already_requested = job.cancelled.swap(true, Ordering::AcqRel); + snapshot.cancellation_requested = true; + Ok((snapshot.clone(), already_requested)) + } +} + +fn cleanup_jobs(jobs: &mut HashMap>, terminal_limit: usize) { + let now = unix_epoch_millis(); + jobs.retain(|_, job| { + let snapshot = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()); + snapshot.status == TableImportStatus::Running + || snapshot.finished_at_ms.is_none_or(|finished| now.saturating_sub(finished) <= JOB_TTL_MS) + }); + let mut terminal = jobs + .iter() + .filter_map(|(id, job)| { + let snapshot = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()); + (snapshot.status != TableImportStatus::Running) + .then_some((id.clone(), snapshot.finished_at_ms.unwrap_or(0))) + }) + .collect::>(); + terminal.sort_by_key(|(_, finished)| *finished); + let remove_count = terminal.len().saturating_sub(terminal_limit); + for (id, _) in terminal.into_iter().take(remove_count) { + jobs.remove(&id); } } @@ -447,7 +521,7 @@ pub fn configured_import_roots() -> Result, EnterpriseToolError> { })?; let roots = std::env::split_paths(&value) .filter_map(|path| std::fs::canonicalize(path).ok()) - .filter(|path| path.parent().is_some()) + .filter(|path| path.parent().is_some() && path.is_dir()) .collect::>(); if roots.is_empty() { return Err(EnterpriseToolError::new("IMPORT_ROOTS_NOT_CONFIGURED", "DBX_MCP_IMPORT_ROOTS 中没有可用目录。")); @@ -482,7 +556,12 @@ pub fn validate_import_file_with_roots( } let canonical = std::fs::canonicalize(path) .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法解析文件路径：{error}")))?; - if !roots.iter().filter_map(|root| std::fs::canonicalize(root).ok()).any(|root| canonical.starts_with(root)) { + if !roots + .iter() + .filter_map(|root| std::fs::canonicalize(root).ok()) + .filter(|root| root.is_dir()) + .any(|root| canonical.starts_with(root)) + { return Err(EnterpriseToolError::new( "IMPORT_PATH_OUTSIDE_ROOTS", "文件不在 DBX_MCP_IMPORT_ROOTS 允许目录中。", @@ -568,86 +647,131 @@ pub fn cell_char_limit(value: Option) -> Result TableImportPreview { for row in &mut preview.rows { for cell in row { - if let Value::String(value) = cell { - let mut chars = value.chars(); - let truncated = chars.by_ref().take(char_limit).collect::(); - if chars.next().is_some() { - *value = format!("{truncated}…"); + sanitize_preview_value(cell, char_limit, 0); + } + } + preview +} + +pub fn validate_preview_headers(preview: &TableImportPreview) -> Result<(), EnterpriseToolError> { + if preview.columns.len() > 1_000 { + return Err(EnterpriseToolError::new("IMPORT_COLUMN_LIMIT_EXCEEDED", "源文件超过 1000 列，拒绝返回或导入。")); + } + if let Some(header) = preview.columns.iter().find(|header| header.chars().count() > MAX_CELL_CHAR_LIMIT) { + return Err(EnterpriseToolError::new( + "IMPORT_HEADER_TOO_LONG", + format!("表头超过 {MAX_CELL_CHAR_LIMIT} 个字符：{}…", truncate_preview_text(header, 80)), + )); + } + Ok(()) +} + +fn sanitize_preview_value(value: &mut Value, char_limit: usize, depth: usize) { + if depth >= 4 { + if value.is_array() || value.is_object() { + *value = Value::String("[嵌套内容已截断]".to_string()); + } + return; + } + match value { + Value::String(text) => *text = truncate_preview_text(text, char_limit), + Value::Array(values) => { + values.truncate(50); + for value in values { + sanitize_preview_value(value, char_limit, depth + 1); + } + } + Value::Object(object) => { + let was_truncated = object.len() > 50; + let original = std::mem::take(object); + for (index, (key, mut value)) in original.into_iter().take(50).enumerate() { + sanitize_preview_value(&mut value, char_limit, depth + 1); + let mut key = truncate_preview_text(&key, char_limit); + if object.contains_key(&key) { + key = format!("{key}__{}", index + 1); } + object.insert(key, value); + } + if was_truncated { + object.insert("_dbx_truncated".to_string(), Value::Bool(true)); } } + _ => {} + } +} + +fn truncate_preview_text(value: &str, char_limit: usize) -> String { + let mut chars = value.chars(); + let truncated = chars.by_ref().take(char_limit).collect::(); + if chars.next().is_some() { + format!("{truncated}…") + } else { + truncated } - preview } pub fn structure_fingerprint(preview: &TableImportPreview, parse_options: &TableImportParseOptions) -> String { let value = json!({ "fileType": preview.file_type, "columns": preview.columns, - "columnTypeHints": column_type_hints(preview), "sheets": preview.sheets, "parseOptions": parse_options, }); sha256_bytes(value.to_string().as_bytes()) } -pub fn column_type_hints(preview: &TableImportPreview) -> Vec { - (0..preview.columns.len()) - .map(|index| { - let kinds = preview - .rows - .iter() - .filter_map(|row| row.get(index)) - .filter_map(|value| match value { - Value::Null => None, - Value::Bool(_) => Some("boolean"), - Value::Number(number) if number.is_i64() || number.is_u64() => Some("integer"), - Value::Number(_) => Some("decimal"), - Value::String(_) => Some("text"), - Value::Array(_) | Value::Object(_) => Some("json"), - }) - .collect::>(); - if kinds.is_empty() { - "empty".to_string() - } else { - kinds.into_iter().collect::>().join("|") - } - }) - .collect() -} - -pub fn validate_staging_target(schema: &str, table: &str) -> Result<(), EnterpriseToolError> { - validate_identifier(schema, "schema")?; - validate_identifier(table, "table")?; +pub fn generated_staging_relation() -> Result<(String, String), EnterpriseToolError> { let allowed = comma_list_env("DBX_MCP_IMPORT_STAGING_SCHEMAS", &["staging"]); - if !allowed.iter().any(|candidate| candidate == schema) { - return Err(EnterpriseToolError::new( - "IMPORT_TARGET_NOT_STAGING", - format!("MCP 导入只能写入隔离 staging schema；当前允许：{}。", allowed.join(", ")), - )); - } - Ok(()) + let schema = allowed + .first() + .cloned() + .ok_or_else(|| EnterpriseToolError::new("IMPORT_STAGING_SCHEMA_REQUIRED", "必须配置一个 staging schema。"))?; + validate_identifier(&schema, "staging_schema")?; + let table = format!("mcp_{}", Uuid::new_v4().simple()); + Ok((schema, table)) } pub fn validate_mappings( - mappings: &[TableImportColumnMapping], + mappings: &[McpImportColumnMapping], source_columns: &[String], -) -> Result<(), EnterpriseToolError> { +) -> Result, EnterpriseToolError> { if mappings.is_empty() { return Err(EnterpriseToolError::new("IMPORT_MAPPING_REQUIRED", "至少需要一个字段映射。")); } - let source = source_columns.iter().collect::>(); let mut targets = HashSet::new(); + let mut positions = HashSet::new(); + let mut normalized = Vec::with_capacity(mappings.len()); for mapping in mappings { - if !source.contains(&mapping.source_column) { + let source_index = mapping.source_position.checked_sub(1).ok_or_else(|| { + EnterpriseToolError::new("IMPORT_SOURCE_POSITION_INVALID", "source_position 从 1 开始，不能为 0。") + })?; + let actual_name = source_columns.get(source_index).ok_or_else(|| { + EnterpriseToolError::new( + "IMPORT_SOURCE_POSITION_INVALID", + format!("源文件没有第 {} 列。", mapping.source_position), + ) + })?; + if actual_name != &mapping.source_name { + return Err(EnterpriseToolError::new( + "IMPORT_SOURCE_NAME_MISMATCH", + format!( + "第 {} 列当前名称为 {}，与请求中的 {} 不一致。", + mapping.source_position, actual_name, mapping.source_name + ), + )); + } + if !positions.insert(mapping.source_position) { return Err(EnterpriseToolError::new( - "IMPORT_SOURCE_COLUMN_NOT_FOUND", - format!("源字段 {} 不存在。", mapping.source_column), + "IMPORT_SOURCE_POSITION_DUPLICATED", + format!("源位置 {} 被重复映射。", mapping.source_position), )); } validate_identifier(&mapping.target_column, "target_column")?; - if let Some(data_type) = mapping.target_data_type.as_deref() { - validate_mcp_postgres_type(data_type)?; + if RESERVED_STAGING_COLUMNS.contains(&mapping.target_column.as_str()) { + return Err(EnterpriseToolError::new( + "IMPORT_TARGET_COLUMN_RESERVED", + format!("目标字段 {} 是 staging 血缘保留字段。", mapping.target_column), + )); } if !targets.insert(mapping.target_column.as_str()) { return Err(EnterpriseToolError::new( @@ -655,8 +779,13 @@ pub fn validate_mappings( format!("目标字段 {} 被重复映射。", mapping.target_column), )); } + normalized.push(TableImportColumnMapping { + source_column: actual_name.clone(), + target_column: mapping.target_column.clone(), + target_data_type: Some("TEXT".to_string()), + }); } - Ok(()) + Ok(normalized) } #[allow(clippy::too_many_arguments)] @@ -735,6 +864,186 @@ pub async fn revalidate_plan_file(plan: &PreparedImportPlan) -> Result<(), Enter Ok(()) } +pub fn validate_governed_source_v1( + file_path: &str, + source_format: Option, + parse_options: &TableImportParseOptions, +) -> Result { + let source_format = dbx_core::table_import::effective_source_format(file_path, source_format) + .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_FORMAT_INVALID", error))?; + if !source_format.is_delimited() { + return Err(EnterpriseToolError::new( + "IMPORT_GOVERNED_FORMAT_UNSUPPORTED_V1", + format!("{} 尚未实现流式有界行血缘转换；v1 治理导入当前仅支持 UTF-8 CSV/TSV。", source_format.label()), + )); + } + if !matches!(parse_options.encoding, None | Some(TableImportTextEncoding::Auto | TableImportTextEncoding::Utf8)) { + return Err(EnterpriseToolError::new( + "IMPORT_GOVERNED_ENCODING_UNSUPPORTED_V1", + "v1 流式治理导入仅支持 UTF-8；其他编码不创建导入计划。", + )); + } + Ok(source_format) +} + +pub async fn build_governed_import_snapshot( + mut request: TableImportRequest, + plan_id: &str, + source_sha: &str, + output_path: &Path, +) -> Result { + let source_format = validate_governed_source_v1(&request.file_path, request.source_format, &request.parse_options)?; + let config = dbx_core::table_import::effective_delimited_config(source_format, &request.parse_options) + .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_PARSE_FAILED", error))?; + let offset = UtcOffset::from_hms(8, 0, 0) + .map_err(|error| EnterpriseToolError::new("IMPORT_LOADED_AT_FAILED", error.to_string()))?; + let loaded_at = OffsetDateTime::now_utc() + .to_offset(offset) + .format(&Rfc3339) + .map_err(|error| EnterpriseToolError::new("IMPORT_LOADED_AT_FAILED", error.to_string()))?; + let mut writer = csv::WriterBuilder::new() + .has_headers(false) + .from_path(output_path) + .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + let mut headers = request.mappings.iter().map(|mapping| mapping.target_column.clone()).collect::>(); + headers.extend(RESERVED_STAGING_COLUMNS.iter().map(|column| column.to_string())); + writer + .write_record(&headers) + .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + let mut reader = csv::ReaderBuilder::new() + .delimiter(config.delimiter) + .has_headers(false) + .flexible(true) + .from_path(&request.file_path) + .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_PARSE_FAILED", error.to_string()))?; + let mut source_columns = Vec::new(); + let mut source_indexes: Option> = None; + let mut source_row_count = 0usize; + for (index, record) in reader.records().enumerate() { + let source_row_number = index + 1; + let record = record.map_err(|error| { + EnterpriseToolError::new( + "IMPORT_SOURCE_PARSE_FAILED", + format!("UTF-8 CSV/TSV 第 {source_row_number} 行解析失败：{error}"), + ) + })?; + if config.row_range.title_row == Some(source_row_number) { + source_columns = unique_staging_headers(record.iter().enumerate().map(|(index, value)| { + let value = if index == 0 { value.trim_start_matches('\u{feff}') } else { value }; + dbx_core::table_import::normalize_header(value, index) + })); + continue; + } + if source_row_number < config.row_range.data_start_row { + continue; + } + if config.row_range.last_data_row.is_some_and(|last| source_row_number > last) { + break; + } + if source_columns.is_empty() { + source_columns = (0..record.len()).map(|index| format!("column_{}", index + 1)).collect(); + } + if source_indexes.is_none() { + source_indexes = Some( + request + .mappings + .iter() + .map(|mapping| { + source_columns.iter().position(|column| column == &mapping.source_column).ok_or_else(|| { + EnterpriseToolError::new( + "IMPORT_SOURCE_COLUMN_NOT_FOUND", + format!("规范化时找不到源字段 {}。", mapping.source_column), + ) + }) + }) + .collect::, _>>()?, + ); + } + let indexes = source_indexes + .as_ref() + .ok_or_else(|| EnterpriseToolError::new("IMPORT_SOURCE_MAPPING_FAILED", "无法建立源字段位置映射。"))?; + let parsed_row = (0..source_columns.len()) + .map(|column| dbx_core::table_import::csv_value_with_config(record.get(column).unwrap_or_default(), config)) + .collect::>(); + let mut output = indexes + .iter() + .map(|index| staging_text_value(parsed_row.get(*index).unwrap_or(&Value::Null))) + .collect::>(); + let row_bytes = serde_json::to_vec(&parsed_row) + .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_ROW_HASH_FAILED", error.to_string()))?; + output.extend([ + request.import_id.clone(), + plan_id.to_string(), + source_sha.to_string(), + source_row_number.to_string(), + sha256_bytes(&row_bytes), + loaded_at.clone(), + ]); + writer + .write_record(&output) + .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + source_row_count += 1; + } + if source_row_count == 0 { + return Err(EnterpriseToolError::new("IMPORT_SOURCE_EMPTY", "选择的源范围没有数据行；未写入数据库。")); + } + writer.flush().map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + + let mut mappings = request + .mappings + .iter() + .map(|mapping| TableImportColumnMapping { + source_column: mapping.target_column.clone(), + target_column: mapping.target_column.clone(), + target_data_type: Some("TEXT".to_string()), + }) + .collect::>(); + mappings.extend(RESERVED_STAGING_COLUMNS.iter().map(|column| TableImportColumnMapping { + source_column: column.to_string(), + target_column: column.to_string(), + target_data_type: Some("TEXT".to_string()), + })); + request.file_path = output_path.to_string_lossy().to_string(); + request.source_format = Some(TableImportSourceFormat::Csv); + request.parse_options = TableImportParseOptions { + has_header: Some(true), + trim_values: Some(false), + empty_string_as_null: Some(false), + ..Default::default() + }; + request.mappings = mappings; + request.mode = TableImportMode::Append; + request.create_table = true; + request.prepared_source = None; + Ok(request) +} + +fn staging_text_value(value: &Value) -> String { + match value { + Value::Null => String::new(), + Value::String(value) => value.clone(), + other => other.to_string(), + } +} + +fn unique_staging_headers(headers: impl IntoIterator) -> Vec { + let mut seen = HashSet::new(); + let mut next_suffix = HashMap::::new(); + headers + .into_iter() + .map(|header| { + let suffix = next_suffix.entry(header.to_lowercase()).or_default(); + loop { + let candidate = if *suffix == 0 { header.clone() } else { format!("{header}_{suffix}") }; + *suffix += 1; + if seen.insert(candidate.to_lowercase()) { + break candidate; + } + } + }) + .collect() +} + pub fn validate_vector_collection(collection: &str) -> Result<(), EnterpriseToolError> { validate_identifier(collection, "collection")?; let allowed = comma_list_env("DBX_MCP_VECTOR_COLLECTIONS", &["semantic_cards"]); @@ -757,7 +1066,7 @@ pub fn vector_top_k(top_k: Option) -> Result } pub fn validate_embedding(embedding: &[f32]) -> Result<(), EnterpriseToolError> { - let dimension = env_usize("DBX_MCP_VECTOR_DIMENSION", DEFAULT_VECTOR_DIMENSION).clamp(1, 65_536); + let dimension = DEFAULT_VECTOR_DIMENSION; if embedding.len() != dimension { return Err(EnterpriseToolError::new( "VECTOR_DIMENSION_MISMATCH", @@ -798,11 +1107,8 @@ pub fn build_milvus_filter( semantic_version: Option<&str>, filters: &BTreeMap, ) -> Result { - if !valid_effective_timestamp(active_at) { - return Err(EnterpriseToolError::new( - "VECTOR_ACTIVE_AT_INVALID", - "active_at 必须是 YYYY-MM-DD 或 RFC 3339 时间。", - )); + if active_at.len() != 10 || !valid_effective_timestamp(active_at) { + return Err(EnterpriseToolError::new("VECTOR_ACTIVE_AT_INVALID", "active_at 必须是合法 YYYY-MM-DD 日期。")); } let allowed = comma_list_env("DBX_MCP_VECTOR_FILTER_FIELDS", DEFAULT_VECTOR_FILTER_FIELDS); let mut clauses = vec![ @@ -828,13 +1134,42 @@ pub fn build_milvus_filter( } pub fn read_semantic_jsonl(path: &Path, semantic_batch_id: &str) -> Result, EnterpriseToolError> { - if semantic_batch_id.trim().is_empty() { - return Err(EnterpriseToolError::new("SEMANTIC_BATCH_ID_REQUIRED", "semantic_batch_id 不能为空。")); + if semantic_batch_id.trim().is_empty() + || semantic_batch_id.len() > 200 + || semantic_batch_id.chars().any(char::is_control) + { + return Err(EnterpriseToolError::new( + "SEMANTIC_BATCH_ID_INVALID", + "semantic_batch_id 必须是 1～200 个非控制字符。", + )); } let source = std::fs::read_to_string(path) .map_err(|error| EnterpriseToolError::new("VECTOR_JSONL_READ_FAILED", format!("读取 JSONL 失败：{error}")))?; - let expected_dimension = env_usize("DBX_MCP_VECTOR_DIMENSION", DEFAULT_VECTOR_DIMENSION); + let expected_dimension = DEFAULT_VECTOR_DIMENSION; let mut records = Vec::new(); + let mut card_ids = HashSet::new(); + let mut file_semantic_version: Option = None; + let allowed_fields = HashSet::from([ + "card_id", + "card_type", + "business_domain", + "dataset_id", + "template_version", + "title", + "content", + "aliases", + "approval_status", + "effective_from", + "effective_to", + "source_uri", + "source_checksum", + "content_checksum", + "semantic_version", + "embedding_model", + "embedding_revision", + "embedding", + "semantic_batch_id", + ]); for (index, line) in source.lines().enumerate() { let line_number = index + 1; let line = line.trim(); @@ -853,7 +1188,33 @@ pub fn read_semantic_jsonl(path: &Path, semantic_batch_id: &str) -> Result Result Result 100 + || aliases.iter().any(|alias| alias.as_str().is_none_or(|value| value.is_empty() || value.len() > 500)) + { return Err(EnterpriseToolError::new( "VECTOR_JSONL_FIELD_INVALID", - format!("JSONL 第 {line_number} 行的 effective_from 不是有效日期或 RFC 3339 时间。"), + format!("JSONL 第 {line_number} 行的 aliases 超出限制或包含非字符串。"), + )); + } + for field in ["source_checksum", "content_checksum"] { + if object.get(field).and_then(Value::as_str).is_none_or(|value| !is_sha256_hex(value)) { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_INVALID", + format!("JSONL 第 {line_number} 行的 {field} 必须是 64 位十六进制 SHA-256。"), + )); + } + } + let content = object.get("content").and_then(Value::as_str).unwrap_or_default(); + if content.chars().count() > 100_000 { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_INVALID", + format!("JSONL 第 {line_number} 行的 content 超过 100000 字节。"), + )); + } + if object.get("content_checksum").and_then(Value::as_str) != Some(sha256_bytes(content.as_bytes()).as_str()) { + return Err(EnterpriseToolError::new( + "VECTOR_CONTENT_CHECKSUM_MISMATCH", + format!("JSONL 第 {line_number} 行的 content_checksum 与 content 不一致。"), + )); + } + let card_id = object.get("card_id").and_then(Value::as_str).unwrap_or_default().to_string(); + if !card_ids.insert(card_id.clone()) { + return Err(EnterpriseToolError::new( + "VECTOR_CARD_ID_DUPLICATED", + format!("JSONL 内重复 card_id：{card_id}。"), + )); + } + let semantic_version = object.get("semantic_version").and_then(Value::as_str).unwrap_or_default().to_string(); + if file_semantic_version.as_ref().is_some_and(|version| version != &semantic_version) { + return Err(EnterpriseToolError::new( + "VECTOR_SEMANTIC_VERSION_MIXED", + "同一 JSONL 只能包含一个 semantic_version。", + )); + } + file_semantic_version.get_or_insert(semantic_version); + if object + .get("effective_from") + .and_then(Value::as_str) + .is_none_or(|value| value.len() != 10 || !valid_effective_timestamp(value)) + { + return Err(EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_INVALID", + format!("JSONL 第 {line_number} 行的 effective_from 不是合法 YYYY-MM-DD 日期。"), )); } match object.get("effective_to") { None | Some(Value::Null) => { object.insert("effective_to".to_string(), Value::String(String::new())); } - Some(Value::String(value)) if value.is_empty() || valid_effective_timestamp(value) => {} + Some(Value::String(value)) + if value.is_empty() || (value.len() == 10 && valid_effective_timestamp(value)) => {} Some(_) => { return Err(EnterpriseToolError::new( "VECTOR_JSONL_FIELD_INVALID", - format!("JSONL 第 {line_number} 行的 effective_to 必须为 null、空字符串或有效日期/RFC 3339 时间。"), + format!("JSONL 第 {line_number} 行的 effective_to 必须为 null、空字符串或合法 YYYY-MM-DD 日期。"), )); } } + let effective_from = object.get("effective_from").and_then(Value::as_str).unwrap_or_default(); + let effective_to = object.get("effective_to").and_then(Value::as_str).unwrap_or_default(); + if !effective_to.is_empty() && effective_to < effective_from { + return Err(EnterpriseToolError::new( + "VECTOR_EFFECTIVE_RANGE_INVALID", + format!("JSONL 第 {line_number} 行的 effective_to 早于 effective_from。"), + )); + } if let Some(existing) = object.get("semantic_batch_id").and_then(Value::as_str) { if existing != semantic_batch_id { return Err(EnterpriseToolError::new( @@ -969,6 +1412,42 @@ pub fn milvus_upsert_query(database: &str, collection: &str, records: &[Value]) ) } +pub fn milvus_existing_cards_query(database: &str, collection: &str, card_ids: &[String]) -> String { + let ids = card_ids.iter().map(|card_id| json_string(card_id)).collect::>().join(", "); + format!( + "POST /v2/vectordb/entities/query\n{}", + json!({ + "dbName": if database.is_empty() { "default" } else { database }, + "collectionName": collection, + "filter": format!("card_id in [{ids}]"), + "limit": card_ids.len(), + "outputFields": ["card_id", "semantic_batch_id", "semantic_version"], + }) + ) +} + +pub fn validate_existing_card_ownership( + rows: &[Value], + semantic_batch_id: &str, + semantic_version: &str, +) -> Result<(), EnterpriseToolError> { + for row in rows { + let object = row.as_object().ok_or_else(|| { + EnterpriseToolError::new("VECTOR_EXISTING_CARD_INVALID", "Milvus 返回了无法验证归属的已有语义卡。") + })?; + let card_id = object.get("card_id").and_then(Value::as_str).unwrap_or(""); + let same_batch = object.get("semantic_batch_id").and_then(Value::as_str) == Some(semantic_batch_id); + let same_version = object.get("semantic_version").and_then(Value::as_str) == Some(semantic_version); + if !same_batch || !same_version { + return Err(EnterpriseToolError::new( + "VECTOR_CARD_OWNERSHIP_CONFLICT", + format!("card_id {card_id} 已属于其他 semantic_batch_id 或 semantic_version，禁止覆盖。"), + )); + } + } + Ok(()) +} + pub fn milvus_delete_batch_query(database: &str, collection: &str, semantic_batch_id: &str) -> String { format!( "POST /v2/vectordb/entities/delete\n{}", @@ -1031,34 +1510,58 @@ fn validate_identifier(value: &str, field: &str) -> Result<(), EnterpriseToolErr Ok(()) } -fn validate_mcp_postgres_type(data_type: &str) -> Result<(), EnterpriseToolError> { - let normalized = data_type.trim().to_ascii_uppercase().replace(' ', ""); - let simple = matches!(normalized.as_str(), "TEXT" | "TIMESTAMPTZ" | "DATE" | "BOOLEAN" | "JSONB"); - let numeric = normalized - .strip_prefix("NUMERIC(") - .and_then(|value| value.strip_suffix(')')) - .and_then(|value| value.split_once(',')) - .and_then(|(precision, scale)| Some((precision.parse::().ok()?, scale.parse::().ok()?))) - .is_some_and(|(precision, scale)| (1..=100).contains(&precision) && scale <= precision); - if !simple && !numeric { - return Err(EnterpriseToolError::new( - "IMPORT_TARGET_TYPE_NOT_ALLOWED", - "v1 MCP 导入类型只允许 TEXT、NUMERIC(p,s)、TIMESTAMPTZ、DATE、BOOLEAN 或 JSONB。", - )); - } - Ok(()) -} - fn json_string(value: &str) -> String { serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string()) } +fn is_sha256_hex(value: &str) -> bool { + value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} + fn valid_effective_timestamp(value: &str) -> bool { let value = value.trim(); - (10..=40).contains(&value.len()) - && value.as_bytes().get(4) == Some(&b'-') - && value.as_bytes().get(7) == Some(&b'-') - && value.chars().all(|ch| ch.is_ascii_digit() || matches!(ch, '-' | ':' | 'T' | 'Z' | '+' | '.')) + if value.len() < 10 || value.as_bytes().get(4) != Some(&b'-') || value.as_bytes().get(7) != Some(&b'-') { + return false; + } + let year = value.get(0..4).and_then(|part| part.parse::().ok()).unwrap_or(0); + let month = value.get(5..7).and_then(|part| part.parse::().ok()).unwrap_or(0); + let day = value.get(8..10).and_then(|part| part.parse::().ok()).unwrap_or(0); + let leap = year.is_multiple_of(4) && (!year.is_multiple_of(100) || year.is_multiple_of(400)); + let max_day = match month { + 1 | 3 | 5 | 7 | 8 | 10 | 12 => 31, + 4 | 6 | 9 | 11 => 30, + 2 if leap => 29, + 2 => 28, + _ => 0, + }; + if year == 0 || day == 0 || day > max_day { + return false; + } + if value.len() == 10 { + return true; + } + if value.len() != 20 && value.len() != 25 { + return false; + } + let bytes = value.as_bytes(); + if bytes.get(10) != Some(&b'T') || bytes.get(13) != Some(&b':') || bytes.get(16) != Some(&b':') { + return false; + } + let hour = value.get(11..13).and_then(|part| part.parse::().ok()).unwrap_or(24); + let minute = value.get(14..16).and_then(|part| part.parse::().ok()).unwrap_or(60); + let second = value.get(17..19).and_then(|part| part.parse::().ok()).unwrap_or(60); + if hour > 23 || minute > 59 || second > 59 { + return false; + } + if value.len() == 20 { + return bytes.get(19) == Some(&b'Z'); + } + if !matches!(bytes.get(19), Some(b'+') | Some(b'-')) || bytes.get(22) != Some(&b':') { + return false; + } + let offset_hour = value.get(20..22).and_then(|part| part.parse::().ok()).unwrap_or(24); + let offset_minute = value.get(23..25).and_then(|part| part.parse::().ok()).unwrap_or(60); + offset_hour <= 23 && offset_minute <= 59 } fn sha256_bytes(value: &[u8]) -> String { @@ -1102,7 +1605,13 @@ mod tests { assert_eq!( validate_import_file_with_roots(source.to_str().unwrap(), &[root.path().to_path_buf()], false, 1024) .unwrap(), - std::fs::canonicalize(source).unwrap() + std::fs::canonicalize(&source).unwrap() + ); + assert_eq!( + validate_import_file_with_roots(source.to_str().unwrap(), &[source.clone()], false, 1024,) + .unwrap_err() + .code, + "IMPORT_PATH_OUTSIDE_ROOTS" ); let outside_file = outside.path().join("outside.csv"); @@ -1127,12 +1636,26 @@ mod tests { } } + #[test] + fn recursive_preview_sanitizer_bounds_nested_values_and_keys() { + let long_key = "键".repeat(20); + let mut value = json!({ + long_key: (0..100).map(|index| json!({ "value": "内容".repeat(20), "index": index })).collect::>() + }); + sanitize_preview_value(&mut value, 8, 0); + let object = value.as_object().unwrap(); + assert!(object.keys().all(|key| key.chars().count() <= 11)); + let array = object.values().next().unwrap().as_array().unwrap(); + assert_eq!(array.len(), 50); + assert!(array[0].get("value").and_then(Value::as_str).unwrap().chars().count() <= 9); + } + #[test] fn milvus_filter_forces_approval_and_semantic_version() { let filters = BTreeMap::from([("business_domain".to_string(), json!("交易"))]); - let filter = build_milvus_filter("2026-08-25T10:00:00+08:00", Some("semantic-v3"), &filters).unwrap(); + let filter = build_milvus_filter("2026-08-25", Some("semantic-v3"), &filters).unwrap(); assert!(filter.contains("approval_status == \"approved\"")); - assert!(filter.contains("effective_from <= \"2026-08-25T10:00:00+08:00\"")); + assert!(filter.contains("effective_from <= \"2026-08-25\"")); assert!(filter.contains("semantic_version == \"semantic-v3\"")); assert!(filter.contains("business_domain == \"交易\"")); @@ -1147,37 +1670,65 @@ mod tests { fn mappings_reject_missing_and_duplicate_targets() { let source = vec!["订单号".to_string(), "金额".to_string()]; let duplicated = vec![ - TableImportColumnMapping { - source_column: "订单号".to_string(), + McpImportColumnMapping { + source_position: 1, + source_name: "订单号".to_string(), target_column: "order_id".to_string(), - target_data_type: None, }, - TableImportColumnMapping { - source_column: "金额".to_string(), + McpImportColumnMapping { + source_position: 2, + source_name: "金额".to_string(), target_column: "order_id".to_string(), - target_data_type: None, }, ]; assert_eq!(validate_mappings(&duplicated, &source).unwrap_err().code, "IMPORT_TARGET_COLUMN_DUPLICATED"); - let unsafe_type = vec![TableImportColumnMapping { - source_column: "订单号".to_string(), - target_column: "order_id".to_string(), - target_data_type: Some("TEXT DEFAULT current_user".to_string()), + let reserved = vec![McpImportColumnMapping { + source_position: 1, + source_name: "订单号".to_string(), + target_column: "source_row_hash".to_string(), }]; - assert_eq!(validate_mappings(&unsafe_type, &source).unwrap_err().code, "IMPORT_TARGET_TYPE_NOT_ALLOWED"); + assert_eq!(validate_mappings(&reserved, &source).unwrap_err().code, "IMPORT_TARGET_COLUMN_RESERVED"); + + let duplicate_headers = vec!["note".to_string(), "note_1".to_string()]; + let second = vec![McpImportColumnMapping { + source_position: 2, + source_name: "note_1".to_string(), + target_column: "second_note".to_string(), + }]; + assert_eq!(validate_mappings(&second, &duplicate_headers).unwrap()[0].source_column, "note_1"); + let wrong_name = vec![McpImportColumnMapping { + source_position: 2, + source_name: "note".to_string(), + target_column: "second_note".to_string(), + }]; + assert_eq!(validate_mappings(&wrong_name, &duplicate_headers).unwrap_err().code, "IMPORT_SOURCE_NAME_MISMATCH"); } #[test] fn semantic_jsonl_normalizes_null_effective_to() { let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("cards.jsonl"); + let content = "成交金额定义"; let record = json!({ "card_id": "metric-order-gmv", + "card_type": "metric", + "business_domain": "交易", + "dataset_id": "orders", + "template_version": "orders-v1", + "title": "成交金额", + "content": content, + "aliases": ["GMV"], "approval_status": "approved", + "semantic_batch_id": "semantic-batch-1", "semantic_version": "semantic-v1", - "effective_from": "2026-08-25T00:00:00+08:00", + "effective_from": "2026-08-25", "effective_to": null, + "source_uri": "s3://semantic/orders.md", + "source_checksum": "a".repeat(64), + "content_checksum": sha256_bytes(content.as_bytes()), + "embedding_model": "BAAI/bge-m3", + "embedding_revision": "revision-1", "embedding": vec![0.0_f32; DEFAULT_VECTOR_DIMENSION], }); std::fs::write(&path, format!("{}\n", record)).unwrap(); @@ -1185,6 +1736,46 @@ mod tests { let records = read_semantic_jsonl(&path, "semantic-batch-1").unwrap(); assert_eq!(records[0].get("effective_to"), Some(&json!(""))); assert_eq!(records[0].get("semantic_batch_id"), Some(&json!("semantic-batch-1"))); + + let mut invalid_date = record.clone(); + invalid_date.as_object_mut().unwrap().insert("effective_from".to_string(), json!("2026-02-30")); + std::fs::write(&path, format!("{}\n", invalid_date)).unwrap(); + assert_eq!(read_semantic_jsonl(&path, "semantic-batch-1").unwrap_err().code, "VECTOR_JSONL_FIELD_INVALID"); + + let mut invalid_checksum = record.clone(); + invalid_checksum.as_object_mut().unwrap().insert("content_checksum".to_string(), json!("b".repeat(64))); + std::fs::write(&path, format!("{}\n", invalid_checksum)).unwrap(); + assert_eq!( + read_semantic_jsonl(&path, "semantic-batch-1").unwrap_err().code, + "VECTOR_CONTENT_CHECKSUM_MISMATCH" + ); + + for forbidden_field in ["chunk_index", "embedding_dimension"] { + let mut forbidden = record.clone(); + forbidden.as_object_mut().unwrap().insert(forbidden_field.to_string(), json!(0)); + std::fs::write(&path, format!("{}\n", forbidden)).unwrap(); + assert_eq!(read_semantic_jsonl(&path, "semantic-batch-1").unwrap_err().code, "VECTOR_JSONL_UNKNOWN_FIELD"); + } + } + + #[test] + fn existing_card_ownership_prevents_cross_batch_overwrite() { + let same = vec![json!({ + "card_id": "metric-order-gmv", + "semantic_batch_id": "batch-1", + "semantic_version": "semantic-v1" + })]; + validate_existing_card_ownership(&same, "batch-1", "semantic-v1").unwrap(); + + let conflict = vec![json!({ + "card_id": "metric-order-gmv", + "semantic_batch_id": "batch-0", + "semantic_version": "semantic-v1" + })]; + assert_eq!( + validate_existing_card_ownership(&conflict, "batch-1", "semantic-v1").unwrap_err().code, + "VECTOR_CARD_OWNERSHIP_CONFLICT" + ); } #[tokio::test] @@ -1217,9 +1808,104 @@ mod tests { .unwrap(); let plan_id = plan.plan_id.clone(); let runtime = EnterpriseRuntime::default(); - runtime.insert_plan(plan).await; + runtime.insert_plan(plan.clone()).await.unwrap(); assert_eq!(runtime.consume_plan(&plan_id).await.unwrap().plan_id, plan_id); assert_eq!(runtime.consume_plan(&plan_id).await.unwrap_err().code, "IMPORT_PLAN_ALREADY_USED"); + let job = runtime.create_job(&plan).await.unwrap(); + let import_id = job.snapshot.lock().unwrap().import_id.clone(); + assert!(!runtime.cancel_job(&import_id).await.unwrap().1); + assert!(runtime.cancel_job(&import_id).await.unwrap().1); + } + + #[tokio::test] + async fn governed_csv_snapshot_preserves_duplicate_position_and_real_row_number() { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("duplicate.csv"); + let output = directory.path().join("normalized.csv"); + std::fs::write(&source, "说明行\nnote,note\nA,B\n").unwrap(); + let request = TableImportRequest { + import_id: "import-1".to_string(), + connection_id: "postgres-1".to_string(), + database: "enterprise".to_string(), + schema: "staging".to_string(), + table: "mcp_test".to_string(), + file_path: source.to_string_lossy().to_string(), + source_ref: Some("a".repeat(64)), + source_format: Some(TableImportSourceFormat::Csv), + parse_options: TableImportParseOptions { + title_row: Some(2), + data_start_row: Some(3), + has_header: Some(true), + ..Default::default() + }, + mappings: vec![ + TableImportColumnMapping { + source_column: "note".to_string(), + target_column: "first_note".to_string(), + target_data_type: Some("TEXT".to_string()), + }, + TableImportColumnMapping { + source_column: "note_1".to_string(), + target_column: "second_note".to_string(), + target_data_type: Some("TEXT".to_string()), + }, + ], + mode: TableImportMode::Append, + create_table: true, + batch_size: 100, + date_time_format: None, + prepared_source: None, + retain_source: true, + }; + + let governed = build_governed_import_snapshot(request, "plan-1", &"a".repeat(64), &output).await.unwrap(); + assert!(governed.create_table); + assert!(governed.mappings.iter().all(|mapping| mapping.target_data_type.as_deref() == Some("TEXT"))); + let mut reader = csv::Reader::from_path(output).unwrap(); + let headers = reader.headers().unwrap().clone(); + let row = reader.records().next().unwrap().unwrap(); + assert_eq!(headers.get(0), Some("first_note")); + assert_eq!(headers.get(1), Some("second_note")); + assert_eq!(row.get(0), Some("A")); + assert_eq!(row.get(1), Some("B")); + assert_eq!(row.get(5), Some("3")); + assert_eq!(row.get(6).map(str::len), Some(64)); + } + + #[tokio::test] + async fn governed_snapshot_blocks_excel_before_database_write() { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("input.xlsx"); + let output = directory.path().join("normalized.csv"); + std::fs::write(&source, b"not-an-xlsx").unwrap(); + let mut request = TableImportRequest { + import_id: "import-unsupported".to_string(), + connection_id: "postgres-1".to_string(), + database: "enterprise".to_string(), + schema: "staging".to_string(), + table: "mcp_test".to_string(), + file_path: source.to_string_lossy().to_string(), + source_ref: Some("b".repeat(64)), + source_format: Some(TableImportSourceFormat::Excel), + parse_options: TableImportParseOptions::default(), + mappings: Vec::new(), + mode: TableImportMode::Append, + create_table: true, + batch_size: 100, + date_time_format: None, + prepared_source: None, + retain_source: true, + }; + request.mappings.push(TableImportColumnMapping { + source_column: "id".to_string(), + target_column: "id".to_string(), + target_data_type: Some("TEXT".to_string()), + }); + assert_eq!( + build_governed_import_snapshot(request, "plan-1", &"b".repeat(64), &output).await.unwrap_err().code, + "IMPORT_GOVERNED_FORMAT_UNSUPPORTED_V1" + ); + assert!(!output.exists()); } } diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index bc8a3f197f..312cb62e29 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -11,13 +11,14 @@ use uuid::Uuid; use crate::backend::{format_query_result, new_connection_config, parse_database_type, ConnectionSummary, DbxBackend}; use crate::enterprise_tools::{ - build_milvus_filter, build_plan, cell_char_limit, column_type_hints, file_identity, milvus_delete_batch_query, - milvus_search_query, milvus_upsert_query, preview_limit, query_result_rows, read_semantic_jsonl, - revalidate_plan_file, sanitize_preview, structure_fingerprint, validate_embedding, validate_import_file, - validate_mappings, validate_staging_target, validate_vector_collection, vector_output_fields, vector_top_k, + build_milvus_filter, build_plan, cell_char_limit, file_identity, generated_staging_relation, + milvus_delete_batch_query, milvus_existing_cards_query, milvus_search_query, milvus_upsert_query, preview_limit, + query_result_rows, read_semantic_jsonl, revalidate_plan_file, sanitize_preview, structure_fingerprint, + validate_embedding, validate_existing_card_ownership, validate_governed_source_v1, validate_import_file, + validate_mappings, validate_preview_headers, validate_vector_collection, vector_output_fields, vector_top_k, vector_upsert_batch_size, EnterpriseRuntime, EnterpriseToolError, ImportStatusRequest, PrepareTableImportRequest, PreviewImportFileRequest, StartTableImportRequest, VectorDeleteByBatchRequest, VectorSearchRequest, - VectorUpsertFileRequest, FORMAT_VERSION, STRUCTURE_PROFILE_ROWS, + VectorUpsertFileRequest, FORMAT_VERSION, }; use crate::mongo::{self, MongoCommand, MongoSafetyError}; use crate::session::{McpSession, McpSessionStore}; @@ -409,22 +410,39 @@ impl DbxMcpServer { source_ref: Some(identity.sha256.clone()), source_format: request.source_format.map(Into::into), parse_options: parse_options.clone(), - preview_limit: Some(STRUCTURE_PROFILE_ROWS), + preview_limit: Some(preview_rows), }, ) .await; - let mut preview = match preview { - Ok(preview) => sanitize_preview(preview, char_limit), + let preview = match preview { + Ok(preview) => preview, Err(error) => { return enterprise_error(EnterpriseToolError::new("IMPORT_PREVIEW_FAILED", error)); } }; + if let Err(error) = validate_preview_headers(&preview) { + return enterprise_error(error); + } let fingerprint = structure_fingerprint(&preview, &parse_options); - let type_hints = column_type_hints(&preview); - preview.rows.truncate(preview_rows); - let title_row = parse_options.title_row.or_else(|| parse_options.has_header.unwrap_or(true).then_some(1)); - let data_start_row = parse_options.data_start_row.unwrap_or_else(|| title_row.map_or(1, |row| row + 1)); - let last_row = data_start_row.saturating_add(preview.total_rows.saturating_sub(1)); + let current_identity = match file_identity(path).await { + Ok(identity) => identity, + Err(error) => return enterprise_error(error), + }; + if current_identity != identity { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_FILE_CHANGED_DURING_PREVIEW", + "文件在预览期间发生变化；结果已丢弃。", + )); + } + let source_columns = preview + .columns + .iter() + .enumerate() + .map(|(index, name)| json!({ "position": index + 1, "name": name })) + .collect::>(); + let used_first_row = preview.source_row_numbers.first().copied(); + let used_last_row = preview.source_row_numbers.last().copied(); + let preview = sanitize_preview(preview, char_limit); let summary = format!( "已只读剖析 {}：{} 列，{} 行，返回 {} 行预览。", preview.file_name, @@ -437,10 +455,10 @@ impl DbxMcpServer { json!({ "file": identity, "structureFingerprint": fingerprint, - "columnTypeHints": type_hints, + "sourceColumns": source_columns, "usedRange": { - "firstRow": title_row.unwrap_or(data_start_row), - "lastRow": last_row, + "firstDataRow": used_first_row, + "lastPreviewDataRow": used_last_row, "firstColumn": if preview.columns.is_empty() { 0 } else { 1 }, "lastColumn": preview.columns.len(), }, @@ -451,7 +469,7 @@ impl DbxMcpServer { #[tool( name = "dbx_prepare_table_import", - description = "Validate an approved-template mapping and immutable local source, then create a single-use import plan valid for 30 minutes. This tool never writes data." + description = "Validate a caller-approved template mapping and immutable local source, then generate a unique PostgreSQL staging relation and a single-use plan valid for 30 minutes. This tool performs no database query or write." )] async fn prepare_table_import(&self, Parameters(request): Parameters) -> CallToolResult { if self.web_mode { @@ -460,9 +478,6 @@ impl DbxMcpServer { "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", )); } - if let Err(error) = validate_staging_target(&request.schema, &request.table) { - return enterprise_error(error); - } let resolved = match self.resolve_connection(&request.selector).await { Ok(resolved) => resolved, Err(error) => return error, @@ -477,9 +492,6 @@ impl DbxMcpServer { "v1 MCP 表导入只允许 PostgreSQL staging 连接。", )); } - if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &database) { - return error; - } let path = match validate_import_file(&request.file_path, false) { Ok(path) => path, Err(error) => return enterprise_error(error), @@ -493,13 +505,16 @@ impl DbxMcpServer { Err(error) => return enterprise_error(error), }; let source_format = request.source_format.map(Into::into); + if let Err(error) = validate_governed_source_v1(&path.to_string_lossy(), source_format, &parse_options) { + return enterprise_error(error); + } let preview = dbx_core::table_import::preview_table_import_file_with_request( dbx_core::table_import::TableImportPreviewRequest { file_path: path.to_string_lossy().to_string(), source_ref: Some(identity.sha256.clone()), source_format, parse_options: parse_options.clone(), - preview_limit: Some(STRUCTURE_PROFILE_ROWS), + preview_limit: Some(1), }, ) .await; @@ -507,30 +522,27 @@ impl DbxMcpServer { Ok(preview) => preview, Err(error) => return enterprise_error(EnterpriseToolError::new("IMPORT_PREVIEW_FAILED", error)), }; - let mappings = request.mappings.into_iter().map(Into::into).collect::>(); - if let Err(error) = validate_mappings(&mappings, &preview.columns) { + if let Err(error) = validate_preview_headers(&preview) { return enterprise_error(error); } - if !request.create_table { - let columns = match self - .backend - .get_columns(&resolved.connection, &database, &request.schema, &request.table) - .await - { - Ok(columns) => columns, - Err(error) => return enterprise_error(EnterpriseToolError::new("IMPORT_TARGET_CHECK_FAILED", error)), - }; - let target_columns = - columns.iter().map(|column| column.name.as_str()).collect::>(); - if let Some(missing) = - mappings.iter().find(|mapping| !target_columns.contains(mapping.target_column.as_str())) - { - return enterprise_error(EnterpriseToolError::new( - "IMPORT_TARGET_COLUMN_NOT_FOUND", - format!("目标表缺少字段 {}。", missing.target_column), - )); - } + let current_identity = match file_identity(path).await { + Ok(identity) => identity, + Err(error) => return enterprise_error(error), + }; + if current_identity != identity { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_FILE_CHANGED_DURING_PREPARE", + "文件在 prepare 剖析期间发生变化；未创建计划。", + )); } + let mappings = match validate_mappings(&request.mappings, &preview.columns) { + Ok(mappings) => mappings, + Err(error) => return enterprise_error(error), + }; + let (schema, table) = match generated_staging_relation() { + Ok(relation) => relation, + Err(error) => return enterprise_error(error), + }; let batch_size = request.batch_size.unwrap_or(1_000); if !(1..=50_000).contains(&batch_size) { return enterprise_error(EnterpriseToolError::new( @@ -543,22 +555,24 @@ impl DbxMcpServer { resolved.connection.id.clone(), resolved.connection.name.clone(), database, - request.schema, - request.table, + schema, + table, request.template_version, identity, fingerprint, source_format, parse_options, mappings, - request.create_table, + true, batch_size, request.date_time_format, ) { Ok(plan) => plan, Err(error) => return enterprise_error(error), }; - self.enterprise.insert_plan(plan.clone()).await; + if let Err(error) = self.enterprise.insert_plan(plan.clone()).await { + return enterprise_error(error); + } structured_success( format!("导入计划 {} 已准备；绑定模板 {}，30 分钟内可启动一次。", plan.plan_id, plan.template_version), json!({ "plan": plan }), @@ -576,6 +590,10 @@ impl DbxMcpServer { "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", )); } + let import_permit = match self.enterprise.try_import_permit().await { + Ok(permit) => permit, + Err(error) => return enterprise_error(error), + }; let plan = match self.enterprise.consume_plan(request.plan_id.trim()).await { Ok(plan) => plan, Err(error) => return enterprise_error(error), @@ -597,13 +615,18 @@ impl DbxMcpServer { if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &plan.database) { return error; } - let job = self.enterprise.create_job(&plan).await; + let job = match self.enterprise.create_job(&plan).await { + Ok(job) => job, + Err(error) => return enterprise_error(error), + }; let import_id = job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).import_id.clone(); let import_request = plan.to_import_request(import_id.clone()); let backend = self.backend.clone(); let connection = resolved.connection.clone(); let task_job = job.clone(); + let plan_id = plan.plan_id.clone(); tokio::spawn(async move { + let _import_permit = import_permit; let progress_job = task_job.clone(); let progress = Arc::new(move |mut update: dbx_core::table_import::TableImportProgress| { if update.status == dbx_core::table_import::TableImportStatus::Done { @@ -613,7 +636,7 @@ impl DbxMcpServer { progress_job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).apply_progress(update); }); let result = backend - .import_table_file_for_mcp(&connection, import_request, task_job.cancelled.clone(), progress) + .import_table_file_for_mcp(&connection, import_request, plan_id, task_job.cancelled.clone(), progress) .await; let mut snapshot = task_job.snapshot.lock().unwrap_or_else(|error| error.into_inner()); match result { @@ -635,10 +658,18 @@ impl DbxMcpServer { snapshot.error = Some(error); } } + snapshot.finished_at_ms = Some( + std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_millis(), + ); }); structured_success( format!("导入任务 {import_id} 已启动；只会写入隔离 staging。"), - json!({ "importId": import_id, "planId": plan.plan_id, "status": "running" }), + json!({ + "importId": import_id, + "planId": plan.plan_id, + "relation": format!("{}.{}", plan.schema, plan.table), + "status": "running" + }), ) } @@ -674,19 +705,28 @@ impl DbxMcpServer { "v1 文件导入仅支持本地 DBX Desktop/MCP 模式。", )); } - let snapshot = match self.enterprise.cancel_job(request.import_id.trim()).await { - Ok(snapshot) => snapshot, + let (snapshot, already_requested) = match self.enterprise.cancel_job(request.import_id.trim()).await { + Ok(outcome) => outcome, Err(error) => return enterprise_error(error), }; structured_success( - format!("已向导入任务 {} 发送取消请求。", snapshot.import_id), - json!({ "importId": snapshot.import_id, "cancellationRequested": true, "lastStatus": snapshot.status }), + if already_requested { + format!("导入任务 {} 已经收到取消请求。", snapshot.import_id) + } else { + format!("已向导入任务 {} 发送取消请求。", snapshot.import_id) + }, + json!({ + "importId": snapshot.import_id, + "cancellationRequested": true, + "alreadyRequested": already_requested, + "lastStatus": snapshot.status + }), ) } #[tool( name = "dbx_vector_search", - description = "Search an allowed Milvus semantic collection with a bounded vector and safe equality filters. Approval status and semantic version are always enforced." + description = "Search an allowed Milvus semantic collection with a bounded vector and safe equality filters. Approval and effective date are always enforced; semantic_version is an optional exact filter." )] async fn vector_search(&self, Parameters(request): Parameters) -> CallToolResult { let resolved = match self.resolve_connection(&request.selector).await { @@ -768,6 +808,10 @@ impl DbxMcpServer { if let Err(error) = validate_vector_collection(&request.collection) { return enterprise_error(error); } + let _semantic_write_permit = match self.enterprise.semantic_write_permit().await { + Ok(permit) => permit, + Err(error) => return enterprise_error(error), + }; let path = match validate_import_file(&request.file_path, true) { Ok(path) => path, Err(error) => return enterprise_error(error), @@ -798,6 +842,38 @@ impl DbxMcpServer { Ok(batch_size) => batch_size, Err(error) => return enterprise_error(error), }; + let semantic_version = records + .first() + .and_then(|record| record.get("semantic_version")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(); + for chunk in records.chunks(batch_size) { + let card_ids = chunk + .iter() + .filter_map(|record| record.get("card_id").and_then(serde_json::Value::as_str)) + .map(str::to_string) + .collect::>(); + let query = milvus_existing_cards_query(&database, &request.collection, &card_ids); + let existing = match self + .backend + .execute_query(&resolved.connection, &database, &query, Some(card_ids.len()), Some(30)) + .await + { + Ok(result) => query_result_rows(result), + Err(error) => { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_OWNERSHIP_CHECK_FAILED", + format!("无法确认现有 card_id 归属，未执行 upsert：{error}"), + )) + } + }; + if let Err(error) = + validate_existing_card_ownership(&existing, &request.semantic_batch_id, &semantic_version) + { + return enterprise_error(error); + } + } let mut upserted = 0usize; for chunk in records.chunks(batch_size) { let query = milvus_upsert_query(&database, &request.collection, chunk); @@ -847,6 +923,12 @@ impl DbxMcpServer { if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &database) { return error; } + if request.published { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_PUBLISHED_BATCH_DELETE_BLOCKED", + "published 必须由调用方显式证明为 false；已发布语义批次不能通过该工具删除。", + )); + } if !resolved.policy.allow_dangerous_sql { return enterprise_error(EnterpriseToolError::new( "VECTOR_HIGH_RISK_WRITE_BLOCKED", @@ -856,6 +938,10 @@ impl DbxMcpServer { if let Err(error) = validate_vector_collection(&request.collection) { return enterprise_error(error); } + let _semantic_write_permit = match self.enterprise.semantic_write_permit().await { + Ok(permit) => permit, + Err(error) => return enterprise_error(error), + }; if request.semantic_batch_id.trim().is_empty() || request.semantic_batch_id.len() > 200 { return enterprise_error(EnterpriseToolError::new( "SEMANTIC_BATCH_ID_INVALID", @@ -894,6 +980,14 @@ impl DbxMcpServer { "Redis connections do not accept SQL through dbx_execute_query. Use dbx_execute_redis_command.", ); } + if is_vector_database(connection.db_type) + && !generic_vector_diagnostic_allowed(connection.db_type, &request.sql) + { + return enterprise_error(EnterpriseToolError::new( + "VECTOR_DEDICATED_TOOL_REQUIRED", + "向量实体检索、写入和删除只能使用 dbx_vector_search、dbx_vector_upsert_file 或 dbx_vector_delete_by_batch；通用查询仅允许明确的管理诊断端点。", + )); + } // Resolve the session before the database so its connection/database // binding is enforced on every stateful query. let session = match request.session_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { @@ -1613,6 +1707,58 @@ fn validate_safe_write_connection( Ok(()) } +fn is_vector_database(database_type: DatabaseType) -> bool { + matches!( + database_type, + DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb + ) +} + +fn generic_vector_diagnostic_allowed(database_type: DatabaseType, source: &str) -> bool { + let request_line = source.lines().find(|line| !line.trim().is_empty()).unwrap_or("").trim(); + let mut parts = request_line.split_whitespace(); + let method = parts.next().unwrap_or("").to_ascii_uppercase(); + let path = parts.next().unwrap_or("").split('?').next().unwrap_or("").trim_end_matches('/').to_ascii_lowercase(); + match database_type { + DatabaseType::Milvus => { + method == "POST" + && matches!( + path.as_str(), + "/v2/vectordb/collections/list" + | "/v2/vectordb/collections/describe" + | "/v2/vectordb/databases/list" + | "/v2/vectordb/indexes/list" + | "/v2/vectordb/indexes/describe" + ) + } + DatabaseType::Qdrant => { + method == "GET" + && (path == "/collections" + || path == "/cluster" + || path == "/healthz" + || path == "/readyz" + || path == "/telemetry" + || (path.starts_with("/collections/") && !path.contains("/points"))) + } + DatabaseType::Weaviate => { + method == "GET" + && (path == "/v1/meta" + || path == "/v1/schema" + || path.starts_with("/v1/schema/") + || path == "/v1/.well-known/live" + || path == "/v1/.well-known/ready") + } + DatabaseType::ChromaDb => { + method == "GET" + && (path.ends_with("/heartbeat") + || path.ends_with("/version") + || (path.contains("/collections") + && !["/get", "/query", "/count", "/peek"].iter().any(|suffix| path.ends_with(suffix)))) + } + _ => false, + } +} + fn backend_tool_error(default_code: &str, error: impl Into) -> CallToolResult { let error = error.into(); for code in [ @@ -2044,6 +2190,68 @@ mod tests { assert!(output.contains("Project\\|A / Staging West")); } + #[test] + fn generic_vector_query_allows_only_management_diagnostics() { + assert!(generic_vector_diagnostic_allowed(DatabaseType::Milvus, "POST /v2/vectordb/collections/describe\n{}")); + for query in [ + "POST /v2/vectordb/entities/search\n{}", + "POST /v2/vectordb/entities/query\n{}", + "POST /v2/vectordb/entities/upsert\n{}", + "POST /v2/vectordb/entities/delete\n{}", + ] { + assert!(!generic_vector_diagnostic_allowed(DatabaseType::Milvus, query)); + } + assert!(generic_vector_diagnostic_allowed(DatabaseType::Qdrant, "GET /collections/cards")); + assert!(!generic_vector_diagnostic_allowed(DatabaseType::Qdrant, "POST /collections/cards/points/search\n{}")); + } + + #[tokio::test] + async fn execute_query_requires_dedicated_vector_entity_tool() { + let milvus = connection("milvus-1", "运营组语义库", "milvus", "default"); + let server = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend { connections: vec![milvus], ..Default::default() }), + McpScope::default(), + false, + ); + let result = server + .execute_query(Parameters(ExecuteQueryRequest { + selector: ConnectionSelector { connection_id: Some("milvus-1".to_string()), connection_name: None }, + database: Some("default".to_string()), + sql: "POST /v2/vectordb/entities/search\n{}".to_string(), + session_id: None, + cell_char_offset: None, + cell_char_limit: None, + })) + .await; + assert_eq!( + result.structured_content.as_ref().and_then(|value| value.pointer("/error/code")), + Some(&json!("VECTOR_DEDICATED_TOOL_REQUIRED")) + ); + } + + #[tokio::test] + async fn vector_batch_delete_requires_explicit_unpublished_attestation() { + let milvus = connection("milvus-1", "运营组语义库", "milvus", "default"); + let server = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend { connections: vec![milvus], ..Default::default() }), + McpScope::default(), + false, + ); + let result = server + .vector_delete_by_batch(Parameters(VectorDeleteByBatchRequest { + selector: ConnectionSelector { connection_id: Some("milvus-1".to_string()), connection_name: None }, + database: Some("default".to_string()), + collection: "semantic_cards".to_string(), + semantic_batch_id: "batch-1".to_string(), + published: true, + })) + .await; + assert_eq!( + result.structured_content.as_ref().and_then(|value| value.pointer("/error/code")), + Some(&json!("VECTOR_PUBLISHED_BATCH_DELETE_BLOCKED")) + ); + } + #[test] fn server_registers_list_connections_tool() { let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); @@ -2104,25 +2312,22 @@ mod tests { .prepare_table_import(Parameters(PrepareTableImportRequest { selector: ConnectionSelector { connection_id: Some("write-pg".to_string()), connection_name: None }, database: Some("enterprise".to_string()), - schema: "staging".to_string(), - table: "orders_batch_1".to_string(), template_version: "orders-v1".to_string(), file_path: source.to_string_lossy().to_string(), source_format: None, parse_options: Default::default(), mappings: vec![ crate::enterprise_tools::McpImportColumnMapping { - source_column: "order_id".to_string(), + source_position: 1, + source_name: "order_id".to_string(), target_column: "order_id".to_string(), - target_data_type: Some("TEXT".to_string()), }, crate::enterprise_tools::McpImportColumnMapping { - source_column: "amount".to_string(), + source_position: 2, + source_name: "amount".to_string(), target_column: "amount".to_string(), - target_data_type: Some("NUMERIC(20,2)".to_string()), }, ], - create_table: true, batch_size: Some(1_000), date_time_format: None, })) @@ -2135,6 +2340,16 @@ mod tests { .and_then(serde_json::Value::as_str) .unwrap() .to_string(); + let plan = prepared.structured_content.as_ref().unwrap().pointer("/data/plan").unwrap(); + assert_eq!(plan.get("schema"), Some(&json!("staging"))); + assert!(plan.get("table").and_then(serde_json::Value::as_str).unwrap().starts_with("mcp_")); + assert_eq!(plan.get("createTable"), Some(&json!(true))); + assert!(plan + .get("mappings") + .and_then(serde_json::Value::as_array) + .unwrap() + .iter() + .all(|mapping| mapping.get("targetDataType") == Some(&json!("TEXT")))); std::fs::write(&source, "order_id,amount\nA001,10.00\nA002,20.00\n").unwrap(); let started = server.start_table_import(Parameters(StartTableImportRequest { plan_id })).await; diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index fde17e844c..33aafa1841 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -117,11 +117,11 @@ DBX MCP 当前提供 21 个工具： ## 本地文件导入与 Milvus 工具 -导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个允许目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。 +导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个真实目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。preview 支持 Excel、JSON、CSV 和 TSV；v1 治理入库有意只开放可流式有界处理的 UTF-8 CSV/TSV。Excel/JSON 或非 UTF-8 文件会在创建计划和访问数据库之前返回稳定的不支持错误。 -`dbx_prepare_table_import` 会把连接、已批准模板版本、规范路径、大小、修改时间、SHA-256、解析参数、字段映射和 staging 目标固化为不可变摘要。计划 30 分钟后过期，且只能启动一次。`dbx_start_table_import` 在后台启动前会重新核对完整文件身份。导入目标仅限 `DBX_MCP_IMPORT_STAGING_SCHEMAS` 中的 PostgreSQL Schema（默认为 `staging`）；发布到业务 Schema 仍由调用方通过独立治理流程完成。 +`dbx_prepare_table_import` 会把连接、调用方已批准模板版本、规范路径、大小、修改时间、SHA-256、解析参数和字段映射固化为不可变摘要。每个映射必须提供 preview 返回的 1 基 `source_position`、精确 `source_name` 和 `target_column`。目标不能由调用方指定：服务端生成唯一 `staging.mcp_`，强制创建新表，所有业务列固定为 `TEXT`。计划 30 分钟后过期，且只能启动一次。start 复验源文件后，先流式写出有界治理快照，再打开数据库连接；每行增加 `import_id`、`plan_id`、`source_sha`、真实 `source_row_number`、`source_row_hash` 和 `loaded_at`。无法形成精确血缘时会在数据库访问前阻断。status 会返回生成的 relation；发布到业务 Schema 仍是独立治理流程。 -Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索请求必须提供 `active_at`；服务端始终追加 `approval_status == "approved"` 以及 `effective_from`/`effective_to` 生效时间过滤，并可选追加精确 `semantic_version`。JSONL upsert 需要“数据读写”权限；删除只能按精确 `semantic_batch_id` 执行，并要求“完全访问”权限。通过 `dbx_execute_query` 提交的通用向量 REST 请求也会按端点分类，因此写入或破坏性 `POST` 不能绕过 MCP 策略。 +Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索必须提供 `active_at`，服务端强制 approved 和生效时间过滤，并可选追加精确 `semantic_version`。upsert 使用严格语义卡 Schema：重算并核对 64 位 `content_checksum`，要求 `embedding_revision`，拒绝 `chunk_index` 与 `embedding_dimension`，文件内 card_id 必须唯一且只能有一个语义版本；写入前还会查询现有 card_id，禁止覆盖其他批次或版本。删除除“完全访问”外还要求调用方显式传入 `published: false`。所有通过 `dbx_execute_query` 提交的向量实体 search/query/mutation 都返回 `VECTOR_DEDICATED_TOOL_REQUIRED`，通用入口只保留明确管理诊断。 向量检索的核心参数形状如下： @@ -130,7 +130,7 @@ Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。 "connection_name": "运营组语义库", "database": "default", "collection": "semantic_cards", - "active_at": "2026-08-25T10:00:00+08:00", + "active_at": "2026-08-25", "semantic_version": "semantic-v3", "embedding": [/* 必须恰好包含 1024 个有限浮点数 */], "top_k": 12, @@ -139,7 +139,7 @@ Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。 } ``` -仅依赖生效日期选择活动语义卡时可以省略 `semantic_version`，但 `active_at` 仍为必填。`embedding` 的实际长度必须等于 `DBX_MCP_VECTOR_DIMENSION`；上面数组内的注释只是文档缩写，不是可直接提交的工具参数。 +仅依赖生效日期选择活动语义卡时可以省略 `semantic_version`，但 `active_at` 仍为必填。v1 的 `embedding` 必须恰好包含 1024 个有限数值；上面数组内的注释只是文档缩写，不是可直接提交的工具参数。 ## 有状态查询会话 @@ -222,9 +222,12 @@ DBX 在 **设置 → MCP** 中保存一份权威策略，并在每次请求时 | `DBX_MCP_IMPORT_ROOTS` | 导入和语义 JSONL 工具可读取的本地目录 allowlist，使用平台路径分隔符 | | `DBX_MCP_IMPORT_STAGING_SCHEMAS` | 允许导入的 PostgreSQL staging Schema，逗号分隔，默认 `staging` | | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Excel/CSV/TSV/JSON 最大字节数，默认 512 MiB | +| `DBX_MCP_IMPORT_CONCURRENCY` | 治理导入并发数，限制为 1～16，默认 `2` | +| `DBX_MCP_IMPORT_PLAN_CAPACITY` | 同时保留的待启动计划上限，默认 `512` | +| `DBX_MCP_IMPORT_JOB_CAPACITY` | 运行中和终态任务总保留上限，默认 `256` | +| `DBX_MCP_IMPORT_TERMINAL_LIMIT` | 24 小时 TTL 内终态任务保留上限，默认 `128` | | `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | 语义 JSONL 最大字节数，默认 64 MiB | | `DBX_MCP_VECTOR_COLLECTIONS` | Milvus 集合 allowlist，逗号分隔，默认 `semantic_cards` | -| `DBX_MCP_VECTOR_DIMENSION` | 固定向量维度，默认 `1024` | | `DBX_MCP_VECTOR_TOP_K_MAX` | Milvus Top K 上限，默认 `20`，硬上限 50 | | `DBX_MCP_VECTOR_FILTER_FIELDS` | 允许用于安全等值过滤的附加字段，逗号分隔 | | `DBX_MCP_VECTOR_OUTPUT_FIELDS` | 向量检索允许返回的标量字段，逗号分隔 | diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index e3a95cacf3..f6a649195d 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -117,11 +117,11 @@ Connection-scoped sessions hide connection-mutating and desktop UI tools. ## Local file imports and Milvus tools -The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. +The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more real directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. Preview supports Excel, JSON, CSV, and TSV; governed database import in v1 is deliberately limited to streaming UTF-8 CSV/TSV. Excel/JSON or non-UTF-8 prepare calls fail with a stable unsupported code before a plan or database write exists. -`dbx_prepare_table_import` binds the connection, approved template version, canonical path, size, modification time, SHA-256, parsing options, mappings, and staging target into an immutable digest. The returned plan expires after 30 minutes and can be started once. `dbx_start_table_import` rechecks the complete file identity before launching a background import. Import targets are limited to PostgreSQL schemas listed in `DBX_MCP_IMPORT_STAGING_SCHEMAS` (default `staging`); publication to business schemas remains the caller's separately governed workflow. +`dbx_prepare_table_import` binds the connection, caller-approved template version, canonical path, size, modification time, SHA-256, parsing options, and mappings into an immutable digest. Each mapping supplies the 1-based `source_position`, the exact `source_name` returned by preview, and `target_column`. The server—not the caller—generates a unique `staging.mcp_` relation, forces a new table, and makes every business column `TEXT`. The returned plan expires after 30 minutes and can be started once. Start rechecks the source and streams it into a bounded governed snapshot before opening a database pool. Every row gains `import_id`, `plan_id`, `source_sha`, the real `source_row_number`, `source_row_hash`, and `loaded_at`; if exact lineage cannot be produced, the import fails before database access. Status returns the generated relation. Publication to business schemas remains a separately governed workflow. -Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires `active_at`, always adds `approval_status == "approved"` plus `effective_from`/`effective_to` activity filters, and optionally applies an exact `semantic_version`. Upsert reads bounded `.jsonl` files and requires Data read/write permission. Delete accepts only an exact `semantic_batch_id` and requires Full access. Generic vector REST passed to `dbx_execute_query` is also classified by endpoint, so write or destructive POST requests cannot bypass MCP policy. +Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires `active_at`, always adds `approval_status == "approved"` plus effective-date filters, and optionally applies an exact `semantic_version`. Upsert accepts a strict semantic-card schema, verifies the 64-hex `content_checksum` against `content`, requires `embedding_revision`, rejects `chunk_index` and `embedding_dimension`, enforces unique card IDs and one semantic version per file, and preflights existing IDs so another batch/version cannot be overwritten. Delete requires Full access and an explicit `published: false` attestation. All entity search/query/mutation REST submitted through generic `dbx_execute_query` returns `VECTOR_DEDICATED_TOOL_REQUIRED`; only explicit management diagnostics remain available there. The core vector-search arguments are: @@ -130,7 +130,7 @@ The core vector-search arguments are: "connection_name": "Operations semantics", "database": "default", "collection": "semantic_cards", - "active_at": "2026-08-25T10:00:00+08:00", + "active_at": "2026-08-25", "semantic_version": "semantic-v3", "embedding": [/* exactly 1024 finite floating-point values */], "top_k": 12, @@ -139,7 +139,7 @@ The core vector-search arguments are: } ``` -Omit `semantic_version` when effective dates alone select the active cards. `active_at` remains required. `embedding` must contain exactly `DBX_MCP_VECTOR_DIMENSION` finite values; the abbreviated comment above is documentation only, not valid tool input. +Omit `semantic_version` when effective dates alone select the active cards. `active_at` remains required. In v1, `embedding` must contain exactly 1024 finite values; the abbreviated comment above is documentation only, not valid tool input. ## Stateful Query Sessions @@ -222,9 +222,12 @@ Updated servers do not let `DBX_MCP_ALLOW_WRITES` or `DBX_MCP_ALLOW_DANGEROUS_SQ | `DBX_MCP_IMPORT_ROOTS` | Platform-separated allowlist of local directories readable by import and semantic JSONL tools | | `DBX_MCP_IMPORT_STAGING_SCHEMAS` | Comma-separated PostgreSQL staging schema allowlist (default `staging`) | | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Maximum Excel/CSV/TSV/JSON source size (default 512 MiB) | +| `DBX_MCP_IMPORT_CONCURRENCY` | Concurrent governed imports, clamped to 1–16 (default `2`) | +| `DBX_MCP_IMPORT_PLAN_CAPACITY` | Maximum live prepared plans (default `512`) | +| `DBX_MCP_IMPORT_JOB_CAPACITY` | Maximum retained running/terminal jobs (default `256`) | +| `DBX_MCP_IMPORT_TERMINAL_LIMIT` | Maximum terminal jobs retained within the 24-hour TTL (default `128`) | | `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | Maximum semantic JSONL size (default 64 MiB) | | `DBX_MCP_VECTOR_COLLECTIONS` | Comma-separated Milvus collection allowlist (default `semantic_cards`) | -| `DBX_MCP_VECTOR_DIMENSION` | Required embedding dimension (default `1024`) | | `DBX_MCP_VECTOR_TOP_K_MAX` | Maximum Milvus Top K, hard-capped at 50 (default `20`) | | `DBX_MCP_VECTOR_FILTER_FIELDS` | Comma-separated additional safe equality-filter fields | | `DBX_MCP_VECTOR_OUTPUT_FIELDS` | Comma-separated scalar fields that vector search may return | diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index 7b7bd23de5..2fcbbf03d6 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -183,7 +183,7 @@ DBX connection storage defaults to: Override the directory with `DBX_DATA_DIR`. -Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Paths must be absolute regular files inside one configured root. Import plans bind SHA-256, parsing, mapping, template version, connection, and staging target; start revalidates the source and imports from a private task snapshot. Import tools support local mode only and return `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` in Web mode. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. +Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Preview accepts Excel/JSON/CSV/TSV, while v1 governed database import is fail-closed to streaming UTF-8 CSV/TSV. Prepare accepts position+name mappings and generates a unique `staging.mcp_` relation; callers cannot append to an existing table. Start revalidates the source and streams a private snapshot containing row numbers, hashes, and batch lineage before opening a database pool. Import tools support local mode only and return `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` in Web mode. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. ### Agent/JDBC databases @@ -294,7 +294,6 @@ SQL text is not included in normal MCP errors or logged by default. Enable tempo | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Maximum tabular import source size (default 512 MiB) | | `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | Maximum semantic JSONL size (default 64 MiB) | | `DBX_MCP_VECTOR_COLLECTIONS` | Comma-separated Milvus collection allowlist (default `semantic_cards`) | -| `DBX_MCP_VECTOR_DIMENSION` | Required embedding dimension (default `1024`) | | `DBX_MCP_VECTOR_TOP_K_MAX` | Maximum Milvus Top K (default `20`, hard cap 50) | | `DBX_MCP_DEBUG_SQL` | Include SQL in temporary diagnostics | | `DBX_MCP_BINARY` | Override the native binary used by the npm launcher | @@ -509,7 +508,7 @@ MCP 配置： 通过 `DBX_DATA_DIR` 覆盖默认目录。Windows 便携版应指向 `DBX.exe` 同级、包含 `dbx.db` 的 `data` 文件夹。 -本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。文件必须是允许目录内的绝对路径普通文件。prepare 会固化 SHA-256、解析参数、映射、模板版本、连接和 staging 目标；start 会重新复验，并从任务私有快照执行导入。导入工具仅支持本地模式，Web 模式稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 +本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。preview 支持 Excel/JSON/CSV/TSV，但 v1 治理入库只允许可流式有界转换的 UTF-8 CSV/TSV。prepare 接受“源位置+源名称”映射并由服务端生成唯一 `staging.mcp_`，调用方不能追加已有表。start 在打开数据库连接前复验文件并流式生成包含真实行号、行哈希和批次血缘的任务快照。导入工具仅支持本地模式，Web 模式稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 ### DBX Web / Docker @@ -597,7 +596,6 @@ MongoDB 更新和删除在未启用完全访问时必须提供可验证有效的 | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | 表格导入源最大字节数，默认 512 MiB | | `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | 语义 JSONL 最大字节数，默认 64 MiB | | `DBX_MCP_VECTOR_COLLECTIONS` | Milvus 集合 allowlist，逗号分隔，默认 `semantic_cards` | -| `DBX_MCP_VECTOR_DIMENSION` | 固定向量维度，默认 `1024` | | `DBX_MCP_VECTOR_TOP_K_MAX` | Milvus Top K 上限，默认 `20`，硬上限 50 | | `DBX_MCP_DEBUG_SQL` | 临时输出 SQL 诊断信息 | | `DBX_MCP_BINARY` | 覆盖 npm 启动器使用的原生文件 | From 75133be12f0b2f6723dc2075048f92cef975248c Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 18:34:09 +0800 Subject: [PATCH 04/10] feat(mcp): add governed XLSX streaming imports --- Cargo.lock | 1 + crates/dbx-core/src/table_import.rs | 64 ++- crates/dbx-mcp/Cargo.toml | 1 + crates/dbx-mcp/src/backend.rs | 1 + crates/dbx-mcp/src/enterprise_tools.rs | 698 +++++++++++++++++++++++-- crates/dbx-mcp/src/server.rs | 48 +- docs/content/docs/mcp.cn.mdx | 4 +- docs/content/docs/mcp.mdx | 4 +- packages/mcp-server/README.md | 4 +- 9 files changed, 756 insertions(+), 69 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9f7af1f1da..0726c380ae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2172,6 +2172,7 @@ dependencies = [ "tokio-rustls 0.26.4", "url", "uuid", + "zip", ] [[package]] diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index aa59fe0626..275ec3c4fb 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -2793,10 +2793,15 @@ fn parse_xlsx_file_with_options_and_text_columns( } #[derive(Debug)] -enum XlsxStreamMessage { +/// XLSX/XLSM 有界流式读取消息，供导入与治理快照复用。 +pub enum XlsxStreamMessage { + /// 已解析并去重的兼容 DBX 表头。 Header(Vec), - Rows(Vec>), + /// 数据批次及每行对应的工作表绝对 1 基行号。 + Rows { rows: Vec>, source_row_numbers: Vec }, + /// 已读取的近似字节进度。 Progress(u64), + /// 工作表已完整解析。 Done, } @@ -2856,6 +2861,7 @@ struct XlsxStreamRowsState { columns: Vec, header_sent: bool, pending_rows: Vec>, + pending_source_row_numbers: Vec, rows_seen: usize, current_row: Option, current_values: Vec, @@ -2881,6 +2887,7 @@ impl XlsxStreamRowsState { columns: expected_columns.unwrap_or_default(), header_sent: false, pending_rows: Vec::with_capacity(batch_size), + pending_source_row_numbers: Vec::with_capacity(batch_size), rows_seen: 0, current_row: None, current_values: Vec::new(), @@ -2923,6 +2930,9 @@ impl XlsxStreamRowsState { absolute_column: usize, text_source_columns: &HashSet, ) -> bool { + if text_source_columns.contains("*") { + return true; + } self.initialize_range(absolute_row, absolute_column); absolute_column .checked_sub(self.start_column) @@ -3007,6 +3017,7 @@ impl XlsxStreamRowsState { values.resize(self.columns.len(), serde_json::Value::Null); values.truncate(self.columns.len()); self.pending_rows.push(values); + self.pending_source_row_numbers.push(absolute_row); self.rows_seen = self.rows_seen.saturating_add(1); if self.pending_rows.len() >= self.batch_size { self.emit_rows(progress)?; @@ -3019,12 +3030,16 @@ impl XlsxStreamRowsState { return Ok(()); } self.sender - .blocking_send(Ok(XlsxStreamMessage::Rows(std::mem::take(&mut self.pending_rows)))) + .blocking_send(Ok(XlsxStreamMessage::Rows { + rows: std::mem::take(&mut self.pending_rows), + source_row_numbers: std::mem::take(&mut self.pending_source_row_numbers), + })) .map_err(|_| "Excel import consumer closed before the stream finished".to_string())?; self.sender .blocking_send(Ok(XlsxStreamMessage::Progress(progress))) .map_err(|_| "Excel import consumer closed before the stream finished".to_string())?; self.pending_rows = Vec::with_capacity(self.batch_size); + self.pending_source_row_numbers = Vec::with_capacity(self.batch_size); Ok(()) } @@ -3063,7 +3078,9 @@ fn stream_xlsx_rows_to_channel( } #[allow(clippy::too_many_arguments)] -fn stream_xlsx_rows_to_channel_with_control( +/// 以有界批次流式读取 XLSX/XLSM，保留绝对源行号并支持外部取消。 +/// 调用方必须消费有界通道；函数不会把整个工作表加载到内存。 +pub fn stream_xlsx_rows_to_channel_with_control( path: &str, options: &TableImportParseOptions, batch_size: usize, @@ -3321,7 +3338,7 @@ async fn validate_xlsx_worksheet_for_import( match message { Ok(XlsxStreamMessage::Header(header)) => columns = Some(header), Ok(XlsxStreamMessage::Progress(bytes_read)) => on_progress(bytes_read), - Ok(XlsxStreamMessage::Rows(_) | XlsxStreamMessage::Done) => {} + Ok(XlsxStreamMessage::Rows { .. } | XlsxStreamMessage::Done) => {} Err(error) => { producer_cancelled.store(true, Ordering::Release); drop(receiver); @@ -6676,7 +6693,7 @@ where }; match message { Ok(XlsxStreamMessage::Header(_)) => {} - Ok(XlsxStreamMessage::Rows(rows)) => { + Ok(XlsxStreamMessage::Rows { rows, .. }) => { if is_cancelled(&request.import_id).await { producer_cancelled.store(true, Ordering::Release); drop(receiver); @@ -7476,7 +7493,7 @@ mod tests { while let Some(message) = receiver.blocking_recv() { match message.unwrap() { XlsxStreamMessage::Header(columns) => streamed_columns = columns, - XlsxStreamMessage::Rows(rows) => streamed_rows.extend(rows), + XlsxStreamMessage::Rows { rows, .. } => streamed_rows.extend(rows), _ => {} } } @@ -7577,7 +7594,7 @@ mod tests { .unwrap(); let mut streamed_rows = Vec::new(); while let Some(message) = receiver.blocking_recv() { - if let XlsxStreamMessage::Rows(rows) = message.unwrap() { + if let XlsxStreamMessage::Rows { rows, .. } = message.unwrap() { streamed_rows.extend(rows); } } @@ -8741,7 +8758,7 @@ mod tests { let streamed_rows = messages .into_iter() .filter_map(|message| match message { - XlsxStreamMessage::Rows(rows) => Some(rows), + XlsxStreamMessage::Rows { rows, .. } => Some(rows), _ => None, }) .flatten() @@ -9020,7 +9037,7 @@ mod tests { let mut streamed_rows = Vec::new(); while let Some(message) = receiver.blocking_recv() { - if let XlsxStreamMessage::Rows(rows) = message.unwrap() { + if let XlsxStreamMessage::Rows { rows, .. } = message.unwrap() { streamed_rows.extend(rows); } } @@ -9051,7 +9068,7 @@ mod tests { let mut streamed_rows = Vec::new(); while let Some(message) = receiver.blocking_recv() { - if let XlsxStreamMessage::Rows(rows) = message.unwrap() { + if let XlsxStreamMessage::Rows { rows, .. } = message.unwrap() { streamed_rows.extend(rows); } } @@ -9059,6 +9076,27 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn streaming_excel_rows_honor_preexisting_cancellation() { + let path = std::env::temp_dir().join(format!("dbx-table-import-stream-cancel-{}.xlsx", uuid::Uuid::new_v4())); + std::fs::write(&path, build_styled_test_xlsx(false, &[("A1", 5, 10.0)])).unwrap(); + let options = TableImportParseOptions { has_header: Some(false), ..TableImportParseOptions::default() }; + let (sender, _receiver) = tokio::sync::mpsc::channel(2); + let error = stream_xlsx_rows_to_channel_with_control( + &path.to_string_lossy(), + &options, + 1, + None, + HashSet::from(["*".to_string()]), + true, + sender, + Arc::new(AtomicBool::new(true)), + ) + .unwrap_err(); + assert_eq!(error, "Import cancelled"); + let _ = std::fs::remove_file(path); + } + #[test] fn streaming_excel_rows_preserve_custom_title_and_data_range() { let path = std::env::temp_dir().join(format!("dbx-table-import-stream-range-{}.xlsx", uuid::Uuid::new_v4())); @@ -9093,7 +9131,7 @@ mod tests { while let Some(message) = receiver.blocking_recv() { match message.unwrap() { XlsxStreamMessage::Header(header) => columns = header, - XlsxStreamMessage::Rows(rows) => streamed_rows.extend(rows), + XlsxStreamMessage::Rows { rows, .. } => streamed_rows.extend(rows), _ => {} } } @@ -9164,7 +9202,7 @@ mod tests { let mut streamed_rows = Vec::new(); while let Some(message) = receiver.blocking_recv() { - if let XlsxStreamMessage::Rows(rows) = message.unwrap() { + if let XlsxStreamMessage::Rows { rows, .. } = message.unwrap() { streamed_rows.extend(rows); } } diff --git a/crates/dbx-mcp/Cargo.toml b/crates/dbx-mcp/Cargo.toml index fd7814b8e2..4114f91a27 100644 --- a/crates/dbx-mcp/Cargo.toml +++ b/crates/dbx-mcp/Cargo.toml @@ -35,3 +35,4 @@ rustls-pemfile = "2.2" tempfile = "3" tokio = { version = "1", features = ["net", "io-util"] } tokio-rustls = "0.26" +zip = "4" diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index ab085eeb84..35ea3da94e 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -727,6 +727,7 @@ impl DbxBackend for LocalBackend { &plan_id, &expected_sha256, &normalized_path, + cancelled.clone(), ) .await .map_err(|error| error.to_string())?; diff --git a/crates/dbx-mcp/src/enterprise_tools.rs b/crates/dbx-mcp/src/enterprise_tools.rs index dffde397f4..9957ff46fd 100644 --- a/crates/dbx-mcp/src/enterprise_tools.rs +++ b/crates/dbx-mcp/src/enterprise_tools.rs @@ -87,7 +87,7 @@ impl std::fmt::Display for EnterpriseToolError { } } -#[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] +#[derive(Debug, Clone, Copy, Deserialize, Serialize, schemars::JsonSchema)] #[serde(rename_all = "snake_case")] pub enum McpImportSourceFormat { Csv, @@ -173,10 +173,20 @@ impl McpImportParseOptions { pub struct McpImportColumnMapping { #[schemars(description = "1-based source column position returned by dbx_preview_import_file")] pub source_position: usize, - pub source_name: String, + pub raw_source_name: String, + pub canonical_source_name: String, pub target_column: String, } +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct McpSourceColumn { + pub source_position: usize, + pub raw_source_name: String, + pub canonical_source_name: String, + pub dbx_source_name: String, +} + #[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)] pub struct PreviewImportFileRequest { pub file_path: String, @@ -710,16 +720,103 @@ fn truncate_preview_text(value: &str, char_limit: usize) -> String { } } -pub fn structure_fingerprint(preview: &TableImportPreview, parse_options: &TableImportParseOptions) -> String { +pub fn structure_fingerprint( + preview: &TableImportPreview, + parse_options: &TableImportParseOptions, + source_columns: &[McpSourceColumn], +) -> String { let value = json!({ "fileType": preview.file_type, "columns": preview.columns, + "sourceColumns": source_columns, "sheets": preview.sheets, "parseOptions": parse_options, }); sha256_bytes(value.to_string().as_bytes()) } +pub async fn source_columns_for_preview( + file_path: &str, + source_format: Option, + parse_options: &TableImportParseOptions, + dbx_columns: &[String], +) -> Result, EnterpriseToolError> { + let row_range = dbx_core::table_import::effective_import_row_range(parse_options) + .map_err(|error| EnterpriseToolError::new("IMPORT_ROW_RANGE_INVALID", error))?; + let raw_names = if let Some(title_row) = row_range.title_row { + let header_options = TableImportParseOptions { + has_header: Some(false), + title_row: Some(0), + data_start_row: Some(title_row), + last_data_row: Some(title_row), + ..parse_options.clone() + }; + let header = dbx_core::table_import::preview_table_import_file_with_request( + dbx_core::table_import::TableImportPreviewRequest { + file_path: file_path.to_string(), + source_ref: None, + source_format, + parse_options: header_options, + preview_limit: Some(1), + }, + ) + .await + .map_err(|error| EnterpriseToolError::new("IMPORT_HEADER_PREVIEW_FAILED", error))?; + header.rows.first().cloned().unwrap_or_default().into_iter().map(source_header_text).collect::>() + } else { + dbx_columns.to_vec() + }; + let normalized = (0..dbx_columns.len()) + .map(|index| { + let raw = raw_names.get(index).cloned().unwrap_or_default(); + dbx_core::table_import::normalize_header(raw.trim_start_matches('\u{feff}'), index) + }) + .collect::>(); + let canonical = canonical_source_names(&normalized); + Ok((0..dbx_columns.len()) + .map(|index| McpSourceColumn { + source_position: index + 1, + raw_source_name: raw_names.get(index).cloned().unwrap_or_default(), + canonical_source_name: canonical[index].clone(), + dbx_source_name: dbx_columns[index].clone(), + }) + .collect()) +} + +fn source_header_text(value: Value) -> String { + match value { + Value::Null => String::new(), + Value::String(value) => value, + other => other.to_string(), + } +} + +fn canonical_source_names(normalized: &[String]) -> Vec { + let counts = normalized.iter().fold(HashMap::::new(), |mut counts, name| { + *counts.entry(name.to_lowercase()).or_default() += 1; + counts + }); + let mut occurrences = HashMap::::new(); + let mut used = HashSet::new(); + normalized + .iter() + .map(|name| { + let key = name.to_lowercase(); + let occurrence = occurrences.entry(key.clone()).or_default(); + *occurrence += 1; + let base = + if counts.get(&key).copied().unwrap_or(0) > 1 { format!("{name}__{occurrence}") } else { name.clone() }; + let mut candidate = base.clone(); + let mut suffix = 1usize; + while !used.insert(candidate.to_lowercase()) { + candidate = format!("{base}__{suffix}"); + suffix += 1; + } + candidate + }) + .collect() +} + pub fn generated_staging_relation() -> Result<(String, String), EnterpriseToolError> { let allowed = comma_list_env("DBX_MCP_IMPORT_STAGING_SCHEMAS", &["staging"]); let schema = allowed @@ -733,7 +830,7 @@ pub fn generated_staging_relation() -> Result<(String, String), EnterpriseToolEr pub fn validate_mappings( mappings: &[McpImportColumnMapping], - source_columns: &[String], + source_columns: &[McpSourceColumn], ) -> Result, EnterpriseToolError> { if mappings.is_empty() { return Err(EnterpriseToolError::new("IMPORT_MAPPING_REQUIRED", "至少需要一个字段映射。")); @@ -745,18 +842,24 @@ pub fn validate_mappings( let source_index = mapping.source_position.checked_sub(1).ok_or_else(|| { EnterpriseToolError::new("IMPORT_SOURCE_POSITION_INVALID", "source_position 从 1 开始，不能为 0。") })?; - let actual_name = source_columns.get(source_index).ok_or_else(|| { + let source = source_columns.get(source_index).ok_or_else(|| { EnterpriseToolError::new( "IMPORT_SOURCE_POSITION_INVALID", format!("源文件没有第 {} 列。", mapping.source_position), ) })?; - if actual_name != &mapping.source_name { + if source.raw_source_name != mapping.raw_source_name + || source.canonical_source_name != mapping.canonical_source_name + { return Err(EnterpriseToolError::new( "IMPORT_SOURCE_NAME_MISMATCH", format!( - "第 {} 列当前名称为 {}，与请求中的 {} 不一致。", - mapping.source_position, actual_name, mapping.source_name + "第 {} 列当前 raw/canonical 名称为 {:?}/{:?}，与请求中的 {:?}/{:?} 不一致。", + mapping.source_position, + source.raw_source_name, + source.canonical_source_name, + mapping.raw_source_name, + mapping.canonical_source_name ), )); } @@ -780,7 +883,7 @@ pub fn validate_mappings( )); } normalized.push(TableImportColumnMapping { - source_column: actual_name.clone(), + source_column: source.dbx_source_name.clone(), target_column: mapping.target_column.clone(), target_data_type: Some("TEXT".to_string()), }); @@ -871,28 +974,60 @@ pub fn validate_governed_source_v1( ) -> Result { let source_format = dbx_core::table_import::effective_source_format(file_path, source_format) .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_FORMAT_INVALID", error))?; - if !source_format.is_delimited() { - return Err(EnterpriseToolError::new( - "IMPORT_GOVERNED_FORMAT_UNSUPPORTED_V1", - format!("{} 尚未实现流式有界行血缘转换；v1 治理导入当前仅支持 UTF-8 CSV/TSV。", source_format.label()), - )); + if source_format.is_delimited() { + if !matches!(parse_options.encoding, None | Some(TableImportTextEncoding::Auto | TableImportTextEncoding::Utf8)) + { + return Err(EnterpriseToolError::new( + "IMPORT_GOVERNED_ENCODING_UNSUPPORTED_V1", + "v1 流式治理导入仅支持 UTF-8；其他编码不创建导入计划。", + )); + } + return Ok(source_format); } - if !matches!(parse_options.encoding, None | Some(TableImportTextEncoding::Auto | TableImportTextEncoding::Utf8)) { + if source_format == TableImportSourceFormat::Excel { + let extension = + Path::new(file_path).extension().and_then(OsStr::to_str).unwrap_or_default().to_ascii_lowercase(); + if matches!(extension.as_str(), "xlsx" | "xlsm") { + return Ok(source_format); + } return Err(EnterpriseToolError::new( - "IMPORT_GOVERNED_ENCODING_UNSUPPORTED_V1", - "v1 流式治理导入仅支持 UTF-8；其他编码不创建导入计划。", + "IMPORT_GOVERNED_XLS_UNSUPPORTED_V1", + "旧版 .xls 仅支持有界 preview；正式治理导入要求另存为 .xlsx。", )); } - Ok(source_format) + Err(EnterpriseToolError::new( + "IMPORT_GOVERNED_FORMAT_UNSUPPORTED_V1", + "JSON 尚未实现流式有界治理快照；不创建导入计划。", + )) } pub async fn build_governed_import_snapshot( - mut request: TableImportRequest, + request: TableImportRequest, plan_id: &str, source_sha: &str, output_path: &Path, + cancelled: Arc, ) -> Result { let source_format = validate_governed_source_v1(&request.file_path, request.source_format, &request.parse_options)?; + let result = if source_format == TableImportSourceFormat::Excel { + build_governed_xlsx_snapshot(request, plan_id, source_sha, output_path, cancelled).await + } else { + build_governed_delimited_snapshot(request, plan_id, source_sha, output_path, source_format, cancelled).await + }; + if result.is_err() { + let _ = tokio::fs::remove_file(output_path).await; + } + result +} + +async fn build_governed_delimited_snapshot( + mut request: TableImportRequest, + plan_id: &str, + source_sha: &str, + output_path: &Path, + source_format: TableImportSourceFormat, + cancelled: Arc, +) -> Result { let config = dbx_core::table_import::effective_delimited_config(source_format, &request.parse_options) .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_PARSE_FAILED", error))?; let offset = UtcOffset::from_hms(8, 0, 0) @@ -920,6 +1055,9 @@ pub async fn build_governed_import_snapshot( let mut source_indexes: Option> = None; let mut source_row_count = 0usize; for (index, record) in reader.records().enumerate() { + if cancelled.load(Ordering::Acquire) { + return Err(EnterpriseToolError::new("IMPORT_CANCELLED", "导入已取消。")); + } let source_row_number = index + 1; let record = record.map_err(|error| { EnterpriseToolError::new( @@ -1018,6 +1156,193 @@ pub async fn build_governed_import_snapshot( Ok(request) } +async fn build_governed_xlsx_snapshot( + mut request: TableImportRequest, + plan_id: &str, + source_sha: &str, + output_path: &Path, + cancelled: Arc, +) -> Result { + if cancelled.load(Ordering::Acquire) { + return Err(EnterpriseToolError::new("IMPORT_CANCELLED", "导入已取消。")); + } + let offset = UtcOffset::from_hms(8, 0, 0) + .map_err(|error| EnterpriseToolError::new("IMPORT_LOADED_AT_FAILED", error.to_string()))?; + let loaded_at = OffsetDateTime::now_utc() + .to_offset(offset) + .format(&Rfc3339) + .map_err(|error| EnterpriseToolError::new("IMPORT_LOADED_AT_FAILED", error.to_string()))?; + let mut writer = csv::WriterBuilder::new() + .has_headers(false) + .from_path(output_path) + .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + let mut headers = request.mappings.iter().map(|mapping| mapping.target_column.clone()).collect::>(); + headers.extend(RESERVED_STAGING_COLUMNS.iter().map(|column| column.to_string())); + writer + .write_record(&headers) + .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + + let (sender, mut receiver) = tokio::sync::mpsc::channel(2); + let path = request.file_path.clone(); + let options = request.parse_options.clone(); + let batch_size = request.batch_size.max(1); + let text_source_columns = HashSet::from(["*".to_string()]); + let producer_cancelled = Arc::new(AtomicBool::new(cancelled.load(Ordering::Acquire))); + let _producer_cancel_guard = CancelOnDrop(producer_cancelled.clone()); + let producer_cancelled_for_task = producer_cancelled.clone(); + let producer = tokio::task::spawn_blocking(move || { + dbx_core::table_import::stream_xlsx_rows_to_channel_with_control( + &path, + &options, + batch_size, + None, + text_source_columns, + true, + sender, + producer_cancelled_for_task, + ) + }); + let user_cancelled = cancelled.clone(); + let monitor_cancelled = producer_cancelled.clone(); + let cancellation_monitor = tokio::spawn(async move { + while !monitor_cancelled.load(Ordering::Acquire) { + if user_cancelled.load(Ordering::Acquire) { + monitor_cancelled.store(true, Ordering::Release); + break; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }); + let mut header_seen = false; + let mut done_seen = false; + let mut source_row_count = 0usize; + let mut source_indexes: Option> = None; + while let Some(message) = receiver.recv().await { + match message.map_err(|error| xlsx_governed_error(error, cancelled.load(Ordering::Acquire)))? { + dbx_core::table_import::XlsxStreamMessage::Header(columns) => { + if columns.len() > 1_000 || columns.iter().any(|column| column.chars().count() > MAX_CELL_CHAR_LIMIT) { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_HEADER_LIMIT_EXCEEDED", + "XLSX 流式表头超过列数或单列表头长度限制。", + )); + } + source_indexes = Some( + request + .mappings + .iter() + .map(|mapping| { + columns.iter().position(|column| column == &mapping.source_column).ok_or_else(|| { + EnterpriseToolError::new( + "IMPORT_XLSX_HEADER_CHANGED", + format!("XLSX 流式读取找不到 prepare 字段 {}。", mapping.source_column), + ) + }) + }) + .collect::, _>>()?, + ); + header_seen = true; + } + dbx_core::table_import::XlsxStreamMessage::Rows { rows, source_row_numbers } => { + if rows.len() != source_row_numbers.len() { + return Err(EnterpriseToolError::new( + "IMPORT_SOURCE_ROW_LINEAGE_UNAVAILABLE", + "XLSX 行批次缺少精确绝对源行号；未访问数据库。", + )); + } + let source_indexes = source_indexes.as_ref().ok_or_else(|| { + EnterpriseToolError::new("IMPORT_XLSX_HEADER_MISSING", "XLSX 数据行先于表头到达。") + })?; + for (row, source_row_number) in rows.into_iter().zip(source_row_numbers) { + if cancelled.load(Ordering::Acquire) { + return Err(EnterpriseToolError::new("IMPORT_CANCELLED", "导入已取消。")); + } + let mut output = source_indexes + .iter() + .map(|index| staging_text_value(row.get(*index).unwrap_or(&Value::Null))) + .collect::>(); + let row_bytes = serde_json::to_vec(&row).map_err(|error| { + EnterpriseToolError::new("IMPORT_SOURCE_ROW_HASH_FAILED", error.to_string()) + })?; + output.extend([ + request.import_id.clone(), + plan_id.to_string(), + source_sha.to_string(), + source_row_number.to_string(), + sha256_bytes(&row_bytes), + loaded_at.clone(), + ]); + writer + .write_record(&output) + .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + source_row_count += 1; + } + } + dbx_core::table_import::XlsxStreamMessage::Progress(_) => {} + dbx_core::table_import::XlsxStreamMessage::Done => done_seen = true, + } + } + match producer.await { + Ok(Ok(())) => {} + Ok(Err(error)) => return Err(xlsx_governed_error(error, cancelled.load(Ordering::Acquire))), + Err(error) => { + return Err(EnterpriseToolError::new("IMPORT_XLSX_STREAM_FAILED", error.to_string())); + } + } + producer_cancelled.store(true, Ordering::Release); + cancellation_monitor.abort(); + if !header_seen || !done_seen || source_row_count == 0 { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_STREAM_INCOMPLETE", + "XLSX 流式治理未完整结束；未访问数据库。", + )); + } + writer.flush().map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + + let mut mappings = request + .mappings + .iter() + .map(|mapping| TableImportColumnMapping { + source_column: mapping.target_column.clone(), + target_column: mapping.target_column.clone(), + target_data_type: Some("TEXT".to_string()), + }) + .collect::>(); + mappings.extend(RESERVED_STAGING_COLUMNS.iter().map(|column| TableImportColumnMapping { + source_column: column.to_string(), + target_column: column.to_string(), + target_data_type: Some("TEXT".to_string()), + })); + request.file_path = output_path.to_string_lossy().to_string(); + request.source_format = Some(TableImportSourceFormat::Csv); + request.parse_options = TableImportParseOptions { + has_header: Some(true), + trim_values: Some(false), + empty_string_as_null: Some(false), + ..Default::default() + }; + request.mappings = mappings; + request.mode = TableImportMode::Append; + request.create_table = true; + request.prepared_source = None; + Ok(request) +} + +fn xlsx_governed_error(error: String, cancelled: bool) -> EnterpriseToolError { + if cancelled || error == "Import cancelled" { + EnterpriseToolError::new("IMPORT_CANCELLED", "导入已取消。") + } else { + EnterpriseToolError::new("IMPORT_XLSX_STREAM_FAILED", error) + } +} + +struct CancelOnDrop(Arc); + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } +} + fn staging_text_value(value: &Value) -> String { match value { Value::Null => String::new(), @@ -1595,6 +1920,72 @@ fn comma_list_env(name: &str, defaults: &[&str]) -> Vec { #[cfg(test)] mod tests { use super::*; + use std::io::Write; + + fn write_xlsx_entry(zip: &mut zip::ZipWriter, path: &str, content: &str) { + zip.start_file( + path, + zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored), + ) + .unwrap(); + zip.write_all(content.as_bytes()).unwrap(); + } + + fn write_governed_test_xlsx(path: &Path) { + let file = std::fs::File::create(path).unwrap(); + let mut zip = zip::ZipWriter::new(file); + write_xlsx_entry( + &mut zip, + "[Content_Types].xml", + r#" + + + + + +"#, + ); + write_xlsx_entry( + &mut zip, + "_rels/.rels", + r#" + + +"#, + ); + write_xlsx_entry( + &mut zip, + "xl/workbook.xml", + r#" + + +"#, + ); + write_xlsx_entry( + &mut zip, + "xl/_rels/workbook.xml.rels", + r#" + + +"#, + ); + write_xlsx_entry( + &mut zip, + "xl/worksheets/sheet1.xml", + r#" + + + + 报告说明 + 商家备注商家备注金额 + AB10 + + CD20 + +"#, + ); + zip.finish().unwrap(); + } #[test] fn file_policy_rejects_outside_root_and_symlink() { @@ -1668,16 +2059,31 @@ mod tests { #[test] fn mappings_reject_missing_and_duplicate_targets() { - let source = vec!["订单号".to_string(), "金额".to_string()]; + let source = vec![ + McpSourceColumn { + source_position: 1, + raw_source_name: "订单号".to_string(), + canonical_source_name: "订单号".to_string(), + dbx_source_name: "订单号".to_string(), + }, + McpSourceColumn { + source_position: 2, + raw_source_name: "金额".to_string(), + canonical_source_name: "金额".to_string(), + dbx_source_name: "金额".to_string(), + }, + ]; let duplicated = vec![ McpImportColumnMapping { source_position: 1, - source_name: "订单号".to_string(), + raw_source_name: "订单号".to_string(), + canonical_source_name: "订单号".to_string(), target_column: "order_id".to_string(), }, McpImportColumnMapping { source_position: 2, - source_name: "金额".to_string(), + raw_source_name: "金额".to_string(), + canonical_source_name: "金额".to_string(), target_column: "order_id".to_string(), }, ]; @@ -1685,21 +2091,37 @@ mod tests { let reserved = vec![McpImportColumnMapping { source_position: 1, - source_name: "订单号".to_string(), + raw_source_name: "订单号".to_string(), + canonical_source_name: "订单号".to_string(), target_column: "source_row_hash".to_string(), }]; assert_eq!(validate_mappings(&reserved, &source).unwrap_err().code, "IMPORT_TARGET_COLUMN_RESERVED"); - let duplicate_headers = vec!["note".to_string(), "note_1".to_string()]; + let duplicate_headers = vec![ + McpSourceColumn { + source_position: 1, + raw_source_name: "note".to_string(), + canonical_source_name: "note__1".to_string(), + dbx_source_name: "note".to_string(), + }, + McpSourceColumn { + source_position: 2, + raw_source_name: "note".to_string(), + canonical_source_name: "note__2".to_string(), + dbx_source_name: "note_1".to_string(), + }, + ]; let second = vec![McpImportColumnMapping { source_position: 2, - source_name: "note_1".to_string(), + raw_source_name: "note".to_string(), + canonical_source_name: "note__2".to_string(), target_column: "second_note".to_string(), }]; assert_eq!(validate_mappings(&second, &duplicate_headers).unwrap()[0].source_column, "note_1"); let wrong_name = vec![McpImportColumnMapping { source_position: 2, - source_name: "note".to_string(), + raw_source_name: "note".to_string(), + canonical_source_name: "note__1".to_string(), target_column: "second_note".to_string(), }]; assert_eq!(validate_mappings(&wrong_name, &duplicate_headers).unwrap_err().code, "IMPORT_SOURCE_NAME_MISMATCH"); @@ -1859,7 +2281,15 @@ mod tests { retain_source: true, }; - let governed = build_governed_import_snapshot(request, "plan-1", &"a".repeat(64), &output).await.unwrap(); + let governed = build_governed_import_snapshot( + request, + "plan-1", + &"a".repeat(64), + &output, + Arc::new(AtomicBool::new(false)), + ) + .await + .unwrap(); assert!(governed.create_table); assert!(governed.mappings.iter().all(|mapping| mapping.target_data_type.as_deref() == Some("TEXT"))); let mut reader = csv::Reader::from_path(output).unwrap(); @@ -1874,9 +2304,9 @@ mod tests { } #[tokio::test] - async fn governed_snapshot_blocks_excel_before_database_write() { + async fn governed_snapshot_blocks_legacy_xls_before_database_write() { let directory = tempfile::tempdir().unwrap(); - let source = directory.path().join("input.xlsx"); + let source = directory.path().join("input.xls"); let output = directory.path().join("normalized.csv"); std::fs::write(&source, b"not-an-xlsx").unwrap(); let mut request = TableImportRequest { @@ -1903,9 +2333,215 @@ mod tests { target_data_type: Some("TEXT".to_string()), }); assert_eq!( - build_governed_import_snapshot(request, "plan-1", &"b".repeat(64), &output).await.unwrap_err().code, - "IMPORT_GOVERNED_FORMAT_UNSUPPORTED_V1" + build_governed_import_snapshot( + request, + "plan-1", + &"b".repeat(64), + &output, + Arc::new(AtomicBool::new(false)), + ) + .await + .unwrap_err() + .code, + "IMPORT_GOVERNED_XLS_UNSUPPORTED_V1" ); assert!(!output.exists()); + assert_eq!( + validate_governed_source_v1( + "input.xlsm", + Some(TableImportSourceFormat::Excel), + &TableImportParseOptions::default(), + ) + .unwrap(), + TableImportSourceFormat::Excel + ); + assert_eq!( + validate_governed_source_v1( + "input.json", + Some(TableImportSourceFormat::Json), + &TableImportParseOptions::default(), + ) + .unwrap_err() + .code, + "IMPORT_GOVERNED_FORMAT_UNSUPPORTED_V1" + ); + } + + #[tokio::test] + async fn governed_xlsx_stream_preserves_absolute_rows_duplicate_positions_and_cancellation() { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("input.xlsx"); + let output = directory.path().join("normalized.csv"); + write_governed_test_xlsx(&source); + let parse_options = + TableImportParseOptions { title_row: Some(2), data_start_row: Some(3), ..Default::default() }; + let preview = dbx_core::table_import::preview_table_import_file_with_request( + dbx_core::table_import::TableImportPreviewRequest { + file_path: source.to_string_lossy().to_string(), + source_ref: None, + source_format: Some(TableImportSourceFormat::Excel), + parse_options: parse_options.clone(), + preview_limit: Some(10), + }, + ) + .await + .unwrap(); + let source_columns = source_columns_for_preview( + &source.to_string_lossy(), + Some(TableImportSourceFormat::Excel), + &parse_options, + &preview.columns, + ) + .await + .unwrap(); + assert_eq!(source_columns[0].raw_source_name, "商家备注"); + assert_eq!(source_columns[0].canonical_source_name, "商家备注__1"); + assert_eq!(source_columns[1].raw_source_name, "商家备注"); + assert_eq!(source_columns[1].canonical_source_name, "商家备注__2"); + assert_eq!(preview.source_row_numbers, vec![3, 4, 5]); + let mappings = validate_mappings( + &[ + McpImportColumnMapping { + source_position: 2, + raw_source_name: "商家备注".to_string(), + canonical_source_name: "商家备注__2".to_string(), + target_column: "second_note".to_string(), + }, + McpImportColumnMapping { + source_position: 3, + raw_source_name: "金额".to_string(), + canonical_source_name: "金额".to_string(), + target_column: "amount".to_string(), + }, + ], + &source_columns, + ) + .unwrap(); + let request = TableImportRequest { + import_id: "xlsx-import".to_string(), + connection_id: "postgres-1".to_string(), + database: "enterprise".to_string(), + schema: "staging".to_string(), + table: "mcp_xlsx".to_string(), + file_path: source.to_string_lossy().to_string(), + source_ref: Some("c".repeat(64)), + source_format: Some(TableImportSourceFormat::Excel), + parse_options, + mappings, + mode: TableImportMode::Append, + create_table: true, + batch_size: 1, + date_time_format: None, + prepared_source: None, + retain_source: true, + }; + build_governed_import_snapshot( + request.clone(), + "plan-xlsx", + &"c".repeat(64), + &output, + Arc::new(AtomicBool::new(false)), + ) + .await + .unwrap(); + let mut reader = csv::Reader::from_path(&output).unwrap(); + let rows = reader.records().map(Result::unwrap).collect::>(); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].get(0), Some("B")); + assert_eq!(rows[1].get(0), Some("D")); + assert_eq!(rows[0].get(5), Some("3")); + assert_eq!(rows[1].get(5), Some("5")); + assert_eq!(rows[0].get(6).map(str::len), Some(64)); + + let cancelled_output = directory.path().join("cancelled.csv"); + assert_eq!( + build_governed_import_snapshot( + request, + "plan-cancelled", + &"c".repeat(64), + &cancelled_output, + Arc::new(AtomicBool::new(true)), + ) + .await + .unwrap_err() + .code, + "IMPORT_CANCELLED" + ); + assert!(!cancelled_output.exists()); + } + + #[tokio::test] + #[ignore = "需要设置 DBX_MCP_REAL_XLSX_FIXTURE，且只执行本地只读治理快照回归"] + async fn real_xlsx_fixture_prepare_and_governed_snapshot_without_database() { + let source = PathBuf::from(std::env::var("DBX_MCP_REAL_XLSX_FIXTURE").expect("fixture path")); + let directory = tempfile::tempdir().unwrap(); + let output = directory.path().join("normalized.csv"); + let identity = file_identity(source.clone()).await.unwrap(); + let parse_options = TableImportParseOptions::default(); + let preview = dbx_core::table_import::preview_table_import_file_with_request( + dbx_core::table_import::TableImportPreviewRequest { + file_path: source.to_string_lossy().to_string(), + source_ref: Some(identity.sha256.clone()), + source_format: Some(TableImportSourceFormat::Excel), + parse_options: parse_options.clone(), + preview_limit: Some(1), + }, + ) + .await + .unwrap(); + validate_preview_headers(&preview).unwrap(); + let source_columns = source_columns_for_preview( + &source.to_string_lossy(), + Some(TableImportSourceFormat::Excel), + &parse_options, + &preview.columns, + ) + .await + .unwrap(); + let requested = source_columns + .iter() + .take(2) + .enumerate() + .map(|(index, column)| McpImportColumnMapping { + source_position: column.source_position, + raw_source_name: column.raw_source_name.clone(), + canonical_source_name: column.canonical_source_name.clone(), + target_column: format!("source_column_{}", index + 1), + }) + .collect::>(); + let mappings = validate_mappings(&requested, &source_columns).unwrap(); + let plan = build_plan( + "fixture-postgres".to_string(), + "fixture".to_string(), + "enterprise".to_string(), + "staging".to_string(), + format!("mcp_{}", Uuid::new_v4().simple()), + "fixture-v1".to_string(), + identity.clone(), + structure_fingerprint(&preview, &parse_options, &source_columns), + Some(TableImportSourceFormat::Excel), + parse_options, + mappings, + true, + 500, + None, + ) + .unwrap(); + let request = plan.to_import_request("fixture-import".to_string()); + build_governed_import_snapshot( + request, + &plan.plan_id, + &identity.sha256, + &output, + Arc::new(AtomicBool::new(false)), + ) + .await + .unwrap(); + let mut reader = csv::Reader::from_path(output).unwrap(); + let headers = reader.headers().unwrap().clone(); + assert!(headers.iter().any(|header| header == "source_row_number")); + assert!(headers.iter().any(|header| header == "source_row_hash")); + let row_count = reader.records().filter_map(Result::ok).count(); + assert!(row_count > 0); } } diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 312cb62e29..f35305054d 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -13,12 +13,12 @@ use crate::backend::{format_query_result, new_connection_config, parse_database_ use crate::enterprise_tools::{ build_milvus_filter, build_plan, cell_char_limit, file_identity, generated_staging_relation, milvus_delete_batch_query, milvus_existing_cards_query, milvus_search_query, milvus_upsert_query, preview_limit, - query_result_rows, read_semantic_jsonl, revalidate_plan_file, sanitize_preview, structure_fingerprint, - validate_embedding, validate_existing_card_ownership, validate_governed_source_v1, validate_import_file, - validate_mappings, validate_preview_headers, validate_vector_collection, vector_output_fields, vector_top_k, - vector_upsert_batch_size, EnterpriseRuntime, EnterpriseToolError, ImportStatusRequest, PrepareTableImportRequest, - PreviewImportFileRequest, StartTableImportRequest, VectorDeleteByBatchRequest, VectorSearchRequest, - VectorUpsertFileRequest, FORMAT_VERSION, + query_result_rows, read_semantic_jsonl, revalidate_plan_file, sanitize_preview, source_columns_for_preview, + structure_fingerprint, validate_embedding, validate_existing_card_ownership, validate_governed_source_v1, + validate_import_file, validate_mappings, validate_preview_headers, validate_vector_collection, + vector_output_fields, vector_top_k, vector_upsert_batch_size, EnterpriseRuntime, EnterpriseToolError, + ImportStatusRequest, PrepareTableImportRequest, PreviewImportFileRequest, StartTableImportRequest, + VectorDeleteByBatchRequest, VectorSearchRequest, VectorUpsertFileRequest, FORMAT_VERSION, }; use crate::mongo::{self, MongoCommand, MongoSafetyError}; use crate::session::{McpSession, McpSessionStore}; @@ -404,11 +404,13 @@ impl DbxMcpServer { Ok(limit) => limit, Err(error) => return enterprise_error(error), }; + let source_format = request.source_format.map(Into::into); + let file_path = path.to_string_lossy().to_string(); let preview = dbx_core::table_import::preview_table_import_file_with_request( dbx_core::table_import::TableImportPreviewRequest { - file_path: path.to_string_lossy().to_string(), + file_path: file_path.clone(), source_ref: Some(identity.sha256.clone()), - source_format: request.source_format.map(Into::into), + source_format, parse_options: parse_options.clone(), preview_limit: Some(preview_rows), }, @@ -423,7 +425,12 @@ impl DbxMcpServer { if let Err(error) = validate_preview_headers(&preview) { return enterprise_error(error); } - let fingerprint = structure_fingerprint(&preview, &parse_options); + let source_columns = + match source_columns_for_preview(&file_path, source_format, &parse_options, &preview.columns).await { + Ok(columns) => columns, + Err(error) => return enterprise_error(error), + }; + let fingerprint = structure_fingerprint(&preview, &parse_options, &source_columns); let current_identity = match file_identity(path).await { Ok(identity) => identity, Err(error) => return enterprise_error(error), @@ -434,12 +441,6 @@ impl DbxMcpServer { "文件在预览期间发生变化；结果已丢弃。", )); } - let source_columns = preview - .columns - .iter() - .enumerate() - .map(|(index, name)| json!({ "position": index + 1, "name": name })) - .collect::>(); let used_first_row = preview.source_row_numbers.first().copied(); let used_last_row = preview.source_row_numbers.last().copied(); let preview = sanitize_preview(preview, char_limit); @@ -525,6 +526,13 @@ impl DbxMcpServer { if let Err(error) = validate_preview_headers(&preview) { return enterprise_error(error); } + let source_columns = + match source_columns_for_preview(&path.to_string_lossy(), source_format, &parse_options, &preview.columns) + .await + { + Ok(columns) => columns, + Err(error) => return enterprise_error(error), + }; let current_identity = match file_identity(path).await { Ok(identity) => identity, Err(error) => return enterprise_error(error), @@ -535,7 +543,7 @@ impl DbxMcpServer { "文件在 prepare 剖析期间发生变化；未创建计划。", )); } - let mappings = match validate_mappings(&request.mappings, &preview.columns) { + let mappings = match validate_mappings(&request.mappings, &source_columns) { Ok(mappings) => mappings, Err(error) => return enterprise_error(error), }; @@ -550,7 +558,7 @@ impl DbxMcpServer { "batch_size 必须在 1 到 50000 之间。", )); } - let fingerprint = structure_fingerprint(&preview, &parse_options); + let fingerprint = structure_fingerprint(&preview, &parse_options, &source_columns); let plan = match build_plan( resolved.connection.id.clone(), resolved.connection.name.clone(), @@ -2319,12 +2327,14 @@ mod tests { mappings: vec![ crate::enterprise_tools::McpImportColumnMapping { source_position: 1, - source_name: "order_id".to_string(), + raw_source_name: "order_id".to_string(), + canonical_source_name: "order_id".to_string(), target_column: "order_id".to_string(), }, crate::enterprise_tools::McpImportColumnMapping { source_position: 2, - source_name: "amount".to_string(), + raw_source_name: "amount".to_string(), + canonical_source_name: "amount".to_string(), target_column: "amount".to_string(), }, ], diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index 33aafa1841..4bf6684525 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -117,9 +117,9 @@ DBX MCP 当前提供 21 个工具： ## 本地文件导入与 Milvus 工具 -导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个真实目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。preview 支持 Excel、JSON、CSV 和 TSV；v1 治理入库有意只开放可流式有界处理的 UTF-8 CSV/TSV。Excel/JSON 或非 UTF-8 文件会在创建计划和访问数据库之前返回稳定的不支持错误。 +导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个真实目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。preview 支持 Excel、JSON、CSV 和 TSV；v1 治理入库使用有界流式处理支持 XLSX/XLSM 和 UTF-8 CSV/TSV。旧版 `.xls`、JSON 和非 UTF-8 分隔文件会在创建计划和访问数据库之前返回稳定的不支持错误。 -`dbx_prepare_table_import` 会把连接、调用方已批准模板版本、规范路径、大小、修改时间、SHA-256、解析参数和字段映射固化为不可变摘要。每个映射必须提供 preview 返回的 1 基 `source_position`、精确 `source_name` 和 `target_column`。目标不能由调用方指定：服务端生成唯一 `staging.mcp_`，强制创建新表，所有业务列固定为 `TEXT`。计划 30 分钟后过期，且只能启动一次。start 复验源文件后，先流式写出有界治理快照，再打开数据库连接；每行增加 `import_id`、`plan_id`、`source_sha`、真实 `source_row_number`、`source_row_hash` 和 `loaded_at`。无法形成精确血缘时会在数据库访问前阻断。status 会返回生成的 relation；发布到业务 Schema 仍是独立治理流程。 +`dbx_prepare_table_import` 会把连接、调用方已批准模板版本、规范路径、大小、修改时间、SHA-256、解析参数和字段映射固化为不可变摘要。preview 对每列返回 1 基 `source_position`、原始 `raw_source_name`、稳定 `canonical_source_name`（重复表头统一为 `名称__1`、`名称__2`）以及兼容旧 DBX 的名称；映射必须复核位置、raw、canonical，并提供 `target_column`。目标不能由调用方指定：服务端生成唯一 `staging.mcp_`，强制创建新表，所有业务列固定为 `TEXT`。计划 30 分钟后过期，且只能启动一次。start 复验源文件后，先流式写出有界治理快照，再打开数据库连接；每行增加 `import_id`、`plan_id`、`source_sha`、真实绝对 `source_row_number`、`source_row_hash` 和 `loaded_at`。无法形成精确血缘时会在数据库访问前阻断。status 会返回生成的 relation；发布到业务 Schema 仍是独立治理流程。 Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索必须提供 `active_at`，服务端强制 approved 和生效时间过滤，并可选追加精确 `semantic_version`。upsert 使用严格语义卡 Schema：重算并核对 64 位 `content_checksum`，要求 `embedding_revision`，拒绝 `chunk_index` 与 `embedding_dimension`，文件内 card_id 必须唯一且只能有一个语义版本；写入前还会查询现有 card_id，禁止覆盖其他批次或版本。删除除“完全访问”外还要求调用方显式传入 `published: false`。所有通过 `dbx_execute_query` 提交的向量实体 search/query/mutation 都返回 `VECTOR_DEDICATED_TOOL_REQUIRED`，通用入口只保留明确管理诊断。 diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index f6a649195d..80683ff290 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -117,9 +117,9 @@ Connection-scoped sessions hide connection-mutating and desktop UI tools. ## Local file imports and Milvus tools -The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more real directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. Preview supports Excel, JSON, CSV, and TSV; governed database import in v1 is deliberately limited to streaming UTF-8 CSV/TSV. Excel/JSON or non-UTF-8 prepare calls fail with a stable unsupported code before a plan or database write exists. +The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more real directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. Preview supports Excel, JSON, CSV, and TSV. Governed v1 database import uses bounded streaming for XLSX/XLSM and UTF-8 CSV/TSV. Legacy `.xls`, JSON, and non-UTF-8 delimited files return stable unsupported errors before a plan or database write exists. -`dbx_prepare_table_import` binds the connection, caller-approved template version, canonical path, size, modification time, SHA-256, parsing options, and mappings into an immutable digest. Each mapping supplies the 1-based `source_position`, the exact `source_name` returned by preview, and `target_column`. The server—not the caller—generates a unique `staging.mcp_` relation, forces a new table, and makes every business column `TEXT`. The returned plan expires after 30 minutes and can be started once. Start rechecks the source and streams it into a bounded governed snapshot before opening a database pool. Every row gains `import_id`, `plan_id`, `source_sha`, the real `source_row_number`, `source_row_hash`, and `loaded_at`; if exact lineage cannot be produced, the import fails before database access. Status returns the generated relation. Publication to business schemas remains a separately governed workflow. +`dbx_prepare_table_import` binds the connection, caller-approved template version, canonical path, size, modification time, SHA-256, parsing options, and mappings into an immutable digest. Preview exposes each source column as 1-based `source_position`, exact `raw_source_name`, stable `canonical_source_name` (`name__1`, `name__2` for duplicates), and the legacy DBX name. Each mapping must repeat the position plus raw/canonical names and supply `target_column`. The server—not the caller—generates a unique `staging.mcp_` relation, forces a new table, and makes every business column `TEXT`. The returned plan expires after 30 minutes and can be started once. Start rechecks the source and streams it into a bounded governed snapshot before opening a database pool. Every row gains `import_id`, `plan_id`, `source_sha`, the real absolute `source_row_number`, `source_row_hash`, and `loaded_at`; if exact lineage cannot be produced, the import fails before database access. Status returns the generated relation. Publication to business schemas remains a separately governed workflow. Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires `active_at`, always adds `approval_status == "approved"` plus effective-date filters, and optionally applies an exact `semantic_version`. Upsert accepts a strict semantic-card schema, verifies the 64-hex `content_checksum` against `content`, requires `embedding_revision`, rejects `chunk_index` and `embedding_dimension`, enforces unique card IDs and one semantic version per file, and preflights existing IDs so another batch/version cannot be overwritten. Delete requires Full access and an explicit `published: false` attestation. All entity search/query/mutation REST submitted through generic `dbx_execute_query` returns `VECTOR_DEDICATED_TOOL_REQUIRED`; only explicit management diagnostics remain available there. diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index 2fcbbf03d6..56742a7fb3 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -183,7 +183,7 @@ DBX connection storage defaults to: Override the directory with `DBX_DATA_DIR`. -Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Preview accepts Excel/JSON/CSV/TSV, while v1 governed database import is fail-closed to streaming UTF-8 CSV/TSV. Prepare accepts position+name mappings and generates a unique `staging.mcp_` relation; callers cannot append to an existing table. Start revalidates the source and streams a private snapshot containing row numbers, hashes, and batch lineage before opening a database pool. Import tools support local mode only and return `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` in Web mode. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. +Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Preview accepts Excel/JSON/CSV/TSV; v1 governed import streams XLSX/XLSM and UTF-8 CSV/TSV, while legacy `.xls` and JSON remain preview-only. Prepare accepts position+raw+canonical-name mappings and generates a unique `staging.mcp_` relation; callers cannot append to an existing table. Start revalidates the source and streams a private snapshot containing absolute row numbers, hashes, and batch lineage before opening a database pool. Import tools support local mode only and return `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` in Web mode. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. ### Agent/JDBC databases @@ -508,7 +508,7 @@ MCP 配置： 通过 `DBX_DATA_DIR` 覆盖默认目录。Windows 便携版应指向 `DBX.exe` 同级、包含 `dbx.db` 的 `data` 文件夹。 -本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。preview 支持 Excel/JSON/CSV/TSV，但 v1 治理入库只允许可流式有界转换的 UTF-8 CSV/TSV。prepare 接受“源位置+源名称”映射并由服务端生成唯一 `staging.mcp_`，调用方不能追加已有表。start 在打开数据库连接前复验文件并流式生成包含真实行号、行哈希和批次血缘的任务快照。导入工具仅支持本地模式，Web 模式稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 +本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。preview 支持 Excel/JSON/CSV/TSV；v1 治理入库流式支持 XLSX/XLSM 和 UTF-8 CSV/TSV，旧版 `.xls` 与 JSON 暂为 preview-only。prepare 接受“源位置+raw+canonical 名称”映射并由服务端生成唯一 `staging.mcp_`，调用方不能追加已有表。start 在打开数据库连接前复验文件并流式生成包含真实绝对行号、行哈希和批次血缘的任务快照。导入工具仅支持本地模式，Web 模式稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 ### DBX Web / Docker From f6cead0e3377a3d90e92386249637dfeab5bba6a Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 19:46:26 +0800 Subject: [PATCH 05/10] fix(core): flatten single Milvus search results --- crates/dbx-core/src/db/vector_driver.rs | 81 ++++++++++++++++++++++--- 1 file changed, 72 insertions(+), 9 deletions(-) diff --git a/crates/dbx-core/src/db/vector_driver.rs b/crates/dbx-core/src/db/vector_driver.rs index 04d9df8958..095f02aaec 100644 --- a/crates/dbx-core/src/db/vector_driver.rs +++ b/crates/dbx-core/src/db/vector_driver.rs @@ -749,14 +749,21 @@ pub async fn find_documents( pub async fn execute_rest_query(client: &VectorClient, input: &str) -> Result { let start = Instant::now(); + let flatten_single_milvus_search = client.kind == VectorDbKind::Milvus && is_milvus_entity_search(input); let request = parse_rest_query(client, input)?; let resp = request.send().await.map_err(|e| format!("{} request failed: {e}", client.kind.label()))?; let status = resp.status().as_u16(); let body = resp.json::().await.unwrap_or(Value::Null); - rest_query_result(client.kind, status, body, start) + rest_query_result(client.kind, status, body, start, flatten_single_milvus_search) } -fn rest_query_result(kind: VectorDbKind, status: u16, body: Value, start: Instant) -> Result { +fn rest_query_result( + kind: VectorDbKind, + status: u16, + body: Value, + start: Instant, + flatten_single_milvus_search: bool, +) -> Result { if !(200..300).contains(&status) { let detail = serde_json::to_string_pretty(&body).unwrap_or_else(|_| body.to_string()); return Err(format!("{} error ({status}): {detail}", kind.label())); @@ -765,10 +772,32 @@ fn rest_query_result(kind: VectorDbKind, status: u16, body: Value, start: Instan if let Some(error) = milvus_business_error(&body) { return Err(error); } + if flatten_single_milvus_search { + if let Some(result) = milvus_single_search_to_query_result(&body, start) { + return Ok(result); + } + } } Ok(json_to_query_result(status, body, start)) } +fn is_milvus_entity_search(input: &str) -> bool { + input + .lines() + .find(|line| !line.trim().is_empty()) + .and_then(|line| line.split_whitespace().nth(1)) + .is_some_and(|path| path.split('?').next() == Some("/v2/vectordb/entities/search")) +} + +fn milvus_single_search_to_query_result(body: &Value, start: Instant) -> Option { + let queries = body.get("data")?.as_array()?; + if queries.len() != 1 { + return None; + } + let hits = queries.first()?.as_array()?.clone(); + Some(values_to_query_result(hits, start)) +} + // Milvus REST uses HTTP-style code 200 for success, while some responses use gRPC-style code 0. fn milvus_business_error(body: &Value) -> Option { let code = body.get("code").and_then(Value::as_i64)?; @@ -1009,9 +1038,9 @@ fn format_reqwest_error(err: &reqwest::Error) -> String { #[cfg(test)] mod tests { use super::{ - chroma_get_response_to_rows, default_collection_query, milvus_collection_schema, milvus_database_names, - rename_collection, rest_query_result, starts_with_http_method, test_connection, test_connection_request, - values_to_query_result, vector_auth, weaviate_collection_names_from_schema, + chroma_get_response_to_rows, default_collection_query, is_milvus_entity_search, milvus_collection_schema, + milvus_database_names, rename_collection, rest_query_result, starts_with_http_method, test_connection, + test_connection_request, values_to_query_result, vector_auth, weaviate_collection_names_from_schema, weaviate_vector_dimension_from_graphql, CollectionInfo, VectorAuth, VectorClient, VectorDbKind, }; use serde_json::{json, Value}; @@ -1097,20 +1126,54 @@ mod tests { VectorDbKind::Milvus, 200, json!({ "code": 1100, "message": "field kind does not exist" }), - Instant::now() + Instant::now(), + false ) .unwrap_err(), "Milvus error (code 1100): field kind does not exist" ); - assert!(rest_query_result(VectorDbKind::Milvus, 200, json!({ "code": 0 }), Instant::now()).is_ok()); + assert!(rest_query_result(VectorDbKind::Milvus, 200, json!({ "code": 0 }), Instant::now(), false).is_ok()); assert!(rest_query_result( VectorDbKind::Milvus, 200, json!({ "code": 200, "data": ["kb_vectors"] }), - Instant::now() + Instant::now(), + false ) .is_ok()); - assert!(rest_query_result(VectorDbKind::Milvus, 200, json!({ "data": [] }), Instant::now()).is_ok()); + assert!(rest_query_result(VectorDbKind::Milvus, 200, json!({ "data": [] }), Instant::now(), false).is_ok()); + } + + #[test] + fn milvus_single_search_flattens_empty_and_multiple_hits_without_affecting_query() { + assert!(is_milvus_entity_search("POST /v2/vectordb/entities/search\n{}")); + assert!(!is_milvus_entity_search("POST /v2/vectordb/entities/query\n{}")); + let empty = + rest_query_result(VectorDbKind::Milvus, 200, json!({ "code": 0, "data": [[]] }), Instant::now(), true) + .unwrap(); + assert!(empty.rows.is_empty()); + + let hits = rest_query_result( + VectorDbKind::Milvus, + 200, + json!({ "code": 0, "data": [[{"card_id": "a"}, {"card_id": "b"}]] }), + Instant::now(), + true, + ) + .unwrap(); + assert_eq!(hits.rows.len(), 2); + assert_eq!(hits.columns, vec!["card_id"]); + + let query = rest_query_result( + VectorDbKind::Milvus, + 200, + json!({ "code": 0, "data": [[{"card_id": "a"}]] }), + Instant::now(), + false, + ) + .unwrap(); + assert_eq!(query.rows.len(), 1); + assert_eq!(query.columns, vec!["value"]); } #[tokio::test] From 3fe42d28da49f405d5f15c80cbbabf84a7c6bfef Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 19:46:46 +0800 Subject: [PATCH 06/10] fix(mcp): close enterprise import release blockers --- Cargo.lock | 2 + crates/dbx-core/src/table_import.rs | 204 +++++- crates/dbx-mcp/Cargo.toml | 6 +- crates/dbx-mcp/src/backend.rs | 77 ++- crates/dbx-mcp/src/enterprise_tools.rs | 895 +++++++++++++++++++++++-- crates/dbx-mcp/src/server.rs | 531 ++++++++++----- crates/dbx-mcp/tests/local.rs | 3 +- crates/dbx-mcp/tests/protocol.rs | 13 +- docs/content/docs/mcp.cn.mdx | 28 +- docs/content/docs/mcp.mdx | 28 +- packages/mcp-server/README.md | 22 +- 11 files changed, 1531 insertions(+), 278 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0726c380ae..5c6c7d1fc1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2158,6 +2158,7 @@ dependencies = [ "csv", "dbx-core", "dirs", + "libc", "reqwest 0.12.28", "rmcp", "rustls 0.23.43", @@ -2166,6 +2167,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", + "sysinfo", "tempfile", "time", "tokio", diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index 275ec3c4fb..ac2fbb069a 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -39,6 +39,7 @@ pub const MAX_LEGACY_XLS_IMPORT_BYTES: u64 = 50 * 1024 * 1024; const IMPORT_PROGRESS_INTERVAL: Duration = Duration::from_millis(100); // Keep preview parsing bounded even when an XLSX dimension declares a huge sparse range. const MAX_FAST_PREVIEW_CELLS: usize = 100_000; +const MAX_XLSX_PREVIEW_CELL_BYTES: usize = 1024 * 1024; // Shared strings stay in memory for small workbooks and spill to an indexed temp file for large ones. const MAX_IN_MEMORY_XLSX_SHARED_STRINGS_BYTES: u64 = 8 * 1024 * 1024; const MAX_XLSX_SHARED_STRINGS_BYTES: u64 = 1024 * 1024 * 1024; @@ -2033,10 +2034,10 @@ fn xlsx_cell_ref_position(reference: &str) -> Option<(usize, usize)> { } if ch.is_ascii_alphabetic() && !saw_row { saw_column = true; - column = column * 26 + (ch.to_ascii_uppercase() as u8 - b'A' + 1) as usize; + column = column.checked_mul(26)?.checked_add((ch.to_ascii_uppercase() as u8 - b'A' + 1) as usize)?; } else if ch.is_ascii_digit() { saw_row = true; - row = row * 10 + ch.to_digit(10)? as usize; + row = row.checked_mul(10)?.checked_add(ch.to_digit(10)? as usize)?; } else { return None; } @@ -2141,6 +2142,11 @@ fn read_xlsx_shared_strings( } Ok(Event::Text(text)) if in_item && in_text => { current.push_str(&text.unescape().map_err(|error| error.to_string())?); + if current.len() > MAX_XLSX_PREVIEW_CELL_BYTES { + return Err(format!( + "Excel shared string exceeds the {MAX_XLSX_PREVIEW_CELL_BYTES} byte preview cell limit" + )); + } } Ok(Event::End(element)) if xml_local_name_eq(element.name().as_ref(), b"t") => { in_text = false; @@ -2257,7 +2263,14 @@ fn create_xlsx_spill_file() -> std::io::Result { #[cfg(test)] fn open_xlsx_shared_strings(zip: &mut zip::ZipArchive, memory_limit: u64) -> Result { - open_xlsx_shared_strings_with_control(zip, memory_limit, &|| false, &mut |_| Ok(())) + open_xlsx_shared_strings_with_control( + zip, + memory_limit, + MAX_XLSX_SHARED_STRINGS_BYTES, + usize::MAX, + &|| false, + &mut |_| Ok(()), + ) } struct XlsxCancellableReader<'a, R> { @@ -2283,6 +2296,8 @@ impl IoRead for XlsxCancellableReader<'_, R> { fn open_xlsx_shared_strings_with_control( zip: &mut zip::ZipArchive, memory_limit: u64, + max_uncompressed_bytes: u64, + max_cell_bytes: usize, is_cancelled: &dyn Fn() -> bool, on_progress: &mut dyn FnMut(u64) -> std::io::Result<()>, ) -> Result { @@ -2291,9 +2306,9 @@ fn open_xlsx_shared_strings_with_control( Err(zip::result::ZipError::FileNotFound) => return Ok(XlsxSharedStrings::Memory(Vec::new())), Err(error) => return Err(error.to_string()), }; - if uncompressed_size > MAX_XLSX_SHARED_STRINGS_BYTES { + if uncompressed_size > max_uncompressed_bytes { return Err(format!( - "Excel shared strings are too large: {uncompressed_size} bytes (max {MAX_XLSX_SHARED_STRINGS_BYTES} bytes)" + "Excel shared strings are too large: {uncompressed_size} bytes (max {max_uncompressed_bytes} bytes)" )); } // A fixed-width offset/length index lets cell parsing seek individual strings without @@ -2334,6 +2349,9 @@ fn open_xlsx_shared_strings_with_control( } Ok(Event::Text(text)) if in_item && in_text => { current.push_str(&text.unescape().map_err(|error| error.to_string())?); + if current.len() > max_cell_bytes { + return Err(format!("Excel shared string exceeds the {max_cell_bytes} byte cell limit")); + } } Ok(Event::End(element)) if xml_local_name_eq(element.name().as_ref(), b"t") => { in_text = false; @@ -2531,9 +2549,15 @@ fn parse_xlsx_preview_file_with_options( } Ok(Event::Text(text)) if in_value => { current_cell.value.push_str(&text.unescape().map_err(|error| error.to_string())?); + if current_cell.value.len() > MAX_XLSX_PREVIEW_CELL_BYTES { + return Err(format!("Excel cell exceeds the {MAX_XLSX_PREVIEW_CELL_BYTES} byte preview limit")); + } } Ok(Event::Text(text)) if in_inline_text => { current_cell.inline_value.push_str(&text.unescape().map_err(|error| error.to_string())?); + if current_cell.inline_value.len() > MAX_XLSX_PREVIEW_CELL_BYTES { + return Err(format!("Excel cell exceeds the {MAX_XLSX_PREVIEW_CELL_BYTES} byte preview limit")); + } } Ok(Event::End(element)) if xml_local_name_eq(element.name().as_ref(), b"v") => { in_value = false; @@ -2578,8 +2602,14 @@ fn parse_xlsx_preview_file_with_options( return Err("Import file has no data rows in the selected row range".to_string()); } let start_row = observed_min_row; - let start_column = observed_min_column; - let observed_end_column = observed_max_column.max(start_column); + let title_column_bounds = row_range.title_row.and_then(|title_row| { + let absolute_title_row = start_row.saturating_add(title_row.saturating_sub(1)); + let mut columns = raw_cells.keys().filter_map(|(row, column)| (*row == absolute_title_row).then_some(*column)); + let first = columns.next()?; + Some(columns.fold((first, first), |(min, max), column| (min.min(column), max.max(column)))) + }); + let start_column = title_column_bounds.map_or(observed_min_column, |(start, _)| start); + let observed_end_column = title_column_bounds.map_or(observed_max_column.max(start_column), |(_, end)| end); let observed_column_count = observed_end_column.saturating_sub(start_column).saturating_add(1); let preview_row_count = requested_last_row .saturating_sub(row_range.data_start_row) @@ -2805,6 +2835,36 @@ pub enum XlsxStreamMessage { Done, } +/// XLSX/XLSM 流式读取的资源上限；企业治理入口使用严格值，普通导入保持兼容默认值。 +#[derive(Debug, Clone, Copy)] +pub struct XlsxStreamLimits { + /// sharedStrings.xml 最大解压字节数。 + pub max_shared_strings_bytes: u64, + /// 选中工作表 XML 最大解压字节数。 + pub max_worksheet_bytes: u64, + /// 工作表最多允许的 row 元素数量。 + pub max_worksheet_rows: usize, + /// 工作表最多允许的 cell 元素数量。 + pub max_worksheet_cells: usize, + /// 单个源单元格解码后的最大 UTF-8 字节数。 + pub max_cell_bytes: usize, + /// 单个待发送行批次的近似 JSON 字节上限。 + pub max_batch_bytes: usize, +} + +impl Default for XlsxStreamLimits { + fn default() -> Self { + Self { + max_shared_strings_bytes: MAX_XLSX_SHARED_STRINGS_BYTES, + max_worksheet_bytes: u64::MAX, + max_worksheet_rows: usize::MAX, + max_worksheet_cells: usize::MAX, + max_cell_bytes: usize::MAX, + max_batch_bytes: usize::MAX, + } + } +} + fn xlsx_stream_cell_value( cell: &XlsxPreviewRawCell, shared_strings: &mut XlsxSharedStrings, @@ -2866,6 +2926,8 @@ struct XlsxStreamRowsState { current_row: Option, current_values: Vec, batch_size: usize, + pending_bytes: usize, + max_batch_bytes: usize, } impl XlsxStreamRowsState { @@ -2875,6 +2937,7 @@ impl XlsxStreamRowsState { dimension: Option<((usize, usize), (usize, usize))>, expected_columns: Option>, batch_size: usize, + max_batch_bytes: usize, ) -> Self { let batch_size = batch_size.max(1); Self { @@ -2892,11 +2955,16 @@ impl XlsxStreamRowsState { current_row: None, current_values: Vec::new(), batch_size, + pending_bytes: 0, + max_batch_bytes: max_batch_bytes.max(1), } } fn initialize_range(&mut self, first_row: usize, first_column: usize) { - if self.start_row.is_some() { + let start_row = *self.start_row.get_or_insert(first_row); + let relative_row = first_row.saturating_sub(start_row).saturating_add(1); + let selected_first_row = self.row_range.title_row.unwrap_or(self.row_range.data_start_row); + if relative_row < selected_first_row || self.start_column > 0 { return; } let expected_column_count = (!self.columns.is_empty()).then_some(self.columns.len()); @@ -2911,7 +2979,6 @@ impl XlsxStreamRowsState { expected == column_count }) }); - self.start_row = Some(first_row); self.start_column = first_column; self.declared_column_count = dimension .map(|((_, start_column), (_, end_column))| end_column.saturating_sub(start_column).saturating_add(1)); @@ -2934,6 +3001,9 @@ impl XlsxStreamRowsState { return true; } self.initialize_range(absolute_row, absolute_column); + if self.start_column == 0 { + return false; + } absolute_column .checked_sub(self.start_column) .and_then(|offset| self.columns.get(offset)) @@ -2948,6 +3018,9 @@ impl XlsxStreamRowsState { progress: u64, ) -> Result<(), String> { self.initialize_range(absolute_row, absolute_column); + if self.start_column == 0 { + return Ok(()); + } if self.current_row != Some(absolute_row) { self.flush_current_row(progress)?; self.current_row = Some(absolute_row); @@ -3016,8 +3089,19 @@ impl XlsxStreamRowsState { } values.resize(self.columns.len(), serde_json::Value::Null); values.truncate(self.columns.len()); + let row_bytes = serde_json::to_vec(&values).map_err(|error| error.to_string())?.len(); + if row_bytes > self.max_batch_bytes { + return Err(format!( + "Excel row {absolute_row} is {row_bytes} bytes after normalization, exceeding the {} byte batch budget", + self.max_batch_bytes + )); + } + if !self.pending_rows.is_empty() && self.pending_bytes.saturating_add(row_bytes) > self.max_batch_bytes { + self.emit_rows(progress)?; + } self.pending_rows.push(values); self.pending_source_row_numbers.push(absolute_row); + self.pending_bytes = self.pending_bytes.saturating_add(row_bytes); self.rows_seen = self.rows_seen.saturating_add(1); if self.pending_rows.len() >= self.batch_size { self.emit_rows(progress)?; @@ -3040,6 +3124,7 @@ impl XlsxStreamRowsState { .map_err(|_| "Excel import consumer closed before the stream finished".to_string())?; self.pending_rows = Vec::with_capacity(self.batch_size); self.pending_source_row_numbers = Vec::with_capacity(self.batch_size); + self.pending_bytes = 0; Ok(()) } @@ -3089,6 +3174,32 @@ pub fn stream_xlsx_rows_to_channel_with_control( scan_full_worksheet: bool, sender: tokio::sync::mpsc::Sender>, cancelled: Arc, +) -> Result<(), String> { + stream_xlsx_rows_to_channel_with_limits( + path, + options, + batch_size, + expected_columns, + text_source_columns, + scan_full_worksheet, + sender, + cancelled, + XlsxStreamLimits::default(), + ) +} + +#[allow(clippy::too_many_arguments)] +/// 以调用方提供的资源预算流式读取 XLSX/XLSM；任何预算不足都会在消费者写入前失败。 +pub fn stream_xlsx_rows_to_channel_with_limits( + path: &str, + options: &TableImportParseOptions, + batch_size: usize, + expected_columns: Option>, + text_source_columns: HashSet, + scan_full_worksheet: bool, + sender: tokio::sync::mpsc::Sender>, + cancelled: Arc, + limits: XlsxStreamLimits, ) -> Result<(), String> { // This producer runs on a blocking thread and communicates in bounded batches. The small // channel capacity applies backpressure when database writes are slower than XML parsing. @@ -3140,16 +3251,26 @@ pub fn stream_xlsx_rows_to_channel_with_control( let is_cancelled = || cancelled.load(Ordering::Acquire); let mut shared_strings = open_xlsx_shared_strings_with_control( &mut zip, - MAX_IN_MEMORY_XLSX_SHARED_STRINGS_BYTES, + MAX_IN_MEMORY_XLSX_SHARED_STRINGS_BYTES.min(limits.max_shared_strings_bytes), + limits.max_shared_strings_bytes, + limits.max_cell_bytes, &is_cancelled, &mut on_shared_progress, )?; let row_range = effective_import_row_range(options)?; let sheet = zip.by_name(&sheet_path).map_err(|error| error.to_string())?; + if sheet.size() > limits.max_worksheet_bytes { + return Err(format!( + "Excel worksheet is too large: {} bytes (max {} bytes)", + sheet.size(), + limits.max_worksheet_bytes + )); + } let uncompressed_sheet_bytes = sheet.size().max(1); let mut reader = XmlReader::from_reader(BufReader::new(sheet)); reader.config_mut().trim_text(false); - let mut rows = XlsxStreamRowsState::new(sender, row_range, None, expected_columns, batch_size); + let mut rows = + XlsxStreamRowsState::new(sender, row_range, None, expected_columns, batch_size, limits.max_batch_bytes); let mut buffer = Vec::new(); let mut current_row = 0usize; let mut current_column = 0usize; @@ -3158,6 +3279,8 @@ pub fn stream_xlsx_rows_to_channel_with_control( let mut in_value = false; let mut in_inline_text = false; let mut inline_phonetic_depth = 0usize; + let mut worksheet_rows = 0usize; + let mut worksheet_cells = 0usize; loop { // Convert the uncompressed worksheet offset into an approximate archive-byte offset so // progress remains monotonic without scanning the ZIP twice. @@ -3179,6 +3302,13 @@ pub fn stream_xlsx_rows_to_channel_with_control( rows.dimension = xml_attr_value(&reader, &element, b"ref").as_deref().and_then(xlsx_dimension_bounds); } Ok(Event::Start(element)) if xml_local_name_eq(element.name().as_ref(), b"row") => { + worksheet_rows = worksheet_rows.saturating_add(1); + if worksheet_rows > limits.max_worksheet_rows { + return Err(format!( + "Excel worksheet row count exceeds the {} row limit", + limits.max_worksheet_rows + )); + } current_row = xml_attr_value(&reader, &element, b"r") .and_then(|value| value.parse::().ok()) .filter(|row| *row > 0) @@ -3189,6 +3319,13 @@ pub fn stream_xlsx_rows_to_channel_with_control( } } Ok(Event::Empty(element)) if xml_local_name_eq(element.name().as_ref(), b"c") => { + worksheet_cells = worksheet_cells.saturating_add(1); + if worksheet_cells > limits.max_worksheet_cells { + return Err(format!( + "Excel worksheet cell count exceeds the {} cell limit", + limits.max_worksheet_cells + )); + } let position = xml_attr_value(&reader, &element, b"r") .as_deref() .and_then(xlsx_cell_ref_position) @@ -3198,6 +3335,13 @@ pub fn stream_xlsx_rows_to_channel_with_control( rows.push_cell(position.0, position.1, serde_json::Value::Null, progress)?; } Ok(Event::Start(element)) if xml_local_name_eq(element.name().as_ref(), b"c") => { + worksheet_cells = worksheet_cells.saturating_add(1); + if worksheet_cells > limits.max_worksheet_cells { + return Err(format!( + "Excel worksheet cell count exceeds the {} cell limit", + limits.max_worksheet_cells + )); + } let position = xml_attr_value(&reader, &element, b"r") .as_deref() .and_then(xlsx_cell_ref_position) @@ -3230,9 +3374,15 @@ pub fn stream_xlsx_rows_to_channel_with_control( } Ok(Event::Text(text)) if in_value => { current_cell.value.push_str(&text.unescape().map_err(|error| error.to_string())?); + if current_cell.value.len() > limits.max_cell_bytes { + return Err(format!("Excel cell exceeds the {} byte limit", limits.max_cell_bytes)); + } } Ok(Event::Text(text)) if in_inline_text => { current_cell.inline_value.push_str(&text.unescape().map_err(|error| error.to_string())?); + if current_cell.inline_value.len() > limits.max_cell_bytes { + return Err(format!("Excel cell exceeds the {} byte limit", limits.max_cell_bytes)); + } } Ok(Event::End(element)) if xml_local_name_eq(element.name().as_ref(), b"v") => in_value = false, Ok(Event::End(element)) if xml_local_name_eq(element.name().as_ref(), b"t") => in_inline_text = false, @@ -3250,6 +3400,9 @@ pub fn stream_xlsx_rows_to_channel_with_control( format_as_text, empty_string_as_null, )?; + if value.as_str().is_some_and(|value| value.len() > limits.max_cell_bytes) { + return Err(format!("Excel cell exceeds the {} byte limit", limits.max_cell_bytes)); + } rows.push_cell(row, column, value, progress)?; current_cell = XlsxPreviewRawCell::default(); } @@ -9097,6 +9250,35 @@ mod tests { let _ = std::fs::remove_file(path); } + #[test] + fn streaming_excel_rows_enforce_row_cell_and_batch_budgets() { + let path = std::env::temp_dir().join(format!("dbx-table-import-stream-budget-{}.xlsx", uuid::Uuid::new_v4())); + std::fs::write(&path, build_styled_test_xlsx(false, &[("A1", 5, 10.0)])).unwrap(); + let options = TableImportParseOptions { has_header: Some(false), ..TableImportParseOptions::default() }; + + for (limits, expected) in [ + (XlsxStreamLimits { max_worksheet_rows: 0, ..Default::default() }, "row count"), + (XlsxStreamLimits { max_worksheet_cells: 0, ..Default::default() }, "cell count"), + (XlsxStreamLimits { max_batch_bytes: 1, ..Default::default() }, "batch budget"), + ] { + let (sender, _receiver) = tokio::sync::mpsc::channel(2); + let error = stream_xlsx_rows_to_channel_with_limits( + &path.to_string_lossy(), + &options, + 1, + None, + HashSet::from(["*".to_string()]), + true, + sender, + Arc::new(AtomicBool::new(false)), + limits, + ) + .unwrap_err(); + assert!(error.contains(expected), "unexpected error: {error}"); + } + let _ = std::fs::remove_file(path); + } + #[test] fn streaming_excel_rows_preserve_custom_title_and_data_range() { let path = std::env::temp_dir().join(format!("dbx-table-import-stream-range-{}.xlsx", uuid::Uuid::new_v4())); diff --git a/crates/dbx-mcp/Cargo.toml b/crates/dbx-mcp/Cargo.toml index 4114f91a27..79a89db9f5 100644 --- a/crates/dbx-mcp/Cargo.toml +++ b/crates/dbx-mcp/Cargo.toml @@ -24,10 +24,15 @@ schemars = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" sha2 = "0.10" +sysinfo = { version = "0.32", features = ["system"] } tokio = { version = "1", features = ["fs", "io-util", "macros", "rt-multi-thread", "sync", "time"] } time = { version = "0.3", features = ["formatting"] } uuid = { version = "1", features = ["v4"] } url = "2" +zip = { version = "4", default-features = false, features = ["deflate"] } + +[target.'cfg(unix)'.dependencies] +libc = "0.2" [dev-dependencies] rustls = { version = "0.23", features = ["aws-lc-rs"] } @@ -35,4 +40,3 @@ rustls-pemfile = "2.2" tempfile = "3" tokio = { version = "1", features = ["net", "io-util"] } tokio-rustls = "0.26" -zip = "4" diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 35ea3da94e..e9f238af04 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -25,6 +25,36 @@ use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use crate::mongo::MongoCommand; +struct ImportSnapshotCleanup { + directory: PathBuf, + active: bool, +} + +impl ImportSnapshotCleanup { + fn new(directory: PathBuf) -> Self { + Self { directory, active: true } + } + + fn cleanup(&mut self) -> std::io::Result<()> { + if self.active { + match std::fs::remove_dir_all(&self.directory) { + Ok(()) => self.active = false, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => self.active = false, + Err(error) => return Err(error), + } + } + Ok(()) + } +} + +impl Drop for ImportSnapshotCleanup { + fn drop(&mut self) { + if self.active { + let _ = std::fs::remove_dir_all(&self.directory); + } + } +} + #[derive(Debug, Clone, Serialize, PartialEq, Eq)] pub struct ConnectionSummary { pub id: String, @@ -706,21 +736,25 @@ impl DbxBackend for LocalBackend { .to_string(); let snapshot_root = self.data_dir.join("tmp").join("mcp_import"); tokio::fs::create_dir_all(&snapshot_root).await.map_err(|error| format!("创建导入快照目录失败：{error}"))?; + let source_size = + std::fs::metadata(&request.file_path).map_err(|error| format!("读取导入源大小失败：{error}"))?.len(); + crate::enterprise_tools::ensure_import_disk_budget(&snapshot_root, source_size) + .map_err(|error| error.to_string())?; let snapshot_dir = snapshot_root.join(&request.import_id); tokio::fs::create_dir(&snapshot_dir).await.map_err(|error| format!("创建任务快照目录失败：{error}"))?; + let mut cleanup_guard = ImportSnapshotCleanup::new(snapshot_dir.clone()); let extension = Path::new(&request.file_path).extension().and_then(|value| value.to_str()).unwrap_or("data"); let snapshot_path = snapshot_dir.join(format!("source.{extension}")); let normalized_path = snapshot_dir.join("normalized.csv"); let snapshot_result = async { - tokio::fs::copy(&request.file_path, &snapshot_path) - .await - .map_err(|error| format!("复制导入源快照失败：{error}"))?; - let identity = crate::enterprise_tools::file_identity(snapshot_path.clone()) - .await - .map_err(|error| error.to_string())?; - if identity.sha256 != expected_sha256 { - return Err("IMPORT_FILE_CHANGED: 源文件在启动导入时发生变化；未写入数据库。".to_string()); - } + crate::enterprise_tools::copy_verified_import_source( + PathBuf::from(&request.file_path), + snapshot_path.clone(), + expected_sha256.clone(), + cancelled.clone(), + ) + .await + .map_err(|error| error.to_string())?; request.file_path = snapshot_path.to_string_lossy().to_string(); request = crate::enterprise_tools::build_governed_import_snapshot( request, @@ -758,10 +792,15 @@ impl DbxBackend for LocalBackend { result } .await; - let _ = tokio::fs::remove_file(&normalized_path).await; - let _ = tokio::fs::remove_file(&snapshot_path).await; - let _ = tokio::fs::remove_dir(&snapshot_dir).await; - snapshot_result + let cleanup_result = cleanup_guard.cleanup(); + match (snapshot_result, cleanup_result) { + (Ok(summary), Ok(())) => Ok(summary), + (Err(error), Ok(())) => Err(error), + (Ok(_), Err(error)) => { + Err(format!("IMPORT_SNAPSHOT_CLEANUP_FAILED: staging 已完成，但清理任务快照失败：{error}")) + } + (Err(error), Err(cleanup)) => Err(format!("{error}; IMPORT_SNAPSHOT_CLEANUP_FAILED: {cleanup}")), + } } async fn execute_redis_command( @@ -2901,6 +2940,18 @@ mod tests { assert_eq!(local_agent_dir(&legacy, data_dir), PathBuf::from("D:/DBX/drivers/agents")); } + #[test] + fn snapshot_cleanup_guard_removes_directory_during_unwind_or_early_return() { + let data_dir = tempfile::tempdir().unwrap(); + let snapshot = data_dir.path().join("mcp-import-test"); + std::fs::create_dir(&snapshot).unwrap(); + std::fs::write(snapshot.join("normalized.csv"), b"partial").unwrap(); + { + let _guard = ImportSnapshotCleanup::new(snapshot.clone()); + } + assert!(!snapshot.exists()); + } + struct StubBackend; #[async_trait] diff --git a/crates/dbx-mcp/src/enterprise_tools.rs b/crates/dbx-mcp/src/enterprise_tools.rs index 9957ff46fd..4f6d260121 100644 --- a/crates/dbx-mcp/src/enterprise_tools.rs +++ b/crates/dbx-mcp/src/enterprise_tools.rs @@ -1,7 +1,8 @@ use std::{ collections::{BTreeMap, HashMap, HashSet}, ffi::OsStr, - io::Read, + fs::{File, OpenOptions}, + io::{Read, Seek, SeekFrom, Write as IoWrite}, path::{Path, PathBuf}, sync::{ atomic::{AtomicBool, Ordering}, @@ -19,6 +20,7 @@ use rmcp::schemars; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use sha2::{Digest, Sha256}; +use sysinfo::Disks; use time::{format_description::well_known::Rfc3339, OffsetDateTime, UtcOffset}; use tokio::sync::{OwnedSemaphorePermit, RwLock, Semaphore}; use uuid::Uuid; @@ -41,6 +43,20 @@ const DEFAULT_PLAN_CAPACITY: usize = 512; const DEFAULT_JOB_CAPACITY: usize = 256; const DEFAULT_TERMINAL_JOB_LIMIT: usize = 128; const JOB_TTL_MS: u128 = 24 * 60 * 60 * 1_000; +const DEFAULT_INSPECTION_CONCURRENCY: usize = 2; +const DEFAULT_INSPECTION_TIMEOUT_SECS: u64 = 30; +const DEFAULT_XLSX_ZIP_ENTRY_LIMIT: usize = 4_096; +const DEFAULT_XLSX_TOTAL_UNCOMPRESSED_BYTES: u64 = 4 * 1024 * 1024 * 1024; +const DEFAULT_XLSX_METADATA_ENTRY_BYTES: u64 = 16 * 1024 * 1024; +const DEFAULT_XLSX_SHARED_STRINGS_BYTES: u64 = 256 * 1024 * 1024; +const DEFAULT_XLSX_WORKSHEET_BYTES: u64 = 2 * 1024 * 1024 * 1024; +const DEFAULT_XLSX_WORKSHEET_ROWS: usize = 5_000_000; +const DEFAULT_XLSX_WORKSHEET_CELLS: usize = 100_000_000; +const DEFAULT_XLSX_CELL_BYTES: usize = 1024 * 1024; +const DEFAULT_XLSX_BATCH_MEMORY_BYTES: usize = 32 * 1024 * 1024; +const DEFAULT_NORMALIZED_OUTPUT_BYTES: u64 = 2 * 1024 * 1024 * 1024; +const DEFAULT_IMPORT_DISK_RESERVE_BYTES: u64 = 2 * 1024 * 1024 * 1024; +const ESTIMATED_IMPORT_CELL_BYTES: usize = 512; const IMPORT_EXTENSIONS: &[&str] = &["xlsx", "xlsm", "xls", "csv", "tsv", "json"]; const SEMANTIC_EXTENSIONS: &[&str] = &["jsonl"]; @@ -68,6 +84,25 @@ const DEFAULT_VECTOR_OUTPUT_FIELDS: &[&str] = &[ ]; const RESERVED_STAGING_COLUMNS: &[&str] = &["import_id", "plan_id", "source_sha", "source_row_number", "source_row_hash", "loaded_at"]; +const SEMANTIC_VARCHAR_LIMITS: &[(&str, usize)] = &[ + ("card_id", 128), + ("card_type", 32), + ("business_domain", 128), + ("dataset_id", 128), + ("template_version", 64), + ("title", 512), + ("content", 8_192), + ("approval_status", 32), + ("effective_from", 10), + ("effective_to", 10), + ("source_uri", 2_048), + ("source_checksum", 64), + ("content_checksum", 64), + ("semantic_version", 128), + ("semantic_batch_id", 128), + ("embedding_model", 128), + ("embedding_revision", 64), +]; #[derive(Debug, Clone, PartialEq, Eq)] pub struct EnterpriseToolError { @@ -269,8 +304,8 @@ pub struct VectorDeleteByBatchRequest { pub database: Option, pub collection: String, pub semantic_batch_id: String, - #[schemars(description = "Must be explicitly false; published semantic batches cannot be deleted by this tool")] - pub published: bool, + #[schemars(extend("type" = "boolean"), description = "Ignored in v1; deletion is disabled server-side")] + pub published: Option, } #[derive(Debug, Clone, Serialize, PartialEq, Eq)] @@ -394,6 +429,7 @@ pub struct ImportJob { pub struct EnterpriseRuntime { plans: RwLock>, jobs: RwLock>>, + inspection_slots: Arc, import_slots: Arc, semantic_write_slots: Arc, plan_capacity: usize, @@ -406,6 +442,9 @@ impl Default for EnterpriseRuntime { Self { plans: RwLock::new(HashMap::new()), jobs: RwLock::new(HashMap::new()), + inspection_slots: Arc::new(Semaphore::new( + env_usize("DBX_MCP_IMPORT_INSPECTION_CONCURRENCY", DEFAULT_INSPECTION_CONCURRENCY).clamp(1, 8), + )), import_slots: Arc::new(Semaphore::new(env_usize("DBX_MCP_IMPORT_CONCURRENCY", 2).clamp(1, 16))), semantic_write_slots: Arc::new(Semaphore::new(1)), plan_capacity: env_usize("DBX_MCP_IMPORT_PLAN_CAPACITY", DEFAULT_PLAN_CAPACITY).clamp(1, 4_096), @@ -416,6 +455,17 @@ impl Default for EnterpriseRuntime { } impl EnterpriseRuntime { + #[cfg(test)] + fn with_inspection_concurrency(concurrency: usize) -> Self { + Self { inspection_slots: Arc::new(Semaphore::new(concurrency.max(1))), ..Self::default() } + } + + pub async fn try_inspection_permit(&self) -> Result { + self.inspection_slots.clone().try_acquire_owned().map_err(|_| { + EnterpriseToolError::new("IMPORT_INSPECTION_CONCURRENCY_LIMIT", "文件剖析并发已达到上限，请稍后重试。") + }) + } + pub async fn insert_plan(&self, plan: PreparedImportPlan) -> Result<(), EnterpriseToolError> { let now = unix_epoch_millis(); let mut plans = self.plans.write().await; @@ -603,7 +653,9 @@ pub async fn file_identity(path: PathBuf) -> Result Result { - let metadata = std::fs::metadata(path) + let mut file = open_regular_file_nofollow(path)?; + let metadata = file + .metadata() .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法读取文件元数据：{error}")))?; let modified_nanos = metadata .modified() @@ -611,8 +663,6 @@ fn file_identity_sync(path: &Path) -> Result .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok()) .map(|duration| duration.as_nanos()) .unwrap_or_default(); - let mut file = std::fs::File::open(path) - .map_err(|error| EnterpriseToolError::new("IMPORT_HASH_FAILED", format!("无法打开文件：{error}")))?; let mut hasher = Sha256::new(); let mut buffer = vec![0_u8; 1024 * 1024]; loop { @@ -624,6 +674,15 @@ fn file_identity_sync(path: &Path) -> Result } hasher.update(&buffer[..read]); } + let final_metadata = file + .metadata() + .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法复核文件元数据：{error}")))?; + if !same_file_metadata(&metadata, &final_metadata) { + return Err(EnterpriseToolError::new( + "IMPORT_FILE_CHANGED_DURING_HASH", + "文件在计算 SHA-256 期间发生变化；结果已丢弃。", + )); + } Ok(FileIdentity { canonical_path: path.to_string_lossy().to_string(), size_bytes: metadata.len(), @@ -632,6 +691,391 @@ fn file_identity_sync(path: &Path) -> Result }) } +fn open_regular_file_nofollow(path: &Path) -> Result { + let mut options = OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW); + } + let file = options + .open(path) + .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法安全打开文件：{error}")))?; + let metadata = file + .metadata() + .map_err(|error| EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法读取文件描述符：{error}")))?; + if !metadata.file_type().is_file() { + return Err(EnterpriseToolError::new("IMPORT_NOT_REGULAR_FILE", "导入源必须是普通文件。")); + } + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + let path_metadata = std::fs::metadata(path).map_err(|error| { + EnterpriseToolError::new("IMPORT_FILE_UNAVAILABLE", format!("无法复核文件路径：{error}")) + })?; + if metadata.dev() != path_metadata.dev() || metadata.ino() != path_metadata.ino() { + return Err(EnterpriseToolError::new( + "IMPORT_FILE_IDENTITY_RACE", + "文件路径在打开期间指向了不同对象；结果已丢弃。", + )); + } + } + Ok(file) +} + +fn same_file_metadata(before: &std::fs::Metadata, after: &std::fs::Metadata) -> bool { + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + before.dev() == after.dev() + && before.ino() == after.ino() + && before.len() == after.len() + && before.mtime() == after.mtime() + && before.mtime_nsec() == after.mtime_nsec() + && before.ctime() == after.ctime() + && before.ctime_nsec() == after.ctime_nsec() + } + #[cfg(not(unix))] + { + before.len() == after.len() + && before.modified().ok().and_then(|value| value.duration_since(UNIX_EPOCH).ok()) + == after.modified().ok().and_then(|value| value.duration_since(UNIX_EPOCH).ok()) + } +} + +pub fn ensure_import_disk_budget(snapshot_root: &Path, source_size: u64) -> Result<(), EnterpriseToolError> { + let canonical = std::fs::canonicalize(snapshot_root).map_err(|error| { + EnterpriseToolError::new("IMPORT_DISK_BUDGET_UNAVAILABLE", format!("无法解析任务磁盘路径：{error}")) + })?; + let disks = Disks::new_with_refreshed_list(); + let disk = disks + .list() + .iter() + .filter(|disk| canonical.starts_with(disk.mount_point())) + .max_by_key(|disk| disk.mount_point().as_os_str().len()) + .ok_or_else(|| EnterpriseToolError::new("IMPORT_DISK_BUDGET_UNAVAILABLE", "无法确定任务磁盘剩余容量。"))?; + let reserve = env_u64("DBX_MCP_IMPORT_DISK_RESERVE_BYTES", DEFAULT_IMPORT_DISK_RESERVE_BYTES) + .clamp(512 * 1024 * 1024, 16 * 1024 * 1024 * 1024); + let required = source_size + .checked_add(normalized_output_limit()) + .and_then(|value| value.checked_add(reserve)) + .ok_or_else(|| EnterpriseToolError::new("IMPORT_DISK_BUDGET_EXCEEDED", "导入磁盘预算累计溢出。"))?; + if disk.available_space() < required { + return Err(EnterpriseToolError::new( + "IMPORT_DISK_BUDGET_EXCEEDED", + format!( + "任务磁盘可用 {} 字节，至少需要 {required} 字节（源文件、规范化上限与安全保留空间）。", + disk.available_space() + ), + )); + } + Ok(()) +} + +pub async fn copy_verified_import_source( + source: PathBuf, + destination: PathBuf, + expected_sha256: String, + cancelled: Arc, +) -> Result<(), EnterpriseToolError> { + tokio::task::spawn_blocking(move || { + let result = copy_verified_import_source_sync(&source, &destination, &expected_sha256, &cancelled); + if result.is_err() { + let _ = std::fs::remove_file(&destination); + } + result + }) + .await + .map_err(|error| EnterpriseToolError::new("IMPORT_COPY_FAILED", error.to_string()))? +} + +fn copy_verified_import_source_sync( + source: &Path, + destination: &Path, + expected_sha256: &str, + cancelled: &AtomicBool, +) -> Result<(), EnterpriseToolError> { + if !is_sha256_hex(expected_sha256) { + return Err(EnterpriseToolError::new("IMPORT_SOURCE_HASH_REQUIRED", "导入计划缺少合法 SHA-256。")); + } + let mut input = open_regular_file_nofollow(source)?; + let initial = input + .metadata() + .map_err(|error| EnterpriseToolError::new("IMPORT_COPY_FAILED", format!("无法读取源文件描述符：{error}")))?; + let mut output = + OpenOptions::new().write(true).create_new(true).open(destination).map_err(|error| { + EnterpriseToolError::new("IMPORT_COPY_FAILED", format!("无法创建任务私有快照：{error}")) + })?; + let mut hasher = Sha256::new(); + let mut buffer = vec![0u8; 1024 * 1024]; + loop { + if cancelled.load(Ordering::Acquire) { + return Err(EnterpriseToolError::new("IMPORT_CANCELLED", "复制任务私有快照时收到取消请求。")); + } + let read = input + .read(&mut buffer) + .map_err(|error| EnterpriseToolError::new("IMPORT_COPY_FAILED", format!("读取源文件失败：{error}")))?; + if read == 0 { + break; + } + output.write_all(&buffer[..read]).map_err(|error| { + EnterpriseToolError::new("IMPORT_COPY_FAILED", format!("写入任务私有快照失败：{error}")) + })?; + hasher.update(&buffer[..read]); + } + output + .flush() + .and_then(|_| output.sync_all()) + .map_err(|error| EnterpriseToolError::new("IMPORT_COPY_FAILED", format!("落盘任务私有快照失败：{error}")))?; + let final_metadata = input + .metadata() + .map_err(|error| EnterpriseToolError::new("IMPORT_COPY_FAILED", format!("无法复核源文件描述符：{error}")))?; + if !same_file_metadata(&initial, &final_metadata) { + return Err(EnterpriseToolError::new( + "IMPORT_FILE_CHANGED", + "源文件在复制任务私有快照期间发生变化；未访问数据库。", + )); + } + let actual_sha256 = format!("{:x}", hasher.finalize()); + if actual_sha256 != expected_sha256 { + return Err(EnterpriseToolError::new("IMPORT_FILE_CHANGED", "源文件在启动导入时发生变化；未访问数据库。")); + } + Ok(()) +} + +#[derive(Debug, Clone)] +pub struct ImportInspection { + pub identity: FileIdentity, + pub preview: TableImportPreview, + pub source_columns: Vec, +} + +pub fn inspection_timeout() -> Duration { + Duration::from_secs( + env_u64("DBX_MCP_IMPORT_INSPECTION_TIMEOUT_SECS", DEFAULT_INSPECTION_TIMEOUT_SECS).clamp(5, 120), + ) +} + +pub async fn inspect_import_source( + path: PathBuf, + source_format: Option, + parse_options: TableImportParseOptions, + preview_rows: usize, +) -> Result { + let limits_path = path.clone(); + tokio::task::spawn_blocking(move || validate_source_archive_limits(&limits_path, source_format)) + .await + .map_err(|error| EnterpriseToolError::new("IMPORT_INSPECTION_FAILED", error.to_string()))??; + let identity = file_identity(path.clone()).await?; + let file_path = path.to_string_lossy().to_string(); + let preview = dbx_core::table_import::preview_table_import_file_with_request( + dbx_core::table_import::TableImportPreviewRequest { + file_path: file_path.clone(), + source_ref: Some(identity.sha256.clone()), + source_format, + parse_options: parse_options.clone(), + preview_limit: Some(preview_rows), + }, + ) + .await + .map_err(|error| EnterpriseToolError::new("IMPORT_PREVIEW_FAILED", error))?; + validate_preview_headers(&preview)?; + let source_columns = + source_columns_for_preview(&file_path, source_format, &parse_options, &preview.columns).await?; + let current_identity = file_identity(path).await?; + if current_identity != identity { + return Err(EnterpriseToolError::new( + "IMPORT_FILE_CHANGED_DURING_INSPECTION", + "文件在剖析期间发生变化；结果已丢弃。", + )); + } + Ok(ImportInspection { identity, preview, source_columns }) +} + +#[derive(Debug, Clone, Copy)] +struct XlsxArchiveLimits { + entry_count: usize, + total_uncompressed_bytes: u64, + metadata_entry_bytes: u64, + shared_strings_bytes: u64, + worksheet_bytes: u64, +} + +fn xlsx_archive_limits() -> XlsxArchiveLimits { + XlsxArchiveLimits { + entry_count: env_usize("DBX_MCP_XLSX_ZIP_ENTRY_LIMIT", DEFAULT_XLSX_ZIP_ENTRY_LIMIT).clamp(16, 16_384), + total_uncompressed_bytes: env_u64( + "DBX_MCP_XLSX_TOTAL_UNCOMPRESSED_MAX_BYTES", + DEFAULT_XLSX_TOTAL_UNCOMPRESSED_BYTES, + ) + .clamp(64 * 1024 * 1024, 8 * 1024 * 1024 * 1024), + metadata_entry_bytes: env_u64("DBX_MCP_XLSX_METADATA_ENTRY_MAX_BYTES", DEFAULT_XLSX_METADATA_ENTRY_BYTES) + .clamp(1024 * 1024, 64 * 1024 * 1024), + shared_strings_bytes: env_u64("DBX_MCP_XLSX_SHARED_STRINGS_MAX_BYTES", DEFAULT_XLSX_SHARED_STRINGS_BYTES) + .clamp(8 * 1024 * 1024, 512 * 1024 * 1024), + worksheet_bytes: env_u64("DBX_MCP_XLSX_WORKSHEET_MAX_BYTES", DEFAULT_XLSX_WORKSHEET_BYTES) + .clamp(64 * 1024 * 1024, 4 * 1024 * 1024 * 1024), + } +} + +fn validate_source_archive_limits( + path: &Path, + source_format: Option, +) -> Result<(), EnterpriseToolError> { + let format = dbx_core::table_import::effective_source_format(&path.to_string_lossy(), source_format) + .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_FORMAT_INVALID", error))?; + let extension = path.extension().and_then(OsStr::to_str).unwrap_or_default(); + if format != TableImportSourceFormat::Excel + || (!extension.eq_ignore_ascii_case("xlsx") && !extension.eq_ignore_ascii_case("xlsm")) + { + return Ok(()); + } + validate_xlsx_archive_limits(path, xlsx_archive_limits()) +} + +fn validate_xlsx_archive_limits(path: &Path, limits: XlsxArchiveLimits) -> Result<(), EnterpriseToolError> { + let mut file = open_regular_file_nofollow(path)?; + let declared_entries = xlsx_eocd_entry_count(&mut file)?; + if declared_entries > limits.entry_count { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_ZIP_ENTRY_LIMIT_EXCEEDED", + format!("XLSX ZIP 包含 {declared_entries} 个条目，超过 {} 个安全上限。", limits.entry_count), + )); + } + file.seek(SeekFrom::Start(0)) + .map_err(|error| EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", error.to_string()))?; + let mut archive = zip::ZipArchive::new(file) + .map_err(|error| EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", error.to_string()))?; + if archive.len() != declared_entries { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_ZIP_ENTRY_COUNT_MISMATCH", + "XLSX ZIP 中央目录条目数不一致。", + )); + } + let mut names = HashSet::new(); + let mut total_uncompressed = 0u64; + for index in 0..archive.len() { + let entry = archive + .by_index(index) + .map_err(|error| EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", error.to_string()))?; + let name = entry.name().replace('\\', "/"); + if !names.insert(name.clone()) { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_ZIP_DUPLICATE_ENTRY", + format!("XLSX ZIP 包含重复条目 {name}。"), + )); + } + total_uncompressed = total_uncompressed.checked_add(entry.size()).ok_or_else(|| { + EnterpriseToolError::new("IMPORT_XLSX_ZIP_BUDGET_EXCEEDED", "XLSX ZIP 解压大小累计溢出。") + })?; + if total_uncompressed > limits.total_uncompressed_bytes { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_ZIP_BUDGET_EXCEEDED", + format!("XLSX ZIP 解压总大小超过 {} 字节。", limits.total_uncompressed_bytes), + )); + } + let metadata_entry = matches!( + name.as_str(), + "xl/workbook.xml" | "xl/_rels/workbook.xml.rels" | "xl/styles.xml" | "[Content_Types].xml" + ); + if metadata_entry && entry.size() > limits.metadata_entry_bytes { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_METADATA_BUDGET_EXCEEDED", + format!("XLSX 元数据条目 {name} 超过 {} 字节。", limits.metadata_entry_bytes), + )); + } + if name == "xl/sharedStrings.xml" && entry.size() > limits.shared_strings_bytes { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_SHARED_STRINGS_BUDGET_EXCEEDED", + format!("XLSX sharedStrings 超过 {} 字节。", limits.shared_strings_bytes), + )); + } + if name.starts_with("xl/worksheets/") && name.ends_with(".xml") && entry.size() > limits.worksheet_bytes { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_WORKSHEET_BUDGET_EXCEEDED", + format!("XLSX 工作表 {name} 超过 {} 字节。", limits.worksheet_bytes), + )); + } + } + Ok(()) +} + +fn xlsx_eocd_entry_count(file: &mut File) -> Result { + let length = + file.metadata().map_err(|error| EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", error.to_string()))?.len(); + let tail_len = length.min(65_557) as usize; + file.seek(SeekFrom::End(-(tail_len as i64))) + .map_err(|error| EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", error.to_string()))?; + let mut tail = vec![0u8; tail_len]; + file.read_exact(&mut tail) + .map_err(|error| EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", error.to_string()))?; + let offset = tail + .windows(4) + .rposition(|window| window == [0x50, 0x4b, 0x05, 0x06]) + .ok_or_else(|| EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", "XLSX ZIP 缺少 EOCD。"))?; + if tail.len().saturating_sub(offset) < 22 { + return Err(EnterpriseToolError::new("IMPORT_XLSX_ZIP_INVALID", "XLSX ZIP EOCD 不完整。")); + } + let disk = u16::from_le_bytes([tail[offset + 4], tail[offset + 5]]); + let central_disk = u16::from_le_bytes([tail[offset + 6], tail[offset + 7]]); + let disk_entries = u16::from_le_bytes([tail[offset + 8], tail[offset + 9]]); + let total_entries = u16::from_le_bytes([tail[offset + 10], tail[offset + 11]]); + if disk != 0 || central_disk != 0 || disk_entries != total_entries { + return Err(EnterpriseToolError::new("IMPORT_XLSX_ZIP_MULTIDISK_UNSUPPORTED", "不支持多磁盘 XLSX ZIP。")); + } + if total_entries == u16::MAX { + return Err(EnterpriseToolError::new( + "IMPORT_XLSX_ZIP64_ENTRY_COUNT_UNSUPPORTED", + "XLSX ZIP64 条目数无法在剖析前安全确定。", + )); + } + Ok(total_entries as usize) +} + +fn governed_xlsx_stream_limits() -> dbx_core::table_import::XlsxStreamLimits { + let archive = xlsx_archive_limits(); + dbx_core::table_import::XlsxStreamLimits { + max_shared_strings_bytes: archive.shared_strings_bytes, + max_worksheet_bytes: archive.worksheet_bytes, + max_worksheet_rows: env_usize("DBX_MCP_XLSX_WORKSHEET_MAX_ROWS", DEFAULT_XLSX_WORKSHEET_ROWS) + .clamp(1_000, 10_000_000), + max_worksheet_cells: env_usize("DBX_MCP_XLSX_WORKSHEET_MAX_CELLS", DEFAULT_XLSX_WORKSHEET_CELLS) + .clamp(1_000, 200_000_000), + max_cell_bytes: env_usize("DBX_MCP_XLSX_CELL_MAX_BYTES", DEFAULT_XLSX_CELL_BYTES) + .clamp(4 * 1024, 4 * 1024 * 1024), + max_batch_bytes: import_batch_memory_bytes(), + } +} + +fn import_batch_memory_bytes() -> usize { + env_usize("DBX_MCP_IMPORT_BATCH_MEMORY_MAX_BYTES", DEFAULT_XLSX_BATCH_MEMORY_BYTES) + .clamp(4 * 1024 * 1024, 128 * 1024 * 1024) +} + +pub fn governed_import_batch_size( + requested: Option, + source_column_count: usize, +) -> Result { + if source_column_count == 0 { + return Err(EnterpriseToolError::new("IMPORT_COLUMN_REQUIRED", "源文件没有可导入字段。")); + } + let estimated_row_bytes = + source_column_count.saturating_mul(ESTIMATED_IMPORT_CELL_BYTES).saturating_add(1024).max(1); + let dynamic_max = import_batch_memory_bytes().checked_div(estimated_row_bytes).unwrap_or(1).clamp(1, 5_000); + let value = requested.unwrap_or(1_000.min(dynamic_max)); + if value == 0 || value > dynamic_max { + return Err(EnterpriseToolError::new( + "IMPORT_BATCH_SIZE_RESOURCE_LIMIT", + format!( + "当前 {} 列结构的 batch_size 必须在 1 到 {dynamic_max} 之间，以满足批次内存预算。", + source_column_count + ), + )); + } + Ok(value) +} + pub fn preview_limit(value: Option) -> Result { let value = value.unwrap_or(DEFAULT_PREVIEW_ROWS); if value == 0 || value > MAX_PREVIEW_ROWS { @@ -975,11 +1419,10 @@ pub fn validate_governed_source_v1( let source_format = dbx_core::table_import::effective_source_format(file_path, source_format) .map_err(|error| EnterpriseToolError::new("IMPORT_SOURCE_FORMAT_INVALID", error))?; if source_format.is_delimited() { - if !matches!(parse_options.encoding, None | Some(TableImportTextEncoding::Auto | TableImportTextEncoding::Utf8)) - { + if !matches!(parse_options.encoding, Some(TableImportTextEncoding::Utf8)) { return Err(EnterpriseToolError::new( - "IMPORT_GOVERNED_ENCODING_UNSUPPORTED_V1", - "v1 流式治理导入仅支持 UTF-8；其他编码不创建导入计划。", + "IMPORT_GOVERNED_EXPLICIT_UTF8_REQUIRED_V1", + "v1 流式治理导入要求 prepare 显式指定 utf-8；auto、GBK 与 UTF-16 不创建导入计划。", )); } return Ok(source_format); @@ -1036,15 +1479,11 @@ async fn build_governed_delimited_snapshot( .to_offset(offset) .format(&Rfc3339) .map_err(|error| EnterpriseToolError::new("IMPORT_LOADED_AT_FAILED", error.to_string()))?; - let mut writer = csv::WriterBuilder::new() - .has_headers(false) - .from_path(output_path) - .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + let output = LimitedWriter::create(output_path, normalized_output_limit())?; + let mut writer = csv::WriterBuilder::new().has_headers(false).from_writer(output); let mut headers = request.mappings.iter().map(|mapping| mapping.target_column.clone()).collect::>(); headers.extend(RESERVED_STAGING_COLUMNS.iter().map(|column| column.to_string())); - writer - .write_record(&headers) - .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + writer.write_record(&headers).map_err(snapshot_write_error)?; let mut reader = csv::ReaderBuilder::new() .delimiter(config.delimiter) .has_headers(false) @@ -1054,6 +1493,7 @@ async fn build_governed_delimited_snapshot( let mut source_columns = Vec::new(); let mut source_indexes: Option> = None; let mut source_row_count = 0usize; + let max_cell_bytes = governed_xlsx_stream_limits().max_cell_bytes; for (index, record) in reader.records().enumerate() { if cancelled.load(Ordering::Acquire) { return Err(EnterpriseToolError::new("IMPORT_CANCELLED", "导入已取消。")); @@ -1081,6 +1521,24 @@ async fn build_governed_delimited_snapshot( if source_columns.is_empty() { source_columns = (0..record.len()).map(|index| format!("column_{}", index + 1)).collect(); } + if record.iter().any(|value| value.len() > max_cell_bytes) { + return Err(EnterpriseToolError::new( + "IMPORT_CELL_BUDGET_EXCEEDED", + format!("UTF-8 CSV/TSV 第 {source_row_number} 行包含超过单元格字节上限的字段。"), + )); + } + if record.len() > source_columns.len() + && record.iter().skip(source_columns.len()).any(|value| !value.is_empty()) + { + return Err(EnterpriseToolError::new( + "IMPORT_SOURCE_ROW_WIDER_THAN_HEADER", + format!( + "UTF-8 CSV/TSV 第 {source_row_number} 行包含表头之外的非空字段；表头 {} 列，数据行 {} 列。", + source_columns.len(), + record.len() + ), + )); + } if source_indexes.is_none() { source_indexes = Some( request @@ -1100,9 +1558,8 @@ async fn build_governed_delimited_snapshot( let indexes = source_indexes .as_ref() .ok_or_else(|| EnterpriseToolError::new("IMPORT_SOURCE_MAPPING_FAILED", "无法建立源字段位置映射。"))?; - let parsed_row = (0..source_columns.len()) - .map(|column| dbx_core::table_import::csv_value_with_config(record.get(column).unwrap_or_default(), config)) - .collect::>(); + let parsed_row = + record.iter().map(|value| dbx_core::table_import::csv_value_with_config(value, config)).collect::>(); let mut output = indexes .iter() .map(|index| staging_text_value(parsed_row.get(*index).unwrap_or(&Value::Null))) @@ -1117,15 +1574,13 @@ async fn build_governed_delimited_snapshot( sha256_bytes(&row_bytes), loaded_at.clone(), ]); - writer - .write_record(&output) - .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + writer.write_record(&output).map_err(snapshot_write_error)?; source_row_count += 1; } if source_row_count == 0 { return Err(EnterpriseToolError::new("IMPORT_SOURCE_EMPTY", "选择的源范围没有数据行；未写入数据库。")); } - writer.flush().map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + writer.flush().map_err(snapshot_write_error)?; let mut mappings = request .mappings @@ -1172,15 +1627,11 @@ async fn build_governed_xlsx_snapshot( .to_offset(offset) .format(&Rfc3339) .map_err(|error| EnterpriseToolError::new("IMPORT_LOADED_AT_FAILED", error.to_string()))?; - let mut writer = csv::WriterBuilder::new() - .has_headers(false) - .from_path(output_path) - .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + let output = LimitedWriter::create(output_path, normalized_output_limit())?; + let mut writer = csv::WriterBuilder::new().has_headers(false).from_writer(output); let mut headers = request.mappings.iter().map(|mapping| mapping.target_column.clone()).collect::>(); headers.extend(RESERVED_STAGING_COLUMNS.iter().map(|column| column.to_string())); - writer - .write_record(&headers) - .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + writer.write_record(&headers).map_err(snapshot_write_error)?; let (sender, mut receiver) = tokio::sync::mpsc::channel(2); let path = request.file_path.clone(); @@ -1191,7 +1642,7 @@ async fn build_governed_xlsx_snapshot( let _producer_cancel_guard = CancelOnDrop(producer_cancelled.clone()); let producer_cancelled_for_task = producer_cancelled.clone(); let producer = tokio::task::spawn_blocking(move || { - dbx_core::table_import::stream_xlsx_rows_to_channel_with_control( + dbx_core::table_import::stream_xlsx_rows_to_channel_with_limits( &path, &options, batch_size, @@ -1200,6 +1651,7 @@ async fn build_governed_xlsx_snapshot( true, sender, producer_cancelled_for_task, + governed_xlsx_stream_limits(), ) }); let user_cancelled = cancelled.clone(); @@ -1271,9 +1723,7 @@ async fn build_governed_xlsx_snapshot( sha256_bytes(&row_bytes), loaded_at.clone(), ]); - writer - .write_record(&output) - .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + writer.write_record(&output).map_err(snapshot_write_error)?; source_row_count += 1; } } @@ -1296,7 +1746,7 @@ async fn build_governed_xlsx_snapshot( "XLSX 流式治理未完整结束；未访问数据库。", )); } - writer.flush().map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + writer.flush().map_err(snapshot_write_error)?; let mut mappings = request .mappings @@ -1343,6 +1793,58 @@ impl Drop for CancelOnDrop { } } +struct LimitedWriter { + file: File, + written: u64, + max_bytes: u64, +} + +impl LimitedWriter { + fn create(path: &Path, max_bytes: u64) -> Result { + let file = OpenOptions::new() + .write(true) + .create_new(true) + .open(path) + .map_err(|error| EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", error.to_string()))?; + Ok(Self { file, written: 0, max_bytes }) + } +} + +impl IoWrite for LimitedWriter { + fn write(&mut self, buffer: &[u8]) -> std::io::Result { + let next = self.written.checked_add(buffer.len() as u64).ok_or_else(|| { + std::io::Error::other("IMPORT_NORMALIZED_OUTPUT_LIMIT_EXCEEDED: 规范化输出大小累计溢出。") + })?; + if next > self.max_bytes { + return Err(std::io::Error::other(format!( + "IMPORT_NORMALIZED_OUTPUT_LIMIT_EXCEEDED: 规范化输出超过 {} 字节上限。", + self.max_bytes + ))); + } + let written = self.file.write(buffer)?; + self.written = self.written.saturating_add(written as u64); + Ok(written) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.file.flush() + } +} + +fn snapshot_write_error(error: impl ToString) -> EnterpriseToolError { + let message = error.to_string(); + if message.contains("IMPORT_NORMALIZED_OUTPUT_LIMIT_EXCEEDED") { + EnterpriseToolError::new("IMPORT_NORMALIZED_OUTPUT_LIMIT_EXCEEDED", message) + } else { + EnterpriseToolError::new("IMPORT_SNAPSHOT_WRITE_FAILED", message) + } +} + +fn normalized_output_limit() -> u64 { + env_u64("DBX_MCP_IMPORT_NORMALIZED_MAX_BYTES", DEFAULT_NORMALIZED_OUTPUT_BYTES) + .clamp(64 * 1024 * 1024, 4 * 1024 * 1024 * 1024) +} + fn staging_text_value(value: &Value) -> String { match value { Value::Null => String::new(), @@ -1460,12 +1962,12 @@ pub fn build_milvus_filter( pub fn read_semantic_jsonl(path: &Path, semantic_batch_id: &str) -> Result, EnterpriseToolError> { if semantic_batch_id.trim().is_empty() - || semantic_batch_id.len() > 200 + || semantic_batch_id.len() > 128 || semantic_batch_id.chars().any(char::is_control) { return Err(EnterpriseToolError::new( "SEMANTIC_BATCH_ID_INVALID", - "semantic_batch_id 必须是 1～200 个非控制字符。", + "semantic_batch_id 必须是 1～128 个 UTF-8 字节且不能包含控制字符。", )); } let source = std::fs::read_to_string(path) @@ -1603,12 +2105,6 @@ pub fn read_semantic_jsonl(path: &Path, semantic_batch_id: &str) -> Result 100_000 { - return Err(EnterpriseToolError::new( - "VECTOR_JSONL_FIELD_INVALID", - format!("JSONL 第 {line_number} 行的 content 超过 100000 字节。"), - )); - } if object.get("content_checksum").and_then(Value::as_str) != Some(sha256_bytes(content.as_bytes()).as_str()) { return Err(EnterpriseToolError::new( "VECTOR_CONTENT_CHECKSUM_MISMATCH", @@ -1671,6 +2167,7 @@ pub fn read_semantic_jsonl(path: &Path, semantic_batch_id: &str) -> Result) -> Result, + line_number: usize, +) -> Result<(), EnterpriseToolError> { + for (field, max_bytes) in SEMANTIC_VARCHAR_LIMITS { + let value = object.get(*field).and_then(Value::as_str).ok_or_else(|| { + EnterpriseToolError::new( + "VECTOR_JSONL_FIELD_INVALID", + format!("JSONL 第 {line_number} 行的 {field} 必须是字符串。"), + ) + })?; + if value.len() > *max_bytes { + return Err(EnterpriseToolError::new( + "VECTOR_VARCHAR_LENGTH_EXCEEDED", + format!( + "JSONL 第 {line_number} 行的 {field} 为 {} 个 UTF-8 字节，超过 semantic_cards VARCHAR({max_bytes})。", + value.len() + ), + )); + } + } + Ok(()) +} + pub fn milvus_search_query( database: &str, collection: &str, @@ -1773,17 +2294,6 @@ pub fn validate_existing_card_ownership( Ok(()) } -pub fn milvus_delete_batch_query(database: &str, collection: &str, semantic_batch_id: &str) -> String { - format!( - "POST /v2/vectordb/entities/delete\n{}", - json!({ - "dbName": if database.is_empty() { "default" } else { database }, - "collectionName": collection, - "filter": format!("semantic_batch_id == {}", json_string(semantic_batch_id)), - }) - ) -} - pub fn query_result_rows(result: dbx_core::db::QueryResult) -> Vec { result .rows @@ -1976,7 +2486,7 @@ mod tests { - 报告说明 + 报告说明 商家备注商家备注金额 AB10 @@ -2027,6 +2537,89 @@ mod tests { } } + #[test] + fn xlsx_archive_preflight_enforces_entry_and_uncompressed_budgets() { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("budget.xlsx"); + write_governed_test_xlsx(&source); + let permissive = XlsxArchiveLimits { + entry_count: 16, + total_uncompressed_bytes: 1024 * 1024, + metadata_entry_bytes: 1024 * 1024, + shared_strings_bytes: 1024 * 1024, + worksheet_bytes: 1024 * 1024, + }; + validate_xlsx_archive_limits(&source, permissive).unwrap(); + assert_eq!( + validate_xlsx_archive_limits(&source, XlsxArchiveLimits { entry_count: 4, ..permissive }).unwrap_err().code, + "IMPORT_XLSX_ZIP_ENTRY_LIMIT_EXCEEDED" + ); + assert_eq!( + validate_xlsx_archive_limits(&source, XlsxArchiveLimits { worksheet_bytes: 16, ..permissive }) + .unwrap_err() + .code, + "IMPORT_XLSX_WORKSHEET_BUDGET_EXCEEDED" + ); + + let shared = directory.path().join("shared.xlsx"); + let file = std::fs::File::create(&shared).unwrap(); + let mut zip = zip::ZipWriter::new(file); + write_xlsx_entry(&mut zip, "xl/sharedStrings.xml", &"x".repeat(128)); + zip.finish().unwrap(); + assert_eq!( + validate_xlsx_archive_limits(&shared, XlsxArchiveLimits { shared_strings_bytes: 64, ..permissive },) + .unwrap_err() + .code, + "IMPORT_XLSX_SHARED_STRINGS_BUDGET_EXCEEDED" + ); + } + + #[test] + fn normalized_writer_and_dynamic_batch_are_bounded() { + let directory = tempfile::tempdir().unwrap(); + let output = directory.path().join("limited.csv"); + let mut writer = LimitedWriter::create(&output, 4).unwrap(); + writer.write_all(b"1234").unwrap(); + assert!(writer.write_all(b"5").is_err()); + + let safe = governed_import_batch_size(None, 145).unwrap(); + assert!(safe < 1_000); + assert_eq!( + governed_import_batch_size(Some(safe + 1), 145).unwrap_err().code, + "IMPORT_BATCH_SIZE_RESOURCE_LIMIT" + ); + } + + #[tokio::test] + async fn private_snapshot_copy_is_hash_verified_and_cancellable() { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("source.csv"); + let destination = directory.path().join("snapshot.csv"); + let bytes = b"id,name\n1,Ada\n"; + std::fs::write(&source, bytes).unwrap(); + copy_verified_import_source( + source.clone(), + destination.clone(), + sha256_bytes(bytes), + Arc::new(AtomicBool::new(false)), + ) + .await + .unwrap(); + assert_eq!(std::fs::read(&destination).unwrap(), bytes); + + std::fs::remove_file(&destination).unwrap(); + let error = copy_verified_import_source( + source, + destination.clone(), + sha256_bytes(bytes), + Arc::new(AtomicBool::new(true)), + ) + .await + .unwrap_err(); + assert_eq!(error.code, "IMPORT_CANCELLED"); + assert!(!destination.exists()); + } + #[test] fn recursive_preview_sanitizer_bounds_nested_values_and_keys() { let long_key = "键".repeat(20); @@ -2178,6 +2771,54 @@ mod tests { std::fs::write(&path, format!("{}\n", forbidden)).unwrap(); assert_eq!(read_semantic_jsonl(&path, "semantic-batch-1").unwrap_err().code, "VECTOR_JSONL_UNKNOWN_FIELD"); } + + for (field, max_bytes) in [ + ("card_id", 128usize), + ("business_domain", 128), + ("dataset_id", 128), + ("template_version", 64), + ("title", 512), + ("content", 8_192), + ("source_uri", 2_048), + ("semantic_version", 128), + ("semantic_batch_id", 128), + ("embedding_model", 128), + ("embedding_revision", 64), + ] { + let mut boundary = record.clone(); + let value = "x".repeat(max_bytes); + boundary.as_object_mut().unwrap().insert(field.to_string(), json!(value.clone())); + if field == "content" { + boundary + .as_object_mut() + .unwrap() + .insert("content_checksum".to_string(), json!(sha256_bytes(value.as_bytes()))); + } + let batch = if field == "semantic_batch_id" { value.as_str() } else { "semantic-batch-1" }; + std::fs::write(&path, format!("{}\n", boundary)).unwrap(); + read_semantic_jsonl(&path, batch).unwrap(); + + let mut overflow = record.clone(); + let overflow_value = "x".repeat(max_bytes + 1); + overflow.as_object_mut().unwrap().insert(field.to_string(), json!(overflow_value.clone())); + if field == "content" { + overflow + .as_object_mut() + .unwrap() + .insert("content_checksum".to_string(), json!(sha256_bytes(overflow_value.as_bytes()))); + } + let overflow_batch = + if field == "semantic_batch_id" { overflow_value.as_str() } else { "semantic-batch-1" }; + std::fs::write(&path, format!("{}\n", overflow)).unwrap(); + let error = read_semantic_jsonl(&path, overflow_batch).unwrap_err(); + assert!(matches!(error.code, "VECTOR_VARCHAR_LENGTH_EXCEEDED" | "SEMANTIC_BATCH_ID_INVALID")); + } + + let mut multibyte_title = record.clone(); + multibyte_title.as_object_mut().unwrap().insert("card_id".to_string(), json!("metric-order-gmv-2")); + multibyte_title.as_object_mut().unwrap().insert("title".to_string(), json!("中".repeat(171))); + std::fs::write(&path, format!("{}\n{}\n", record, multibyte_title)).unwrap(); + assert_eq!(read_semantic_jsonl(&path, "semantic-batch-1").unwrap_err().code, "VECTOR_VARCHAR_LENGTH_EXCEEDED"); } #[test] @@ -2240,6 +2881,15 @@ mod tests { assert!(runtime.cancel_job(&import_id).await.unwrap().1); } + #[tokio::test] + async fn inspection_semaphore_fails_closed_at_capacity() { + let runtime = EnterpriseRuntime::with_inspection_concurrency(1); + let permit = runtime.try_inspection_permit().await.unwrap(); + assert_eq!(runtime.try_inspection_permit().await.unwrap_err().code, "IMPORT_INSPECTION_CONCURRENCY_LIMIT"); + drop(permit); + let _permit = runtime.try_inspection_permit().await.unwrap(); + } + #[tokio::test] async fn governed_csv_snapshot_preserves_duplicate_position_and_real_row_number() { let directory = tempfile::tempdir().unwrap(); @@ -2256,6 +2906,7 @@ mod tests { source_ref: Some("a".repeat(64)), source_format: Some(TableImportSourceFormat::Csv), parse_options: TableImportParseOptions { + encoding: Some(TableImportTextEncoding::Utf8), title_row: Some(2), data_start_row: Some(3), has_header: Some(true), @@ -2303,6 +2954,77 @@ mod tests { assert_eq!(row.get(6).map(str::len), Some(64)); } + #[tokio::test] + async fn governed_csv_distinguishes_short_and_trailing_empty_rows_and_blocks_nonempty_overflow() { + let directory = tempfile::tempdir().unwrap(); + let source = directory.path().join("width.csv"); + let output = directory.path().join("normalized.csv"); + std::fs::write(&source, "a,b\nshort\nexact,value\ntail,empty,\n").unwrap(); + let request = TableImportRequest { + import_id: "import-width".to_string(), + connection_id: "postgres-1".to_string(), + database: "enterprise".to_string(), + schema: "staging".to_string(), + table: "mcp_width".to_string(), + file_path: source.to_string_lossy().to_string(), + source_ref: Some("d".repeat(64)), + source_format: Some(TableImportSourceFormat::Csv), + parse_options: TableImportParseOptions { + encoding: Some(TableImportTextEncoding::Utf8), + ..Default::default() + }, + mappings: vec![ + TableImportColumnMapping { + source_column: "a".to_string(), + target_column: "a".to_string(), + target_data_type: Some("TEXT".to_string()), + }, + TableImportColumnMapping { + source_column: "b".to_string(), + target_column: "b".to_string(), + target_data_type: Some("TEXT".to_string()), + }, + ], + mode: TableImportMode::Append, + create_table: true, + batch_size: 10, + date_time_format: None, + prepared_source: None, + retain_source: true, + }; + build_governed_import_snapshot( + request.clone(), + "plan-width", + &"d".repeat(64), + &output, + Arc::new(AtomicBool::new(false)), + ) + .await + .unwrap(); + let mut reader = csv::Reader::from_path(&output).unwrap(); + let rows = reader.records().map(Result::unwrap).collect::>(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].get(5), Some("2")); + assert_eq!(rows[2].get(5), Some("4")); + assert_ne!(rows[0].get(6), rows[2].get(6)); + assert!(rows.iter().all(|row| row.get(6).is_some_and(|hash| hash.len() == 64))); + + std::fs::write(&source, "a,b\nwide,value,unexpected\n").unwrap(); + let overflow_output = directory.path().join("overflow.csv"); + let error = build_governed_import_snapshot( + request, + "plan-overflow", + &"d".repeat(64), + &overflow_output, + Arc::new(AtomicBool::new(false)), + ) + .await + .unwrap_err(); + assert_eq!(error.code, "IMPORT_SOURCE_ROW_WIDER_THAN_HEADER"); + assert!(error.message.contains("第 2 行")); + assert!(!overflow_output.exists()); + } + #[tokio::test] async fn governed_snapshot_blocks_legacy_xls_before_database_write() { let directory = tempfile::tempdir().unwrap(); @@ -2365,6 +3087,18 @@ mod tests { .code, "IMPORT_GOVERNED_FORMAT_UNSUPPORTED_V1" ); + for encoding in [TableImportTextEncoding::Auto, TableImportTextEncoding::Gbk] { + assert_eq!( + validate_governed_source_v1( + "input.csv", + Some(TableImportSourceFormat::Csv), + &TableImportParseOptions { encoding: Some(encoding), ..Default::default() }, + ) + .unwrap_err() + .code, + "IMPORT_GOVERNED_EXPLICIT_UTF8_REQUIRED_V1" + ); + } } #[tokio::test] @@ -2476,6 +3210,7 @@ mod tests { let source = PathBuf::from(std::env::var("DBX_MCP_REAL_XLSX_FIXTURE").expect("fixture path")); let directory = tempfile::tempdir().unwrap(); let output = directory.path().join("normalized.csv"); + validate_source_archive_limits(&source, Some(TableImportSourceFormat::Excel)).unwrap(); let identity = file_identity(source.clone()).await.unwrap(); let parse_options = TableImportParseOptions::default(); let preview = dbx_core::table_import::preview_table_import_file_with_request( @@ -2490,6 +3225,7 @@ mod tests { .await .unwrap(); validate_preview_headers(&preview).unwrap(); + assert_eq!(preview.columns.len(), 145); let source_columns = source_columns_for_preview( &source.to_string_lossy(), Some(TableImportSourceFormat::Excel), @@ -2498,9 +3234,15 @@ mod tests { ) .await .unwrap(); + assert_eq!(source_columns.len(), 145); + let merchant_notes = source_columns + .iter() + .filter(|column| column.raw_source_name == "商家备注") + .map(|column| column.canonical_source_name.as_str()) + .collect::>(); + assert_eq!(merchant_notes, vec!["商家备注__1", "商家备注__2"]); let requested = source_columns .iter() - .take(2) .enumerate() .map(|(index, column)| McpImportColumnMapping { source_position: column.source_position, @@ -2510,6 +3252,7 @@ mod tests { }) .collect::>(); let mappings = validate_mappings(&requested, &source_columns).unwrap(); + let batch_size = governed_import_batch_size(None, source_columns.len()).unwrap(); let plan = build_plan( "fixture-postgres".to_string(), "fixture".to_string(), @@ -2523,12 +3266,12 @@ mod tests { parse_options, mappings, true, - 500, + batch_size, None, ) .unwrap(); let request = plan.to_import_request("fixture-import".to_string()); - build_governed_import_snapshot( + let governed = build_governed_import_snapshot( request, &plan.plan_id, &identity.sha256, @@ -2537,11 +3280,43 @@ mod tests { ) .await .unwrap(); + assert!(governed.mappings.iter().all(|mapping| mapping.target_data_type.as_deref() == Some("TEXT"))); let mut reader = csv::Reader::from_path(output).unwrap(); let headers = reader.headers().unwrap().clone(); + assert_eq!(headers.len(), source_columns.len() + RESERVED_STAGING_COLUMNS.len()); assert!(headers.iter().any(|header| header == "source_row_number")); assert!(headers.iter().any(|header| header == "source_row_hash")); - let row_count = reader.records().filter_map(Result::ok).count(); + let mut non_empty_columns = vec![false; source_columns.len()]; + let mut previous_source_row = 0usize; + let mut row_count = 0usize; + for record in reader.records() { + let record = record.unwrap(); + assert_eq!(record.len(), source_columns.len() + RESERVED_STAGING_COLUMNS.len()); + for (index, value) in record.iter().take(source_columns.len()).enumerate() { + let normalized = value.trim(); + non_empty_columns[index] |= !normalized.is_empty() && !matches!(normalized, "-" | "--"); + } + let source_row = record.get(source_columns.len() + 3).unwrap().parse::().unwrap(); + assert!(source_row > previous_source_row); + previous_source_row = source_row; + let row_hash = record.get(source_columns.len() + 4).unwrap(); + assert!(is_sha256_hex(row_hash)); + row_count += 1; + } assert!(row_count > 0); + let empty_positions = non_empty_columns + .iter() + .enumerate() + .filter_map(|(index, non_empty)| (!*non_empty).then_some(index + 1)) + .collect::>(); + assert_eq!( + empty_positions, + vec![ + 2, 3, 8, 9, 13, 14, 20, 22, 25, 26, 28, 29, 30, 31, 36, 37, 38, 39, 40, 41, 42, 46, 47, 48, 53, 54, 58, + 59, 60, 61, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, + 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, + 134, 135, 136, 137, 140, 141, 145, + ] + ); } } diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index f35305054d..4dcef15197 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -12,13 +12,13 @@ use uuid::Uuid; use crate::backend::{format_query_result, new_connection_config, parse_database_type, ConnectionSummary, DbxBackend}; use crate::enterprise_tools::{ build_milvus_filter, build_plan, cell_char_limit, file_identity, generated_staging_relation, - milvus_delete_batch_query, milvus_existing_cards_query, milvus_search_query, milvus_upsert_query, preview_limit, - query_result_rows, read_semantic_jsonl, revalidate_plan_file, sanitize_preview, source_columns_for_preview, - structure_fingerprint, validate_embedding, validate_existing_card_ownership, validate_governed_source_v1, - validate_import_file, validate_mappings, validate_preview_headers, validate_vector_collection, - vector_output_fields, vector_top_k, vector_upsert_batch_size, EnterpriseRuntime, EnterpriseToolError, - ImportStatusRequest, PrepareTableImportRequest, PreviewImportFileRequest, StartTableImportRequest, - VectorDeleteByBatchRequest, VectorSearchRequest, VectorUpsertFileRequest, FORMAT_VERSION, + governed_import_batch_size, inspect_import_source, inspection_timeout, milvus_existing_cards_query, + milvus_search_query, milvus_upsert_query, preview_limit, query_result_rows, read_semantic_jsonl, + revalidate_plan_file, sanitize_preview, structure_fingerprint, validate_embedding, + validate_existing_card_ownership, validate_governed_source_v1, validate_import_file, validate_mappings, + validate_vector_collection, vector_output_fields, vector_top_k, vector_upsert_batch_size, EnterpriseRuntime, + EnterpriseToolError, ImportStatusRequest, PrepareTableImportRequest, PreviewImportFileRequest, + StartTableImportRequest, VectorDeleteByBatchRequest, VectorSearchRequest, VectorUpsertFileRequest, FORMAT_VERSION, }; use crate::mongo::{self, MongoCommand, MongoSafetyError}; use crate::session::{McpSession, McpSessionStore}; @@ -205,6 +205,7 @@ pub struct DbxMcpServer { sessions: Arc, enterprise: Arc, web_mode: bool, + connection_management_enabled: bool, tool_router: ToolRouter, } @@ -259,8 +260,22 @@ impl DbxMcpServer { } pub fn with_runtime_options(backend: Arc, scope: McpScope, web_mode: bool) -> Self { + let connection_management_requested = std::env::var("DBX_MCP_ENABLE_CONNECTION_MANAGEMENT") + .ok() + .is_some_and(|value| matches!(value.trim().to_ascii_lowercase().as_str(), "1" | "true")); + Self::with_runtime_options_and_connection_management(backend, scope, web_mode, connection_management_requested) + } + + /// 构造带显式连接维护权限的服务器；任何作用域仍会强制关闭连接管理。 + pub fn with_runtime_options_and_connection_management( + backend: Arc, + scope: McpScope, + web_mode: bool, + connection_management_requested: bool, + ) -> Self { let mut tool_router = Self::tool_router(); - if scope.enabled() { + let connection_management_enabled = connection_management_requested && !scope.enabled(); + if !connection_management_enabled { tool_router.disable_route("dbx_add_connection"); tool_router.disable_route("dbx_duplicate_connection"); tool_router.disable_route("dbx_remove_connection"); @@ -276,10 +291,28 @@ impl DbxMcpServer { sessions: McpSessionStore::new(), enterprise: Arc::new(EnterpriseRuntime::default()), web_mode, + connection_management_enabled, tool_router, } } + #[allow(clippy::result_large_err)] + fn connection_management_guard(&self) -> Result<(), CallToolResult> { + if self.scope.enabled() { + return Err(tool_error( + "CONNECTION_MANAGEMENT_DISABLED_IN_SCOPED_MODE", + "设置任何 DBX_MCP_SCOPE_* 后，连接新增、复制和删除都会被禁用。", + )); + } + if !self.connection_management_enabled { + return Err(tool_error( + "CONNECTION_MANAGEMENT_DISABLED", + "连接管理默认关闭；仅安装维护时可临时设置 DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1。", + )); + } + Ok(()) + } + async fn close_backend_sessions_best_effort(&self, sessions: Vec) { for session in sessions { let _ = self @@ -375,7 +408,14 @@ impl DbxMcpServer { #[tool( name = "dbx_preview_import_file", - description = "Preview an allowed local Excel/CSV/TSV/JSON file without writing data. Returns bounded rows, sheets, hashes, columns, row counts and a structure fingerprint. Local DBX Desktop/MCP mode only." + description = "Preview an allowed local Excel/CSV/TSV/JSON file without writing data. Returns bounded rows, sheets, hashes, columns, row counts and a structure fingerprint. Local DBX Desktop/MCP mode only.", + annotations( + title = "预览本地导入文件", + read_only_hint = true, + destructive_hint = false, + idempotent_hint = true, + open_world_hint = false + ) )] async fn preview_import_file(&self, Parameters(request): Parameters) -> CallToolResult { if self.web_mode { @@ -388,10 +428,6 @@ impl DbxMcpServer { Ok(path) => path, Err(error) => return enterprise_error(error), }; - let identity = match file_identity(path.clone()).await { - Ok(identity) => identity, - Err(error) => return enterprise_error(error), - }; let parse_options = match request.parse_options.into_core() { Ok(options) => options, Err(error) => return enterprise_error(error), @@ -405,50 +441,46 @@ impl DbxMcpServer { Err(error) => return enterprise_error(error), }; let source_format = request.source_format.map(Into::into); - let file_path = path.to_string_lossy().to_string(); - let preview = dbx_core::table_import::preview_table_import_file_with_request( - dbx_core::table_import::TableImportPreviewRequest { - file_path: file_path.clone(), - source_ref: Some(identity.sha256.clone()), - source_format, - parse_options: parse_options.clone(), - preview_limit: Some(preview_rows), - }, - ) - .await; - let preview = match preview { - Ok(preview) => preview, - Err(error) => { - return enterprise_error(EnterpriseToolError::new("IMPORT_PREVIEW_FAILED", error)); + let inspection_permit = match self.enterprise.try_inspection_permit().await { + Ok(permit) => permit, + Err(error) => return enterprise_error(error), + }; + let inspection_parse_options = parse_options.clone(); + let mut inspection_task = tokio::spawn(async move { + let _inspection_permit = inspection_permit; + inspect_import_source(path, source_format, inspection_parse_options, preview_rows).await + }); + let inspection = match tokio::time::timeout(inspection_timeout(), &mut inspection_task).await { + Ok(Ok(Ok(inspection))) => inspection, + Ok(Ok(Err(error))) => return enterprise_error(error), + Ok(Err(error)) => { + return enterprise_error(EnterpriseToolError::new("IMPORT_INSPECTION_FAILED", error.to_string())) + } + Err(_) => { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_INSPECTION_TIMEOUT", + "文件剖析超过时间上限；后台检查会继续占用受控槽位直到安全结束。", + )) } }; - if let Err(error) = validate_preview_headers(&preview) { - return enterprise_error(error); - } - let source_columns = - match source_columns_for_preview(&file_path, source_format, &parse_options, &preview.columns).await { - Ok(columns) => columns, - Err(error) => return enterprise_error(error), - }; + let identity = inspection.identity; + let mut preview = inspection.preview; + let source_columns = inspection.source_columns; let fingerprint = structure_fingerprint(&preview, &parse_options, &source_columns); - let current_identity = match file_identity(path).await { - Ok(identity) => identity, - Err(error) => return enterprise_error(error), - }; - if current_identity != identity { - return enterprise_error(EnterpriseToolError::new( - "IMPORT_FILE_CHANGED_DURING_PREVIEW", - "文件在预览期间发生变化；结果已丢弃。", - )); - } let used_first_row = preview.source_row_numbers.first().copied(); let used_last_row = preview.source_row_numbers.last().copied(); - let preview = sanitize_preview(preview, char_limit); + let used_range_exact = preview.total_rows_exact; + preview = sanitize_preview(preview, char_limit); + let row_summary = if preview.total_rows_exact { + format!("精确共 {} 行", preview.total_rows) + } else { + format!("受限扫描到 {} 行，文件总行数尚未精确统计", preview.total_rows) + }; let summary = format!( - "已只读剖析 {}：{} 列，{} 行，返回 {} 行预览。", + "已只读剖析 {}：{} 列，{}，返回 {} 行预览。", preview.file_name, preview.columns.len(), - preview.total_rows, + row_summary, preview.rows.len() ); structured_success( @@ -462,7 +494,10 @@ impl DbxMcpServer { "lastPreviewDataRow": used_last_row, "firstColumn": if preview.columns.is_empty() { 0 } else { 1 }, "lastColumn": preview.columns.len(), + "exact": used_range_exact, }, + "usedRangeExact": used_range_exact, + "totalRowsExact": preview.total_rows_exact, "preview": preview, }), ) @@ -470,7 +505,14 @@ impl DbxMcpServer { #[tool( name = "dbx_prepare_table_import", - description = "Validate a caller-approved template mapping and immutable local source, then generate a unique PostgreSQL staging relation and a single-use plan valid for 30 minutes. This tool performs no database query or write." + description = "Validate a caller-approved template mapping and immutable local source, then generate a unique PostgreSQL staging relation and a single-use plan valid for 30 minutes. This tool performs no database query or write.", + annotations( + title = "准备隔离导入计划", + read_only_hint = false, + destructive_hint = false, + idempotent_hint = false, + open_world_hint = false + ) )] async fn prepare_table_import(&self, Parameters(request): Parameters) -> CallToolResult { if self.web_mode { @@ -497,10 +539,6 @@ impl DbxMcpServer { Ok(path) => path, Err(error) => return enterprise_error(error), }; - let identity = match file_identity(path.clone()).await { - Ok(identity) => identity, - Err(error) => return enterprise_error(error), - }; let parse_options = match request.parse_options.into_core() { Ok(options) => options, Err(error) => return enterprise_error(error), @@ -509,40 +547,31 @@ impl DbxMcpServer { if let Err(error) = validate_governed_source_v1(&path.to_string_lossy(), source_format, &parse_options) { return enterprise_error(error); } - let preview = dbx_core::table_import::preview_table_import_file_with_request( - dbx_core::table_import::TableImportPreviewRequest { - file_path: path.to_string_lossy().to_string(), - source_ref: Some(identity.sha256.clone()), - source_format, - parse_options: parse_options.clone(), - preview_limit: Some(1), - }, - ) - .await; - let preview = match preview { - Ok(preview) => preview, - Err(error) => return enterprise_error(EnterpriseToolError::new("IMPORT_PREVIEW_FAILED", error)), - }; - if let Err(error) = validate_preview_headers(&preview) { - return enterprise_error(error); - } - let source_columns = - match source_columns_for_preview(&path.to_string_lossy(), source_format, &parse_options, &preview.columns) - .await - { - Ok(columns) => columns, - Err(error) => return enterprise_error(error), - }; - let current_identity = match file_identity(path).await { - Ok(identity) => identity, + let inspection_permit = match self.enterprise.try_inspection_permit().await { + Ok(permit) => permit, Err(error) => return enterprise_error(error), }; - if current_identity != identity { - return enterprise_error(EnterpriseToolError::new( - "IMPORT_FILE_CHANGED_DURING_PREPARE", - "文件在 prepare 剖析期间发生变化；未创建计划。", - )); - } + let inspection_parse_options = parse_options.clone(); + let mut inspection_task = tokio::spawn(async move { + let _inspection_permit = inspection_permit; + inspect_import_source(path, source_format, inspection_parse_options, 1).await + }); + let inspection = match tokio::time::timeout(inspection_timeout(), &mut inspection_task).await { + Ok(Ok(Ok(inspection))) => inspection, + Ok(Ok(Err(error))) => return enterprise_error(error), + Ok(Err(error)) => { + return enterprise_error(EnterpriseToolError::new("IMPORT_INSPECTION_FAILED", error.to_string())) + } + Err(_) => { + return enterprise_error(EnterpriseToolError::new( + "IMPORT_INSPECTION_TIMEOUT", + "prepare 剖析超过时间上限；后台检查会继续占用受控槽位直到安全结束。", + )) + } + }; + let identity = inspection.identity; + let preview = inspection.preview; + let source_columns = inspection.source_columns; let mappings = match validate_mappings(&request.mappings, &source_columns) { Ok(mappings) => mappings, Err(error) => return enterprise_error(error), @@ -551,13 +580,10 @@ impl DbxMcpServer { Ok(relation) => relation, Err(error) => return enterprise_error(error), }; - let batch_size = request.batch_size.unwrap_or(1_000); - if !(1..=50_000).contains(&batch_size) { - return enterprise_error(EnterpriseToolError::new( - "IMPORT_BATCH_SIZE_INVALID", - "batch_size 必须在 1 到 50000 之间。", - )); - } + let batch_size = match governed_import_batch_size(request.batch_size, source_columns.len()) { + Ok(value) => value, + Err(error) => return enterprise_error(error), + }; let fingerprint = structure_fingerprint(&preview, &parse_options, &source_columns); let plan = match build_plan( resolved.connection.id.clone(), @@ -589,7 +615,14 @@ impl DbxMcpServer { #[tool( name = "dbx_start_table_import", - description = "Start a previously prepared single-use PostgreSQL staging import. The file, target, parsing options and mappings are revalidated and cannot be replaced." + description = "Start a previously prepared single-use PostgreSQL staging import. The file, target, parsing options and mappings are revalidated and cannot be replaced.", + annotations( + title = "启动隔离表导入", + read_only_hint = false, + destructive_hint = false, + idempotent_hint = false, + open_world_hint = false + ) )] async fn start_table_import(&self, Parameters(request): Parameters) -> CallToolResult { if self.web_mode { @@ -643,9 +676,16 @@ impl DbxMcpServer { } progress_job.snapshot.lock().unwrap_or_else(|error| error.into_inner()).apply_progress(update); }); - let result = backend - .import_table_file_for_mcp(&connection, import_request, plan_id, task_job.cancelled.clone(), progress) - .await; + let worker_cancelled = task_job.cancelled.clone(); + let worker = tokio::spawn(async move { + backend + .import_table_file_for_mcp(&connection, import_request, plan_id, worker_cancelled, progress) + .await + }); + let result = match worker.await { + Ok(result) => result, + Err(error) => Err(format!("IMPORT_TASK_PANICKED: 导入后台任务异常结束：{error}")), + }; let mut snapshot = task_job.snapshot.lock().unwrap_or_else(|error| error.into_inner()); match result { Ok(summary) => { @@ -683,7 +723,14 @@ impl DbxMcpServer { #[tool( name = "dbx_get_import_status", - description = "Get machine-readable progress, row counts, bytes, elapsed time, errors and final summary for an MCP table import." + description = "Get machine-readable progress, row counts, bytes, elapsed time, errors and final summary for an MCP table import.", + annotations( + title = "查询导入状态", + read_only_hint = true, + destructive_hint = false, + idempotent_hint = true, + open_world_hint = false + ) )] async fn get_import_status(&self, Parameters(request): Parameters) -> CallToolResult { if self.web_mode { @@ -704,7 +751,14 @@ impl DbxMcpServer { #[tool( name = "dbx_cancel_import", - description = "Request cancellation of an MCP staging import. Cancellation stops subsequent batches and never publishes staging data to core or mart." + description = "Request cancellation of an MCP staging import. Cancellation stops subsequent batches and never publishes staging data to core or mart.", + annotations( + title = "取消隔离导入", + read_only_hint = false, + destructive_hint = true, + idempotent_hint = true, + open_world_hint = false + ) )] async fn cancel_import(&self, Parameters(request): Parameters) -> CallToolResult { if self.web_mode { @@ -734,7 +788,14 @@ impl DbxMcpServer { #[tool( name = "dbx_vector_search", - description = "Search an allowed Milvus semantic collection with a bounded vector and safe equality filters. Approval and effective date are always enforced; semantic_version is an optional exact filter." + description = "Search an allowed Milvus semantic collection with a bounded vector and safe equality filters. Approval and effective date are always enforced; semantic_version is an optional exact filter.", + annotations( + title = "检索获批语义卡", + read_only_hint = true, + destructive_hint = false, + idempotent_hint = true, + open_world_hint = false + ) )] async fn vector_search(&self, Parameters(request): Parameters) -> CallToolResult { let resolved = match self.resolve_connection(&request.selector).await { @@ -793,7 +854,14 @@ impl DbxMcpServer { #[tool( name = "dbx_vector_upsert_file", - description = "Upsert approved semantic cards from one allowed JSONL file into an allowed Milvus collection. Requires MCP safe-write permission." + description = "Upsert approved semantic cards from one allowed JSONL file into an allowed Milvus collection. Requires MCP safe-write permission.", + annotations( + title = "写入获批语义卡", + read_only_hint = false, + destructive_hint = true, + idempotent_hint = true, + open_world_hint = false + ) )] async fn vector_upsert_file(&self, Parameters(request): Parameters) -> CallToolResult { let resolved = match self.resolve_connection(&request.selector).await { @@ -890,7 +958,9 @@ impl DbxMcpServer { { return enterprise_error(EnterpriseToolError::new( "VECTOR_UPSERT_FAILED", - format!("已提交 {upserted} 条后失败；语义版本尚未激活，可按 semantic_batch_id 撤销：{error}"), + format!( + "已提交 {upserted} 条后失败；语义版本尚未激活。v1 禁用 MCP 删除，请由管理员按 semantic_batch_id 执行受审计恢复：{error}" + ), )); } upserted += chunk.len(); @@ -908,66 +978,23 @@ impl DbxMcpServer { #[tool( name = "dbx_vector_delete_by_batch", - description = "Delete one unpublished Milvus semantic batch by exact semantic_batch_id. Requires high-risk MCP permission; arbitrary filter expressions are never accepted." + description = "Disabled in v1 because publication state cannot be authoritatively proven by the MCP server. Always returns VECTOR_DELETE_DISABLED_V1 and performs no connection or Milvus call.", + annotations( + title = "语义批次删除已禁用", + read_only_hint = true, + destructive_hint = false, + idempotent_hint = true, + open_world_hint = false + ) )] async fn vector_delete_by_batch( &self, - Parameters(request): Parameters, + Parameters(_request): Parameters, ) -> CallToolResult { - let resolved = match self.resolve_connection(&request.selector).await { - Ok(resolved) => resolved, - Err(error) => return error, - }; - if resolved.connection.db_type != DatabaseType::Milvus { - return enterprise_error(EnterpriseToolError::new( - "VECTOR_REQUIRES_MILVUS", - "dbx_vector_delete_by_batch 只接受 Milvus 连接。", - )); - } - let database = match self.resolve_database(request.database, &resolved.connection) { - Ok(database) => database, - Err(error) => return error, - }; - if let Err(error) = validate_safe_write_connection(&resolved.connection, &resolved.policy, &database) { - return error; - } - if request.published { - return enterprise_error(EnterpriseToolError::new( - "VECTOR_PUBLISHED_BATCH_DELETE_BLOCKED", - "published 必须由调用方显式证明为 false；已发布语义批次不能通过该工具删除。", - )); - } - if !resolved.policy.allow_dangerous_sql { - return enterprise_error(EnterpriseToolError::new( - "VECTOR_HIGH_RISK_WRITE_BLOCKED", - "按批删除语义向量需要开启 DBX MCP 高风险写入权限。", - )); - } - if let Err(error) = validate_vector_collection(&request.collection) { - return enterprise_error(error); - } - let _semantic_write_permit = match self.enterprise.semantic_write_permit().await { - Ok(permit) => permit, - Err(error) => return enterprise_error(error), - }; - if request.semantic_batch_id.trim().is_empty() || request.semantic_batch_id.len() > 200 { - return enterprise_error(EnterpriseToolError::new( - "SEMANTIC_BATCH_ID_INVALID", - "semantic_batch_id 必须是 1～200 个字符。", - )); - } - let query = milvus_delete_batch_query(&database, &request.collection, &request.semantic_batch_id); - match self.backend.execute_query(&resolved.connection, &database, &query, Some(1), Some(30)).await { - Ok(_) => structured_success( - format!("已删除未发布语义批次 {}。", request.semantic_batch_id), - json!({ - "collection": request.collection, - "semanticBatchId": request.semantic_batch_id, - "deleted": true, - }), - ), - Err(error) => enterprise_error(EnterpriseToolError::new("VECTOR_DELETE_FAILED", error)), - } + enterprise_error(EnterpriseToolError::new( + "VECTOR_DELETE_DISABLED_V1", + "v1 禁用 MCP 语义批次删除：服务端无法权威证明该批次尚未发布。请使用受审计的管理员恢复流程。", + )) } #[tool( @@ -993,7 +1020,7 @@ impl DbxMcpServer { { return enterprise_error(EnterpriseToolError::new( "VECTOR_DEDICATED_TOOL_REQUIRED", - "向量实体检索、写入和删除只能使用 dbx_vector_search、dbx_vector_upsert_file 或 dbx_vector_delete_by_batch；通用查询仅允许明确的管理诊断端点。", + "向量实体检索和写入只能使用 dbx_vector_search 与 dbx_vector_upsert_file；v1 禁止 MCP 删除，通用查询仅允许明确的管理诊断端点。", )); } // Resolve the session before the database so its connection/database @@ -1259,6 +1286,9 @@ impl DbxMcpServer { #[tool(name = "dbx_add_connection", description = "Add a new database connection to DBX")] async fn add_connection(&self, Parameters(request): Parameters) -> CallToolResult { + if let Err(error) = self.connection_management_guard() { + return error; + } let policy = match self.load_policy().await { Ok(policy) => policy, Err(error) => return error, @@ -1313,6 +1343,9 @@ impl DbxMcpServer { &self, Parameters(request): Parameters, ) -> CallToolResult { + if let Err(error) = self.connection_management_guard() { + return error; + } let policy = match self.load_policy().await { Ok(policy) => policy, Err(error) => return error, @@ -1368,6 +1401,9 @@ impl DbxMcpServer { #[tool(name = "dbx_remove_connection", description = "Remove a database connection from DBX")] async fn remove_connection(&self, Parameters(request): Parameters) -> CallToolResult { + if let Err(error) = self.connection_management_guard() { + return error; + } let policy = match self.load_policy().await { Ok(policy) => policy, Err(error) => return error, @@ -1382,10 +1418,15 @@ impl DbxMcpServer { Ok(connections) => connections, Err(error) => return tool_error("CONNECTION_LOAD_ERROR", error), }; + let allowed = connections + .iter() + .filter(|connection| policy_allows_connection(&policy, connection)) + .cloned() + .collect::>(); let target = if let Some(id) = request.connection_id.as_deref().map(str::trim).filter(|id| !id.is_empty()) { - connections.iter().find(|connection| connection.id == id).cloned() + allowed.iter().find(|connection| connection.id == id).cloned() } else { - let matching = connections + let matching = allowed .iter() .filter(|connection| connection.name.eq_ignore_ascii_case(&request.connection_name)) .cloned() @@ -2061,6 +2102,7 @@ mod tests { struct FakeBackend { connections: Vec, + policy: McpGlobalPolicy, recorded_arguments: std::sync::Mutex>, closed_sessions: std::sync::Mutex>, pinned_sessions: std::sync::Mutex>, @@ -2071,6 +2113,7 @@ mod tests { fn default() -> Self { Self { connections: Vec::new(), + policy: McpGlobalPolicy::default(), recorded_arguments: std::sync::Mutex::new(Vec::new()), closed_sessions: std::sync::Mutex::new(Vec::new()), pinned_sessions: std::sync::Mutex::new(HashSet::new()), @@ -2109,7 +2152,7 @@ mod tests { #[async_trait] impl DbxBackend for FakeBackend { async fn load_mcp_global_policy(&self) -> Result { - Ok(McpGlobalPolicy::default()) + Ok(self.policy.clone()) } async fn load_connections(&self) -> Result, String> { @@ -2238,7 +2281,7 @@ mod tests { } #[tokio::test] - async fn vector_batch_delete_requires_explicit_unpublished_attestation() { + async fn vector_batch_delete_is_fail_closed_in_v1() { let milvus = connection("milvus-1", "运营组语义库", "milvus", "default"); let server = DbxMcpServer::with_runtime_options( Arc::new(FakeBackend { connections: vec![milvus], ..Default::default() }), @@ -2251,18 +2294,23 @@ mod tests { database: Some("default".to_string()), collection: "semantic_cards".to_string(), semantic_batch_id: "batch-1".to_string(), - published: true, + published: None, })) .await; assert_eq!( result.structured_content.as_ref().and_then(|value| value.pointer("/error/code")), - Some(&json!("VECTOR_PUBLISHED_BATCH_DELETE_BLOCKED")) + Some(&json!("VECTOR_DELETE_DISABLED_V1")) ); } #[test] fn server_registers_list_connections_tool() { - let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(FakeBackend::default()), + McpScope::default(), + false, + true, + ); let tools = server.tool_router.list_all(); let names = tools.iter().map(|tool| tool.name.as_ref()).collect::>(); assert_eq!(tools.len(), 21); @@ -2287,6 +2335,30 @@ mod tests { assert!(names.contains(&"dbx_vector_search")); assert!(names.contains(&"dbx_vector_upsert_file")); assert!(names.contains(&"dbx_vector_delete_by_batch")); + for name in [ + "dbx_preview_import_file", + "dbx_prepare_table_import", + "dbx_start_table_import", + "dbx_get_import_status", + "dbx_cancel_import", + "dbx_vector_search", + "dbx_vector_upsert_file", + "dbx_vector_delete_by_batch", + ] { + let annotations = tools + .iter() + .find(|tool| tool.name == name) + .and_then(|tool| tool.annotations.as_ref()) + .unwrap_or_else(|| panic!("{name} 必须声明 MCP annotations")); + assert_eq!(annotations.open_world_hint, Some(false)); + } + let delete_annotations = tools + .iter() + .find(|tool| tool.name == "dbx_vector_delete_by_batch") + .and_then(|tool| tool.annotations.as_ref()) + .unwrap(); + assert_eq!(delete_annotations.read_only_hint, Some(true)); + assert_eq!(delete_annotations.destructive_hint, Some(false)); } #[tokio::test] @@ -2308,13 +2380,24 @@ mod tests { .preview_import_file(Parameters(PreviewImportFileRequest { file_path: source.to_string_lossy().to_string(), source_format: None, - parse_options: Default::default(), + parse_options: crate::enterprise_tools::McpImportParseOptions { + encoding: Some("utf-8".to_string()), + ..Default::default() + }, preview_rows: Some(5), cell_char_limit: Some(100), })) .await; assert_eq!(preview.is_error, Some(false)); assert_eq!(preview.structured_content.as_ref().and_then(|value| value.get("formatVersion")), Some(&json!(1))); + assert_eq!( + preview.structured_content.as_ref().and_then(|value| value.pointer("/data/usedRangeExact")), + Some(&json!(false)) + ); + assert_eq!( + preview.structured_content.as_ref().and_then(|value| value.pointer("/data/totalRowsExact")), + Some(&json!(false)) + ); let prepared = server .prepare_table_import(Parameters(PrepareTableImportRequest { @@ -2323,7 +2406,10 @@ mod tests { template_version: "orders-v1".to_string(), file_path: source.to_string_lossy().to_string(), source_format: None, - parse_options: Default::default(), + parse_options: crate::enterprise_tools::McpImportParseOptions { + encoding: Some("utf-8".to_string()), + ..Default::default() + }, mappings: vec![ crate::enterprise_tools::McpImportColumnMapping { source_position: 1, @@ -2377,7 +2463,12 @@ mod tests { #[test] fn schema_context_tables_schema_is_gemini_compatible() { - let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(FakeBackend::default()), + McpScope::default(), + false, + true, + ); let tool = server .tool_router .list_all() @@ -2402,7 +2493,12 @@ mod tests { #[test] fn connection_selector_schema_uses_optional_strings() { - let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(FakeBackend::default()), + McpScope::default(), + false, + true, + ); let tools = server.tool_router.list_all(); for tool_name in [ @@ -2445,7 +2541,12 @@ mod tests { // field on these request structs emitting the same union shape. Every optional field // must instead publish a single concrete `type`, relying on omission from `required` // (not a `"null"` union member) to signal optionality. - let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(FakeBackend::default()), + McpScope::default(), + false, + true, + ); let tools = server.tool_router.list_all(); let checks: &[(&str, &[&str])] = &[ @@ -2557,6 +2658,94 @@ mod tests { assert!(names.iter().any(|name| name == "dbx_close_session")); } + #[test] + fn connection_management_is_default_closed_without_explicit_maintenance_flag() { + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(FakeBackend::default()), + McpScope::default(), + false, + false, + ); + let names = server.tool_router.list_all().into_iter().map(|tool| tool.name).collect::>(); + assert_eq!(names.len(), 18); + for name in ["dbx_add_connection", "dbx_duplicate_connection", "dbx_remove_connection"] { + assert!(!names.iter().any(|candidate| candidate == name)); + } + } + + #[tokio::test] + async fn daily_scope_cannot_add_duplicate_or_remove_connections_even_when_management_is_requested() { + let daily = connection("daily", "运营组数据查询", "postgres", "enterprise"); + let management = connection("management", "运营组数据管理", "postgres", "enterprise"); + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(FakeBackend { connections: vec![daily, management], ..Default::default() }), + McpScope { connection_ids: vec!["daily".to_string()], ..Default::default() }, + false, + true, + ); + let names = server.tool_router.list_all().into_iter().map(|tool| tool.name).collect::>(); + for name in ["dbx_add_connection", "dbx_duplicate_connection", "dbx_remove_connection"] { + assert!(!names.iter().any(|candidate| candidate == name)); + } + + let add = server + .add_connection(Parameters(AddConnectionRequest { + name: "forbidden".to_string(), + db_type: "postgres".to_string(), + host: "127.0.0.1".to_string(), + port: Some(5432), + username: String::new(), + password: String::new(), + database: Some("enterprise".to_string()), + ssl: false, + driver_profile: None, + })) + .await; + let duplicate = server + .duplicate_connection(Parameters(DuplicateConnectionRequest { + selector: ConnectionSelector { connection_id: Some("management".to_string()), connection_name: None }, + new_name: "forbidden-copy".to_string(), + })) + .await; + let remove = server + .remove_connection(Parameters(RemoveConnectionRequest { + connection_name: "运营组数据管理".to_string(), + connection_id: Some("management".to_string()), + })) + .await; + for result in [add, duplicate, remove] { + assert!(result_text(&result).contains("CONNECTION_MANAGEMENT_DISABLED_IN_SCOPED_MODE")); + } + } + + #[tokio::test] + async fn remove_connection_respects_global_allowed_connection_ids() { + let daily = connection("daily", "运营组数据查询", "postgres", "enterprise"); + let management = connection("management", "运营组数据管理", "postgres", "enterprise"); + let backend = FakeBackend { + connections: vec![daily, management], + policy: McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: Some(vec!["daily".to_string()]), + }, + ..Default::default() + }; + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(backend), + McpScope::default(), + false, + true, + ); + let result = server + .remove_connection(Parameters(RemoveConnectionRequest { + connection_name: "运营组数据管理".to_string(), + connection_id: Some("management".to_string()), + })) + .await; + assert!(result_text(&result).contains("CONNECTION_NOT_FOUND")); + } + #[test] fn scoped_connection_ids_are_deduplicated_and_take_precedence_over_name() { assert_eq!(scoped_connection_ids(Some(" first, second,first ,, ")), vec!["first", "second"]); diff --git a/crates/dbx-mcp/tests/local.rs b/crates/dbx-mcp/tests/local.rs index 45676bd438..74779ca03c 100644 --- a/crates/dbx-mcp/tests/local.rs +++ b/crates/dbx-mcp/tests/local.rs @@ -169,7 +169,8 @@ async fn duplicate_connection_preserves_secrets_ssh_and_sidebar_group() { let backend = Arc::new(LocalBackend::open(&db_path).await.expect("open local backend")); let policy = backend.load_mcp_global_policy().await.expect("load configured policy"); assert!(!policy.read_only); - let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + let server = + DbxMcpServer::with_runtime_options_and_connection_management(backend, McpScope::default(), false, true); let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); let server_task = tokio::spawn(async move { server.serve(server_transport).await }); let client = ().serve(client_transport).await.expect("initialize client"); diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index d364d6f9bb..64c1f7909b 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -165,9 +165,11 @@ async fn initializes_lists_tools_and_calls_a_tool() { let tools = client.peer().list_tools(None).await.expect("list tools"); let names = tools.tools.iter().map(|tool| tool.name.as_ref()).collect::>(); - assert_eq!(names.len(), 21); + assert_eq!(names.len(), 18); assert!(names.contains(&"dbx_list_connections")); - assert!(names.contains(&"dbx_duplicate_connection")); + assert!(!names.contains(&"dbx_add_connection")); + assert!(!names.contains(&"dbx_duplicate_connection")); + assert!(!names.contains(&"dbx_remove_connection")); assert!(names.contains(&"dbx_execute_redis_command")); assert!(names.contains(&"dbx_execute_and_show")); assert!(names.contains(&"dbx_open_session")); @@ -362,7 +364,12 @@ async fn duplicate_connection_rejects_ambiguous_source_names() { group_paths: Ok(HashMap::new()), }; let (server_transport, client_transport) = tokio::io::duplex(16 * 1024); - let server = DbxMcpServer::with_runtime_options(Arc::new(backend), McpScope::default(), false); + let server = DbxMcpServer::with_runtime_options_and_connection_management( + Arc::new(backend), + McpScope::default(), + false, + true, + ); let server_task = tokio::spawn(async move { server.serve(server_transport).await }); let client = ().serve(client_transport).await.expect("initialize MCP client"); let result = client diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index 4bf6684525..aba586ab3f 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -87,7 +87,7 @@ DeepSeek Harness 通过 Cordis 插件条目加载 MCP Server，不读取 `mcpSer ## 工具列表 -DBX MCP 当前提供 21 个工具： +DBX MCP 定义 21 个工具；日常默认注册 18 个，只有显式无作用域维护模式才额外注册 3 个连接管理工具： | 工具 | 说明 | | --- | --- | @@ -111,17 +111,19 @@ DBX MCP 当前提供 21 个工具： | `dbx_cancel_import` | 请求在后续 staging 批次写入前取消任务 | | `dbx_vector_search` | 在允许的 Milvus 集合中按批准状态和生效时间检索语义卡 | | `dbx_vector_upsert_file` | 从允许目录中的 JSONL upsert 已批准语义卡 | -| `dbx_vector_delete_by_batch` | 在完全访问权限下删除一个明确的未发布语义批次 | +| `dbx_vector_delete_by_batch` | v1 稳定返回 `VECTOR_DELETE_DISABLED_V1`，不连接 Milvus、不执行删除 | -启用连接作用域后，修改连接和桌面 UI 工具会被隐藏。 +连接管理默认关闭；只有无任何 `DBX_MCP_SCOPE_*` 且安装维护进程显式设置 `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1` 时，新增、复制和删除连接工具才会注册。启用任一作用域后，即使设置该开关也会隐藏并在直接调用时阻断连接管理。删除还会复核全局 `allowed_connection_ids`。桌面 UI 工具在 scoped 模式下同样隐藏。 ## 本地文件导入与 Milvus 工具 -导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个真实目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。preview 支持 Excel、JSON、CSV 和 TSV；v1 治理入库使用有界流式处理支持 XLSX/XLSM 和 UTF-8 CSV/TSV。旧版 `.xls`、JSON 和非 UTF-8 分隔文件会在创建计划和访问数据库之前返回稳定的不支持错误。 +导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个真实目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。preview 支持 Excel、JSON、CSV 和 TSV，并明确返回 `usedRangeExact` 与 `totalRowsExact`；受限预览不会把已扫描行数伪装成全文件总行数。v1 治理入库使用有界流式处理支持 XLSX/XLSM 和显式指定 `utf-8` 的 CSV/TSV。旧版 `.xls`、JSON、`auto`、GBK 和 UTF-16 分隔文件会在创建计划和访问数据库之前返回稳定的不支持错误。 -`dbx_prepare_table_import` 会把连接、调用方已批准模板版本、规范路径、大小、修改时间、SHA-256、解析参数和字段映射固化为不可变摘要。preview 对每列返回 1 基 `source_position`、原始 `raw_source_name`、稳定 `canonical_source_name`（重复表头统一为 `名称__1`、`名称__2`）以及兼容旧 DBX 的名称；映射必须复核位置、raw、canonical，并提供 `target_column`。目标不能由调用方指定：服务端生成唯一 `staging.mcp_`，强制创建新表，所有业务列固定为 `TEXT`。计划 30 分钟后过期，且只能启动一次。start 复验源文件后，先流式写出有界治理快照，再打开数据库连接；每行增加 `import_id`、`plan_id`、`source_sha`、真实绝对 `source_row_number`、`source_row_hash` 和 `loaded_at`。无法形成精确血缘时会在数据库访问前阻断。status 会返回生成的 relation；发布到业务 Schema 仍是独立治理流程。 +`dbx_preview_import_file` 与 `dbx_prepare_table_import` 共用受控 inspection semaphore 和超时。XLSX 会在解析前核验 ZIP 条目数、总解压大小、元数据、sharedStrings 与 worksheet 解压预算；完整治理流继续限制工作表行数、单元格数、单元格字节数和批次内存。`dbx_prepare_table_import` 会把连接、调用方已批准模板版本、规范路径、大小、修改时间、SHA-256、解析参数和字段映射固化为不可变摘要。preview 对每列返回 1 基 `source_position`、原始 `raw_source_name`、稳定 `canonical_source_name`（重复表头统一为 `名称__1`、`名称__2`）以及兼容旧 DBX 的名称；映射必须复核位置、raw、canonical，并提供 `target_column`。batch 行数上限根据列数和内存预算动态计算。目标不能由调用方指定：服务端生成唯一 `staging.mcp_`，强制创建新表，所有业务列固定为 `TEXT`。计划 30 分钟后过期，且只能启动一次。 -Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索必须提供 `active_at`，服务端强制 approved 和生效时间过滤，并可选追加精确 `semantic_version`。upsert 使用严格语义卡 Schema：重算并核对 64 位 `content_checksum`，要求 `embedding_revision`，拒绝 `chunk_index` 与 `embedding_dimension`，文件内 card_id 必须唯一且只能有一个语义版本；写入前还会查询现有 card_id，禁止覆盖其他批次或版本。删除除“完全访问”外还要求调用方显式传入 `published: false`。所有通过 `dbx_execute_query` 提交的向量实体 search/query/mutation 都返回 `VECTOR_DEDICATED_TOOL_REQUIRED`，通用入口只保留明确管理诊断。 +start 复验源文件后，用 `O_NOFOLLOW`/文件描述符元数据打开源文件，在复制任务私有快照时一遍完成可取消复制与 SHA-256；复制期间身份漂移会阻断。服务端先检查任务磁盘是否能容纳源文件、规范化输出上限和保留空间，再流式写出有大小上限的治理快照，成功后才打开数据库连接。CSV/TSV 短行按缺失值处理；超出表头的尾部空字段会保留在完整行哈希中，任一尾部非空字段按真实源行号阻断。每行增加 `import_id`、`plan_id`、`source_sha`、真实绝对 `source_row_number`、覆盖完整规范源记录的 `source_row_hash` 和 `loaded_at`。无法形成精确血缘时会在数据库访问前阻断。status 会返回生成的 relation；发布到业务 Schema 仍是独立治理流程。 + +Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索必须提供 `active_at`，服务端强制 approved 和生效时间过滤，并可选追加精确 `semantic_version`。单查询 search 响应会正确展开 Milvus 的外层 query 数组：`data:[[]]` 返回零行，多命中返回语义卡行；query/list 解析不受影响。upsert 使用严格语义卡 Schema：全部 VARCHAR 在任何 Milvus 调用前按当前集合的 UTF-8 字节长度核验，`content` 上限为 8192 字节；重算并核对 64 位 `content_checksum`，要求 `embedding_revision`，拒绝 `chunk_index` 与 `embedding_dimension`，文件内 card_id 必须唯一且只能有一个语义版本。整份 JSONL 和全部已有 card_id 归属检查通过后才开始分批 upsert。v1 不再信任调用方提供的 `published:false`，`dbx_vector_delete_by_batch` 始终返回 `VECTOR_DELETE_DISABLED_V1`；撤销残留批次必须走受审计管理员流程。所有通过 `dbx_execute_query` 提交的向量实体 search/query/mutation 都返回 `VECTOR_DEDICATED_TOOL_REQUIRED`，通用入口只保留明确管理诊断。 向量检索的核心参数形状如下： @@ -219,10 +221,24 @@ DBX 在 **设置 → MCP** 中保存一份权威策略，并在每次请求时 | `DBX_MCP_SCOPE_CONNECTION_IDS` | 兼容旧配置：限制为多个连接 ID | | `DBX_MCP_SCOPE_CONNECTION_NAME` | 限制为一个连接名称 | | `DBX_MCP_SCOPE_DATABASE` | 限制为一个数据库 | +| `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT` | 仅无作用域安装维护进程临时设置为 `1`；默认关闭连接新增、复制和删除 | | `DBX_MCP_IMPORT_ROOTS` | 导入和语义 JSONL 工具可读取的本地目录 allowlist，使用平台路径分隔符 | | `DBX_MCP_IMPORT_STAGING_SCHEMAS` | 允许导入的 PostgreSQL staging Schema，逗号分隔，默认 `staging` | | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Excel/CSV/TSV/JSON 最大字节数，默认 512 MiB | | `DBX_MCP_IMPORT_CONCURRENCY` | 治理导入并发数，限制为 1～16，默认 `2` | +| `DBX_MCP_IMPORT_INSPECTION_CONCURRENCY` | preview/prepare 文件剖析并发，限制为 1～8，默认 `2` | +| `DBX_MCP_IMPORT_INSPECTION_TIMEOUT_SECS` | 单次文件剖析等待上限，限制为 5～120 秒，默认 `30` | +| `DBX_MCP_IMPORT_BATCH_MEMORY_MAX_BYTES` | XLSX 单个流式行批次近似 JSON 内存上限，默认 32 MiB | +| `DBX_MCP_IMPORT_NORMALIZED_MAX_BYTES` | 单个规范化治理 CSV 硬上限，默认 2 GiB | +| `DBX_MCP_IMPORT_DISK_RESERVE_BYTES` | 任务磁盘除源文件和规范化上限外必须保留的空间，默认 2 GiB | +| `DBX_MCP_XLSX_ZIP_ENTRY_LIMIT` | XLSX ZIP 条目数上限，默认 `4096` | +| `DBX_MCP_XLSX_TOTAL_UNCOMPRESSED_MAX_BYTES` | XLSX ZIP 解压总字节上限，默认 4 GiB | +| `DBX_MCP_XLSX_METADATA_ENTRY_MAX_BYTES` | 单个 workbook/rels/styles 元数据条目上限，默认 16 MiB | +| `DBX_MCP_XLSX_SHARED_STRINGS_MAX_BYTES` | sharedStrings 解压上限，默认 256 MiB | +| `DBX_MCP_XLSX_WORKSHEET_MAX_BYTES` | 单个 worksheet XML 解压上限，默认 2 GiB | +| `DBX_MCP_XLSX_WORKSHEET_MAX_ROWS` | 单个 worksheet row 元素上限，默认 500 万 | +| `DBX_MCP_XLSX_WORKSHEET_MAX_CELLS` | 单个 worksheet cell 元素上限，默认 1 亿 | +| `DBX_MCP_XLSX_CELL_MAX_BYTES` | 单个解码单元格 UTF-8 字节上限，默认 1 MiB | | `DBX_MCP_IMPORT_PLAN_CAPACITY` | 同时保留的待启动计划上限，默认 `512` | | `DBX_MCP_IMPORT_JOB_CAPACITY` | 运行中和终态任务总保留上限，默认 `256` | | `DBX_MCP_IMPORT_TERMINAL_LIMIT` | 24 小时 TTL 内终态任务保留上限，默认 `128` | diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index 80683ff290..4b606f11e1 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -87,7 +87,7 @@ Run `dsh web --dump-config` to verify the composed configuration, then restart D ## Tools -DBX MCP currently provides 21 tools: +DBX MCP defines 21 tools. Daily mode registers 18 by default; an explicit unscoped maintenance process adds the three connection-management tools: | Tool | Description | | --- | --- | @@ -111,17 +111,19 @@ DBX MCP currently provides 21 tools: | `dbx_cancel_import` | Request cancellation before subsequent staging batches are written | | `dbx_vector_search` | Search an allowed Milvus collection with bounded output and enforced approval/effective-time filters | | `dbx_vector_upsert_file` | Upsert approved semantic cards from one allowed JSONL file | -| `dbx_vector_delete_by_batch` | Delete one exact unpublished semantic batch with full-access permission | +| `dbx_vector_delete_by_batch` | Always return `VECTOR_DELETE_DISABLED_V1` in v1 without connecting to Milvus | -Connection-scoped sessions hide connection-mutating and desktop UI tools. +Connection management is default-closed. Add, duplicate, and remove tools are registered only when an unscoped maintenance process explicitly sets `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1`. Any `DBX_MCP_SCOPE_*` setting disables these tools even when the maintenance flag is present, and remove also enforces the global `allowed_connection_ids`. Scoped sessions also hide desktop UI tools. ## Local file imports and Milvus tools -The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more real directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. Preview supports Excel, JSON, CSV, and TSV. Governed v1 database import uses bounded streaming for XLSX/XLSM and UTF-8 CSV/TSV. Legacy `.xls`, JSON, and non-UTF-8 delimited files return stable unsupported errors before a plan or database write exists. +The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more real directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. Preview supports Excel, JSON, CSV, and TSV and explicitly reports `usedRangeExact` and `totalRowsExact`; a bounded scan is never presented as a complete row count. Governed v1 database import uses bounded streaming for XLSX/XLSM and CSV/TSV whose prepare request explicitly selects `utf-8`. Legacy `.xls`, JSON, `auto`, GBK, and UTF-16 delimited files return stable unsupported errors before a plan or database write exists. -`dbx_prepare_table_import` binds the connection, caller-approved template version, canonical path, size, modification time, SHA-256, parsing options, and mappings into an immutable digest. Preview exposes each source column as 1-based `source_position`, exact `raw_source_name`, stable `canonical_source_name` (`name__1`, `name__2` for duplicates), and the legacy DBX name. Each mapping must repeat the position plus raw/canonical names and supply `target_column`. The server—not the caller—generates a unique `staging.mcp_` relation, forces a new table, and makes every business column `TEXT`. The returned plan expires after 30 minutes and can be started once. Start rechecks the source and streams it into a bounded governed snapshot before opening a database pool. Every row gains `import_id`, `plan_id`, `source_sha`, the real absolute `source_row_number`, `source_row_hash`, and `loaded_at`; if exact lineage cannot be produced, the import fails before database access. Status returns the generated relation. Publication to business schemas remains a separately governed workflow. +Preview and prepare share a bounded inspection semaphore and timeout. XLSX inspection checks ZIP entry count, total uncompressed bytes, metadata entries, shared strings, and worksheet sizes before parsing; the full governed stream additionally bounds worksheet rows, cells, decoded cell bytes, and batch memory. `dbx_prepare_table_import` binds the connection, caller-approved template version, canonical path, size, modification time, SHA-256, parsing options, and mappings into an immutable digest. Preview exposes each source column as 1-based `source_position`, exact `raw_source_name`, stable `canonical_source_name` (`name__1`, `name__2` for duplicates), and the legacy DBX name. Each mapping must repeat the position plus raw/canonical names and supply `target_column`. The maximum row batch is derived from column count and the memory budget. The server—not the caller—generates a unique `staging.mcp_` relation, forces a new table, and makes every business column `TEXT`. -Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires `active_at`, always adds `approval_status == "approved"` plus effective-date filters, and optionally applies an exact `semantic_version`. Upsert accepts a strict semantic-card schema, verifies the 64-hex `content_checksum` against `content`, requires `embedding_revision`, rejects `chunk_index` and `embedding_dimension`, enforces unique card IDs and one semantic version per file, and preflights existing IDs so another batch/version cannot be overwritten. Delete requires Full access and an explicit `published: false` attestation. All entity search/query/mutation REST submitted through generic `dbx_execute_query` returns `VECTOR_DEDICATED_TOOL_REQUIRED`; only explicit management diagnostics remain available there. +Start rechecks the source, opens it with `O_NOFOLLOW` plus descriptor metadata checks, and performs cancellable copy plus SHA-256 in one pass into a task-private snapshot. It verifies that the task disk can hold the source, normalized-output ceiling, and reserve, then writes a size-limited governed CSV before opening a database pool. Short CSV/TSV rows retain missing values; extra empty trailing fields remain part of the full-record row hash, while any extra non-empty field fails with its real source row number. Every row gains `import_id`, `plan_id`, `source_sha`, the real absolute `source_row_number`, a `source_row_hash` over the complete normalized source record, and `loaded_at`. If exact lineage cannot be produced, the import fails before database access. Status returns the generated relation. Publication to business schemas remains separately governed. + +Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires `active_at`, always adds `approval_status == "approved"` plus effective-date filters, and optionally applies an exact `semantic_version`. A single-query search flattens Milvus's outer query array: `data:[[]]` becomes zero rows and multiple hits become semantic-card rows, without changing query/list parsing. Upsert validates every VARCHAR against the current collection's UTF-8 byte length before any Milvus call (`content` is limited to 8192 bytes), verifies the 64-hex `content_checksum`, requires `embedding_revision`, rejects `chunk_index` and `embedding_dimension`, and validates the entire file plus all existing card-ID ownership before the first batch is written. V1 no longer trusts caller-supplied `published:false`; `dbx_vector_delete_by_batch` always returns `VECTOR_DELETE_DISABLED_V1`, and cleanup requires an audited administrator workflow. All entity search/query/mutation REST submitted through generic `dbx_execute_query` returns `VECTOR_DEDICATED_TOOL_REQUIRED`; only explicit management diagnostics remain available there. The core vector-search arguments are: @@ -219,10 +221,24 @@ Updated servers do not let `DBX_MCP_ALLOW_WRITES` or `DBX_MCP_ALLOW_DANGEROUS_SQ | `DBX_MCP_SCOPE_CONNECTION_IDS` | Compatibility scope for multiple connection IDs | | `DBX_MCP_SCOPE_CONNECTION_NAME` | Restrict the session to one connection name | | `DBX_MCP_SCOPE_DATABASE` | Restrict the session to one database | +| `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT` | Set to `1` only in an unscoped installation/maintenance process; add, duplicate, and remove are disabled by default | | `DBX_MCP_IMPORT_ROOTS` | Platform-separated allowlist of local directories readable by import and semantic JSONL tools | | `DBX_MCP_IMPORT_STAGING_SCHEMAS` | Comma-separated PostgreSQL staging schema allowlist (default `staging`) | | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Maximum Excel/CSV/TSV/JSON source size (default 512 MiB) | | `DBX_MCP_IMPORT_CONCURRENCY` | Concurrent governed imports, clamped to 1–16 (default `2`) | +| `DBX_MCP_IMPORT_INSPECTION_CONCURRENCY` | Concurrent preview/prepare inspections, clamped to 1–8 (default `2`) | +| `DBX_MCP_IMPORT_INSPECTION_TIMEOUT_SECS` | Inspection wait timeout, clamped to 5–120 seconds (default `30`) | +| `DBX_MCP_IMPORT_BATCH_MEMORY_MAX_BYTES` | Approximate JSON memory ceiling for one XLSX row batch (default 32 MiB) | +| `DBX_MCP_IMPORT_NORMALIZED_MAX_BYTES` | Hard limit for one normalized governed CSV (default 2 GiB) | +| `DBX_MCP_IMPORT_DISK_RESERVE_BYTES` | Required free space beyond source plus normalized ceiling (default 2 GiB) | +| `DBX_MCP_XLSX_ZIP_ENTRY_LIMIT` | Maximum XLSX ZIP entries (default `4096`) | +| `DBX_MCP_XLSX_TOTAL_UNCOMPRESSED_MAX_BYTES` | Maximum total XLSX uncompressed bytes (default 4 GiB) | +| `DBX_MCP_XLSX_METADATA_ENTRY_MAX_BYTES` | Maximum workbook/rels/styles metadata entry size (default 16 MiB) | +| `DBX_MCP_XLSX_SHARED_STRINGS_MAX_BYTES` | Maximum uncompressed sharedStrings size (default 256 MiB) | +| `DBX_MCP_XLSX_WORKSHEET_MAX_BYTES` | Maximum uncompressed worksheet XML size (default 2 GiB) | +| `DBX_MCP_XLSX_WORKSHEET_MAX_ROWS` | Maximum worksheet row elements (default 5 million) | +| `DBX_MCP_XLSX_WORKSHEET_MAX_CELLS` | Maximum worksheet cell elements (default 100 million) | +| `DBX_MCP_XLSX_CELL_MAX_BYTES` | Maximum decoded UTF-8 bytes per cell (default 1 MiB) | | `DBX_MCP_IMPORT_PLAN_CAPACITY` | Maximum live prepared plans (default `512`) | | `DBX_MCP_IMPORT_JOB_CAPACITY` | Maximum retained running/terminal jobs (default `256`) | | `DBX_MCP_IMPORT_TERMINAL_LIMIT` | Maximum terminal jobs retained within the 24-hour TTL (default `128`) | diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index 56742a7fb3..c580a4056e 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -161,9 +161,9 @@ Ask the MCP client to: | `dbx_cancel_import` | Request cancellation of a staging import | | `dbx_vector_search` | Search approved, active semantic cards in an allowed Milvus collection | | `dbx_vector_upsert_file` | Upsert approved semantic cards from an allowed JSONL file | -| `dbx_vector_delete_by_batch` | Delete one exact unpublished semantic batch with Full access | +| `dbx_vector_delete_by_batch` | Return `VECTOR_DELETE_DISABLED_V1`; v1 performs no MCP deletion | -When connection scoping is enabled, mutating connection tools and desktop UI tools are hidden. +Connection management is default-closed. An unscoped installation/maintenance process must temporarily set `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1`; any `DBX_MCP_SCOPE_*` setting still disables add, duplicate, and remove. Remove also enforces the global connection allowlist. Scoped sessions hide desktop UI tools too. ## Execution Modes @@ -183,7 +183,7 @@ DBX connection storage defaults to: Override the directory with `DBX_DATA_DIR`. -Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Preview accepts Excel/JSON/CSV/TSV; v1 governed import streams XLSX/XLSM and UTF-8 CSV/TSV, while legacy `.xls` and JSON remain preview-only. Prepare accepts position+raw+canonical-name mappings and generates a unique `staging.mcp_` relation; callers cannot append to an existing table. Start revalidates the source and streams a private snapshot containing absolute row numbers, hashes, and batch lineage before opening a database pool. Import tools support local mode only and return `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` in Web mode. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. +Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Preview accepts Excel/JSON/CSV/TSV and reports whether row/range counts are exact. V1 governed import streams XLSX/XLSM and CSV/TSV with explicit `utf-8`; legacy `.xls`, JSON, `auto`, GBK, and UTF-16 remain preview-only. Inspection, ZIP expansion, shared strings, worksheets, cells, row batches, normalized output, and task disk all have hard budgets. Prepare accepts position+raw+canonical-name mappings and generates a unique `staging.mcp_` relation; callers cannot append to an existing table. Start performs cancellable copy+SHA into a private snapshot, hashes the complete source record with its absolute row number, and opens the database only after normalization succeeds. Semantic VARCHAR lengths are fully prevalidated, Milvus's single-search outer result array is flattened, and v1 batch deletion is disabled. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. ### Agent/JDBC databases @@ -289,9 +289,14 @@ SQL text is not included in normal MCP errors or logged by default. Enable tempo | `DBX_MCP_SCOPE_CONNECTION_IDS` | Compatibility scope for multiple connection IDs | | `DBX_MCP_SCOPE_CONNECTION_NAME` | Restrict tools to one connection name | | `DBX_MCP_SCOPE_DATABASE` | Restrict tools to one database | +| `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT` | Set to `1` only for an unscoped installation/maintenance process; default disabled | | `DBX_MCP_IMPORT_ROOTS` | Platform-separated allowlist of local import/semantic-file directories | | `DBX_MCP_IMPORT_STAGING_SCHEMAS` | Comma-separated PostgreSQL staging schema allowlist (default `staging`) | | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | Maximum tabular import source size (default 512 MiB) | +| `DBX_MCP_IMPORT_INSPECTION_CONCURRENCY` | Concurrent preview/prepare inspections (default `2`) | +| `DBX_MCP_IMPORT_INSPECTION_TIMEOUT_SECS` | Inspection wait timeout (default `30`) | +| `DBX_MCP_IMPORT_NORMALIZED_MAX_BYTES` | Maximum normalized governed CSV size (default 2 GiB) | +| `DBX_MCP_IMPORT_DISK_RESERVE_BYTES` | Required free task-disk reserve (default 2 GiB) | | `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | Maximum semantic JSONL size (default 64 MiB) | | `DBX_MCP_VECTOR_COLLECTIONS` | Comma-separated Milvus collection allowlist (default `semantic_cards`) | | `DBX_MCP_VECTOR_TOP_K_MAX` | Maximum Milvus Top K (default `20`, hard cap 50) | @@ -397,7 +402,7 @@ MCP 协议、连接读取、SQL 安全检查、Schema、Redis、MongoDB、Web ### 主要能力 -- 21 个 MCP 工具，覆盖连接、Schema、SQL、会话、本地受控导入、Milvus 语义、Redis 和 DBX 桌面操作 +- 定义 21 个 MCP 工具；日常默认注册 18 个，显式无作用域维护模式再开放 3 个连接管理工具 - 不依赖 `better-sqlite3`，没有 Node 原生模块 ABI 问题 - 支持本地 DBX、DBX Web 和 Docker - 支持预编译原生二进制和离线运行 @@ -498,7 +503,7 @@ MCP 配置： | `dbx_cancel_import` | 请求取消 staging 导入 | | `dbx_vector_search` | 在允许的 Milvus 集合中检索已批准且生效的语义卡 | | `dbx_vector_upsert_file` | 从允许目录中的 JSONL upsert 已批准语义卡 | -| `dbx_vector_delete_by_batch` | 在完全访问权限下删除一个明确的未发布语义批次 | +| `dbx_vector_delete_by_batch` | v1 返回 `VECTOR_DELETE_DISABLED_V1`，不执行 MCP 删除 | ### 本地数据目录 @@ -508,7 +513,7 @@ MCP 配置： 通过 `DBX_DATA_DIR` 覆盖默认目录。Windows 便携版应指向 `DBX.exe` 同级、包含 `dbx.db` 的 `data` 文件夹。 -本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。preview 支持 Excel/JSON/CSV/TSV；v1 治理入库流式支持 XLSX/XLSM 和 UTF-8 CSV/TSV，旧版 `.xls` 与 JSON 暂为 preview-only。prepare 接受“源位置+raw+canonical 名称”映射并由服务端生成唯一 `staging.mcp_`，调用方不能追加已有表。start 在打开数据库连接前复验文件并流式生成包含真实绝对行号、行哈希和批次血缘的任务快照。导入工具仅支持本地模式，Web 模式稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 +连接管理默认关闭；仅无作用域安装维护进程临时设置 `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1`，任何 `DBX_MCP_SCOPE_*` 都继续禁用新增、复制和删除。本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。preview 支持 Excel/JSON/CSV/TSV 并标明行数/范围是否精确；v1 治理入库流式支持 XLSX/XLSM 和显式 `utf-8` 的 CSV/TSV，旧版 `.xls`、JSON、`auto`、GBK 与 UTF-16 暂为 preview-only。inspection、ZIP 解压、sharedStrings、worksheet、单元格、批次、规范化输出与任务磁盘均有硬预算。start 以可取消复制+SHA 生成任务私有快照，保留绝对行号和完整源记录哈希，规范化成功后才打开数据库。语义 VARCHAR 会整文件预校验，Milvus 单查询 search 会展开外层数组，v1 批次删除禁用。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 ### DBX Web / Docker @@ -591,9 +596,14 @@ MongoDB 更新和删除在未启用完全访问时必须提供可验证有效的 | `DBX_MCP_SCOPE_CONNECTION_IDS` | 兼容旧配置：限制到多个连接 ID | | `DBX_MCP_SCOPE_CONNECTION_NAME` | 限制到指定连接名称 | | `DBX_MCP_SCOPE_DATABASE` | 限制到指定数据库 | +| `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT` | 仅无作用域安装维护进程临时设为 `1`；默认禁用 | | `DBX_MCP_IMPORT_ROOTS` | 本地导入/语义文件目录 allowlist，使用平台路径分隔符 | | `DBX_MCP_IMPORT_STAGING_SCHEMAS` | PostgreSQL staging Schema allowlist，逗号分隔，默认 `staging` | | `DBX_MCP_IMPORT_FILE_MAX_BYTES` | 表格导入源最大字节数，默认 512 MiB | +| `DBX_MCP_IMPORT_INSPECTION_CONCURRENCY` | preview/prepare 剖析并发，默认 `2` | +| `DBX_MCP_IMPORT_INSPECTION_TIMEOUT_SECS` | 文件剖析等待上限，默认 `30` 秒 | +| `DBX_MCP_IMPORT_NORMALIZED_MAX_BYTES` | 单个规范化治理 CSV 上限，默认 2 GiB | +| `DBX_MCP_IMPORT_DISK_RESERVE_BYTES` | 任务磁盘必须额外保留的空间，默认 2 GiB | | `DBX_MCP_SEMANTIC_FILE_MAX_BYTES` | 语义 JSONL 最大字节数，默认 64 MiB | | `DBX_MCP_VECTOR_COLLECTIONS` | Milvus 集合 allowlist，逗号分隔，默认 `semantic_cards` | | `DBX_MCP_VECTOR_TOP_K_MAX` | Milvus Top K 上限，默认 `20`，硬上限 50 | From 3fc7828838ab26a80f5c410678dadabc6b2c6a5a Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 20:05:40 +0800 Subject: [PATCH 07/10] fix(mcp): require exact semantic search version --- crates/dbx-mcp/src/enterprise_tools.rs | 57 +++++++++++++++++++++---- crates/dbx-mcp/src/server.rs | 58 +++++++++++++++++++++----- crates/dbx-mcp/tests/local.rs | 53 +++++++++++++++++++++++ crates/dbx-mcp/tests/protocol.rs | 6 +++ docs/content/docs/mcp.cn.mdx | 4 +- docs/content/docs/mcp.mdx | 4 +- packages/mcp-server/README.md | 4 +- 7 files changed, 161 insertions(+), 25 deletions(-) diff --git a/crates/dbx-mcp/src/enterprise_tools.rs b/crates/dbx-mcp/src/enterprise_tools.rs index 4f6d260121..3f311e26fd 100644 --- a/crates/dbx-mcp/src/enterprise_tools.rs +++ b/crates/dbx-mcp/src/enterprise_tools.rs @@ -271,9 +271,16 @@ pub struct VectorSearchRequest { #[schemars(extend("type" = "string"))] pub database: Option, pub collection: String, + #[schemars( + pattern(r"^\d{4}-\d{2}-\d{2}$"), + description = "Required Asia/Shanghai business date in YYYY-MM-DD format" + )] pub active_at: String, - #[schemars(extend("type" = "string"))] - pub semantic_version: Option, + #[schemars( + length(min = 1, max = 128), + description = "Required exact semantic version; UTF-8 byte length must be 1 to 128" + )] + pub semantic_version: String, pub embedding: Vec, #[schemars(extend("type" = "integer"))] pub top_k: Option, @@ -1931,21 +1938,20 @@ pub fn vector_output_fields(requested: Option>) -> Result, + semantic_version: &str, filters: &BTreeMap, ) -> Result { if active_at.len() != 10 || !valid_effective_timestamp(active_at) { return Err(EnterpriseToolError::new("VECTOR_ACTIVE_AT_INVALID", "active_at 必须是合法 YYYY-MM-DD 日期。")); } let allowed = comma_list_env("DBX_MCP_VECTOR_FILTER_FIELDS", DEFAULT_VECTOR_FILTER_FIELDS); + let semantic_version = validate_semantic_version(semantic_version)?; let mut clauses = vec![ format!("approval_status == {}", json_string("approved")), format!("effective_from <= {}", json_string(active_at)), format!("(effective_to == \"\" or effective_to >= {})", json_string(active_at)), + format!("semantic_version == {}", json_string(semantic_version)), ]; - if let Some(semantic_version) = semantic_version.map(str::trim).filter(|value| !value.is_empty()) { - clauses.push(format!("semantic_version == {}", json_string(semantic_version))); - } for (field, value) in filters { validate_identifier(field, "filter_field")?; if field == "approval_status" || field == "semantic_version" || !allowed.iter().any(|allowed| allowed == field) @@ -1960,6 +1966,23 @@ pub fn build_milvus_filter( Ok(clauses.join(" and ")) } +pub fn validate_semantic_version(value: &str) -> Result<&str, EnterpriseToolError> { + let value = value.trim(); + if value.is_empty() { + return Err(EnterpriseToolError::new( + "VECTOR_SEMANTIC_VERSION_REQUIRED", + "semantic_version 是 v1 必填的精确语义版本。", + )); + } + if value.len() > 128 || value.chars().any(char::is_control) { + return Err(EnterpriseToolError::new( + "VECTOR_SEMANTIC_VERSION_INVALID", + "semantic_version 必须是 1～128 个 UTF-8 字节且不能包含控制字符。", + )); + } + Ok(value) +} + pub fn read_semantic_jsonl(path: &Path, semantic_batch_id: &str) -> Result, EnterpriseToolError> { if semantic_batch_id.trim().is_empty() || semantic_batch_id.len() > 128 @@ -2637,17 +2660,35 @@ mod tests { #[test] fn milvus_filter_forces_approval_and_semantic_version() { let filters = BTreeMap::from([("business_domain".to_string(), json!("交易"))]); - let filter = build_milvus_filter("2026-08-25", Some("semantic-v3"), &filters).unwrap(); + let filter = build_milvus_filter("2026-08-25", "semantic-v3", &filters).unwrap(); assert!(filter.contains("approval_status == \"approved\"")); assert!(filter.contains("effective_from <= \"2026-08-25\"")); + assert!(filter.contains("effective_to == \"\" or effective_to >= \"2026-08-25\"")); assert!(filter.contains("semantic_version == \"semantic-v3\"")); assert!(filter.contains("business_domain == \"交易\"")); let forbidden = BTreeMap::from([("approval_status".to_string(), json!("draft"))]); assert_eq!( - build_milvus_filter("2026-08-25", None, &forbidden).unwrap_err().code, + build_milvus_filter("2026-08-25", "semantic-v3", &forbidden).unwrap_err().code, + "VECTOR_FILTER_FIELD_NOT_ALLOWED" + ); + let forbidden_version = BTreeMap::from([("semantic_version".to_string(), json!("semantic-v2"))]); + assert_eq!( + build_milvus_filter("2026-08-25", "semantic-v3", &forbidden_version).unwrap_err().code, "VECTOR_FILTER_FIELD_NOT_ALLOWED" ); + assert_eq!( + build_milvus_filter("2026-08-25", " ", &BTreeMap::new()).unwrap_err().code, + "VECTOR_SEMANTIC_VERSION_REQUIRED" + ); + assert_eq!( + build_milvus_filter("2026-08-25", &"x".repeat(129), &BTreeMap::new()).unwrap_err().code, + "VECTOR_SEMANTIC_VERSION_INVALID" + ); + assert_eq!( + build_milvus_filter("2026-08-25T00:00:00+08:00", "semantic-v3", &BTreeMap::new()).unwrap_err().code, + "VECTOR_ACTIVE_AT_INVALID" + ); } #[test] diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 4dcef15197..949a465656 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -788,7 +788,7 @@ impl DbxMcpServer { #[tool( name = "dbx_vector_search", - description = "Search an allowed Milvus semantic collection with a bounded vector and safe equality filters. Approval and effective date are always enforced; semantic_version is an optional exact filter.", + description = "Search one exact semantic_version in an allowed Milvus collection. Approval, Asia/Shanghai business date (YYYY-MM-DD), and semantic version are always enforced together.", annotations( title = "检索获批语义卡", read_only_hint = true, @@ -822,11 +822,11 @@ impl DbxMcpServer { Ok(fields) => fields, Err(error) => return enterprise_error(error), }; - let filter = - match build_milvus_filter(&request.active_at, request.semantic_version.as_deref(), &request.filters) { - Ok(filter) => filter, - Err(error) => return enterprise_error(error), - }; + let semantic_version = request.semantic_version.trim().to_string(); + let filter = match build_milvus_filter(&request.active_at, &semantic_version, &request.filters) { + Ok(filter) => filter, + Err(error) => return enterprise_error(error), + }; let database = match self.resolve_database(request.database, &resolved.connection) { Ok(database) => database, Err(error) => return error, @@ -841,7 +841,7 @@ impl DbxMcpServer { json!({ "collection": request.collection, "activeAt": request.active_at, - "semanticVersion": request.semantic_version, + "semanticVersion": semantic_version, "topK": top_k, "filter": filter, "rows": rows, @@ -2532,6 +2532,45 @@ mod tests { } } + #[test] + fn vector_search_schema_requires_exact_semantic_version_and_shanghai_business_date() { + let server = DbxMcpServer::with_runtime_options(Arc::new(FakeBackend::default()), McpScope::default(), false); + let tool = server + .tool_router + .list_all() + .into_iter() + .find(|tool| tool.name == "dbx_vector_search") + .expect("vector search tool should be registered"); + let required = tool + .input_schema + .get("required") + .and_then(serde_json::Value::as_array) + .expect("vector search should publish required fields"); + for field in ["active_at", "semantic_version"] { + assert!(required.iter().any(|required| required == field), "{field} 必须是 inputSchema 必填字段"); + } + let properties = tool.input_schema.get("properties").and_then(serde_json::Value::as_object).unwrap(); + assert!(properties + .get("active_at") + .and_then(|schema| schema.get("description")) + .and_then(serde_json::Value::as_str) + .is_some_and(|description| description.contains("Asia/Shanghai") && description.contains("YYYY-MM-DD"))); + assert_eq!( + properties.get("active_at").and_then(|schema| schema.get("pattern")), + Some(&json!(r"^\d{4}-\d{2}-\d{2}$")) + ); + assert_eq!(properties.get("semantic_version").and_then(|schema| schema.get("minLength")), Some(&json!(1))); + assert_eq!(properties.get("semantic_version").and_then(|schema| schema.get("maxLength")), Some(&json!(128))); + + let missing = serde_json::from_value::(json!({ + "collection": "semantic_cards", + "active_at": "2026-08-25", + "embedding": vec![0.0_f32; 1024] + })) + .unwrap_err(); + assert!(missing.to_string().contains("semantic_version")); + } + #[test] fn optional_fields_never_publish_nullable_union_types() { // Some MCP clients (e.g. OpenCode, see #6344) cannot resolve a JSON Schema @@ -2565,10 +2604,7 @@ mod tests { "dbx_prepare_table_import", &["connection_id", "connection_name", "database", "source_format", "batch_size", "date_time_format"], ), - ( - "dbx_vector_search", - &["connection_id", "connection_name", "database", "semantic_version", "top_k", "output_fields"], - ), + ("dbx_vector_search", &["connection_id", "connection_name", "database", "top_k", "output_fields"]), ("dbx_vector_upsert_file", &["connection_id", "connection_name", "database", "batch_size"]), ("dbx_vector_delete_by_batch", &["connection_id", "connection_name", "database"]), ]; diff --git a/crates/dbx-mcp/tests/local.rs b/crates/dbx-mcp/tests/local.rs index 74779ca03c..ce3ec7b888 100644 --- a/crates/dbx-mcp/tests/local.rs +++ b/crates/dbx-mcp/tests/local.rs @@ -89,6 +89,59 @@ async fn local_backend_reads_dbx_storage_without_desktop_process() { server_task.abort(); } +#[tokio::test] +async fn local_vector_search_rejects_empty_semantic_version_before_network_access() { + let directory = tempdir().expect("temporary data directory"); + let db_path = directory.path().join("dbx.db"); + let storage = Storage::open(&db_path).await.expect("open storage"); + let connection: ConnectionConfig = serde_json::from_value(json!({ + "id": "milvus-local", + "name": "semantic-local", + "db_type": "milvus", + "host": "127.0.0.1", + "port": 19530, + "username": "", + "password": "", + "database": "default", + "ssl": false + })) + .expect("Milvus connection"); + storage.save_connections(&[connection]).await.expect("save Milvus connection"); + drop(storage); + + let backend = Arc::new(LocalBackend::open(&db_path).await.expect("open local backend")); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + let (server_transport, client_transport) = tokio::io::duplex(64 * 1024); + let server_task = tokio::spawn(async move { server.serve(server_transport).await }); + let client = ().serve(client_transport).await.expect("initialize client"); + let result = client + .peer() + .call_tool( + CallToolRequestParams::new("dbx_vector_search").with_arguments( + json!({ + "connection_id": "milvus-local", + "database": "default", + "collection": "semantic_cards", + "active_at": "2026-08-25", + "semantic_version": " ", + "embedding": vec![0.0_f32; 1024] + }) + .as_object() + .cloned() + .unwrap_or_else(Map::new), + ), + ) + .await + .expect("call vector search"); + assert_eq!(result.is_error, Some(true)); + assert_eq!( + result.structured_content.as_ref().and_then(|value| value.pointer("/error/code")), + Some(&json!("VECTOR_SEMANTIC_VERSION_REQUIRED")) + ); + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} + #[tokio::test] async fn duplicate_connection_preserves_secrets_ssh_and_sidebar_group() { let directory = tempdir().expect("temporary data directory"); diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index 64c1f7909b..06ce1da003 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -182,6 +182,12 @@ async fn initializes_lists_tools_and_calls_a_tool() { assert!(names.contains(&"dbx_vector_search")); assert!(names.contains(&"dbx_vector_upsert_file")); assert!(names.contains(&"dbx_vector_delete_by_batch")); + let vector_search = + tools.tools.iter().find(|tool| tool.name == "dbx_vector_search").expect("vector search tool schema"); + let required = + vector_search.input_schema.get("required").and_then(Value::as_array).expect("vector search required fields"); + assert!(required.iter().any(|field| field == "active_at")); + assert!(required.iter().any(|field| field == "semantic_version")); let result = client.peer().call_tool(CallToolRequestParams::new("dbx_list_connections")).await.expect("call tool"); let response = result.content[0].as_text().expect("text response"); diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index aba586ab3f..60e0a932bb 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -123,7 +123,7 @@ DBX MCP 定义 21 个工具；日常默认注册 18 个，只有显式无作用 start 复验源文件后，用 `O_NOFOLLOW`/文件描述符元数据打开源文件，在复制任务私有快照时一遍完成可取消复制与 SHA-256；复制期间身份漂移会阻断。服务端先检查任务磁盘是否能容纳源文件、规范化输出上限和保留空间，再流式写出有大小上限的治理快照，成功后才打开数据库连接。CSV/TSV 短行按缺失值处理；超出表头的尾部空字段会保留在完整行哈希中，任一尾部非空字段按真实源行号阻断。每行增加 `import_id`、`plan_id`、`source_sha`、真实绝对 `source_row_number`、覆盖完整规范源记录的 `source_row_hash` 和 `loaded_at`。无法形成精确血缘时会在数据库访问前阻断。status 会返回生成的 relation；发布到业务 Schema 仍是独立治理流程。 -Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索必须提供 `active_at`，服务端强制 approved 和生效时间过滤，并可选追加精确 `semantic_version`。单查询 search 响应会正确展开 Milvus 的外层 query 数组：`data:[[]]` 返回零行，多命中返回语义卡行；query/list 解析不受影响。upsert 使用严格语义卡 Schema：全部 VARCHAR 在任何 Milvus 调用前按当前集合的 UTF-8 字节长度核验，`content` 上限为 8192 字节；重算并核对 64 位 `content_checksum`，要求 `embedding_revision`，拒绝 `chunk_index` 与 `embedding_dimension`，文件内 card_id 必须唯一且只能有一个语义版本。整份 JSONL 和全部已有 card_id 归属检查通过后才开始分批 upsert。v1 不再信任调用方提供的 `published:false`，`dbx_vector_delete_by_batch` 始终返回 `VECTOR_DELETE_DISABLED_V1`；撤销残留批次必须走受审计管理员流程。所有通过 `dbx_execute_query` 提交的向量实体 search/query/mutation 都返回 `VECTOR_DEDICATED_TOOL_REQUIRED`，通用入口只保留明确管理诊断。 +Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。检索必须同时提供 `active_at` 与 `semantic_version`：`active_at` 是将查询时点换算到 `Asia/Shanghai` 后得到的 `YYYY-MM-DD` 业务日期；`semantic_version` 是 1～128 个 UTF-8 字节的非空精确版本。服务端始终组合 `approval_status=approved`、生效日期区间和精确语义版本，`filters` 不能提供或覆盖 `approval_status`/`semantic_version`。单查询 search 响应会正确展开 Milvus 的外层 query 数组：`data:[[]]` 返回零行，多命中返回语义卡行；query/list 解析不受影响。upsert 使用严格语义卡 Schema：全部 VARCHAR 在任何 Milvus 调用前按当前集合的 UTF-8 字节长度核验，`content` 上限为 8192 字节；重算并核对 64 位 `content_checksum`，要求 `embedding_revision`，拒绝 `chunk_index` 与 `embedding_dimension`，文件内 card_id 必须唯一且只能有一个语义版本。整份 JSONL 和全部已有 card_id 归属检查通过后才开始分批 upsert。v1 不再信任调用方提供的 `published:false`，`dbx_vector_delete_by_batch` 始终返回 `VECTOR_DELETE_DISABLED_V1`；撤销残留批次必须走受审计管理员流程。所有通过 `dbx_execute_query` 提交的向量实体 search/query/mutation 都返回 `VECTOR_DEDICATED_TOOL_REQUIRED`，通用入口只保留明确管理诊断。 向量检索的核心参数形状如下： @@ -141,7 +141,7 @@ Milvus 工具只接受 allowlist 中的集合、过滤字段和输出字段。 } ``` -仅依赖生效日期选择活动语义卡时可以省略 `semantic_version`，但 `active_at` 仍为必填。v1 的 `embedding` 必须恰好包含 1024 个有限数值；上面数组内的注释只是文档缩写，不是可直接提交的工具参数。 +`semantic_version` 与 `active_at` 均为必填，不能只依赖生效日期选择版本；调用方应先从已发布的新鲜度目录取得活动 `dataset_id + semantic_version`，再逐对检索。`active_at` 只接受 `Asia/Shanghai` 业务日期 `YYYY-MM-DD`，不接受带时区 datetime。v1 的 `embedding` 必须恰好包含 1024 个有限数值；上面数组内的注释只是文档缩写，不是可直接提交的工具参数。 ## 有状态查询会话 diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index 4b606f11e1..15fcf671e6 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -123,7 +123,7 @@ Preview and prepare share a bounded inspection semaphore and timeout. XLSX inspe Start rechecks the source, opens it with `O_NOFOLLOW` plus descriptor metadata checks, and performs cancellable copy plus SHA-256 in one pass into a task-private snapshot. It verifies that the task disk can hold the source, normalized-output ceiling, and reserve, then writes a size-limited governed CSV before opening a database pool. Short CSV/TSV rows retain missing values; extra empty trailing fields remain part of the full-record row hash, while any extra non-empty field fails with its real source row number. Every row gains `import_id`, `plan_id`, `source_sha`, the real absolute `source_row_number`, a `source_row_hash` over the complete normalized source record, and `loaded_at`. If exact lineage cannot be produced, the import fails before database access. Status returns the generated relation. Publication to business schemas remains separately governed. -Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires `active_at`, always adds `approval_status == "approved"` plus effective-date filters, and optionally applies an exact `semantic_version`. A single-query search flattens Milvus's outer query array: `data:[[]]` becomes zero rows and multiple hits become semantic-card rows, without changing query/list parsing. Upsert validates every VARCHAR against the current collection's UTF-8 byte length before any Milvus call (`content` is limited to 8192 bytes), verifies the 64-hex `content_checksum`, requires `embedding_revision`, rejects `chunk_index` and `embedding_dimension`, and validates the entire file plus all existing card-ID ownership before the first batch is written. V1 no longer trusts caller-supplied `published:false`; `dbx_vector_delete_by_batch` always returns `VECTOR_DELETE_DISABLED_V1`, and cleanup requires an audited administrator workflow. All entity search/query/mutation REST submitted through generic `dbx_execute_query` returns `VECTOR_DEDICATED_TOOL_REQUIRED`; only explicit management diagnostics remain available there. +Milvus tools accept only collections, filter fields, and output fields on their corresponding allowlists. Search requires both `active_at` and `semantic_version`. `active_at` is the `YYYY-MM-DD` business date after converting the query instant to `Asia/Shanghai`; `semantic_version` is a non-empty exact version of 1–128 UTF-8 bytes. The server always combines `approval_status=approved`, effective-date bounds, and exact semantic version; `filters` cannot supply or override `approval_status` or `semantic_version`. A single-query search flattens Milvus's outer query array: `data:[[]]` becomes zero rows and multiple hits become semantic-card rows, without changing query/list parsing. Upsert validates every VARCHAR against the current collection's UTF-8 byte length before any Milvus call (`content` is limited to 8192 bytes), verifies the 64-hex `content_checksum`, requires `embedding_revision`, rejects `chunk_index` and `embedding_dimension`, and validates the entire file plus all existing card-ID ownership before the first batch is written. V1 no longer trusts caller-supplied `published:false`; `dbx_vector_delete_by_batch` always returns `VECTOR_DELETE_DISABLED_V1`, and cleanup requires an audited administrator workflow. All entity search/query/mutation REST submitted through generic `dbx_execute_query` returns `VECTOR_DEDICATED_TOOL_REQUIRED`; only explicit management diagnostics remain available there. The core vector-search arguments are: @@ -141,7 +141,7 @@ The core vector-search arguments are: } ``` -Omit `semantic_version` when effective dates alone select the active cards. `active_at` remains required. In v1, `embedding` must contain exactly 1024 finite values; the abbreviated comment above is documentation only, not valid tool input. +Both `semantic_version` and `active_at` are required; effective dates alone must never select a version. The caller should first obtain active `dataset_id + semantic_version` pairs from the published freshness catalog, then search each pair. `active_at` accepts only an `Asia/Shanghai` business date in `YYYY-MM-DD`, not a timezone-bearing datetime. In v1, `embedding` must contain exactly 1024 finite values; the abbreviated comment above is documentation only, not valid tool input. ## Stateful Query Sessions diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index c580a4056e..e0dd995976 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -183,7 +183,7 @@ DBX connection storage defaults to: Override the directory with `DBX_DATA_DIR`. -Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Preview accepts Excel/JSON/CSV/TSV and reports whether row/range counts are exact. V1 governed import streams XLSX/XLSM and CSV/TSV with explicit `utf-8`; legacy `.xls`, JSON, `auto`, GBK, and UTF-16 remain preview-only. Inspection, ZIP expansion, shared strings, worksheets, cells, row batches, normalized output, and task disk all have hard budgets. Prepare accepts position+raw+canonical-name mappings and generates a unique `staging.mcp_` relation; callers cannot append to an existing table. Start performs cancellable copy+SHA into a private snapshot, hashes the complete source record with its absolute row number, and opens the database only after normalization succeeds. Semantic VARCHAR lengths are fully prevalidated, Milvus's single-search outer result array is flattened, and v1 batch deletion is disabled. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. +Local file imports additionally require `DBX_MCP_IMPORT_ROOTS`. Preview accepts Excel/JSON/CSV/TSV and reports whether row/range counts are exact. V1 governed import streams XLSX/XLSM and CSV/TSV with explicit `utf-8`; legacy `.xls`, JSON, `auto`, GBK, and UTF-16 remain preview-only. Inspection, ZIP expansion, shared strings, worksheets, cells, row batches, normalized output, and task disk all have hard budgets. Prepare accepts position+raw+canonical-name mappings and generates a unique `staging.mcp_` relation; callers cannot append to an existing table. Start performs cancellable copy+SHA into a private snapshot, hashes the complete source record with its absolute row number, and opens the database only after normalization succeeds. Semantic VARCHAR lengths are fully prevalidated, Milvus's single-search outer result array is flattened, and v1 batch deletion is disabled. Vector search requires both an exact `semantic_version` and an `Asia/Shanghai` business date `active_at` in `YYYY-MM-DD`; approved status, effective dates, and version are always enforced together. See the [MCP documentation](../../docs/content/docs/mcp.mdx#local-file-imports-and-milvus-tools) for the full contract. ### Agent/JDBC databases @@ -513,7 +513,7 @@ MCP 配置： 通过 `DBX_DATA_DIR` 覆盖默认目录。Windows 便携版应指向 `DBX.exe` 同级、包含 `dbx.db` 的 `data` 文件夹。 -连接管理默认关闭；仅无作用域安装维护进程临时设置 `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1`，任何 `DBX_MCP_SCOPE_*` 都继续禁用新增、复制和删除。本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。preview 支持 Excel/JSON/CSV/TSV 并标明行数/范围是否精确；v1 治理入库流式支持 XLSX/XLSM 和显式 `utf-8` 的 CSV/TSV，旧版 `.xls`、JSON、`auto`、GBK 与 UTF-16 暂为 preview-only。inspection、ZIP 解压、sharedStrings、worksheet、单元格、批次、规范化输出与任务磁盘均有硬预算。start 以可取消复制+SHA 生成任务私有快照，保留绝对行号和完整源记录哈希，规范化成功后才打开数据库。语义 VARCHAR 会整文件预校验，Milvus 单查询 search 会展开外层数组，v1 批次删除禁用。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 +连接管理默认关闭；仅无作用域安装维护进程临时设置 `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1`，任何 `DBX_MCP_SCOPE_*` 都继续禁用新增、复制和删除。本地文件导入还必须配置 `DBX_MCP_IMPORT_ROOTS`。preview 支持 Excel/JSON/CSV/TSV 并标明行数/范围是否精确；v1 治理入库流式支持 XLSX/XLSM 和显式 `utf-8` 的 CSV/TSV，旧版 `.xls`、JSON、`auto`、GBK 与 UTF-16 暂为 preview-only。inspection、ZIP 解压、sharedStrings、worksheet、单元格、批次、规范化输出与任务磁盘均有硬预算。start 以可取消复制+SHA 生成任务私有快照，保留绝对行号和完整源记录哈希，规范化成功后才打开数据库。语义 VARCHAR 会整文件预校验，Milvus 单查询 search 会展开外层数组，v1 批次删除禁用。向量检索必须同时提供精确 `semantic_version` 与 `Asia/Shanghai` 的 `YYYY-MM-DD` 业务日期 `active_at`，服务端始终组合 approved、有效期和版本过滤。完整契约见 [MCP 中文文档](../../docs/content/docs/mcp.cn.mdx#本地文件导入与-milvus-工具)。 ### DBX Web / Docker From 72e931a79a7dd69a07981dba939a70824f8eff4e Mon Sep 17 00:00:00 2001 From: xingranya Date: Tue, 25 Aug 2026 21:00:17 +0800 Subject: [PATCH 08/10] fix(mcp): resolve all scoped connections --- crates/dbx-mcp/src/server.rs | 177 +++++++++++++++++++++++++------ crates/dbx-mcp/tests/protocol.rs | 81 ++++++++++++++ docs/content/docs/mcp.cn.mdx | 2 + docs/content/docs/mcp.mdx | 2 + packages/mcp-server/README.md | 4 +- 5 files changed, 229 insertions(+), 37 deletions(-) diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 949a465656..9a5163886d 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -1640,48 +1640,69 @@ impl DbxMcpServer { } return Ok(ResolvedConnection { connection, policy }); } - if self.scope.connection_scope_enabled() { - let connection = connections - .into_iter() - .find(|connection| self.scope.matches(connection)) - .ok_or_else(|| tool_error("CONNECTION_NOT_FOUND", "Scoped DBX connection was not found."))?; - if let Some(name) = selector.connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) { - if name != connection.name && name != connection.id { - return Err(tool_error( - "CONNECTION_OUT_OF_SCOPE", - format!("Connection \"{name}\" is outside this DBX AI session scope."), - )); + let scoped = connections + .iter() + .filter(|connection| !self.scope.connection_scope_enabled() || self.scope.matches(connection)) + .collect::>(); + let allowed = scoped + .iter() + .copied() + .filter(|connection| policy_allows_connection(&policy, connection)) + .collect::>(); + + if let Some(name) = selector.connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) { + if let Some(exact_id) = connections.iter().find(|connection| connection.id == name) { + if allowed.iter().any(|connection| connection.id == exact_id.id) { + return Ok(ResolvedConnection { connection: exact_id.clone(), policy }); } - } - if !policy_allows_connection(&policy, &connection) { return Err(tool_error( "CONNECTION_OUT_OF_SCOPE", - "The DBX AI session scope is outside the global MCP connection allowlist.", + format!("Connection \"{name}\" is outside the effective DBX MCP scope or allowlist."), )); } - return Ok(ResolvedConnection { connection, policy }); + let matching = allowed + .iter() + .copied() + .filter(|connection| connection.name.eq_ignore_ascii_case(name)) + .cloned() + .collect::>(); + return match matching.as_slice() { + [connection] => Ok(ResolvedConnection { connection: connection.clone(), policy }), + [_, _, ..] => Err(tool_error("AMBIGUOUS_CONNECTION", ambiguous_connections(name, &matching))), + [] => { + let exists = connections + .iter() + .any(|connection| connection.id == name || connection.name.eq_ignore_ascii_case(name)); + if !exists { + Err(tool_error("CONNECTION_NOT_FOUND", format!("Connection \"{name}\" not found."))) + } else { + Err(tool_error( + "CONNECTION_OUT_OF_SCOPE", + format!("Connection \"{name}\" is outside the effective DBX MCP scope or allowlist."), + )) + } + } + }; } - let Some(name) = selector.connection_name.as_deref().map(str::trim).filter(|name| !name.is_empty()) else { - return Err(tool_error("CONNECTION_NOT_FOUND", "Either connection_id or connection_name is required.")); - }; - let matching = - connections.into_iter().filter(|connection| connection.name.eq_ignore_ascii_case(name)).collect::>(); - let allowed = matching - .iter() - .filter(|connection| policy_allows_connection(&policy, connection)) - .cloned() - .collect::>(); - match allowed.as_slice() { - [] if matching.is_empty() => { - Err(tool_error("CONNECTION_NOT_FOUND", format!("Connection \"{name}\" not found."))) - } - [] => Err(tool_error( - "CONNECTION_OUT_OF_SCOPE", - format!("Connection \"{name}\" is not allowed by DBX MCP settings."), - )), - [connection] => Ok(ResolvedConnection { connection: connection.clone(), policy }), - _ => Err(tool_error("AMBIGUOUS_CONNECTION", ambiguous_connections(name, &allowed))), + + if self.scope.connection_scope_enabled() { + return match allowed.as_slice() { + [connection] => Ok(ResolvedConnection { connection: (*connection).clone(), policy }), + [] if scoped.is_empty() => { + Err(tool_error("CONNECTION_NOT_FOUND", "Scoped DBX connection was not found.")) + } + [] => Err(tool_error( + "CONNECTION_OUT_OF_SCOPE", + "The DBX AI session scope is outside the global MCP connection allowlist.", + )), + _ => Err(tool_error( + "CONNECTION_SELECTOR_REQUIRED", + "Multiple scoped DBX connections are available; provide connection_id or connection_name.", + )), + }; } + + Err(tool_error("CONNECTION_NOT_FOUND", "Either connection_id or connection_name is required.")) } } @@ -2782,6 +2803,92 @@ mod tests { assert!(result_text(&result).contains("CONNECTION_NOT_FOUND")); } + #[tokio::test] + async fn multi_connection_scope_resolves_all_allowed_selectors_and_fails_closed_without_one() { + let query = connection("query-id", "Operations Query", "postgres", "enterprise"); + let semantic = connection("semantic-id", "Operations Semantics", "milvus", "default"); + let trace = connection("trace-id", "Operations Trace", "postgres", "enterprise"); + let scope_only = connection("scope-only", "Scope Only", "postgres", "enterprise"); + let outside = connection("outside", "Outside", "postgres", "enterprise"); + let connections = vec![query.clone(), semantic.clone(), trace.clone(), scope_only, outside]; + let policy = McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: Some(vec![ + "query-id".to_string(), + "semantic-id".to_string(), + "trace-id".to_string(), + "outside".to_string(), + ]), + }; + let server = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend { connections: connections.clone(), policy: policy.clone(), ..Default::default() }), + McpScope { + connection_ids: vec![ + "query-id".to_string(), + "semantic-id".to_string(), + "trace-id".to_string(), + "scope-only".to_string(), + ], + ..Default::default() + }, + false, + ); + + for (selector, expected_id) in [ + ( + ConnectionSelector { connection_id: None, connection_name: Some("operations query".to_string()) }, + "query-id", + ), + ( + ConnectionSelector { connection_id: None, connection_name: Some("OPERATIONS SEMANTICS".to_string()) }, + "semantic-id", + ), + ( + ConnectionSelector { connection_id: None, connection_name: Some("semantic-id".to_string()) }, + "semantic-id", + ), + (ConnectionSelector { connection_id: Some("trace-id".to_string()), connection_name: None }, "trace-id"), + ] { + assert_eq!(server.resolve_connection(&selector).await.unwrap().connection.id, expected_id); + } + + let empty_selector = ConnectionSelector { connection_id: None, connection_name: None }; + let missing_selector = match server.resolve_connection(&empty_selector).await { + Err(error) => error, + Ok(_) => panic!("多个 scoped connection 不得隐式选择"), + }; + assert!(result_text(&missing_selector).contains("CONNECTION_SELECTOR_REQUIRED")); + for name in ["Scope Only", "Outside"] { + let error = match server + .resolve_connection(&ConnectionSelector { + connection_id: None, + connection_name: Some(name.to_string()), + }) + .await + { + Err(error) => error, + Ok(_) => panic!("scope 或 allowlist 外连接不得解析：{name}"), + }; + assert!(result_text(&error).contains("CONNECTION_OUT_OF_SCOPE")); + } + + let single_id = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend { connections: connections.clone(), policy: policy.clone(), ..Default::default() }), + McpScope { connection_ids: vec!["query-id".to_string()], ..Default::default() }, + false, + ); + let empty_selector = ConnectionSelector { connection_id: None, connection_name: None }; + assert_eq!(single_id.resolve_connection(&empty_selector).await.unwrap().connection.id, query.id); + let single_name = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend { connections, policy, ..Default::default() }), + McpScope { connection_name: Some("Operations Trace".to_string()), ..Default::default() }, + false, + ); + let empty_selector = ConnectionSelector { connection_id: None, connection_name: None }; + assert_eq!(single_name.resolve_connection(&empty_selector).await.unwrap().connection.id, trace.id); + } + #[test] fn scoped_connection_ids_are_deduplicated_and_take_precedence_over_name() { assert_eq!(scoped_connection_ids(Some(" first, second,first ,, ")), vec!["first", "second"]); diff --git a/crates/dbx-mcp/tests/protocol.rs b/crates/dbx-mcp/tests/protocol.rs index 06ce1da003..ac8a07812f 100644 --- a/crates/dbx-mcp/tests/protocol.rs +++ b/crates/dbx-mcp/tests/protocol.rs @@ -443,3 +443,84 @@ async fn runtime_connection_scope_preserves_group_paths() { client.cancel().await.expect("close MCP client"); server_task.abort(); } + +#[tokio::test] +async fn multi_connection_runtime_scope_routes_each_name_and_requires_selector() { + let backend = PolicyBackend { + policy: McpGlobalPolicy { + read_only: false, + allow_dangerous_sql: false, + allowed_connection_ids: Some(vec![ + "query-id".to_string(), + "semantic-id".to_string(), + "trace-id".to_string(), + "outside".to_string(), + ]), + }, + connections: vec![ + test_connection("query-id", "Operations Query"), + test_connection("semantic-id", "Operations Semantics"), + test_connection("trace-id", "Operations Trace"), + test_connection("scope-only", "Scope Only"), + test_connection("outside", "Outside"), + ], + group_paths: Ok(HashMap::new()), + }; + let (server_transport, client_transport) = tokio::io::duplex(32 * 1024); + let server = DbxMcpServer::with_runtime_options( + Arc::new(backend), + McpScope { + connection_ids: vec![ + "query-id".to_string(), + "semantic-id".to_string(), + "trace-id".to_string(), + "scope-only".to_string(), + ], + ..Default::default() + }, + false, + ); + let server_task = tokio::spawn(async move { server.serve(server_transport).await }); + let client = ().serve(client_transport).await.expect("initialize MCP client"); + + for name in ["operations query", "OPERATIONS SEMANTICS", "Operations Trace"] { + let result = client + .peer() + .call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments( + json!({ "connection_name": name, "sql": "SELECT 1" }).as_object().cloned().unwrap_or_else(Map::new), + )) + .await + .expect("route scoped connection by name"); + let text = result.content[0].as_text().expect("query result").text.as_str(); + assert!(text.contains("query should have been blocked"), "{name} did not reach selected backend: {text}"); + assert!(!text.contains("CONNECTION_OUT_OF_SCOPE")); + } + + let no_selector = client + .peer() + .call_tool( + CallToolRequestParams::new("dbx_execute_query") + .with_arguments(json!({ "sql": "SELECT 1" }).as_object().cloned().unwrap_or_else(Map::new)), + ) + .await + .expect("reject missing selector"); + assert!(no_selector.content[0] + .as_text() + .expect("missing selector result") + .text + .contains("CONNECTION_SELECTOR_REQUIRED")); + + for name in ["Scope Only", "Outside"] { + let result = client + .peer() + .call_tool(CallToolRequestParams::new("dbx_execute_query").with_arguments( + json!({ "connection_name": name, "sql": "SELECT 1" }).as_object().cloned().unwrap_or_else(Map::new), + )) + .await + .expect("reject out-of-scope selector"); + assert!(result.content[0].as_text().expect("out-of-scope result").text.contains("CONNECTION_OUT_OF_SCOPE")); + } + + client.cancel().await.expect("close MCP client"); + server_task.abort(); +} diff --git a/docs/content/docs/mcp.cn.mdx b/docs/content/docs/mcp.cn.mdx index 60e0a932bb..87a9fd2f9a 100644 --- a/docs/content/docs/mcp.cn.mdx +++ b/docs/content/docs/mcp.cn.mdx @@ -115,6 +115,8 @@ DBX MCP 定义 21 个工具；日常默认注册 18 个，只有显式无作用 连接管理默认关闭；只有无任何 `DBX_MCP_SCOPE_*` 且安装维护进程显式设置 `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1` 时，新增、复制和删除连接工具才会注册。启用任一作用域后，即使设置该开关也会隐藏并在直接调用时阻断连接管理。删除还会复核全局 `allowed_connection_ids`。桌面 UI 工具在 scoped 模式下同样隐藏。 +连接解析始终先取运行时 scope 与全局 `allowed_connection_ids` 的交集。`connection_id` 按 ID 精确解析；`connection_name` 可按 ID 精确解析，也可在全部有效候选中按名称大小写不敏感解析。只有交集内恰好一个连接且没有 selector 时才允许隐式使用；多个 `DBX_MCP_SCOPE_CONNECTION_IDS` 同时有效时，省略 selector 会返回 `CONNECTION_SELECTOR_REQUIRED`。scope 或全局 allowlist 之外的连接继续返回 `CONNECTION_OUT_OF_SCOPE`。 + ## 本地文件导入与 Milvus 工具 导入工具只支持本地 DBX Desktop/MCP 模式；在 Web 模式中会稳定返回 `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1`。必须用平台路径分隔符在 `DBX_MCP_IMPORT_ROOTS` 中配置一个或多个真实目录。文件路径必须是绝对路径，解析后仍位于允许目录内，是普通且非符号链接的文件，扩展名受支持，并且没有超过大小限制。preview 支持 Excel、JSON、CSV 和 TSV，并明确返回 `usedRangeExact` 与 `totalRowsExact`；受限预览不会把已扫描行数伪装成全文件总行数。v1 治理入库使用有界流式处理支持 XLSX/XLSM 和显式指定 `utf-8` 的 CSV/TSV。旧版 `.xls`、JSON、`auto`、GBK 和 UTF-16 分隔文件会在创建计划和访问数据库之前返回稳定的不支持错误。 diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index 15fcf671e6..f26373018b 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -115,6 +115,8 @@ DBX MCP defines 21 tools. Daily mode registers 18 by default; an explicit unscop Connection management is default-closed. Add, duplicate, and remove tools are registered only when an unscoped maintenance process explicitly sets `DBX_MCP_ENABLE_CONNECTION_MANAGEMENT=1`. Any `DBX_MCP_SCOPE_*` setting disables these tools even when the maintenance flag is present, and remove also enforces the global `allowed_connection_ids`. Scoped sessions also hide desktop UI tools. +Connection resolution always starts from the intersection of the runtime scope and global `allowed_connection_ids`. `connection_id` resolves an exact ID. `connection_name` may resolve an exact ID or a case-insensitive name across all effective candidates. An omitted selector is accepted only when that intersection contains exactly one connection; multiple effective `DBX_MCP_SCOPE_CONNECTION_IDS` return `CONNECTION_SELECTOR_REQUIRED`. Connections outside either boundary continue to return `CONNECTION_OUT_OF_SCOPE`. + ## Local file imports and Milvus tools The import tools are available only in local DBX Desktop/MCP mode. In Web mode they return the stable `IMPORT_UNSUPPORTED_IN_WEB_MODE_V1` code. Configure `DBX_MCP_IMPORT_ROOTS` with one or more real directories using the platform path separator. Every file path must be absolute, resolve inside one configured root, be a regular non-symlink file, use an allowed extension, and stay below the configured size limit. Preview supports Excel, JSON, CSV, and TSV and explicitly reports `usedRangeExact` and `totalRowsExact`; a bounded scan is never presented as a complete row count. Governed v1 database import uses bounded streaming for XLSX/XLSM and CSV/TSV whose prepare request explicitly selects `utf-8`. Legacy `.xls`, JSON, `auto`, GBK, and UTF-16 delimited files return stable unsupported errors before a plan or database write exists. diff --git a/packages/mcp-server/README.md b/packages/mcp-server/README.md index e0dd995976..2f24213af8 100644 --- a/packages/mcp-server/README.md +++ b/packages/mcp-server/README.md @@ -262,7 +262,7 @@ Legacy connection scope variables can still narrow the DBX allowlist for existin } ``` -Use `DBX_MCP_SCOPE_CONNECTION_ID`, comma-separated `DBX_MCP_SCOPE_CONNECTION_IDS`, or `DBX_MCP_SCOPE_CONNECTION_NAME`. ID scopes take precedence over the name scope. The scoped database is optional. +Use `DBX_MCP_SCOPE_CONNECTION_ID`, comma-separated `DBX_MCP_SCOPE_CONNECTION_IDS`, or `DBX_MCP_SCOPE_CONNECTION_NAME`. ID scopes take precedence over the name scope. The scoped database is optional. Resolution uses the intersection with the global connection allowlist. An exact `connection_id`, exact ID passed through `connection_name`, or case-insensitive name can select any effective candidate. Omit the selector only when exactly one candidate remains; multiple scoped IDs return `CONNECTION_SELECTOR_REQUIRED`. ## Safety @@ -571,7 +571,7 @@ DBX 在 **设置 → MCP** 中保存一份权威策略，并在每次请求时 } ``` -可使用 `DBX_MCP_SCOPE_CONNECTION_ID`、逗号分隔的 `DBX_MCP_SCOPE_CONNECTION_IDS` 或 `DBX_MCP_SCOPE_CONNECTION_NAME`。ID scope 优先于名称 scope；作用域模式会隐藏连接增删和桌面 UI 工具。 +可使用 `DBX_MCP_SCOPE_CONNECTION_ID`、逗号分隔的 `DBX_MCP_SCOPE_CONNECTION_IDS` 或 `DBX_MCP_SCOPE_CONNECTION_NAME`。ID scope 优先于名称 scope；作用域模式会隐藏连接增删和桌面 UI 工具。连接候选取 scope 与全局 allowlist 的交集，可按精确 ID 或大小写不敏感名称选择；只有唯一候选时才能省略 selector，多个 scoped ID 会返回 `CONNECTION_SELECTOR_REQUIRED`。 ### SQL 和命令安全 From 86e5eeb13b2bd0a0546687185a6fb6fdd5c184b8 Mon Sep 17 00:00:00 2001 From: xingranya Date: Thu, 27 Aug 2026 21:52:30 +0800 Subject: [PATCH 09/10] fix(core): initialize benchmark source row numbers --- crates/dbx-core/examples/data_transfer_bench.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/dbx-core/examples/data_transfer_bench.rs b/crates/dbx-core/examples/data_transfer_bench.rs index ac82902943..a70013d96d 100644 --- a/crates/dbx-core/examples/data_transfer_bench.rs +++ b/crates/dbx-core/examples/data_transfer_bench.rs @@ -70,6 +70,7 @@ fn run() -> Result<(), String> { let parsed = ParsedImportFile { columns: columns.clone(), rows: rows.clone(), + source_row_numbers: (1..=rows.len()).collect(), total_rows: rows.len(), effective_encoding: None, }; From 57478532c063e58ab3ea8230f9789090c7773be1 Mon Sep 17 00:00:00 2001 From: xingranya Date: Thu, 27 Aug 2026 22:09:38 +0800 Subject: [PATCH 10/10] fix(core): preserve SQL import source row numbers --- crates/dbx-core/src/table_import.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/crates/dbx-core/src/table_import.rs b/crates/dbx-core/src/table_import.rs index ac2fbb069a..97ce384ab2 100644 --- a/crates/dbx-core/src/table_import.rs +++ b/crates/dbx-core/src/table_import.rs @@ -1588,7 +1588,14 @@ pub fn parse_sql_bytes_with_options( } let target = target.ok_or_else(|| "No INSERT statements found in SQL file".to_string())?; - Ok(ParsedImportFile { columns: target.columns, rows, total_rows, effective_encoding: Some(encoding) }) + let source_row_numbers = (1..=rows.len()).collect(); + Ok(ParsedImportFile { + columns: target.columns, + rows, + source_row_numbers, + total_rows, + effective_encoding: Some(encoding), + }) } pub fn parse_sql_bytes(bytes: &[u8], preview_limit: usize) -> Result { @@ -8404,6 +8411,7 @@ mod tests { assert_eq!(parsed.columns, vec!["id", "name"]); assert_eq!(parsed.total_rows, 3); + assert_eq!(parsed.source_row_numbers, vec![1, 2, 3]); assert_eq!(parsed.rows[0], vec![serde_json::json!(1), serde_json::json!("Ada")]); assert_eq!(parsed.rows[2], vec![serde_json::json!(3), serde_json::json!("Cathy")]); } @@ -8577,6 +8585,7 @@ mod tests { assert_eq!(parsed.total_rows, 5); assert_eq!(parsed.rows.len(), 2); + assert_eq!(parsed.source_row_numbers, vec![1, 2]); } #[test] @@ -10146,6 +10155,7 @@ mod tests { let data = ParsedImportFile { columns: vec!["name".to_string()], rows: vec![vec![serde_json::json!("Ada")]], + source_row_numbers: Vec::new(), total_rows: 1, effective_encoding: None, };