diff --git a/apps/desktop/src/components/structure/TableStructureEditor.vue b/apps/desktop/src/components/structure/TableStructureEditor.vue index 6b8e45e041..7d3f3a6ad4 100644 --- a/apps/desktop/src/components/structure/TableStructureEditor.vue +++ b/apps/desktop/src/components/structure/TableStructureEditor.vue @@ -2301,6 +2301,7 @@ async function revalidateCachedStructureMetadata(loadRequestId: number, scope: { const revalidationId = ++structureMetadataRevalidationId; const metadataRequest = { connectionId, database, schema, tableName, catalog }; try { + await store.ensureConnected(connectionId); // Force alone only clears this facet's own key; the web backend keeps its // own backend-columns/backend-comment entries under the same table prefix // and would serve them to the forced re-fetch. Drop the whole table scope @@ -3923,10 +3924,18 @@ onMounted(() => { void loadTableOwnerRoles(); void loadMysqlTableEngine(props.draft?.mysqlTableEngine !== undefined); if (props.draft?.initialized) { - void hydrateRestoredDraftFromDatabase().then(() => { + // A clean persisted editor snapshot is not a live schema cache. After an + // MCP DDL, restoring its loaded-facet flags would otherwise bypass the + // invalidated backend cache entirely. Legacy/dirty drafts remain intact. + const revalidateRestoredColumns = props.draft.dirty === false && !isCreateMode.value && loadedMetadataFacets.has("columns") && !(databaseType.value === "manticoresearch" && tableMetadataCapabilities.value.ddl); + void hydrateRestoredDraftFromDatabase().then(async () => { applyInitialStructureTarget(); void loadMysqlAutoIncrementCounter(true); - void loadActiveTableStructureMetadataIfNeeded(); + await loadActiveTableStructureMetadataIfNeeded(); + if (revalidateRestoredColumns) { + // The existing revalidation checks again for edits made while loading. + void revalidateCachedStructureMetadata(structureLoadRequestId, { columns: true, tableComment: false }, undefined); + } }); } else if (isCreateMode.value) { markDraftHydratedAndSync(); diff --git a/apps/desktop/src/components/structure/__tests__/TableStructureEditor.metadataRevalidation.spec.ts b/apps/desktop/src/components/structure/__tests__/TableStructureEditor.metadataRevalidation.spec.ts index bfe0761650..c3f3aa3aa4 100644 --- a/apps/desktop/src/components/structure/__tests__/TableStructureEditor.metadataRevalidation.spec.ts +++ b/apps/desktop/src/components/structure/__tests__/TableStructureEditor.metadataRevalidation.spec.ts @@ -252,6 +252,8 @@ vi.mock("@/lib/backend/api", () => ({ })); import TableStructureEditor from "@/components/structure/TableStructureEditor.vue"; +import { createColumnDrafts } from "@/lib/table/tableStructureEditorState"; +import type { TableStructureEditorDraft } from "@/types/database"; const mountedApps: App[] = []; @@ -272,13 +274,14 @@ async function settle() { } } -async function mountStructureEditor() { +async function mountStructureEditor(draft?: TableStructureEditorDraft) { const root = document.createElement("div"); document.body.append(root); const app = createApp(TableStructureEditor, { connectionId: mocks.connection.id, database: "test", tableName: "users", + draft, }); mountedApps.push(app); app.mount(root); @@ -404,3 +407,61 @@ describe("TableStructureEditor metadata revalidation (#8816)", () => { expect(root.textContent).toContain("email"); }); }); + +describe("restored structure snapshots after external DDL", () => { + function restoredDraft(dirty: boolean | undefined): TableStructureEditorDraft { + return { + initialized: true, + dirty, + activeTab: "columns", + newTableName: "", + tableComment: "", + originalTableComment: "", + columns: createColumnDrafts( + initialColumns.map((column) => ({ name: column.name, data_type: column.data_type, is_nullable: column.nullable, column_default: column.default_value, is_primary_key: false, extra: null, comment: column.comment })), + "mysql", + ), + indexes: [], + foreignKeys: [], + triggers: [], + loadedMetadataFacets: ["columns", "comment"], + }; + } + + it("revalidates an explicitly clean restored snapshot", async () => { + mocks.loadObjectMetadataFacet.mockImplementation(async (_request, facet) => ({ value: facet === "columns" ? rebuiltColumns : facet === "comment" ? "" : [], cacheStatus: "remote" })); + const root = await mountStructureEditor(restoredDraft(false)); + await vi.waitFor(() => expect(Array.from(root.querySelectorAll("[data-column-name-input]")).map((input) => input.value)).toContain("email")); + expect(columnsFacetCalls()).toHaveLength(1); + }); + + it.each([true, undefined])("preserves restored edits when dirty is %s", async (dirty) => { + mocks.loadObjectMetadataFacet.mockImplementation(async (_request, facet) => ({ value: facet === "columns" ? rebuiltColumns : facet === "comment" ? "" : [], cacheStatus: "remote" })); + const draft = restoredDraft(dirty); + draft.columns[0].name = "local_edit"; + const root = await mountStructureEditor(draft); + await settle(); + expect(root.querySelector("[data-column-name-input]")?.value).toBe("local_edit"); + expect(columnsFacetCalls()).toHaveLength(0); + }); + + it("preserves an edit started while restored snapshot validation is pending", async () => { + let resolveColumns!: (value: unknown) => void; + mocks.loadObjectMetadataFacet.mockImplementation((_request, facet) => + facet === "columns" + ? new Promise((resolve) => { + resolveColumns = resolve; + }) + : Promise.resolve({ value: facet === "comment" ? "" : [], cacheStatus: "remote" }), + ); + const root = await mountStructureEditor(restoredDraft(false)); + await vi.waitFor(() => expect(resolveColumns).toBeDefined()); + const input = root.querySelector("[data-column-name-input]")!; + input.value = "local_edit"; + input.dispatchEvent(new Event("input")); + await settle(); + resolveColumns({ value: rebuiltColumns, cacheStatus: "remote" }); + await settle(); + expect(root.querySelector("[data-column-name-input]")?.value).toBe("local_edit"); + }); +}); diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index 322452d211..3ebce7875c 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -68,6 +68,7 @@ pub mod mysql_event_sql; pub mod nacos; #[cfg(all(target_os = "windows", target_env = "gnu"))] mod nanosleep_stub; +pub mod object_cache; pub mod object_source_sql; pub mod path_utils; pub mod plugins; diff --git a/crates/dbx-core/src/object_cache.rs b/crates/dbx-core/src/object_cache.rs new file mode 100644 index 0000000000..8ec76dca5b --- /dev/null +++ b/crates/dbx-core/src/object_cache.rs @@ -0,0 +1,136 @@ +//! Persisted object metadata shared by Web and SQL execution entry points. +//! +//! These keys are also produced by the frontend's encodeURIComponent-based +//! caches. Keep the encoding and namespace stable when invalidating snapshots. + +use crate::models::connection::DatabaseType; +use crate::storage::Storage; + +pub const OBJECT_METADATA_CACHE_PREFIX: &str = "object-meta:v1"; +const OBJECT_DDL_CACHE_PREFIX: &str = "object-ddl:v1"; + +pub fn metadata_cache_segment(value: &str) -> String { + const HEX: &[u8; 16] = b"0123456789ABCDEF"; + let mut encoded = String::with_capacity(value.len()); + for byte in value.bytes() { + match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'!' | b'~' | b'*' | b'\'' | b'(' | b')' => { + encoded.push(byte as char) + } + _ => { + encoded.push('%'); + encoded.push(HEX[(byte >> 4) as usize] as char); + encoded.push(HEX[(byte & 0x0f) as usize] as char); + } + } + } + encoded +} + +pub fn object_metadata_cache_prefix(connection_id: &str, database: &str) -> String { + format!( + "{}:{}:{}:", + OBJECT_METADATA_CACHE_PREFIX, + metadata_cache_segment(connection_id), + metadata_cache_segment(database) + ) +} + +/// Conservative cache policy, separate from SQL authorization. A failed DDL +/// attempt may already have committed, so callers apply this even on errors. +pub(crate) fn sql_may_change_object_metadata(sql: &str, db_type: Option) -> bool { + let statements = match db_type { + Some(db_type) => crate::sql::split_sql_statements_for_database(sql, db_type), + None => crate::sql::split_sql_statements(sql), + }; + statements.iter().any(|statement| { + let risk = match db_type { + Some(db_type) => crate::sql_risk::classify_sql_risk_for_database(statement, db_type), + None => crate::sql_risk::classify_sql_risk(statement, "generic"), + }; + if matches!(risk, Ok(crate::sql_risk::SqlRisk::Ddl)) { + return true; + } + // Include transaction completion: a previous DDL call may have run + // before commit, while another connection cached the old structure. + let keywords = [ + "CREATE", "ALTER", "DROP", "TRUNCATE", "COMMENT", "GRANT", "REVOKE", "RENAME", "COMMIT", "END", "ROLLBACK", + "ABORT", + ]; + match db_type { + Some(db_type) => crate::sql::starts_with_executable_sql_keyword_for_database(statement, &keywords, db_type), + None => crate::sql::starts_with_executable_sql_keyword(statement, &keywords), + } + }) +} + +/// Connection scope also covers qualified cross-database DDL and session USE. +/// Await before returning the SQL result; cache failures must not turn an +/// already committed database operation into an apparent execution failure. +pub(crate) async fn invalidate_connection_object_cache(storage: &Storage, connection_id: &str) { + let connection = metadata_cache_segment(connection_id); + for namespace in [OBJECT_METADATA_CACHE_PREFIX, OBJECT_DDL_CACHE_PREFIX] { + let prefix = format!("{namespace}:{connection}:"); + if let Err(error) = storage.delete_schema_cache_prefix(&prefix).await { + log::warn!( + "failed to invalidate object cache: connection_id={connection_id} namespace={namespace} error={error}" + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ddl_schema_cache_classifies_scripts_without_changing_permissions() { + for sql in [ + "ALTER TABLE users ADD COLUMN test4 INTEGER", + "BEGIN; ALTER TABLE users ADD COLUMN test4 INTEGER; COMMIT;", + "-- leading comment\nCOMMENT ON COLUMN users.test4 IS 'new'", + "CREATE SOMETHING dialect_specific", + "COMMIT", + "END", + "ROLLBACK", + "ABORT", + ] { + assert!(sql_may_change_object_metadata(sql, Some(DatabaseType::Postgres)), "{sql}"); + } + for sql in [ + "SELECT 'ALTER TABLE users'", + "-- DROP TABLE users\nSELECT 1", + "INSERT INTO users VALUES (1)", + "UPDATE users SET id = 1", + "DELETE FROM users", + "BEGIN", + "", + "/* ALTER TABLE */", + ] { + assert!(!sql_may_change_object_metadata(sql, Some(DatabaseType::Postgres)), "{sql}"); + } + } + + #[tokio::test] + async fn ddl_schema_cache_encoding_and_connection_isolation() { + let dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&dir.path().join("storage.db")).await.unwrap(); + let id = "conn:% 中文"; + assert_eq!(metadata_cache_segment(id), "conn%3A%25%20%E4%B8%AD%E6%96%87"); + let prefix = object_metadata_cache_prefix(id, "db% name"); + assert_eq!(prefix, "object-meta:v1:conn%3A%25%20%E4%B8%AD%E6%96%87:db%25%20name:"); + let removed = [ + format!("{prefix}public:users::backend-columns:"), + format!("object-ddl:v1:{}:other-db:public:users::TABLE:", metadata_cache_segment(id)), + ]; + let kept = format!("{}columns:", object_metadata_cache_prefix(&format!("{id}2"), "db% name")); + for key in removed.iter().chain(std::iter::once(&kept)) { + storage.save_schema_cache(key, &serde_json::json!([])).await.unwrap(); + } + invalidate_connection_object_cache(&storage, id).await; + for key in removed { + assert!(storage.load_schema_cache(&key).await.unwrap().is_none()); + } + assert!(storage.load_schema_cache(&kept).await.unwrap().is_some()); + } +} diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 9696eb78d9..8977fe0a78 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -2774,6 +2774,33 @@ pub async fn execute_sql_statement_with_options_typed( schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, +) -> Result { + let db_type = connection_database_type(state, connection_id).await; + let invalidate = crate::object_cache::sql_may_change_object_metadata(sql, db_type); + let result = execute_sql_statement_with_options_typed_inner( + state, + connection_id, + database, + sql, + schema, + cancel_token, + options, + ) + .await; + if invalidate { + crate::object_cache::invalidate_connection_object_cache(&state.storage, connection_id).await; + } + result +} + +async fn execute_sql_statement_with_options_typed_inner( + state: &AppState, + connection_id: &str, + database: &str, + sql: &str, + schema: Option<&str>, + cancel_token: Option, + options: QueryExecutionOptions, ) -> Result { // MongoDB connections use shell-style commands dispatched through the // frontend parser. Queries that fall through to the generic SQL executor @@ -3093,6 +3120,35 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( cancel_token: Option, options: QueryExecutionOptions, progress: Option, +) -> Result, QueryExecutionError> { + let db_type = connection_database_type(state, connection_id).await; + let invalidate = crate::object_cache::sql_may_change_object_metadata(sql, db_type); + let result = execute_multi_core_with_options_for_client_and_progress_typed_inner( + state, + connection_id, + database, + sql, + schema, + cancel_token, + options, + progress, + ) + .await; + if invalidate { + crate::object_cache::invalidate_connection_object_cache(&state.storage, connection_id).await; + } + result +} + +async fn execute_multi_core_with_options_for_client_and_progress_typed_inner( + state: &AppState, + connection_id: &str, + database: &str, + sql: &str, + schema: Option<&str>, + cancel_token: Option, + options: QueryExecutionOptions, + progress: Option, ) -> Result, QueryExecutionError> { let pool_database = query_pool_database(database, options.catalog.as_deref()); // Reject MongoDB queries that fall through to the generic executor. @@ -3931,6 +3987,23 @@ pub async fn execute_statements( statements: &[String], schema: Option<&str>, timeout_secs: Option, +) -> Result { + let db_type = connection_database_type(state, connection_id).await; + let invalidate = statements.iter().any(|sql| crate::object_cache::sql_may_change_object_metadata(sql, db_type)); + let result = execute_statements_inner(state, connection_id, database, statements, schema, timeout_secs).await; + if invalidate { + crate::object_cache::invalidate_connection_object_cache(&state.storage, connection_id).await; + } + result +} + +async fn execute_statements_inner( + state: &AppState, + connection_id: &str, + database: &str, + statements: &[String], + schema: Option<&str>, + timeout_secs: Option, ) -> Result { let sql_ctx = statements.first().map(|s| s.as_str()).unwrap_or(""); let pool_key = if database.is_empty() { @@ -7486,6 +7559,136 @@ for line in sys.stdin: } } + #[tokio::test] + async fn ddl_schema_cache_invalidates_persisted_object_snapshots() { + let dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&dir.path().join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "schema-cache"; + let sqlite = + db::sqlite::connect_path_create_if_missing(dir.path().join("query.db").to_str().unwrap()).await.unwrap(); + state + .update_connection_pools(|pools| { + pools.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); + }) + .await; + state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Sqlite)); + let keys = [ + "object-meta:v1:schema-cache:db:public:users::TABLE:columns:", + "object-meta:v1:schema-cache:db:public:users::backend-columns:", + "object-ddl:v1:schema-cache:db:public:users::TABLE:", + ]; + for key in keys { + state.storage.save_schema_cache(key, &serde_json::json!({"old": true})).await.unwrap(); + } + let unrelated = "object-meta:v1:schema-cache-other:db:public:users::backend-columns:"; + state.storage.save_schema_cache(unrelated, &serde_json::json!([])).await.unwrap(); + execute_sql_statement(&state, connection_id, "", "SELECT 1", None, None).await.unwrap(); + for key in keys { + assert!(state.storage.load_schema_cache(key).await.unwrap().is_some()); + } + execute_sql_statement(&state, connection_id, "", "CREATE TABLE users (id INTEGER)", None, None).await.unwrap(); + for key in keys { + assert!(state.storage.load_schema_cache(key).await.unwrap().is_none(), "stale cache: {key}"); + } + assert!(state.storage.load_schema_cache(unrelated).await.unwrap().is_some()); + } + + #[tokio::test] + async fn ddl_schema_cache_batch_outcomes_and_transaction_completion() { + for (sql, use_transaction, expected_error, expected_column) in [ + ("ALTER TABLE users ADD COLUMN added INTEGER; SELECT 1", false, false, true), + ("ALTER TABLE users ADD COLUMN added INTEGER; SELECT * FROM missing_table", false, true, true), + ("ALTER TABLE users ADD COLUMN added INTEGER; SELECT 1", true, false, true), + ("ALTER TABLE users ADD COLUMN added INTEGER; SELECT * FROM missing_table", true, true, false), + ("BEGIN; ALTER TABLE users ADD COLUMN added INTEGER; COMMIT", false, false, true), + ] { + let dir = tempfile::tempdir().unwrap(); + let storage = Storage::open(&dir.path().join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let sqlite = db::sqlite::connect_path_create_if_missing(dir.path().join("query.db").to_str().unwrap()) + .await + .unwrap(); + state + .update_connection_pools(|pools| { + pools.insert("cache".into(), PoolKind::Sqlite(sqlite)); + }) + .await; + state.configs.write().await.insert("cache".into(), test_connection_config(DatabaseType::Sqlite)); + execute_sql_statement(&state, "cache", "", "CREATE TABLE users (id INTEGER)", None, None).await.unwrap(); + let key = "object-meta:v1:cache:db:public:users::backend-columns:"; + state.storage.save_schema_cache(key, &serde_json::json!([])).await.unwrap(); + let result = execute_multi_core_with_options( + &state, + "cache", + "", + sql, + None, + None, + QueryExecutionOptions { use_transaction: Some(use_transaction), ..Default::default() }, + ) + .await; + let failed = match &result { + Err(_) => true, + Ok(results) => results.iter().any(|r| r.columns == vec!["Error"]), + }; + assert_eq!(failed, expected_error, "{sql}"); + assert!(state.storage.load_schema_cache(key).await.unwrap().is_none(), "{sql}"); + let columns = + execute_sql_statement(&state, "cache", "", "PRAGMA table_info(users)", None, None).await.unwrap(); + assert_eq!(columns.rows.iter().any(|row| row[1] == "added"), expected_column, "{sql}"); + // A failed DDL still keeps its execution error while clearing cache. + state.storage.save_schema_cache(key, &serde_json::json!([])).await.unwrap(); + let error = execute_sql_statement( + &state, + "cache", + "", + "ALTER TABLE missing_table ADD COLUMN x INTEGER", + None, + None, + ) + .await + .unwrap_err(); + assert!(error.contains("missing_table")); + assert!(state.storage.load_schema_cache(key).await.unwrap().is_none()); + for sql in ["BEGIN", "COMMIT", "BEGIN", "ROLLBACK"] { + state.storage.save_schema_cache(key, &serde_json::json!([])).await.unwrap(); + execute_sql_statement(&state, "cache", "", sql, None, None).await.unwrap(); + assert_eq!(state.storage.load_schema_cache(key).await.unwrap().is_some(), sql == "BEGIN"); + } + } + } + + #[tokio::test] + async fn ddl_schema_cache_storage_failure_preserves_sql_outcome() { + let dir = tempfile::tempdir().unwrap(); + let storage_path = dir.path().join("storage.db"); + let state = AppState::new(Storage::open(&storage_path).await.unwrap()); + let sqlite = + db::sqlite::connect_path_create_if_missing(dir.path().join("query.db").to_str().unwrap()).await.unwrap(); + state + .update_connection_pools(|pools| { + pools.insert("cache".into(), PoolKind::Sqlite(sqlite)); + }) + .await; + state.configs.write().await.insert("cache".into(), test_connection_config(DatabaseType::Sqlite)); + let key = "object-meta:v1:cache:db:public:users::backend-columns:"; + state.storage.save_schema_cache(key, &serde_json::json!([])).await.unwrap(); + // Fail only deletion; SQL execution and cache reads remain available. + let fault = rusqlite::Connection::open(&storage_path).unwrap(); + fault.execute_batch("CREATE TRIGGER reject_cache_delete BEFORE DELETE ON schema_cache BEGIN SELECT RAISE(FAIL, 'injected cache delete failure'); END;").unwrap(); + execute_sql_statement(&state, "cache", "", "CREATE TABLE users (id INTEGER)", None, None).await.unwrap(); + assert!(state.storage.load_schema_cache(key).await.unwrap().is_some()); + let result = execute_sql_statement(&state, "cache", "", "SELECT id FROM users", None, None).await.unwrap(); + assert_eq!(result.columns, vec!["id"]); + let error = + execute_sql_statement(&state, "cache", "", "ALTER TABLE missing_table ADD COLUMN x INTEGER", None, None) + .await + .unwrap_err(); + assert!(error.contains("missing_table")); + assert!(!error.contains("injected cache delete failure")); + } + async fn assert_sqlite_batch_error_behavior(failure_first: bool, continue_on_error: bool) { let dir = std::env::temp_dir().join(format!("dbx-query-batch-error-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); diff --git a/crates/dbx-mcp/tests/schema_cache.rs b/crates/dbx-mcp/tests/schema_cache.rs new file mode 100644 index 0000000000..740c7c7809 --- /dev/null +++ b/crates/dbx-mcp/tests/schema_cache.rs @@ -0,0 +1,57 @@ +use dbx_core::{ + models::connection::ConnectionConfig, + storage::{McpGlobalPolicy, Storage}, +}; +use dbx_mcp::{DbxMcpServer, LocalBackend, McpScope}; +use rmcp::{model::CallToolRequestParams, ServiceExt}; +use serde_json::json; +use std::sync::Arc; + +#[tokio::test] +async fn ddl_schema_cache_mcp_query_and_automatic_batch_dispatch() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("query.db"); + dbx_core::db::sqlite::connect_path_create_if_missing(path.to_str().unwrap()).await.unwrap(); + let storage = Storage::open(&dir.path().join("storage.db")).await.unwrap(); + let config: ConnectionConfig = serde_json::from_value(json!({ + "id": "cache-mcp", "name": "cache-mcp", "db_type": "sqlite", "host": path, + "port": 0, "username": "", "password": "", "database": "main" + })) + .unwrap(); + storage.add_connection_for_mcp(config).await.unwrap(); + storage + .save_mcp_global_policy(&McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, ..Default::default() }) + .await + .unwrap(); + let backend = Arc::new(LocalBackend::open(&dir.path().join("storage.db")).await.unwrap()); + let server = DbxMcpServer::with_runtime_options(backend, McpScope::default(), false); + 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.unwrap(); + for (tool, sql) in [ + ("dbx_execute_query", "CREATE TABLE users (id INTEGER)"), + ("dbx_execute_query", "ALTER TABLE users ADD COLUMN from_query INTEGER; SELECT 1"), + ("dbx_execute_batch", "ALTER TABLE users ADD COLUMN from_batch INTEGER; SELECT 1"), + ] { + let key = "object-meta:v1:cache-mcp:main:main:users::backend-columns:"; + storage.save_schema_cache(key, &json!([])).await.unwrap(); + let result = client + .peer() + .call_tool( + CallToolRequestParams::new(tool).with_arguments( + json!({ + "connection_id": "cache-mcp", "database": "main", "sql": sql + }) + .as_object() + .unwrap() + .clone(), + ), + ) + .await + .unwrap(); + assert_ne!(result.is_error, Some(true), "{result:?}"); + assert!(storage.load_schema_cache(key).await.unwrap().is_none(), "{tool}: {sql}"); + } + client.cancel().await.unwrap(); + server_task.await.unwrap().unwrap().cancel().await.unwrap(); +} diff --git a/crates/dbx-web/src/routes/mod.rs b/crates/dbx-web/src/routes/mod.rs index d3b308cb5a..81e19d5623 100644 --- a/crates/dbx-web/src/routes/mod.rs +++ b/crates/dbx-web/src/routes/mod.rs @@ -43,3 +43,6 @@ pub mod tunnel_profiles; pub mod update; pub mod vector; pub mod zookeeper; + +#[cfg(test)] +mod object_cache_tests; diff --git a/crates/dbx-web/src/routes/object_cache_tests.rs b/crates/dbx-web/src/routes/object_cache_tests.rs new file mode 100644 index 0000000000..52eba664a2 --- /dev/null +++ b/crates/dbx-web/src/routes/object_cache_tests.rs @@ -0,0 +1,214 @@ +//! Regression coverage across real Web handlers and MCP backends. +use super::{connection, query, schema, schema_cache}; +use crate::state::WebState; +use axum::{ + routing::{get, post}, + Router, +}; +use dbx_core::{ + agent_tools::AgentSqlPermissions, + connection::AppState, + models::connection::ConnectionConfig, + query::QueryExecutionOptions, + storage::{McpGlobalPolicy, Storage}, +}; +use dbx_mcp::{DbxBackend, LocalBackend, WebBackend}; +use serde_json::{json, Value}; +use std::sync::Arc; + +async fn run_object_cache_regression(postgres: bool) { + let dir = std::env::temp_dir().join(format!("dbx-object-cache-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let database = + if postgres { std::env::var("DBX_LIVE_POSTGRES_DATABASE").expect("test database") } else { "main".into() }; + let config: ConnectionConfig = serde_json::from_value(json!({ + "id": "cache:中文%", "name": "cache regression", "db_type": if postgres { "postgres" } else { "sqlite" }, + "host": if postgres { std::env::var("DBX_LIVE_POSTGRES_HOST").expect("test host") } else { dir.join("query.db").to_str().unwrap().to_string() }, + "port": if postgres { std::env::var("DBX_LIVE_POSTGRES_PORT").expect("test port").parse::().unwrap() } else { 0 }, + "username": if postgres { std::env::var("DBX_LIVE_POSTGRES_USER").expect("test user") } else { String::new() }, + "password": if postgres { std::env::var("DBX_LIVE_POSTGRES_PASSWORD").expect("test password") } else { String::new() }, + "database": database, "save_password": true + })).unwrap(); + if !postgres { + dbx_core::db::sqlite::connect_path_create_if_missing(&config.host).await.unwrap(); + } + storage.save_connections(std::slice::from_ref(&config)).await.unwrap(); + storage + .save_mcp_global_policy(&McpGlobalPolicy { read_only: false, allow_dangerous_sql: true, ..Default::default() }) + .await + .unwrap(); + let app = Arc::new(AppState::new_with_plugin_dir(storage, dir.join("plugins"))); + app.configs.write().await.insert(config.id.clone(), config.clone()); + let mut state = WebState::for_tests(app.clone(), dir.clone()); + state.password_disabled = true; + let state = Arc::new(state); + let router = Router::new() + .route("/api/auth/check", get(crate::auth::check)) + .route("/api/connection/connect", post(connection::connect_db)) + .route("/api/query/execute", post(query::execute_query)) + .route("/api/query/execute-multi", post(query::execute_multi)) + .route("/api/query/execute-batch", post(query::execute_batch)) + .route("/api/schema/columns", get(schema::list_columns)) + .route("/api/schema/ddl", get(schema::get_ddl)) + .route("/api/schema/cache", get(schema_cache::load_schema_cache).post(schema_cache::save_schema_cache)) + .with_state(state); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + let local = LocalBackend::from_app_state(app.clone(), dir.clone()); + let web = WebBackend::new(base.clone(), String::new()).unwrap(); + let client = reqwest::Client::new(); + let schema_name = if postgres { format!("dbx_cache_{}", uuid::Uuid::new_v4().simple()) } else { "main".into() }; + if postgres { + local.execute_query(&config, &database, &format!("CREATE SCHEMA {schema_name}"), None, None).await.unwrap(); + } + let table = format!("{schema_name}.users"); + local.execute_query(&config, &database, &format!("CREATE TABLE {table} (id INTEGER)"), None, None).await.unwrap(); + let prefix = dbx_core::object_cache::object_metadata_cache_prefix(&config.id, &database); + let schema_segment = dbx_core::object_cache::metadata_cache_segment(&schema_name); + let front_key = format!("{prefix}{schema_segment}:users::TABLE:columns:"); + let ddl_key = format!( + "object-ddl:v1:{}:{}:{schema_segment}:users::TABLE:", + dbx_core::object_cache::metadata_cache_segment(&config.id), + dbx_core::object_cache::metadata_cache_segment(&database) + ); + let backend_key = format!("{prefix}{schema_segment}:users::backend-columns:"); + for mode in 0..5 { + let parameters = [ + ("connection_id", config.id.as_str()), + ("database", database.as_str()), + ("schema", schema_name.as_str()), + ("table", "users"), + ]; + let old_columns: Value = client + .get(format!("{base}/api/schema/columns")) + .query(¶meters) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + let old_ddl: Value = client + .get(format!("{base}/api/schema/ddl")) + .query(¶meters) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + for (key, payload) in [ + (&front_key, json!({"version": 1, "cachedAt": chrono::Utc::now().to_rfc3339(), "value": old_columns})), + (&ddl_key, json!({"version": 1, "cachedAt": chrono::Utc::now().to_rfc3339(), "ddl": old_ddl})), + ] { + client + .post(format!("{base}/api/schema/cache")) + .json(&json!({"cacheKey": key, "payload": payload})) + .send() + .await + .unwrap() + .error_for_status() + .unwrap(); + } + assert!(app.storage.load_schema_cache(&backend_key).await.unwrap().is_some()); + let column = format!("added_{mode}"); + let sql = format!("ALTER TABLE {table} ADD COLUMN {column} INTEGER"); + if mode == 0 || mode == 2 { + let backend: &dyn DbxBackend = if mode == 0 { &local } else { &web }; + let result = backend + .execute_agent_tool( + &config, + &database, + "execute_query", + json!({"sql": sql}), + AgentSqlPermissions { allow_writes: true, allow_dangerous: true, confirmed_write_sql: None }, + ) + .await; + assert!(!result.is_error, "{}", result.content); + } else if mode == 1 || mode == 3 { + let backend: &dyn DbxBackend = if mode == 1 { &local } else { &web }; + let results = backend + .execute_batch( + &config, + &database, + Some(&schema_name), + &format!("{sql}; SELECT 1"), + QueryExecutionOptions::default(), + ) + .await + .unwrap(); + assert_eq!(results.len(), 2); + assert!(results.iter().all(|result| !result.execution_error), "{results:?}"); + } else { + client.post(format!("{base}/api/query/execute-batch")).json(&json!({"connectionId": config.id, "database": database, "statements": [sql], "schema": schema_name})).send().await.unwrap().error_for_status().unwrap(); + } + for key in [&front_key, &ddl_key, &backend_key] { + assert!(app.storage.load_schema_cache(key).await.unwrap().is_none(), "mode {mode}: {key}"); + } + // A fresh frontend read must miss persisted snapshots, then fetch the + // actual updated columns and DDL through the real Web handlers. + let cached: Value = client + .get(format!("{base}/api/schema/cache")) + .query(&[("cache_key", &front_key)]) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert!(cached.is_null()); + let columns: Value = client + .get(format!("{base}/api/schema/columns")) + .query(¶meters) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + assert!(columns.as_array().unwrap().iter().any(|c| c["name"] == column), "{columns}"); + let ddl: Value = client + .get(format!("{base}/api/schema/ddl")) + .query(¶meters) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + assert!(ddl.to_string().contains(&column), "{ddl}"); + } + if postgres { + local + .execute_query(&config, &database, &format!("DROP SCHEMA {schema_name} CASCADE"), None, None) + .await + .unwrap(); + } + server.abort(); + let _ = server.await; + drop(web); + drop(local); + drop(app); + std::fs::remove_dir_all(dir).unwrap(); +} + +#[tokio::test] +async fn ddl_schema_cache_real_web_handlers_and_mcp_backends() { + run_object_cache_regression(false).await; +} + +#[tokio::test] +#[ignore = "requires isolated writable DBX_LIVE_POSTGRES_HOST/PORT/USER/PASSWORD/DATABASE"] +async fn ddl_schema_cache_live_postgres_web_handlers_and_mcp_backends() { + run_object_cache_regression(true).await; +} diff --git a/crates/dbx-web/src/routes/schema.rs b/crates/dbx-web/src/routes/schema.rs index 7c63f6fffe..45975f8b19 100644 --- a/crates/dbx-web/src/routes/schema.rs +++ b/crates/dbx-web/src/routes/schema.rs @@ -402,25 +402,8 @@ pub async fn get_custom_type_details( Ok(Json(result)) } -const OBJECT_METADATA_CACHE_PREFIX: &str = "object-meta:v1"; - -fn metadata_cache_segment(value: &str) -> String { - const HEX: &[u8; 16] = b"0123456789ABCDEF"; - let mut encoded = String::with_capacity(value.len()); - for byte in value.bytes() { - match byte { - b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'!' | b'~' | b'*' | b'\'' | b'(' | b')' => { - encoded.push(byte as char) - } - _ => { - encoded.push('%'); - encoded.push(HEX[(byte >> 4) as usize] as char); - encoded.push(HEX[(byte & 0x0f) as usize] as char); - } - } - } - encoded -} +pub(crate) use dbx_core::object_cache::object_metadata_cache_prefix; +use dbx_core::object_cache::{metadata_cache_segment, OBJECT_METADATA_CACHE_PREFIX}; fn metadata_cache_key( connection_id: &str, @@ -443,15 +426,6 @@ fn metadata_cache_key( .join(":") } -pub(crate) fn object_metadata_cache_prefix(connection_id: &str, database: &str) -> String { - format!( - "{}:{}:{}:", - OBJECT_METADATA_CACHE_PREFIX, - metadata_cache_segment(connection_id), - metadata_cache_segment(database) - ) -} - fn decode_metadata_cache(value: serde_json::Value) -> Option { serde_json::from_value(value).ok() }