From 84f1498a4b3b3dfe7d5cd41798a4525d11ebad9c Mon Sep 17 00:00:00 2001 From: Xinyu Kuo Date: Sat, 15 Aug 2026 14:41:16 +0800 Subject: [PATCH] fix(duckdb): restore parquet preview and reclaim connections --- apps/desktop/src/composables/useFileDrop.ts | 9 +- apps/desktop/src/i18n/locales/en.ts | 1 + apps/desktop/src/i18n/locales/es.ts | 1 + apps/desktop/src/i18n/locales/it.ts | 1 + apps/desktop/src/i18n/locales/ja.ts | 1 + apps/desktop/src/i18n/locales/ko.ts | 1 + apps/desktop/src/i18n/locales/pt-BR.ts | 1 + apps/desktop/src/i18n/locales/zh-CN.ts | 1 + apps/desktop/src/i18n/locales/zh-TW.ts | 1 + .../connectionStore.oneTimeCleanup.spec.ts | 122 +++++++++++++ apps/desktop/src/stores/connectionStore.ts | 41 +++++ crates/dbx-core/src/connection.rs | 2 +- crates/dbx-core/src/lib.rs | 1 + crates/dbx-core/src/runtime_config.rs | 171 ++++++++++++++++++ crates/dbx-web/src/routes/connection.rs | 60 +++++- src-tauri/src/commands/connection.rs | 75 +++++++- 16 files changed, 470 insertions(+), 19 deletions(-) create mode 100644 apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts create mode 100644 crates/dbx-core/src/runtime_config.rs diff --git a/apps/desktop/src/composables/useFileDrop.ts b/apps/desktop/src/composables/useFileDrop.ts index 53584a98f2..eb823a51b1 100644 --- a/apps/desktop/src/composables/useFileDrop.ts +++ b/apps/desktop/src/composables/useFileDrop.ts @@ -60,8 +60,15 @@ export function useFileDrop() { port: 0, username: "", password: "", + one_time: true, }; - const connectionId = await api.connectDb(config); + let connectionId: string; + try { + connectionId = await api.connectDb(config); + } catch (e: any) { + toast(t("welcome.fileOpenFailed", { name, message: e?.message || String(e) }), 5000); + continue; + } connectionStore.addEphemeralConnection({ ...config, id: connectionId }); const tabId = queryStore.createTab(connectionId, "", name, "query"); queryStore.updateSql(tabId, dataQuery); diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 2074126b9a..79076184cb 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -2033,6 +2033,7 @@ export default { tipSidebar: "Click a table on the left to view data", tipExecute: "to execute query", fileOpened: "Opened {name}", + fileOpenFailed: "Failed to open {name}: {message}", mcpTitle: "AI Agent Integration", mcpDescription: "Beyond the built-in AI assistant, you can also use Claude Code, Cursor, and other coding agents to query your databases via MCP.", mcpLearnMore: "Learn more", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index fcb32051a1..b27b86c6b9 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -1967,6 +1967,7 @@ export default withEnglishFallback({ tipSidebar: "Haz clic en una tabla a la izquierda para ver los datos", tipExecute: "para ejecutar la consulta", fileOpened: "Abierto {name}", + fileOpenFailed: "No se pudo abrir {name}: {message}", mcpTitle: "Integración con agentes de IA", mcpDescription: "Además del asistente de IA integrado, puedes usar Claude Code, Cursor y otros agentes de programación para consultar tus bases de datos mediante MCP.", mcpLearnMore: "Más información", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index 6184de8103..964477e7a2 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -1965,6 +1965,7 @@ export default withEnglishFallback({ tipSidebar: "Fai clic su una tabella a sinistra per visualizzarne i dati", tipExecute: "per eseguire la query", fileOpened: "Aperto {name}", + fileOpenFailed: "Impossibile aprire {name}: {message}", mcpTitle: "Integrazione Agente AI", mcpDescription: "Oltre all'assistente AI integrato, puoi utilizzare Claude Code, Cursor e altri agenti di sviluppo per interrogare i tuoi database tramite MCP.", mcpLearnMore: "Scopri di più", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index e7b7ec70cd..da6ff5de51 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -1992,6 +1992,7 @@ export default withEnglishFallback({ tipSidebar: "左側のテーブルをクリックしてデータを表示", tipExecute: "でクエリを実行", fileOpened: "{name} を開きました", + fileOpenFailed: "{name} を開けませんでした：{message}", mcpTitle: "AIエージェント連携", mcpDescription: "内蔵AIアシスタントに加えて、Claude Code、Cursor、その他のコーディングエージェントをMCP経由でデータベースクエリに使用できます。", mcpLearnMore: "詳細を見る", diff --git a/apps/desktop/src/i18n/locales/ko.ts b/apps/desktop/src/i18n/locales/ko.ts index 19314e7b02..a1fe196aec 100644 --- a/apps/desktop/src/i18n/locales/ko.ts +++ b/apps/desktop/src/i18n/locales/ko.ts @@ -1894,6 +1894,7 @@ export default withEnglishFallback({ tipSidebar: "데이터를 보려면 왼쪽의 테이블을 클릭하세요", tipExecute: "쿼리를 실행하려면", fileOpened: "{name}을(를) 열었습니다", + fileOpenFailed: "{name}을(를) 열지 못했습니다: {message}", mcpTitle: "AI 에이전트 연동", mcpDescription: "내장된 AI 어시스턴트 외에도 Claude Code, Cursor 등의 코딩 에이전트를 MCP로 연결하여 데이터베이스를 쿼리할 수 있습니다.", mcpLearnMore: "더 알아보기", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index d7a3a48470..5d1cd5288f 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -1967,6 +1967,7 @@ export default withEnglishFallback({ tipSidebar: "Clique em uma tabela à esquerda para visualizar os dados", tipExecute: "para executar a consulta", fileOpened: "{name} aberto", + fileOpenFailed: "Falha ao abrir {name}: {message}", mcpTitle: "Integração com Agente de IA", mcpDescription: "Além do assistente de IA integrado, você também pode usar Claude Code, Cursor e outros agentes de codificação para consultar seus bancos de dados via MCP.", mcpLearnMore: "Saiba mais", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 5d5b6cf845..64e4f24e08 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -1957,6 +1957,7 @@ export default withEnglishFallback({ tipSidebar: "点击左侧表名查看数据", tipExecute: "执行查询", fileOpened: "已打开 {name}", + fileOpenFailed: "打开 {name} 失败：{message}", mcpTitle: "AI 编程助手集成", mcpDescription: "除了内置 AI 助手，还可以通过 MCP 让 Claude Code、Cursor 等编程助手直接使用这里的数据库连接。", mcpLearnMore: "了解更多", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index aa4d4bc294..7163f5c2f5 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -1966,6 +1966,7 @@ export default withEnglishFallback({ tipSidebar: "點選左側資料表以檢視資料", tipExecute: "執行查詢", fileOpened: "已開啟 {name}", + fileOpenFailed: "開啟 {name} 失敗：{message}", mcpTitle: "AI 程式設計助理整合", mcpDescription: "除了內建 AI 助理，也可以藉由 MCP 讓 Claude Code、Cursor 等程式設計助理查詢你的資料庫。", mcpLearnMore: "了解更多", diff --git a/apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts new file mode 100644 index 0000000000..dffd076b27 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/connectionStore.oneTimeCleanup.spec.ts @@ -0,0 +1,122 @@ +import { createPinia, setActivePinia } from "pinia"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { ConnectionConfig } from "@/types/database"; + +function installLocalStorage() { + const data = new Map(); + vi.stubGlobal("localStorage", { + getItem: vi.fn((key: string) => data.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => data.set(key, value)), + removeItem: vi.fn((key: string) => data.delete(key)), + }); +} + +function installApiMocks() { + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + connectDb: vi.fn().mockResolvedValue("preview-1"), + disconnectDb: vi.fn().mockResolvedValue(undefined), + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + loadSchemaCache: vi.fn().mockResolvedValue(null), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + connectionDatabaseInfo: vi.fn().mockResolvedValue(undefined), + listInstalledAgents: vi.fn().mockResolvedValue([]), + sessionCredentialStatus: vi.fn().mockResolvedValue(false), + forgetSessionCredential: vi.fn().mockResolvedValue(undefined), + })); + // Result-snapshot cleanup hits a relative URL, which node's fetch rejects outright. + // Unrelated to what these cases assert. + vi.doMock("@/lib/tabs/tabResultCache", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, deleteTabResultSnapshotsForOwner: vi.fn().mockResolvedValue(undefined) }; + }); +} + +function previewConnection(overrides: Partial = {}): ConnectionConfig { + return { + id: "preview-1", + name: "[Preview] sales.parquet", + db_type: "duckdb", + host: ":memory:", + port: 0, + username: "", + password: "", + one_time: true, + ...overrides, + } as ConnectionConfig; +} + +describe("connectionStore one_time runtime cleanup", () => { + beforeEach(() => { + vi.resetModules(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + // One-time connections are never persisted, so the backend's save_connections + // sync never reclaims them; disconnect_db is the only reclaim point, which makes + // an explicit disconnect on removal mandatory. + it("removeConnection disconnects a one_time connection so the backend can reclaim it", async () => { + installApiMocks(); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [previewConnection()]; + + await store.removeConnection("preview-1"); + + const { disconnectDb } = await import("@/lib/backend/api"); + // No clientAttempt: removal is terminal, so a superseded attempt number must not + // skip the cleanup. + expect(disconnectDb).toHaveBeenCalledWith("preview-1"); + }); + + it("removeConnection leaves saved connections to the save_connections sync", async () => { + installApiMocks(); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [previewConnection({ id: "saved-1", one_time: false })]; + + await store.removeConnection("saved-1"); + + const { disconnectDb } = await import("@/lib/backend/api"); + expect(disconnectDb).not.toHaveBeenCalled(); + }); + + // The connection is gone for good, so its tabs can never reconnect: executing in + // one would fail with "Connection config not found". + it("removeConnection closes the tabs of a one_time connection", async () => { + installApiMocks(); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useConnectionStore(); + const queryStore = useQueryStore(); + store.connections = [previewConnection()]; + queryStore.createTab("preview-1", "", "sales.parquet", "query"); + expect(queryStore.tabs.some((tab) => tab.connectionId === "preview-1")).toBe(true); + + await store.removeConnection("preview-1"); + + expect(queryStore.tabs.some((tab) => tab.connectionId === "preview-1")).toBe(false); + }); + + it("removeConnection keeps the tabs of a saved connection for the disconnect handling mode", async () => { + installApiMocks(); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useConnectionStore(); + const queryStore = useQueryStore(); + store.connections = [previewConnection({ id: "saved-1", one_time: false })]; + queryStore.createTab("saved-1", "", "query.sql", "query"); + + await store.removeConnection("saved-1"); + + expect(queryStore.tabs.some((tab) => tab.connectionId === "saved-1")).toBe(true); + }); +}); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index ad52d9fda0..292aeff54c 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -711,6 +711,44 @@ export const useConnectionStore = defineStore("connection", () => { return trackDisconnectRequest(connectionId, request, clientAttempt != null); } + /** + * One-time connections are never persisted, so the backend's "not in the saved + * list, so reclaim it" branch in `sync_connection_configs` never fires for them + * (see dbx-core `should_retain_runtime_config`) and `disconnect_db` is the only + * reclaim point. Removing one must disconnect it explicitly, or its runtime + * config, pool, and tunnel live until the process exits. + * + * `clientAttempt` is deliberately omitted: removal is terminal, so a superseded + * attempt number must not skip the cleanup. + */ + function releaseOneTimeRuntimeConnections(connectionIds: string[]) { + for (const connectionId of connectionIds) { + let request: Promise; + try { + request = api.disconnectDb(connectionId); + } catch (error) { + request = Promise.reject(error); + } + void trackDisconnectRequest(connectionId, request, false).catch(() => {}); + } + } + + /** + * A one-time connection is deleted outright rather than left behind for a + * reconnect, so its tabs have nothing to point at once it is gone: executing in + * one would fail with "Connection config not found". Close them regardless of + * `disconnectTabHandlingMode`, which only makes sense for connections that can + * still be reconnected. + */ + async function closeOneTimeConnectionTabs(connectionIds: string[]) { + if (!connectionIds.length) return; + const { useQueryStore } = await import("@/stores/queryStore"); + const queryStore = useQueryStore(); + for (const connectionId of connectionIds) { + queryStore.closeConnectionTabs(connectionId); + } + } + function cancelDisconnectKey(connectionId: string, attempt: number): string { return `${connectionId}:${attempt}`; } @@ -2700,8 +2738,11 @@ export const useConnectionStore = defineStore("connection", () => { if (!connectionIds.length) return; const removedIds = new Set(connectionIds); + const oneTimeIds = connectionIds.filter((id) => getConfig(id)?.one_time === true); const nextConnections = connections.value.filter((c) => !removedIds.has(c.id)); await persistConnections(nextConnections); + await closeOneTimeConnectionTabs(oneTimeIds); + releaseOneTimeRuntimeConnections(oneTimeIds); await persistTimeoutInheritanceIds( settingsStore.editorSettings.connectTimeoutInheritConnectionIds.filter((id) => !removedIds.has(id)), settingsStore.editorSettings.queryTimeoutInheritConnectionIds.filter((id) => !removedIds.has(id)), diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index a4bb07434a..35745c5a59 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -2354,7 +2354,7 @@ impl AppState { // Temporary "__test_*" probes must not retain agents in the registry. // reconnect fast-path caching only applies to durable connection ids; // drain_connection_pools no longer drops MQ adapters (reconnect reuse). - if connection_id.starts_with("__test_") { + if connection_id.starts_with(crate::runtime_config::TEST_PROBE_ID_PREFIX) { let adapter = self.mq_registry.build_transient_config(mqc, agent_launch).await?; adapter.test_connection().await?; if let Err(err) = self.ensure_current_connection_attempt(connection_id, connection_attempt).await { diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index c8ace7f311..3c88121ddc 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -70,6 +70,7 @@ pub mod query_result_export; pub mod query_result_sql; pub mod redis_ops; pub mod risk_metrics; +pub mod runtime_config; pub mod saved_sql; pub mod schema; pub mod schema_diff; diff --git a/crates/dbx-core/src/runtime_config.rs b/crates/dbx-core/src/runtime_config.rs new file mode 100644 index 0000000000..1d78e37eab --- /dev/null +++ b/crates/dbx-core/src/runtime_config.rs @@ -0,0 +1,171 @@ +//! Lifecycle rules for runtime connection configs (`AppState::configs`). +//! +//! `save_connections` reclaims runtime configs by asking "is this id in the saved +//! list?", but two kinds of config are never in that list and must be judged +//! separately, or they get evicted by mistake: +//! +//! - **Transient probe ids**: `test_connection`, the visible-database and +//! visible-schema drafts, and other runtime ids generated by prefix convention. +//! Their creators are responsible for cleaning them up. +//! - **One-time connections (`one_time`)**: connections the user explicitly marked +//! as one-time, plus the dropped-file preview connections. They are never +//! persisted (the frontend `persistConnections` filters them out), so they can +//! never appear in the saved list. +//! +//! The desktop backend (`src-tauri`) and the web backend (`dbx-web`) each have +//! their own `sync_connection_configs` and `disconnect_db`; the rules must agree +//! across both, so they live here. + +use crate::connection::AppState; +use crate::models::connection::ConnectionConfig; + +/// Prefix for the throwaway probe connection created by `test_connection`. +pub const TEST_PROBE_ID_PREFIX: &str = "__test_"; +/// Prefix for the connection dialog's "visible databases" picker draft. +pub const VISIBLE_DRAFT_ID_PREFIX: &str = "__visible_draft_"; +/// Prefix for the connection dialog's "visible schemas" picker draft. +pub const VISIBLE_SCHEMA_DRAFT_ID_PREFIX: &str = "__visible_schema_draft_"; + +/// Runtime configs generated by the transient-id prefix convention. +/// +/// These ids never appear in the saved connection list, so `save_connections` +/// must skip them; otherwise one concurrent save kills an in-flight probe. +pub fn is_transient_runtime_config_id(id: &str) -> bool { + id.starts_with(TEST_PROBE_ID_PREFIX) + || id.starts_with(VISIBLE_DRAFT_ID_PREFIX) + || id.starts_with(VISIBLE_SCHEMA_DRAFT_ID_PREFIX) +} + +/// Whether `sync_connection_configs` must keep this runtime config. The caller +/// separately checks whether the id is in the saved list. +pub fn should_retain_runtime_config(id: &str, config: &ConnectionConfig) -> bool { + config.one_time || is_transient_runtime_config_id(id) +} + +/// Whether `disconnect_db` should drop the runtime config along with the session. +/// +/// The asymmetry with [`should_retain_runtime_config`] is deliberate: +/// `__test_*` probes are removed synchronously by `run_temporary_connection_test` +/// in its own cleanup and never reach `disconnect_db`, so that prefix is not +/// matched here. +pub fn should_release_runtime_config_on_disconnect(id: &str, config: Option<&ConnectionConfig>) -> bool { + id.starts_with(VISIBLE_DRAFT_ID_PREFIX) + || id.starts_with(VISIBLE_SCHEMA_DRAFT_ID_PREFIX) + || config.is_some_and(|config| config.one_time) +} + +/// Reclaim a transient runtime config after its session is torn down. +/// +/// These configs are not in the saved list, so the reclaim branch of +/// `sync_connection_configs` never fires for them; this is the only reclaim +/// point. Both the config itself and the session credential (a plaintext +/// password) must be released here, or they live until the process exits. +/// +/// The check and the removal share one write lock so that a concurrent +/// `connect_db` cannot have a freshly inserted config deleted out from under it. +pub async fn release_runtime_config_on_disconnect(state: &AppState, connection_id: &str) { + let released = { + let mut configs = state.configs.write().await; + if should_release_runtime_config_on_disconnect(connection_id, configs.get(connection_id)) { + configs.remove(connection_id); + true + } else { + false + } + }; + if released { + // Matches how `sync_connection_configs` treats a deleted connection: once + // the config is gone the credential has no owner, and leaving it behind + // keeps a plaintext password in memory. + state.session_credentials.clear_connection(connection_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config_with_one_time(one_time: bool) -> ConnectionConfig { + serde_json::from_value(serde_json::json!({ + "id": "cfg", + "name": "cfg", + "db_type": "mysql", + "host": "127.0.0.1", + "port": 3306, + "username": "root", + "password": "", + "database": null, + "one_time": one_time, + })) + .expect("connection config fixture") + } + + #[test] + fn transient_ids_are_recognized_by_prefix() { + assert!(is_transient_runtime_config_id("__test_abc")); + assert!(is_transient_runtime_config_id("__visible_draft_abc")); + assert!(is_transient_runtime_config_id("__visible_schema_draft_abc")); + assert!(!is_transient_runtime_config_id("saved-connection")); + assert!(!is_transient_runtime_config_id("__production_database_draft_abc")); + } + + #[test] + fn retain_covers_both_transient_ids_and_one_time_configs() { + assert!(should_retain_runtime_config("__test_abc", &config_with_one_time(false))); + assert!(should_retain_runtime_config("preview-duckdb", &config_with_one_time(true))); + // Production-database drafts have no dedicated prefix; only `one_time` keeps them alive. + assert!(should_retain_runtime_config("__production_database_draft_abc", &config_with_one_time(true))); + assert!(!should_retain_runtime_config("saved-connection", &config_with_one_time(false))); + } + + #[test] + fn disconnect_releases_drafts_and_one_time_but_not_test_probes() { + assert!(should_release_runtime_config_on_disconnect("__visible_draft_abc", None)); + assert!(should_release_runtime_config_on_disconnect("__visible_schema_draft_abc", None)); + assert!(should_release_runtime_config_on_disconnect("preview-duckdb", Some(&config_with_one_time(true)))); + // `__test_*` cleans up after itself in run_temporary_connection_test, not here. + assert!(!should_release_runtime_config_on_disconnect("__test_abc", Some(&config_with_one_time(false)))); + assert!(!should_release_runtime_config_on_disconnect("saved-connection", Some(&config_with_one_time(false)))); + assert!(!should_release_runtime_config_on_disconnect("saved-connection", None)); + } + + async fn app_state() -> (AppState, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let storage = crate::storage::Storage::open(&dir.path().join("storage.db")).await.unwrap(); + (AppState::new(storage), dir) + } + + #[tokio::test] + async fn disconnect_release_drops_one_time_config_and_its_session_credential() { + let (state, _dir) = app_state().await; + let mut config = config_with_one_time(true); + config.id = "preview-duckdb".to_string(); + state.configs.write().await.insert(config.id.clone(), config.clone()); + state.session_credentials.set("", &config.id, "secret").expect("session credential fixture"); + + release_runtime_config_on_disconnect(&state, &config.id).await; + + assert!(!state.configs.read().await.contains_key(&config.id)); + assert!( + !state.session_credentials.has("", &config.id), + "the plaintext session password must be cleared along with the config" + ); + } + + #[tokio::test] + async fn disconnect_release_keeps_saved_config_and_its_session_credential() { + let (state, _dir) = app_state().await; + let mut config = config_with_one_time(false); + config.id = "saved-connection".to_string(); + state.configs.write().await.insert(config.id.clone(), config.clone()); + state.session_credentials.set("", &config.id, "secret").expect("session credential fixture"); + + release_runtime_config_on_disconnect(&state, &config.id).await; + + // Disconnecting a normal connection only closes the session: the config stays + // so it can reconnect, and the credential must survive too, or every + // save_password=false connection would re-prompt after each disconnect. + assert!(state.configs.read().await.contains_key(&config.id)); + assert!(state.session_credentials.has("", &config.id)); + } +} diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index 683124aa19..ab18d3470d 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -6,6 +6,9 @@ use axum::http::HeaderMap; use axum::Json; use dbx_core::connection::{connection_configs_pool_equivalent, AppState, PoolKind}; use dbx_core::models::connection::{ConnectionConfig, ConnectionTestResult, DatabaseConnectionInfo, DatabaseType}; +use dbx_core::runtime_config::{ + release_runtime_config_on_disconnect, should_retain_runtime_config, TEST_PROBE_ID_PREFIX, +}; use serde::Deserialize; use crate::auth::session_token_from_headers; @@ -173,7 +176,7 @@ async fn run_temporary_connection_test( config: ConnectionConfig, include_database_info: bool, ) -> Result { - let temp_id = format!("__test_{}", uuid::Uuid::new_v4()); + let temp_id = format!("{TEST_PROBE_ID_PREFIX}{}", uuid::Uuid::new_v4()); app.configs.write().await.insert(temp_id.clone(), config.clone()); let pool_result = app.get_or_create_pool(&temp_id, config.database.as_deref()).await; @@ -401,9 +404,7 @@ pub async fn disconnect_db( #[cfg(feature = "mq-admin")] app.mq_registry.drop_connection(&body.connection_id).await; app.reset_connection_transport(&body.connection_id).await; - if body.connection_id.starts_with("__visible_draft_") || body.connection_id.starts_with("__visible_schema_draft_") { - app.configs.write().await.remove(&body.connection_id); - } + release_runtime_config_on_disconnect(app, &body.connection_id).await; Ok(Json(())) } @@ -560,7 +561,7 @@ async fn sync_connection_configs(state: &WebState, configs: &[ConnectionConfig]) let mut connection_pool_ids_to_drop = HashSet::new(); let mut runtime_configs = state.app.configs.write().await; runtime_configs.retain(|id, existing| { - if saved_ids.contains(id.as_str()) || is_transient_runtime_config_id(id) { + if saved_ids.contains(id.as_str()) || should_retain_runtime_config(id, existing) { true } else { connection_pool_ids_to_drop.insert(id.clone()); @@ -606,10 +607,6 @@ async fn sync_connection_configs(state: &WebState, configs: &[ConnectionConfig]) } } -fn is_transient_runtime_config_id(id: &str) -> bool { - id.starts_with("__test_") || id.starts_with("__visible_draft_") || id.starts_with("__visible_schema_draft_") -} - async fn drop_nacos_adapters_for_connection_ids(state: &WebState, connection_ids: &[String]) { for connection_id in connection_ids { state.app.nacos_registry.drop_connection(connection_id).await; @@ -1680,4 +1677,49 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + + #[tokio::test] + async fn sync_connection_configs_retains_one_time_runtime_config_and_its_pool() { + let (state, dir) = test_web_state().await; + let mut one_time = sqlite_config("preview-file", &dir.join("preview.db").to_string_lossy()); + one_time.one_time = true; + state.app.configs.write().await.insert(one_time.id.clone(), one_time.clone()); + + // Any other browser tab saving any connection syncs a list that excludes + // one-time connections. + let sync = sync_connection_configs(&state, &[]).await; + + assert!(state.app.configs.read().await.contains_key(&one_time.id)); + assert!( + !sync.connection_pool_ids_to_drop.contains(&one_time.id), + "one_time connection pool must not be torn down by save sync" + ); + + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn disconnect_db_removes_one_time_config_and_its_session_credential() { + let (state, dir) = test_web_state().await; + let mut one_time = sqlite_config("preview-file", &dir.join("preview.db").to_string_lossy()); + one_time.one_time = true; + one_time.save_password = false; + state.app.configs.write().await.insert(one_time.id.clone(), one_time.clone()); + let _ = state.app.session_credentials.set("token-a", &one_time.id, "owner-a-secret"); + + let result = disconnect_db( + State(state.clone()), + Json(DisconnectRequest { connection_id: one_time.id.clone(), client_attempt: None }), + ) + .await; + assert!(result.is_ok()); + + assert!(!state.app.configs.read().await.contains_key(&one_time.id)); + assert!( + !state.app.session_credentials.has("token-a", &one_time.id), + "the plaintext session password must be cleared along with the config" + ); + + let _ = std::fs::remove_dir_all(dir); + } } diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index a425fc9ec3..479cdce876 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -20,6 +20,7 @@ use dbx_core::models::connection::{ DatabaseConnectionInfo, DatabaseType, }; pub use dbx_core::path_utils::expand_tilde; +use dbx_core::runtime_config::{release_runtime_config_on_disconnect, should_retain_runtime_config}; const MONGO_LEGACY_DRIVER_PROFILE: &str = "mongodb-legacy"; const MONGO_LEGACY_DRIVER_LABEL: &str = "MongoDB (Legacy)"; @@ -625,6 +626,70 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + /// Dropped-file preview connection: in-memory DuckDB, `one_time`, never in the saved list. + fn duckdb_preview_config() -> ConnectionConfig { + ConnectionConfig { + id: "preview-duckdb".to_string(), + name: "[Preview] sales.parquet".to_string(), + db_type: DatabaseType::DuckDb, + driver_profile: Some("duckdb".to_string()), + driver_label: Some("DuckDB".to_string()), + url_params: Some(String::new()), + host: ":memory:".to_string(), + port: 0, + username: String::new(), + password: String::new(), + database: None, + one_time: true, + ..mongodb_config() + } + } + + #[tokio::test] + async fn save_connection_configs_retains_one_time_runtime_config_and_its_pool() { + let dir = std::env::temp_dir().join(format!("dbx-tauri-conn-test-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new_with_plugin_dir(storage, dir.join("plugins")); + let persisted = mongodb_config(); + let preview = duckdb_preview_config(); + state.configs.write().await.insert(preview.id.clone(), preview.clone()); + + let sync = sync_connection_configs(&state, std::slice::from_ref(&persisted)).await; + + let configs = state.configs.read().await; + assert!(configs.contains_key(&persisted.id)); + assert!(configs.contains_key(&preview.id), "one_time runtime config must survive save sync"); + // The preview broke because the sync tore its pool down; asserting only that + // the config survives would miss the actual regression. + assert!( + !sync.connection_pool_ids_to_drop.contains(&preview.id), + "one_time connection pool must not be torn down by save sync" + ); + + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn save_connection_configs_keeps_session_credential_of_one_time_config() { + let dir = std::env::temp_dir().join(format!("dbx-tauri-conn-test-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new_with_plugin_dir(storage, dir.join("plugins")); + let persisted = mongodb_config(); + let preview = duckdb_preview_config(); + state.configs.write().await.insert(preview.id.clone(), preview.clone()); + state.session_credentials.set("", &preview.id, "secret").expect("session credential fixture"); + + save_connection_configs(&state, std::slice::from_ref(&persisted)).await.unwrap(); + + // If the config is retained the credential must be retained with it, or the + // next query re-prompts for a password that was already entered. + assert!(state.session_credentials.has("", &preview.id)); + + let _ = std::fs::remove_dir_all(dir); + } + #[cfg(feature = "mq-admin")] #[tokio::test] async fn save_connection_configs_removes_deleted_connection_pools() { @@ -719,7 +784,7 @@ async fn sync_connection_configs(state: &AppState, configs: &[ConnectionConfig]) let mut connection_pool_ids_to_drop = HashSet::new(); let mut runtime_configs = state.configs.write().await; runtime_configs.retain(|id, existing| { - if saved_ids.contains(id.as_str()) || is_transient_runtime_config_id(id) { + if saved_ids.contains(id.as_str()) || should_retain_runtime_config(id, existing) { true } else { connection_pool_ids_to_drop.insert(id.clone()); @@ -765,10 +830,6 @@ async fn sync_connection_configs(state: &AppState, configs: &[ConnectionConfig]) } } -fn is_transient_runtime_config_id(id: &str) -> bool { - id.starts_with("__test_") || id.starts_with("__visible_draft_") || id.starts_with("__visible_schema_draft_") -} - async fn drop_nacos_adapters_for_connection_ids(state: &AppState, connection_ids: &[String]) { for connection_id in connection_ids { state.nacos_registry.drop_connection(connection_id).await; @@ -1728,9 +1789,7 @@ pub async fn disconnect_db( drop_nacos_adapters_for_connection_ids(state.inner(), std::slice::from_ref(&connection_id)).await; drop_mq_adapters_for_connection_ids(state.inner(), std::slice::from_ref(&connection_id)).await; state.reset_connection_transport(&connection_id).await; - if connection_id.starts_with("__visible_draft_") || connection_id.starts_with("__visible_schema_draft_") { - state.configs.write().await.remove(&connection_id); - } + release_runtime_config_on_disconnect(state.inner(), &connection_id).await; Ok(()) }