mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-19 16:16:45 +00:00
fix(duckdb): restore parquet preview and reclaim connections
This commit is contained in:
@@ -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);
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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ù",
|
||||
|
||||
@@ -1992,6 +1992,7 @@ export default withEnglishFallback({
|
||||
tipSidebar: "左側のテーブルをクリックしてデータを表示",
|
||||
tipExecute: "でクエリを実行",
|
||||
fileOpened: "{name} を開きました",
|
||||
fileOpenFailed: "{name} を開けませんでした:{message}",
|
||||
mcpTitle: "AIエージェント連携",
|
||||
mcpDescription: "内蔵AIアシスタントに加えて、Claude Code、Cursor、その他のコーディングエージェントをMCP経由でデータベースクエリに使用できます。",
|
||||
mcpLearnMore: "詳細を見る",
|
||||
|
||||
@@ -1894,6 +1894,7 @@ export default withEnglishFallback({
|
||||
tipSidebar: "데이터를 보려면 왼쪽의 테이블을 클릭하세요",
|
||||
tipExecute: "쿼리를 실행하려면",
|
||||
fileOpened: "{name}을(를) 열었습니다",
|
||||
fileOpenFailed: "{name}을(를) 열지 못했습니다: {message}",
|
||||
mcpTitle: "AI 에이전트 연동",
|
||||
mcpDescription: "내장된 AI 어시스턴트 외에도 Claude Code, Cursor 등의 코딩 에이전트를 MCP로 연결하여 데이터베이스를 쿼리할 수 있습니다.",
|
||||
mcpLearnMore: "더 알아보기",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1957,6 +1957,7 @@ export default withEnglishFallback({
|
||||
tipSidebar: "点击左侧表名查看数据",
|
||||
tipExecute: "执行查询",
|
||||
fileOpened: "已打开 {name}",
|
||||
fileOpenFailed: "打开 {name} 失败:{message}",
|
||||
mcpTitle: "AI 编程助手集成",
|
||||
mcpDescription: "除了内置 AI 助手,还可以通过 MCP 让 Claude Code、Cursor 等编程助手直接使用这里的数据库连接。",
|
||||
mcpLearnMore: "了解更多",
|
||||
|
||||
@@ -1966,6 +1966,7 @@ export default withEnglishFallback({
|
||||
tipSidebar: "點選左側資料表以檢視資料",
|
||||
tipExecute: "執行查詢",
|
||||
fileOpened: "已開啟 {name}",
|
||||
fileOpenFailed: "開啟 {name} 失敗:{message}",
|
||||
mcpTitle: "AI 程式設計助理整合",
|
||||
mcpDescription: "除了內建 AI 助理,也可以藉由 MCP 讓 Claude Code、Cursor 等程式設計助理查詢你的資料庫。",
|
||||
mcpLearnMore: "了解更多",
|
||||
|
||||
@@ -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<string, string>();
|
||||
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<typeof import("@/lib/tabs/tabResultCache")>();
|
||||
return { ...actual, deleteTabResultSnapshotsForOwner: vi.fn().mockResolvedValue(undefined) };
|
||||
});
|
||||
}
|
||||
|
||||
function previewConnection(overrides: Partial<ConnectionConfig> = {}): 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);
|
||||
});
|
||||
});
|
||||
@@ -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<void>;
|
||||
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)),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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<ConnectionTestResult, String> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user