diff --git a/apps/desktop/src/components/editor/EditorSettingsDialog.vue b/apps/desktop/src/components/editor/EditorSettingsDialog.vue index 076ce39ecc..0b319e015d 100644 --- a/apps/desktop/src/components/editor/EditorSettingsDialog.vue +++ b/apps/desktop/src/components/editor/EditorSettingsDialog.vue @@ -74,6 +74,7 @@ import { EMPTY_TABLE_COLUMN_TEMPLATE_DATA_TYPE, parseTableColumnTemplateFields, import { buildMcpCodexConfig, buildMcpJsonConfig, buildMcpOpenCodeConfig, buildMcpVsCodeConfig, type McpEnvEntry, type McpLaunchConfig } from "@/lib/mcp/mcpConfigTemplates"; import { isWindows } from "@/lib/backend/platform"; import { combineDataTypeForDatabase, dataTypeLengthInputValue, getDataTypeOptions, getDefaultLengthForType, isDataTypeLengthDisabled, splitDataType } from "@/lib/table/tableStructureEditorState"; +import { useToast } from "@/composables/useToast"; import type { DatabaseType, SqlSnippet } from "@/types/database"; import { uuid } from "@/lib/common/utils"; import { DEFAULT_SQL_SNIPPETS } from "@/lib/sql/sqlCompletion"; @@ -90,6 +91,7 @@ import { apiUrl } from "@/lib/common/webPath"; import { DEFAULT_UI_FONT_FAMILY, SYSTEM_UI_FONT_FAMILY } from "@/lib/app/appFonts"; const { t } = useI18n(); +const { toast } = useToast(); const settingsStore = useSettingsStore(); const connectionStore = useConnectionStore(); const savedSqlStore = useSavedSqlStore(); @@ -227,6 +229,10 @@ const editQuitOnClose = ref(settingsStore.desktopSettings.quit_on_close); const desktopCloseBehaviorResetPending = ref(false); const editIconTheme = ref(settingsStore.desktopSettings.icon_theme); const editDebugLoggingEnabled = ref(settingsStore.desktopSettings.debug_logging_enabled); +const editDuckDbWorkerProcessIsolation = ref(settingsStore.desktopSettings.duckdb_worker_process_isolation); +const startupDuckDbWorkerProcessIsolation = ref(settingsStore.desktopSettings.duckdb_worker_process_isolation); +const duckDbWorkerStartupCaptured = ref(false); +const duckDbRestarting = ref(false); const editSidebarTablePageSize = ref(settingsStore.desktopSettings.sidebar_table_page_size ?? DEFAULT_SIDEBAR_TABLE_PAGE_SIZE); const debugLogCopied = ref(false); const debugLogDownloaded = ref(false); @@ -526,6 +532,7 @@ watch( editQuitOnClose.value = settingsStore.desktopSettings.quit_on_close; editIconTheme.value = settingsStore.desktopSettings.icon_theme; editDebugLoggingEnabled.value = settingsStore.desktopSettings.debug_logging_enabled; + editDuckDbWorkerProcessIsolation.value = settingsStore.desktopSettings.duckdb_worker_process_isolation; editSidebarTablePageSize.value = settingsStore.desktopSettings.sidebar_table_page_size ?? DEFAULT_SIDEBAR_TABLE_PAGE_SIZE; editShowColumnCommentsInHeader.value = settingsStore.editorSettings.showColumnCommentsInHeader; editShowColumnTypesInHeader.value = settingsStore.editorSettings.showColumnTypesInHeader; @@ -569,6 +576,7 @@ const formatterEditorShortcutIds: ShortcutActionId[] = ["formatSql", "find", "re const formatterEditorShortcutDefinitions = computed(() => formatterEditorShortcutIds.map((id) => SHORTCUT_DEFINITIONS.find((definition) => definition.id === id)).filter((definition): definition is (typeof SHORTCUT_DEFINITIONS)[number] => !!definition)); const hasShortcutConflicts = computed(() => shortcutConflicts.value.length > 0); const shortcutsChanged = computed(() => JSON.stringify(editShortcuts.value) !== JSON.stringify(settingsStore.editorSettings.shortcuts)); +const duckDbWorkerProcessIsolationRequiresRestart = computed(() => editDuckDbWorkerProcessIsolation.value !== startupDuckDbWorkerProcessIsolation.value); const hasBlockingShortcutConflicts = computed(() => shortcutsChanged.value && hasShortcutConflicts.value); const hasBlockingFormatterConfig = computed(() => activeSettingsTab.value === "formatter" && !sqlFormatterConfigValid.value); const hasApplyBlocker = computed(() => hasBlockingShortcutConflicts.value || hasBlockingFormatterConfig.value); @@ -596,6 +604,7 @@ function hasChanges(): boolean { editQuitOnClose.value !== settingsStore.desktopSettings.quit_on_close || editIconTheme.value !== settingsStore.desktopSettings.icon_theme || editDebugLoggingEnabled.value !== settingsStore.desktopSettings.debug_logging_enabled || + editDuckDbWorkerProcessIsolation.value !== settingsStore.desktopSettings.duckdb_worker_process_isolation || editSidebarTablePageSize.value !== (settingsStore.desktopSettings.sidebar_table_page_size ?? DEFAULT_SIDEBAR_TABLE_PAGE_SIZE) || editShowColumnCommentsInHeader.value !== settingsStore.editorSettings.showColumnCommentsInHeader || editShowColumnTypesInHeader.value !== settingsStore.editorSettings.showColumnTypesInHeader || @@ -680,6 +689,7 @@ async function persistSettings() { close_action_prompted: desktopCloseBehaviorResetPending.value ? false : true, icon_theme: editIconTheme.value, debug_logging_enabled: editDebugLoggingEnabled.value, + duckdb_worker_process_isolation: editDuckDbWorkerProcessIsolation.value, sidebar_table_page_size: editSidebarTablePageSize.value, }); desktopCloseBehaviorResetPending.value = false; @@ -699,6 +709,20 @@ async function applySettingsAndClose() { closeSettings(); } +async function restartDbxForDuckDbIsolation() { + if (duckDbRestarting.value || hasApplyBlocker.value || isWeb) return; + duckDbRestarting.value = true; + try { + await persistSettings(); + const { relaunch } = await import("@tauri-apps/plugin-process"); + await relaunch(); + } catch (e: any) { + toast(t("settings.restartDbxFailed", { error: e?.message || String(e) }), 5000); + } finally { + duckDbRestarting.value = false; + } +} + function resetDefaultsForTab(tab: SettingsCategory) { if (tab === "editor") { editFontFamily.value = DEFAULT_EDITOR_SETTINGS.fontFamily; @@ -747,6 +771,7 @@ function resetDefaultsForTab(tab: SettingsCategory) { editDataGridQuickEntry.value = DEFAULT_EDITOR_SETTINGS.dataGridQuickEntry; editInfiniteScroll.value = DEFAULT_EDITOR_SETTINGS.infiniteScroll; editInfiniteScrollMaxRows.value = DEFAULT_EDITOR_SETTINGS.infiniteScrollMaxRows; + editDuckDbWorkerProcessIsolation.value = DEFAULT_DESKTOP_SETTINGS.duckdb_worker_process_isolation; editTableColumnTemplateRows.value = tableColumnTemplateRowsFromSettings(DEFAULT_EDITOR_SETTINGS.tableColumnTemplateFields); editExportBatchSize.value = DEFAULT_EDITOR_SETTINGS.exportBatchSize; editExportRowLimitEnabled.value = DEFAULT_EDITOR_SETTINGS.exportRowLimitEnabled; @@ -784,6 +809,7 @@ function resetAllDefaults() { desktopCloseBehaviorResetPending.value = true; editIconTheme.value = DEFAULT_DESKTOP_SETTINGS.icon_theme; editDebugLoggingEnabled.value = DEFAULT_DESKTOP_SETTINGS.debug_logging_enabled; + editDuckDbWorkerProcessIsolation.value = DEFAULT_DESKTOP_SETTINGS.duckdb_worker_process_isolation; editSidebarTablePageSize.value = DEFAULT_SIDEBAR_TABLE_PAGE_SIZE; editShowColumnCommentsInHeader.value = DEFAULT_EDITOR_SETTINGS.showColumnCommentsInHeader; editShowColumnTypesInHeader.value = DEFAULT_EDITOR_SETTINGS.showColumnTypesInHeader; @@ -1462,6 +1488,11 @@ watch( editQuitOnClose.value = settingsStore.desktopSettings.quit_on_close; editIconTheme.value = settingsStore.desktopSettings.icon_theme; editDebugLoggingEnabled.value = settingsStore.desktopSettings.debug_logging_enabled; + editDuckDbWorkerProcessIsolation.value = settingsStore.desktopSettings.duckdb_worker_process_isolation; + if (!duckDbWorkerStartupCaptured.value) { + startupDuckDbWorkerProcessIsolation.value = settingsStore.desktopSettings.duckdb_worker_process_isolation; + duckDbWorkerStartupCaptured.value = true; + } editSidebarTablePageSize.value = settingsStore.desktopSettings.sidebar_table_page_size ?? DEFAULT_SIDEBAR_TABLE_PAGE_SIZE; webdavPassword.value = ""; webdavSecretsPassphrase.value = ""; @@ -2873,6 +2904,35 @@ onUnmounted(cleanupPreviewEditor);
+ +
{{ t("settings.exportSection") }}
diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index b16673ddf2..251319a5f5 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -2705,6 +2705,11 @@ export default { updateDownloadSourceAtomgit: "AtomGit", debugLoggingEnabled: "Enable debug logs", debugLoggingEnabledDescription: "Record more detailed local user-side diagnostic logs while troubleshooting reports. Disabled by default.", + duckDbWorkerProcessIsolation: "DuckDB process isolation mode (Experimental)", + duckDbWorkerProcessIsolationDescription: "Runs DuckDB queries in a separate worker process. If a long query does not respond to cancellation, DBX can stop the worker process to release the database file lock.", + duckDbWorkerProcessIsolationRestartRequired: "Restart DBX for this change to take effect.", + restartDbx: "Restart DBX", + restartDbxFailed: "Failed to restart DBX: {error}", debugLogsCopy: "Copy logs", debugLogsCopied: "Copied", debugLogsDownload: "Download logs", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index a1df642ff7..7b0600d3f6 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -2661,6 +2661,11 @@ export default withEnglishFallback({ updateDownloadSourceAtomgit: "AtomGit", debugLoggingEnabled: "Activar logs de depuración", debugLoggingEnabledDescription: "Registra logs de diagnóstico locales del lado del usuario al investigar reportes. Desactivado por defecto.", + duckDbWorkerProcessIsolation: "Modo de aislamiento de proceso de DuckDB (experimental)", + duckDbWorkerProcessIsolationDescription: "Ejecuta las consultas de DuckDB en un proceso worker independiente. Si una consulta larga no responde a la cancelación, DBX puede detener el proceso worker para liberar el bloqueo del archivo de base de datos.", + duckDbWorkerProcessIsolationRestartRequired: "Reinicia DBX para que este cambio surta efecto.", + restartDbx: "Reiniciar DBX", + restartDbxFailed: "No se pudo reiniciar DBX: {error}", debugLogsCopy: "Copiar logs", debugLogsCopied: "Copiado", debugLogsDownload: "Descargar logs", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index f3960d9d3f..f09f8cf007 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -2659,6 +2659,11 @@ export default withEnglishFallback({ updateDownloadSourceAtomgit: "AtomGit", debugLoggingEnabled: "Abilita log di debug", debugLoggingEnabledDescription: "Registra log diagnostici locali lato utente durante l'analisi delle segnalazioni. Disattivato per impostazione predefinita.", + duckDbWorkerProcessIsolation: "Modalità di isolamento del processo DuckDB (sperimentale)", + duckDbWorkerProcessIsolationDescription: "Esegue le query DuckDB in un processo worker separato. Se una query lunga non risponde all'annullamento, DBX può arrestare il processo worker per rilasciare il blocco del file di database.", + duckDbWorkerProcessIsolationRestartRequired: "Riavvia DBX affinché questa modifica abbia effetto.", + restartDbx: "Riavvia DBX", + restartDbxFailed: "Impossibile riavviare DBX: {error}", debugLogsCopy: "Copia log", debugLogsCopied: "Copiati", debugLogsDownload: "Scarica log", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index c6e2cc0a29..7925095211 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -2654,6 +2654,11 @@ export default withEnglishFallback({ updateNotificationsEnabledDescription: "無効時、DBXはアプリやドライバーのアップデートを自動確認せず、アップデートバッジも表示しません。手動確認は引き続き利用可能です。", debugLoggingEnabled: "デバッグログを有効にする", debugLoggingEnabledDescription: "トラブルシューティング時に、より詳細なローカル診断ログを記録します。デフォルトでは無効です。", + duckDbWorkerProcessIsolation: "DuckDB プロセス分離モード（実験的）", + duckDbWorkerProcessIsolationDescription: "DuckDB クエリを独立した worker プロセスで実行します。長時間クエリがキャンセルに応答しない場合、DBX は worker プロセスを停止してデータベースファイルのロックを解放できます。", + duckDbWorkerProcessIsolationRestartRequired: "この変更を有効にするには DBX を再起動してください。", + restartDbx: "DBX を再起動", + restartDbxFailed: "DBX の再起動に失敗しました: {error}", debugLogsCopy: "ログをコピー", debugLogsCopied: "コピーしました", debugLogsDownload: "ログをダウンロード", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index a550206d59..9fe248abdf 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -2660,6 +2660,11 @@ export default withEnglishFallback({ updateDownloadSourceAtomgit: "AtomGit", debugLoggingEnabled: "Ativar logs de depuração", debugLoggingEnabledDescription: "Registra logs locais de diagnóstico do lado do usuário ao investigar relatos. Desativado por padrão.", + duckDbWorkerProcessIsolation: "Modo de isolamento de processo do DuckDB (experimental)", + duckDbWorkerProcessIsolationDescription: "Executa consultas DuckDB em um processo worker separado. Se uma consulta longa não responder ao cancelamento, o DBX pode encerrar o processo worker para liberar o bloqueio do arquivo de banco de dados.", + duckDbWorkerProcessIsolationRestartRequired: "Reinicie o DBX para que esta alteração entre em vigor.", + restartDbx: "Reiniciar DBX", + restartDbxFailed: "Falha ao reiniciar o DBX: {error}", debugLogsCopy: "Copiar logs", debugLogsCopied: "Copiado", debugLogsDownload: "Baixar logs", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 4502157374..8d01dc901c 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -2705,6 +2705,11 @@ export default withEnglishFallback({ updateDownloadSourceAtomgit: "AtomGit", debugLoggingEnabled: "启用调试日志", debugLoggingEnabledDescription: "开启后在本机记录更详细的用户侧诊断日志，反馈问题时可临时开启；默认关闭。", + duckDbWorkerProcessIsolation: "DuckDB 进程隔离模式（实验性）", + duckDbWorkerProcessIsolationDescription: "将 DuckDB 查询放到独立 worker 进程中运行。停止长查询时，如果 DuckDB 未及时响应取消请求，DBX 会结束 worker 进程以释放数据库文件锁。", + duckDbWorkerProcessIsolationRestartRequired: "重启 DBX 后生效。", + restartDbx: "重启 DBX", + restartDbxFailed: "重启 DBX 失败：{error}", debugLogsCopy: "复制日志", debugLogsCopied: "已复制", debugLogsDownload: "下载日志", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index 0a327aa8c4..7a093b181d 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -2557,6 +2557,11 @@ export default withEnglishFallback({ updateNotificationsEnabledDescription: "關閉後，DBX 不會自動檢查應用程式和驅動程式更新，也不會顯示更新紅點；仍可手動檢查更新。", debugLoggingEnabled: "啟用偵錯日誌", debugLoggingEnabledDescription: "開啟後在本機記錄更詳細的使用者側診斷日誌，回報問題時可暫時開啟；預設關閉。", + duckDbWorkerProcessIsolation: "DuckDB 進程隔離模式（實驗性）", + duckDbWorkerProcessIsolationDescription: "將 DuckDB 查詢放到獨立 worker 進程中執行。停止長查詢時，如果 DuckDB 未及時回應取消請求，DBX 會結束 worker 進程以釋放資料庫檔案鎖。", + duckDbWorkerProcessIsolationRestartRequired: "重啟 DBX 後生效。", + restartDbx: "重啟 DBX", + restartDbxFailed: "重啟 DBX 失敗：{error}", debugLogsCopy: "複製日誌", debugLogsCopied: "已複製", debugLogsDownload: "下載日誌", diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index f689c1992b..1016504294 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -151,6 +151,7 @@ const DEFAULT_DESKTOP_SETTINGS: DesktopSettings = { quit_on_close: false, close_action_prompted: false, debug_logging_enabled: false, + duckdb_worker_process_isolation: false, saved_sql_sync_dir: null, driver_store_dir: null, plugin_store_dir: null, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index e3aad86622..9eb7ea0e23 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -144,6 +144,7 @@ export interface DesktopSettings { quit_on_close: boolean; close_action_prompted: boolean; debug_logging_enabled: boolean; + duckdb_worker_process_isolation: boolean; saved_sql_sync_dir?: string | null; driver_store_dir?: string | null; plugin_store_dir?: string | null; diff --git a/apps/desktop/src/stores/__tests__/settingsStore.spec.ts b/apps/desktop/src/stores/__tests__/settingsStore.spec.ts index 5edacea695..de5dda0ec1 100644 --- a/apps/desktop/src/stores/__tests__/settingsStore.spec.ts +++ b/apps/desktop/src/stores/__tests__/settingsStore.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { normalizeEditorSettings } from "@/stores/settingsStore"; +import { normalizeDesktopSettings, normalizeEditorSettings } from "@/stores/settingsStore"; describe("normalizeEditorSettings", () => { it("enables automatic table aliases by default", () => { @@ -57,3 +57,9 @@ describe("normalizeEditorSettings", () => { expect(normalizeEditorSettings({ dataGridSearchMode: "invalid" as any }).dataGridSearchMode).toBe("filter"); }); }); + +describe("normalizeDesktopSettings", () => { + it("defaults DuckDB worker process isolation to disabled for old settings", () => { + expect(normalizeDesktopSettings({}).duckdb_worker_process_isolation).toBe(false); + }); +}); diff --git a/apps/desktop/src/stores/settingsStore.ts b/apps/desktop/src/stores/settingsStore.ts index 003d2ec4f2..9cab2293b2 100644 --- a/apps/desktop/src/stores/settingsStore.ts +++ b/apps/desktop/src/stores/settingsStore.ts @@ -48,6 +48,7 @@ export interface DesktopSettings { quit_on_close: boolean; close_action_prompted: boolean; debug_logging_enabled: boolean; + duckdb_worker_process_isolation: boolean; saved_sql_sync_dir?: string | null; driver_store_dir?: string | null; plugin_store_dir?: string | null; @@ -72,6 +73,7 @@ export const DEFAULT_DESKTOP_SETTINGS: DesktopSettings = { quit_on_close: false, close_action_prompted: false, debug_logging_enabled: false, + duckdb_worker_process_isolation: false, saved_sql_sync_dir: null, driver_store_dir: null, plugin_store_dir: null, @@ -79,7 +81,7 @@ export const DEFAULT_DESKTOP_SETTINGS: DesktopSettings = { sidebar_table_page_size: DEFAULT_SIDEBAR_TABLE_PAGE_SIZE, }; -function normalizeDesktopSettings(settings: Partial | null | undefined): DesktopSettings { +export function normalizeDesktopSettings(settings: Partial | null | undefined): DesktopSettings { const iconTheme = settings?.icon_theme === "black" ? "black" : DEFAULT_DESKTOP_SETTINGS.icon_theme; const sidebarTablePageSize = typeof settings?.sidebar_table_page_size === "number" && settings.sidebar_table_page_size > 0 ? settings.sidebar_table_page_size : DEFAULT_DESKTOP_SETTINGS.sidebar_table_page_size; return { @@ -88,6 +90,7 @@ function normalizeDesktopSettings(settings: Partial | null | un quit_on_close: settings?.quit_on_close ?? DEFAULT_DESKTOP_SETTINGS.quit_on_close, close_action_prompted: settings?.close_action_prompted ?? DEFAULT_DESKTOP_SETTINGS.close_action_prompted, debug_logging_enabled: settings?.debug_logging_enabled ?? DEFAULT_DESKTOP_SETTINGS.debug_logging_enabled, + duckdb_worker_process_isolation: settings?.duckdb_worker_process_isolation ?? DEFAULT_DESKTOP_SETTINGS.duckdb_worker_process_isolation, saved_sql_sync_dir: settings?.saved_sql_sync_dir?.trim() || DEFAULT_DESKTOP_SETTINGS.saved_sql_sync_dir, driver_store_dir: settings?.driver_store_dir?.trim() || DEFAULT_DESKTOP_SETTINGS.driver_store_dir, plugin_store_dir: settings?.plugin_store_dir?.trim() || DEFAULT_DESKTOP_SETTINGS.plugin_store_dir, diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index 27aff4f428..4a334ddc5c 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -4,6 +4,13 @@ version = "0.1.0" edition = "2021" license = "Apache-2.0" +[[bin]] +name = "duckdb-worker-test-host" +path = "tests/support/duckdb_worker_test_host.rs" +required-features = ["duckdb-bundled"] +test = false +bench = false + [features] default = ["duckdb-bundled", "mq-admin"] duckdb-bundled = ["duckdb/bundled"] diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 9358de5ede..cd11e529a9 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -1,5 +1,6 @@ use std::collections::{HashMap, HashSet}; use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::{Mutex, RwLock}; @@ -39,15 +40,17 @@ const POOL_CLOSE_TIMEOUT_SECS: u64 = 3; mod duckdb_types { use std::sync::Arc; pub type DuckDbHandle = Arc; + pub type DuckDbWorkerHandle = Arc; pub type ExternalTabularHandle = Arc; } #[cfg(not(feature = "duckdb-bundled"))] mod duckdb_types { pub type DuckDbHandle = (); + pub type DuckDbWorkerHandle = (); pub type ExternalTabularHandle = (); } -use duckdb_types::{DuckDbHandle, ExternalTabularHandle}; +use duckdb_types::{DuckDbHandle, DuckDbWorkerHandle, ExternalTabularHandle}; #[derive(Debug, Clone, Copy, PartialEq)] pub enum MysqlMode { @@ -77,6 +80,7 @@ pub enum PoolKind { Turso(db::turso_driver::TursoClient), Redis(db::redis_driver::RedisConnection), DuckDb(DuckDbHandle), + DuckDbWorker(DuckDbWorkerHandle), MongoDb(mongodb::Client), ClickHouse(db::clickhouse_driver::ChClient), SqlServer(Arc>), @@ -167,6 +171,7 @@ pub struct AppState { pub plugins: PluginRegistry, pub agent_manager: crate::agent_manager::AgentManager, pub nacos_registry: crate::nacos::NacosAdminRegistry, + duckdb_worker_process_isolation: AtomicBool, /// PostgreSQL TLS cancel context, keyed by pool_key. /// Used to reconstruct a TLS connector compatible with the original connection when cancelling. postgres_cancel_contexts: Arc>>, @@ -467,6 +472,7 @@ impl AppState { app_version, ), nacos_registry: crate::nacos::NacosAdminRegistry::new(), + duckdb_worker_process_isolation: AtomicBool::new(false), postgres_cancel_contexts: Arc::new(RwLock::new(HashMap::new())), transaction_sessions: Arc::new(RwLock::new(HashMap::new())), #[cfg(feature = "mq-admin")] @@ -482,6 +488,21 @@ impl AppState { } } + pub fn set_duckdb_worker_process_isolation_enabled(&self, enabled: bool) { + self.duckdb_worker_process_isolation.store(enabled, Ordering::Relaxed); + } + + pub async fn apply_duckdb_worker_process_isolation(&self, enabled: bool) { + let previous = self.duckdb_worker_process_isolation.swap(enabled, Ordering::Relaxed); + if previous != enabled { + self.remove_duckdb_pools_detached().await; + } + } + + fn duckdb_worker_process_isolation_enabled(&self) -> bool { + self.duckdb_worker_process_isolation.load(Ordering::Relaxed) + } + pub async fn test_external_driver(&self, driver_id: &str, config: &ConnectionConfig) -> Result { let params = serde_json::json!({ "connection": config }); let env = self.external_driver_runtime_env(driver_id)?; @@ -744,7 +765,9 @@ impl AppState { let conns = self.connections.read().await; if conns.contains_key(&pool_key) { drop(conns); - if !self.remove_stale_connection_pool(&pool_key).await { + if self.remove_pool_if_duckdb_isolation_mismatch(&pool_key).await { + // Recreate below using the current DuckDB isolation mode. + } else if !self.remove_stale_connection_pool(&pool_key).await { self.touch_pool_activity(&pool_key).await; return Ok(pool_key); } @@ -859,14 +882,35 @@ impl AppState { } #[cfg(feature = "duckdb-bundled")] DatabaseType::DuckDb => { - let con = db::duckdb_driver::connect_path(&expand_tilde(&db_config.host))?; - { - let locked = con.lock().map_err(|e| e.to_string())?; - for attached in &db_config.attached_databases { - crate::schema::duckdb_attach_database(&locked, &attached.name, &expand_tilde(&attached.path))?; + if self.duckdb_worker_process_isolation_enabled() { + let attached_databases = db_config + .attached_databases + .iter() + .map(|attached| crate::models::connection::AttachedDatabaseConfig { + name: attached.name.clone(), + path: expand_tilde(&attached.path), + }) + .collect(); + let client = db::duckdb_worker_process::DuckDbWorkerClient::open( + expand_tilde(&db_config.host), + attached_databases, + ) + .await?; + PoolKind::DuckDbWorker(Arc::new(client)) + } else { + let con = db::duckdb_driver::connect_path(&expand_tilde(&db_config.host))?; + { + let locked = con.lock().map_err(|e| e.to_string())?; + for attached in &db_config.attached_databases { + crate::schema::duckdb_attach_database( + &locked, + &attached.name, + &expand_tilde(&attached.path), + )?; + } } + PoolKind::DuckDb(con) } - PoolKind::DuckDb(con) } #[cfg(not(feature = "duckdb-bundled"))] DatabaseType::DuckDb => { @@ -1532,6 +1576,7 @@ impl AppState { } PoolKind::Sqlite(_) | PoolKind::DuckDb(_) + | PoolKind::DuckDbWorker(_) | PoolKind::ExternalTabular(_) | PoolKind::ExternalDriver { .. } | PoolKind::MessageQueue @@ -1630,6 +1675,43 @@ impl AppState { } } + #[cfg(feature = "duckdb-bundled")] + async fn remove_pool_by_key_detached(&self, pool_key: &str) -> bool { + self.stop_keepalive_task(pool_key).await; + self.pool_activity.write().await.remove(pool_key); + self.postgres_cancel_contexts.write().await.remove(pool_key); + let removed = self.connections.write().await.remove(pool_key); + if let Some(pool) = removed { + close_removed_pools_in_background(vec![(pool_key.to_string(), pool)]); + true + } else { + false + } + } + + #[cfg(feature = "duckdb-bundled")] + async fn remove_pool_if_duckdb_isolation_mismatch(&self, pool_key: &str) -> bool { + let isolation_enabled = self.duckdb_worker_process_isolation_enabled(); + let mismatch = { + let connections = self.connections.read().await; + match connections.get(pool_key) { + Some(PoolKind::DuckDb(_)) => isolation_enabled, + Some(PoolKind::DuckDbWorker(_)) => !isolation_enabled, + _ => false, + } + }; + if mismatch { + self.remove_pool_by_key_detached(pool_key).await + } else { + false + } + } + + #[cfg(not(feature = "duckdb-bundled"))] + async fn remove_pool_if_duckdb_isolation_mismatch(&self, _pool_key: &str) -> bool { + false + } + #[cfg(feature = "duckdb-bundled")] pub fn spawn_duckdb_pool_cleanup(&self, pool_key: String, con: DuckDbHandle) { let connections = self.connections.clone(); @@ -1995,6 +2077,7 @@ impl AppState { } PoolKind::Sqlite(_) | PoolKind::DuckDb(_) + | PoolKind::DuckDbWorker(_) | PoolKind::ExternalTabular(_) | PoolKind::ExternalDriver { .. } | PoolKind::MessageQueue @@ -2063,6 +2146,15 @@ impl AppState { close_removed_pools_in_background(removed); } + #[cfg(feature = "duckdb-bundled")] + async fn remove_duckdb_pools_detached(&self) { + let removed = self.drain_duckdb_pools().await; + close_removed_pools_in_background(removed); + } + + #[cfg(not(feature = "duckdb-bundled"))] + async fn remove_duckdb_pools_detached(&self) {} + pub async fn remove_external_driver_pools(&self, driver_id: &str) { let removed = self.drain_external_driver_pools(driver_id).await; close_removed_pools(removed).await; @@ -2101,6 +2193,37 @@ impl AppState { removed } + #[cfg(feature = "duckdb-bundled")] + async fn drain_duckdb_pools(&self) -> Vec<(String, PoolKind)> { + let keys_to_remove: Vec = self + .connections + .read() + .await + .iter() + .filter_map(|(key, pool)| match pool { + PoolKind::DuckDb(_) | PoolKind::DuckDbWorker(_) => Some(key.clone()), + _ => None, + }) + .collect(); + self.stop_keepalive_tasks(&keys_to_remove).await; + { + let mut activity = self.pool_activity.write().await; + let mut cancel_contexts = self.postgres_cancel_contexts.write().await; + for key in &keys_to_remove { + activity.remove(key); + cancel_contexts.remove(key); + } + } + let mut conns = self.connections.write().await; + let mut removed = Vec::with_capacity(keys_to_remove.len()); + for key in keys_to_remove { + if let Some(pool) = conns.remove(&key) { + removed.push((key, pool)); + } + } + removed + } + async fn drain_external_driver_pools(&self, driver_id: &str) -> Vec<(String, PoolKind)> { let keys_to_remove: Vec = self .connections @@ -2360,8 +2483,12 @@ fn clone_pool_kind(pool: &PoolKind) -> PoolKind { PoolKind::Turso(client) => PoolKind::Turso(client.clone()), #[cfg(feature = "duckdb-bundled")] PoolKind::DuckDb(con) => PoolKind::DuckDb(con.clone()), + #[cfg(feature = "duckdb-bundled")] + PoolKind::DuckDbWorker(client) => PoolKind::DuckDbWorker(client.clone()), #[cfg(not(feature = "duckdb-bundled"))] PoolKind::DuckDb(con) => PoolKind::DuckDb(con.clone()), + #[cfg(not(feature = "duckdb-bundled"))] + PoolKind::DuckDbWorker(client) => PoolKind::DuckDbWorker(client.clone()), PoolKind::MongoDb(client) => PoolKind::MongoDb(client.clone()), PoolKind::ClickHouse(client) => PoolKind::ClickHouse(client.clone()), PoolKind::SqlServer(client) => PoolKind::SqlServer(client.clone()), @@ -2395,8 +2522,14 @@ pub async fn close_pool_kind(pool: PoolKind) { PoolKind::DuckDb(con) => { crate::db::duckdb_driver::close_connection(con); } + #[cfg(feature = "duckdb-bundled")] + PoolKind::DuckDbWorker(client) => { + client.shutdown().await; + } #[cfg(not(feature = "duckdb-bundled"))] PoolKind::DuckDb(_) => {} + #[cfg(not(feature = "duckdb-bundled"))] + PoolKind::DuckDbWorker(_) => {} PoolKind::MongoDb(client) => { drop(client); } @@ -3612,6 +3745,55 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + #[cfg(feature = "duckdb-bundled")] + #[tokio::test] + async fn applying_duckdb_worker_isolation_drops_existing_duckdb_pools() { + let (state, dir) = test_app_state().await; + let db_path = dir.join("app.duckdb"); + duckdb::Connection::open(&db_path).unwrap(); + let mut config = mysql_config(None); + config.id = "duckdb-conn".to_string(); + config.name = "DuckDB".to_string(); + config.db_type = DatabaseType::DuckDb; + config.host = db_path.to_string_lossy().to_string(); + config.port = 0; + + state.configs.write().await.insert(config.id.clone(), config); + state.get_or_create_pool("duckdb-conn", None).await.unwrap(); + assert!(matches!(state.connections.read().await.get("duckdb-conn"), Some(PoolKind::DuckDb(_)))); + + state.apply_duckdb_worker_process_isolation(true).await; + + assert!(!state.connections.read().await.contains_key("duckdb-conn")); + + let _ = std::fs::remove_dir_all(dir); + } + + #[cfg(feature = "duckdb-bundled")] + #[tokio::test] + async fn duckdb_pool_mode_mismatch_removes_existing_pool() { + let (state, dir) = test_app_state().await; + let db_path = dir.join("app.duckdb"); + duckdb::Connection::open(&db_path).unwrap(); + let mut config = mysql_config(None); + config.id = "duckdb-conn".to_string(); + config.name = "DuckDB".to_string(); + config.db_type = DatabaseType::DuckDb; + config.host = db_path.to_string_lossy().to_string(); + config.port = 0; + + state.configs.write().await.insert(config.id.clone(), config); + state.get_or_create_pool("duckdb-conn", None).await.unwrap(); + assert!(matches!(state.connections.read().await.get("duckdb-conn"), Some(PoolKind::DuckDb(_)))); + + state.set_duckdb_worker_process_isolation_enabled(true); + + assert!(state.remove_pool_if_duckdb_isolation_mismatch("duckdb-conn").await); + assert!(!state.connections.read().await.contains_key("duckdb-conn")); + + let _ = std::fs::remove_dir_all(dir); + } + #[tokio::test] async fn remove_connection_pools_clears_base_and_database_scoped_pools() { let (state, dir) = test_app_state().await; diff --git a/crates/dbx-core/src/db/duckdb_worker_process.rs b/crates/dbx-core/src/db/duckdb_worker_process.rs new file mode 100644 index 0000000000..f628b6b467 --- /dev/null +++ b/crates/dbx-core/src/db/duckdb_worker_process.rs @@ -0,0 +1,456 @@ +#![cfg(feature = "duckdb-bundled")] + +use std::collections::HashMap; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::process::{Child, ChildStdin, Command}; +use tokio::sync::{oneshot, Mutex, OwnedSemaphorePermit, Semaphore}; +use tokio_util::sync::CancellationToken; + +use crate::db; +use crate::db::duckdb_worker_protocol::{ + DuckDbWorkerConnectParams, DuckDbWorkerError, DuckDbWorkerExecuteParams, DuckDbWorkerMethod, DuckDbWorkerRequest, + DuckDbWorkerResponse, +}; +use crate::models::connection::AttachedDatabaseConfig; + +const DEFAULT_WORKER_REQUEST_TIMEOUT: Duration = Duration::from_secs(60); +const DEFAULT_WORKER_KILL_WAIT: Duration = Duration::from_secs(3); +const DEFAULT_WORKER_START_WAIT: Duration = Duration::from_secs(5); +const DEFAULT_DUCKDB_WORKER_PROCESS_LIMIT: usize = 4; + +type PendingRequests = Arc>>; + +struct PendingRequest { + generation: u64, + sender: oneshot::Sender, +} + +#[derive(Clone)] +pub struct DuckDbWorkerClient { + inner: Arc, +} + +struct DuckDbWorkerClientInner { + state: Mutex, + pending: PendingRequests, + connect_lock: Mutex<()>, + query_lock: Mutex<()>, + process_limiter: Arc, + executable: PathBuf, + connect_params: DuckDbWorkerConnectParams, + request_timeout: Duration, + worker_start_timeout: Duration, + next_id: AtomicU64, +} + +#[derive(Default)] +struct WorkerProcessState { + child: Option, + stdin: Option, + connected: bool, + generation: u64, +} + +impl DuckDbWorkerClient { + pub async fn open(path: String, attached_databases: Vec) -> Result { + let executable = std::env::current_exe().map_err(|e| e.to_string())?; + Self::open_with_executable(executable, path, attached_databases).await + } + + pub async fn open_with_executable( + executable: PathBuf, + path: String, + attached_databases: Vec, + ) -> Result { + let client = Self { + inner: Arc::new(DuckDbWorkerClientInner { + state: Mutex::new(WorkerProcessState::default()), + pending: Arc::new(Mutex::new(HashMap::new())), + connect_lock: Mutex::new(()), + query_lock: Mutex::new(()), + process_limiter: duckdb_worker_process_limiter(), + executable, + connect_params: DuckDbWorkerConnectParams { path, attached_databases }, + request_timeout: DEFAULT_WORKER_REQUEST_TIMEOUT, + worker_start_timeout: DEFAULT_WORKER_START_WAIT, + next_id: AtomicU64::new(1), + }), + }; + client.ensure_connected().await?; + Ok(client) + } + + pub async fn execute( + &self, + database: Option, + sql: String, + max_rows: Option, + cancel_token: Option, + query_timeout: Option, + ) -> Result { + let _query_guard = self.inner.query_lock.lock().await; + let client = self.clone(); + let future = async move { + client + .request::( + DuckDbWorkerMethod::Execute, + DuckDbWorkerExecuteParams { sql, database, max_rows }, + None, + ) + .await + }; + tokio::pin!(future); + + match (cancel_token, query_timeout) { + (Some(token), Some(duration)) => { + tokio::select! { + biased; + _ = token.cancelled() => self.cancel_or_kill(crate::query::canceled_error()).await, + result = &mut future => result, + _ = tokio::time::sleep(duration) => self.cancel_or_kill(crate::query::timeout_error()).await, + } + } + (Some(token), None) => { + tokio::select! { + biased; + _ = token.cancelled() => self.cancel_or_kill(crate::query::canceled_error()).await, + result = &mut future => result, + } + } + (None, Some(duration)) => { + tokio::select! { + result = &mut future => result, + _ = tokio::time::sleep(duration) => self.cancel_or_kill(crate::query::timeout_error()).await, + } + } + (None, None) => future.await, + } + } + + async fn cancel_or_kill(&self, final_error: String) -> Result { + let _ = self.cancel().await; + Err(final_error) + } + + pub async fn list_databases(&self) -> Result, String> { + self.request(DuckDbWorkerMethod::ListDatabases, serde_json::json!({}), Some(self.inner.request_timeout)).await + } + + pub async fn list_schemas(&self, database: String) -> Result, String> { + self.request( + DuckDbWorkerMethod::ListSchemas, + serde_json::json!({ "database": database }), + Some(self.inner.request_timeout), + ) + .await + } + + pub async fn list_tables(&self, database: String, schema: String) -> Result, String> { + self.request( + DuckDbWorkerMethod::ListTables, + serde_json::json!({ "database": database, "schema": schema }), + Some(self.inner.request_timeout), + ) + .await + } + + pub async fn list_columns( + &self, + database: String, + schema: String, + table: String, + ) -> Result, String> { + self.request( + DuckDbWorkerMethod::ListColumns, + serde_json::json!({ "database": database, "schema": schema, "table": table }), + Some(self.inner.request_timeout), + ) + .await + } + + pub async fn attach_database(&self, attached: AttachedDatabaseConfig) -> Result<(), String> { + self.request::( + DuckDbWorkerMethod::AttachDatabase, + attached, + Some(self.inner.request_timeout), + ) + .await?; + Ok(()) + } + + pub async fn cancel(&self) -> Result<(), String> { + self.kill().await; + Ok(()) + } + + pub async fn shutdown(&self) { + let _ = self + .send_request::( + DuckDbWorkerMethod::Shutdown, + serde_json::json!({}), + Some(self.inner.request_timeout), + ) + .await; + self.kill().await; + } + + pub async fn kill(&self) { + let (child, generation) = { + let mut state = self.inner.state.lock().await; + let child = state.child.take(); + let generation = state.generation; + state.stdin = None; + state.connected = false; + (child, generation) + }; + + if let Some(mut child) = child { + let _ = child.start_kill(); + match tokio::time::timeout(DEFAULT_WORKER_KILL_WAIT, child.wait()).await { + Ok(Ok(status)) => log::info!("[duckdb-worker:kill:exit] status={status}"), + Ok(Err(err)) => log::warn!("[duckdb-worker:kill:wait-failed] error={err}"), + Err(_) => { + log::warn!("[duckdb-worker:kill:wait-timeout] wait_ms={}", DEFAULT_WORKER_KILL_WAIT.as_millis()) + } + } + } + self.fail_pending_for_generation(generation, "duckdb_worker_killed", "DuckDB worker process was killed").await; + } + + async fn request( + &self, + method: DuckDbWorkerMethod, + params: impl serde::Serialize, + timeout: Option, + ) -> Result + where + T: serde::de::DeserializeOwned, + { + self.ensure_connected().await?; + self.send_request(method, params, timeout).await + } + + async fn ensure_connected(&self) -> Result<(), String> { + let _guard = self.inner.connect_lock.lock().await; + { + let mut state = self.inner.state.lock().await; + self.ensure_started_locked(&mut state).await?; + if state.connected { + return Ok(()); + } + } + + self.send_request::( + DuckDbWorkerMethod::Connect, + self.inner.connect_params.clone(), + Some(self.inner.request_timeout), + ) + .await?; + + let mut state = self.inner.state.lock().await; + state.connected = true; + Ok(()) + } + + async fn ensure_started_locked(&self, state: &mut WorkerProcessState) -> Result<(), String> { + let should_restart = match state.child.as_mut() { + Some(child) => match child.try_wait() { + Ok(Some(status)) => { + log::warn!("[duckdb-worker:exit] status={status}"); + true + } + Ok(None) => return Ok(()), + Err(err) => { + log::warn!("[duckdb-worker:exit-check-failed] error={err}"); + true + } + }, + None => true, + }; + + if should_restart { + state.child = None; + state.stdin = None; + state.connected = false; + } + state.generation = state.generation.wrapping_add(1); + let generation = state.generation; + + let permit = + tokio::time::timeout(self.inner.worker_start_timeout, self.inner.process_limiter.clone().acquire_owned()) + .await + .map_err(|_| duckdb_worker_process_limit_error())? + .map_err(|_| "DuckDB worker process limiter is closed".to_string())?; + + log::info!("[duckdb-worker:start] executable={}", self.inner.executable.display()); + let mut child = Command::new(&self.inner.executable) + .arg("--duckdb-worker") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .kill_on_drop(true) + .spawn() + .map_err(|e| format!("Failed to start DuckDB worker: {e}"))?; + + let stdin = child.stdin.take().ok_or("DuckDB worker stdin unavailable")?; + let stdout = child.stdout.take().ok_or("DuckDB worker stdout unavailable")?; + spawn_stdout_reader(stdout, self.inner.pending.clone(), generation, permit); + + state.child = Some(child); + state.stdin = Some(stdin); + state.connected = false; + Ok(()) + } + + async fn send_request( + &self, + method: DuckDbWorkerMethod, + params: impl serde::Serialize, + timeout: Option, + ) -> Result + where + T: serde::de::DeserializeOwned, + { + let id = format!("duckdb-worker-{}", self.inner.next_id.fetch_add(1, Ordering::Relaxed)); + let request = DuckDbWorkerRequest::new(&id, method, params)?; + let line = serde_json::to_string(&request).map_err(|e| e.to_string())?; + let (tx, rx) = oneshot::channel(); + + let write_result = async { + let mut state = self.inner.state.lock().await; + self.ensure_started_locked(&mut state).await?; + if method != DuckDbWorkerMethod::Connect && !state.connected { + return Err("DuckDB worker is not connected".to_string()); + } + let generation = state.generation; + self.inner.pending.lock().await.insert(id.clone(), PendingRequest { generation, sender: tx }); + let stdin = state.stdin.as_mut().ok_or("DuckDB worker stdin unavailable")?; + stdin.write_all(line.as_bytes()).await.map_err(|e| e.to_string())?; + stdin.write_all(b"\n").await.map_err(|e| e.to_string())?; + stdin.flush().await.map_err(|e| e.to_string()) + } + .await; + + if let Err(err) = write_result { + self.inner.pending.lock().await.remove(&id); + return Err(err); + } + + let response = match timeout { + Some(timeout) => match tokio::time::timeout(timeout, rx).await { + Ok(Ok(response)) => response, + Ok(Err(_)) => return Err("DuckDB worker response channel closed".to_string()), + Err(_) => { + self.inner.pending.lock().await.remove(&id); + return Err(format!("DuckDB worker request timed out after {}s", timeout.as_secs())); + } + }, + None => rx.await.map_err(|_| "DuckDB worker response channel closed".to_string())?, + }; + + if !response.ok { + let error = response + .error + .unwrap_or_else(|| DuckDbWorkerError::new("duckdb_worker_error", "DuckDB worker request failed")); + return Err(error.message); + } + + let result = response.result.unwrap_or(serde_json::Value::Null); + serde_json::from_value(result).map_err(|e| e.to_string()) + } + + async fn fail_pending_for_generation(&self, generation: u64, code: &'static str, message: &'static str) { + let pending = { + let mut pending = self.inner.pending.lock().await; + let ids = pending + .iter() + .filter_map(|(id, request)| (request.generation == generation).then(|| id.clone())) + .collect::>(); + ids.into_iter().filter_map(|id| pending.remove(&id).map(|request| (id, request.sender))).collect::>() + }; + for (id, sender) in pending { + let _ = sender.send(DuckDbWorkerResponse::err(id, DuckDbWorkerError::new(code, message))); + } + } +} + +fn duckdb_worker_process_limiter() -> Arc { + static LIMITER: OnceLock> = OnceLock::new(); + LIMITER.get_or_init(|| Arc::new(Semaphore::new(DEFAULT_DUCKDB_WORKER_PROCESS_LIMIT))).clone() +} + +fn duckdb_worker_process_limit_error() -> String { + format!( + "DuckDB worker process limit reached (max {}). Close another DuckDB worker connection or retry later.", + DEFAULT_DUCKDB_WORKER_PROCESS_LIMIT + ) +} + +fn spawn_stdout_reader( + stdout: tokio::process::ChildStdout, + pending: PendingRequests, + generation: u64, + permit: OwnedSemaphorePermit, +) { + tokio::spawn(async move { + let _permit = permit; + let mut lines = BufReader::new(stdout).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => { + if line.trim().is_empty() { + continue; + } + match serde_json::from_str::(&line) { + Ok(response) => { + let sender = { + let mut pending = pending.lock().await; + let matches_generation = pending + .get(&response.id) + .map(|request| request.generation == generation) + .unwrap_or(false); + if matches_generation { + pending.remove(&response.id).map(|request| request.sender) + } else { + None + } + }; + if let Some(sender) = sender { + let _ = sender.send(response); + } + } + Err(err) => { + log::warn!("[duckdb-worker:invalid-response] error={err} line={line}"); + } + } + } + Ok(None) => break, + Err(err) => { + log::warn!("[duckdb-worker:stdout-error] error={err}"); + break; + } + } + } + + let pending = { + let mut pending = pending.lock().await; + let ids = pending + .iter() + .filter_map(|(id, request)| (request.generation == generation).then(|| id.clone())) + .collect::>(); + ids.into_iter().filter_map(|id| pending.remove(&id).map(|request| (id, request.sender))).collect::>() + }; + for (id, sender) in pending { + let _ = sender.send(DuckDbWorkerResponse::err( + id, + DuckDbWorkerError::new("duckdb_worker_exited", "DuckDB worker process exited"), + )); + } + }); +} diff --git a/crates/dbx-core/src/db/duckdb_worker_protocol.rs b/crates/dbx-core/src/db/duckdb_worker_protocol.rs new file mode 100644 index 0000000000..868607be3b --- /dev/null +++ b/crates/dbx-core/src/db/duckdb_worker_protocol.rs @@ -0,0 +1,178 @@ +#![cfg(feature = "duckdb-bundled")] + +use serde::{de::DeserializeOwned, Deserialize, Serialize}; + +use crate::models::connection::AttachedDatabaseConfig; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct DuckDbWorkerRequest { + pub id: String, + pub method: DuckDbWorkerMethod, + #[serde(default)] + pub params: serde_json::Value, +} + +impl DuckDbWorkerRequest { + pub fn new(id: impl Into, method: DuckDbWorkerMethod, params: impl Serialize) -> Result { + let params = serde_json::to_value(params).map_err(|e| e.to_string())?; + Ok(Self { id: id.into(), method, params }) + } + + pub fn parse_params(&self) -> Result { + serde_json::from_value(self.params.clone()).map_err(|err| { + DuckDbWorkerError::new("invalid_params", format!("Invalid params for {:?}: {err}", self.method)) + }) + } +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum DuckDbWorkerMethod { + Connect, + Execute, + ListDatabases, + ListSchemas, + ListTables, + ListColumns, + AttachDatabase, + Cancel, + Shutdown, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct DuckDbWorkerResponse { + pub id: String, + pub ok: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +impl DuckDbWorkerResponse { + pub fn ok(id: impl Into, result: impl Serialize) -> Self { + match serde_json::to_value(result) { + Ok(result) => Self { id: id.into(), ok: true, result: Some(result), error: None }, + Err(err) => Self::err(id, DuckDbWorkerError::new("serialization_failed", err.to_string())), + } + } + + pub fn ok_empty(id: impl Into) -> Self { + Self { id: id.into(), ok: true, result: None, error: None } + } + + pub fn err(id: impl Into, error: DuckDbWorkerError) -> Self { + Self { id: id.into(), ok: false, result: None, error: Some(error) } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct DuckDbWorkerError { + pub code: String, + pub message: String, +} + +impl DuckDbWorkerError { + pub fn new(code: impl Into, message: impl Into) -> Self { + Self { code: code.into(), message: message.into() } + } + + pub fn from_message(code: &'static str, error: impl ToString) -> Self { + Self::new(code, error.to_string()) + } +} + +impl From for DuckDbWorkerError { + fn from(message: String) -> Self { + Self::new("duckdb_worker_error", message) + } +} + +impl From<&str> for DuckDbWorkerError { + fn from(message: &str) -> Self { + Self::new("duckdb_worker_error", message) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DuckDbWorkerConnectParams { + pub path: String, + #[serde(default)] + pub attached_databases: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DuckDbWorkerExecuteParams { + pub sql: String, + #[serde(default)] + pub database: Option, + #[serde(default)] + pub max_rows: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DuckDbWorkerDatabaseParams { + pub database: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DuckDbWorkerTableParams { + pub database: String, + pub schema: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct DuckDbWorkerColumnParams { + pub database: String, + pub schema: String, + pub table: String, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn protocol_round_trips_connect_request() { + let request = DuckDbWorkerRequest::new( + "req-1", + DuckDbWorkerMethod::Connect, + DuckDbWorkerConnectParams { + path: "D:/tmp/app.duckdb".to_string(), + attached_databases: vec![AttachedDatabaseConfig { + name: "analytics".to_string(), + path: "D:/tmp/analytics.duckdb".to_string(), + }], + }, + ) + .expect("serialize request"); + + let json = serde_json::to_string(&request).expect("request json"); + let parsed: DuckDbWorkerRequest = serde_json::from_str(&json).expect("parse request"); + let params: DuckDbWorkerConnectParams = parsed.parse_params().expect("parse params"); + + assert_eq!(parsed.id, "req-1"); + assert_eq!(parsed.method, DuckDbWorkerMethod::Connect); + assert_eq!(params.path, "D:/tmp/app.duckdb"); + assert_eq!(params.attached_databases[0].name, "analytics"); + } + + #[test] + fn protocol_round_trips_error_response() { + let response = DuckDbWorkerResponse::err( + "req-2", + DuckDbWorkerError::new("duckdb_open_failed", "DuckDB connection failed"), + ); + + let json = serde_json::to_string(&response).expect("response json"); + let parsed: DuckDbWorkerResponse = serde_json::from_str(&json).expect("parse response"); + + assert!(!parsed.ok); + assert_eq!(parsed.error.expect("error").code, "duckdb_open_failed"); + } +} diff --git a/crates/dbx-core/src/db/duckdb_worker_runtime.rs b/crates/dbx-core/src/db/duckdb_worker_runtime.rs new file mode 100644 index 0000000000..bf28129c23 --- /dev/null +++ b/crates/dbx-core/src/db/duckdb_worker_runtime.rs @@ -0,0 +1,525 @@ +#![cfg(feature = "duckdb-bundled")] + +use std::sync::{Arc, Mutex}; + +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Stdout}; + +use crate::db; +use crate::db::duckdb_worker_protocol::{ + DuckDbWorkerColumnParams, DuckDbWorkerConnectParams, DuckDbWorkerDatabaseParams, DuckDbWorkerError, + DuckDbWorkerExecuteParams, DuckDbWorkerMethod, DuckDbWorkerRequest, DuckDbWorkerResponse, DuckDbWorkerTableParams, +}; +use crate::models::connection::AttachedDatabaseConfig; +use crate::path_utils::expand_tilde; + +#[derive(Default)] +pub struct DuckDbWorkerSession { + connection: Option>, + attached_names: Vec, +} + +impl DuckDbWorkerSession { + pub fn connect(&mut self, params: DuckDbWorkerConnectParams) -> Result<(), String> { + let path = expand_tilde(¶ms.path); + let connection = db::duckdb_driver::connect_path(&path)?; + let mut attached_names = Vec::new(); + { + let locked = connection.lock().map_err(|e| e.to_string())?; + for attached in ¶ms.attached_databases { + let path = expand_tilde(&attached.path); + crate::schema::duckdb_attach_database(&locked, &attached.name, &path)?; + attached_names.push(attached.name.clone()); + } + } + self.connection = Some(connection); + self.attached_names = attached_names; + Ok(()) + } + + pub fn execute(&mut self, params: DuckDbWorkerExecuteParams) -> Result { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + let locked = connection.lock().map_err(|e| e.to_string())?; + let result = crate::query::duckdb_execute_for_database( + &locked, + &self.attached_names, + params.database.as_deref(), + ¶ms.sql, + params.max_rows, + )?; + if let Some(name) = duckdb_attached_name_from_attach_sql(¶ms.sql) { + if !self.attached_names.iter().any(|attached| attached.eq_ignore_ascii_case(&name)) { + self.attached_names.push(name); + } + } + Ok(result) + } + + pub fn list_databases(&self) -> Result, String> { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + let locked = connection.lock().map_err(|e| e.to_string())?; + crate::schema::duckdb_list_databases_with_attached(&locked, &self.attached_names) + } + + pub fn list_schemas(&self, params: DuckDbWorkerDatabaseParams) -> Result, String> { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + let locked = connection.lock().map_err(|e| e.to_string())?; + crate::schema::duckdb_list_schemas_with_attached(&locked, ¶ms.database, &self.attached_names) + } + + pub fn list_tables(&self, params: DuckDbWorkerTableParams) -> Result, String> { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + let locked = connection.lock().map_err(|e| e.to_string())?; + crate::schema::duckdb_query_tables_in_database_with_attached( + &locked, + ¶ms.database, + ¶ms.schema, + &self.attached_names, + ) + } + + pub fn list_columns(&self, params: DuckDbWorkerColumnParams) -> Result, String> { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + let locked = connection.lock().map_err(|e| e.to_string())?; + crate::schema::duckdb_query_columns_in_database_with_attached( + &locked, + ¶ms.database, + ¶ms.schema, + ¶ms.table, + &self.attached_names, + ) + } + + pub fn attach_database(&mut self, attached: AttachedDatabaseConfig) -> Result<(), String> { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + let locked = connection.lock().map_err(|e| e.to_string())?; + let path = expand_tilde(&attached.path); + crate::schema::duckdb_attach_database(&locked, &attached.name, &path)?; + self.attached_names.push(attached.name); + Ok(()) + } + + fn interrupt_handle(&self) -> Result, String> { + let connection = self.connection.as_ref().ok_or("DuckDB worker is not connected")?.clone(); + Ok(connection.interrupt_handle()) + } +} + +#[derive(Clone, Default)] +struct DuckDbWorkerRuntime { + session: Arc>, + active_interrupt: Arc>>>, +} + +struct WorkerHandleResult { + response: Option, + shutdown: bool, +} + +impl DuckDbWorkerRuntime { + async fn handle_request( + &self, + request: DuckDbWorkerRequest, + stdout: Arc>, + ) -> WorkerHandleResult { + match request.method { + DuckDbWorkerMethod::Connect => self.handle_session_request(request, |session, request| { + session.connect(request.parse_params()?)?; + Ok(DuckDbWorkerResponse::ok(request.id, serde_json::json!({ "connected": true }))) + }), + DuckDbWorkerMethod::ListDatabases => self.handle_session_request(request, |session, request| { + Ok(DuckDbWorkerResponse::ok(request.id, session.list_databases()?)) + }), + DuckDbWorkerMethod::ListSchemas => self.handle_session_request(request, |session, request| { + let params = request.parse_params()?; + Ok(DuckDbWorkerResponse::ok(request.id, session.list_schemas(params)?)) + }), + DuckDbWorkerMethod::ListTables => self.handle_session_request(request, |session, request| { + let params = request.parse_params()?; + Ok(DuckDbWorkerResponse::ok(request.id, session.list_tables(params)?)) + }), + DuckDbWorkerMethod::ListColumns => self.handle_session_request(request, |session, request| { + let params = request.parse_params()?; + Ok(DuckDbWorkerResponse::ok(request.id, session.list_columns(params)?)) + }), + DuckDbWorkerMethod::AttachDatabase => self.handle_session_request(request, |session, request| { + session.attach_database(request.parse_params()?)?; + Ok(DuckDbWorkerResponse::ok_empty(request.id)) + }), + DuckDbWorkerMethod::Execute => { + if self.active_interrupt.lock().unwrap_or_else(|e| e.into_inner()).is_some() { + return WorkerHandleResult { + response: Some(DuckDbWorkerResponse::err( + request.id, + DuckDbWorkerError::new("duckdb_worker_busy", "DuckDB worker is already executing a query"), + )), + shutdown: false, + }; + } + + let interrupt_handle = match self.session.lock().unwrap_or_else(|e| e.into_inner()).interrupt_handle() { + Ok(handle) => handle, + Err(err) => { + return WorkerHandleResult { + response: Some(DuckDbWorkerResponse::err( + request.id, + DuckDbWorkerError::from_message("duckdb_not_connected", err), + )), + shutdown: false, + } + } + }; + *self.active_interrupt.lock().unwrap_or_else(|e| e.into_inner()) = Some(interrupt_handle); + + let session = self.session.clone(); + let active_interrupt = self.active_interrupt.clone(); + tokio::spawn(async move { + let id = request.id.clone(); + let params = request.parse_params::(); + let result = match params { + Ok(params) => tokio::task::spawn_blocking(move || { + let mut session = session.lock().unwrap_or_else(|e| e.into_inner()); + session.execute(params) + }) + .await + .map_err(|e| e.to_string()) + .and_then(|result| result), + Err(err) => Err(err.message), + }; + *active_interrupt.lock().unwrap_or_else(|e| e.into_inner()) = None; + let response = match result { + Ok(result) => DuckDbWorkerResponse::ok(id, result), + Err(err) => { + DuckDbWorkerResponse::err(id, DuckDbWorkerError::from_message("duckdb_execute_failed", err)) + } + }; + write_response(stdout, &response).await; + }); + + WorkerHandleResult { response: None, shutdown: false } + } + DuckDbWorkerMethod::Cancel => { + if let Some(handle) = self.active_interrupt.lock().unwrap_or_else(|e| e.into_inner()).as_ref() { + handle.interrupt(); + } + WorkerHandleResult { + response: Some(DuckDbWorkerResponse::ok(request.id, serde_json::json!({ "cancelled": true }))), + shutdown: false, + } + } + DuckDbWorkerMethod::Shutdown => WorkerHandleResult { + response: Some(DuckDbWorkerResponse::ok(request.id, serde_json::json!({ "shutdown": true }))), + shutdown: true, + }, + } + } + + fn handle_session_request( + &self, + request: DuckDbWorkerRequest, + handler: impl FnOnce( + &mut DuckDbWorkerSession, + DuckDbWorkerRequest, + ) -> Result, + ) -> WorkerHandleResult { + let id = request.id.clone(); + if self.active_interrupt.lock().unwrap_or_else(|e| e.into_inner()).is_some() { + return WorkerHandleResult { + response: Some(DuckDbWorkerResponse::err( + id, + DuckDbWorkerError::new("duckdb_worker_busy", "DuckDB worker is executing a query"), + )), + shutdown: false, + }; + } + let response = { + match self.session.try_lock() { + Ok(mut session) => match handler(&mut session, request) { + Ok(response) => response, + Err(err) => DuckDbWorkerResponse::err(id, err), + }, + Err(_) => DuckDbWorkerResponse::err( + id, + DuckDbWorkerError::new("duckdb_worker_busy", "DuckDB worker session is busy"), + ), + } + }; + WorkerHandleResult { response: Some(response), shutdown: false } + } +} + +pub async fn run_stdio_worker() -> Result<(), String> { + let runtime = DuckDbWorkerRuntime::default(); + let stdout = Arc::new(tokio::sync::Mutex::new(tokio::io::stdout())); + let stdin = BufReader::new(tokio::io::stdin()); + let mut lines = stdin.lines(); + + loop { + let line = lines.next_line().await.map_err(|e| e.to_string())?; + let Some(line) = line else { + break; + }; + if line.trim().is_empty() { + continue; + } + let request: DuckDbWorkerRequest = match serde_json::from_str(&line) { + Ok(request) => request, + Err(err) => { + let response = DuckDbWorkerResponse::err( + "", + DuckDbWorkerError::new("invalid_json", format!("Invalid worker request JSON: {err}")), + ); + write_response(stdout.clone(), &response).await; + continue; + } + }; + let result = runtime.handle_request(request, stdout.clone()).await; + if let Some(response) = result.response { + write_response(stdout.clone(), &response).await; + } + if result.shutdown { + break; + } + } + Ok(()) +} + +async fn write_response(stdout: Arc>, response: &DuckDbWorkerResponse) { + let Ok(line) = serde_json::to_string(response) else { + return; + }; + let mut stdout = stdout.lock().await; + let _ = stdout.write_all(line.as_bytes()).await; + let _ = stdout.write_all(b"\n").await; + let _ = stdout.flush().await; +} + +fn duckdb_attached_name_from_attach_sql(sql: &str) -> Option { + let trimmed = sql.trim_start(); + let first_word = trimmed.split(|ch: char| ch.is_whitespace() || ch == ';').next().unwrap_or_default(); + if !first_word.eq_ignore_ascii_case("ATTACH") { + return None; + } + + let as_index = find_as_keyword_outside_quotes(trimmed)?; + parse_identifier_after_as(&trimmed[as_index + 2..]) +} + +fn find_as_keyword_outside_quotes(sql: &str) -> Option { + let bytes = sql.as_bytes(); + let mut i = 0; + let mut in_single = false; + let mut in_double = false; + while i < bytes.len() { + match bytes[i] { + b'\'' if !in_double => { + if in_single && i + 1 < bytes.len() && bytes[i + 1] == b'\'' { + i += 2; + continue; + } + in_single = !in_single; + } + b'"' if !in_single => { + if in_double && i + 1 < bytes.len() && bytes[i + 1] == b'"' { + i += 2; + continue; + } + in_double = !in_double; + } + b'a' | b'A' if !in_single && !in_double && i + 1 < bytes.len() => { + if (bytes[i + 1] == b's' || bytes[i + 1] == b'S') + && is_sql_word_boundary(bytes.get(i.wrapping_sub(1)).copied()) + && is_sql_word_boundary(bytes.get(i + 2).copied()) + { + return Some(i); + } + } + _ => {} + } + i += 1; + } + None +} + +fn is_sql_word_boundary(byte: Option) -> bool { + !matches!(byte, Some(b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_')) +} + +fn parse_identifier_after_as(input: &str) -> Option { + let input = input.trim_start(); + if input.is_empty() { + return None; + } + if let Some(rest) = input.strip_prefix('"') { + let mut name = String::new(); + let mut chars = rest.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == '"' { + if chars.peek() == Some(&'"') { + name.push('"'); + chars.next(); + continue; + } + return (!name.trim().is_empty()).then_some(name); + } + name.push(ch); + } + return None; + } + + let name = input.split(|ch: char| ch.is_whitespace() || ch == ';').next().unwrap_or_default().trim(); + (!name.is_empty()).then(|| name.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn worker_session_connects_to_memory_duckdb() { + let mut session = DuckDbWorkerSession::default(); + + session + .connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() }) + .expect("connect"); + + assert_eq!(session.list_databases().expect("list databases")[0].name, "main"); + } + + #[test] + fn worker_session_executes_select_query() { + let mut session = DuckDbWorkerSession::default(); + session + .connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() }) + .expect("connect"); + + let result = session + .execute(DuckDbWorkerExecuteParams { sql: "SELECT 1 AS value".to_string(), database: None, max_rows: None }) + .expect("execute"); + + assert_eq!(result.columns, vec!["value"]); + assert_eq!(result.rows, vec![vec![serde_json::json!(1)]]); + } + + #[test] + fn worker_session_lists_tables_and_columns() { + let mut session = DuckDbWorkerSession::default(); + session + .connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() }) + .expect("connect"); + session + .execute(DuckDbWorkerExecuteParams { + sql: "CREATE TABLE users (id INTEGER PRIMARY KEY, name VARCHAR)".to_string(), + database: None, + max_rows: None, + }) + .expect("create table"); + + let schemas = + session.list_schemas(DuckDbWorkerDatabaseParams { database: "main".to_string() }).expect("list schemas"); + let tables = session + .list_tables(DuckDbWorkerTableParams { database: "main".to_string(), schema: "main".to_string() }) + .expect("list tables"); + let columns = session + .list_columns(DuckDbWorkerColumnParams { + database: "main".to_string(), + schema: "main".to_string(), + table: "users".to_string(), + }) + .expect("list columns"); + + assert!(schemas.iter().any(|schema| schema == "main")); + assert!(tables.iter().any(|table| table.name == "users")); + assert!(columns.iter().any(|column| column.name == "id" && column.is_primary_key)); + } + + #[test] + fn worker_session_attaches_database() { + let dir = std::env::temp_dir().join(format!("dbx-duckdb-worker-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).expect("create temp dir"); + let attached_path = dir.join("analytics.duckdb"); + { + let con = duckdb::Connection::open(&attached_path).expect("create attached db"); + con.execute_batch("CREATE TABLE events (id INTEGER);").expect("create attached table"); + } + let mut session = DuckDbWorkerSession::default(); + session + .connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() }) + .expect("connect"); + + session + .attach_database(AttachedDatabaseConfig { + name: "analytics".to_string(), + path: attached_path.to_string_lossy().to_string(), + }) + .expect("attach"); + + let databases = session.list_databases().expect("list databases"); + let tables = session + .list_tables(DuckDbWorkerTableParams { database: "analytics".to_string(), schema: "main".to_string() }) + .expect("list tables"); + assert!(databases.iter().any(|database| database.name == "analytics")); + assert!(tables.iter().any(|table| table.name == "events")); + + let _ = std::fs::remove_dir_all(dir); + } + + #[test] + fn worker_session_tracks_attach_sql_alias() { + let dir = std::env::temp_dir().join(format!("dbx-duckdb-worker-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).expect("create temp dir"); + let attached_path = dir.join("sales.duckdb"); + { + let con = duckdb::Connection::open(&attached_path).expect("create attached db"); + con.execute_batch("CREATE TABLE orders (id INTEGER);").expect("create attached table"); + } + let mut session = DuckDbWorkerSession::default(); + session + .connect(DuckDbWorkerConnectParams { path: ":memory:".to_string(), attached_databases: Vec::new() }) + .expect("connect"); + + session + .execute(DuckDbWorkerExecuteParams { + sql: format!("ATTACH '{}' AS \"sales db\";", attached_path.to_string_lossy().replace('\'', "''")), + database: None, + max_rows: None, + }) + .expect("attach sql"); + + assert!(session.attached_names.iter().any(|name| name == "sales db")); + let tables = session + .list_tables(DuckDbWorkerTableParams { database: "sales db".to_string(), schema: "main".to_string() }) + .expect("list tables"); + assert!(tables.iter().any(|table| table.name == "orders")); + + let _ = std::fs::remove_dir_all(dir); + } + + #[tokio::test] + async fn metadata_returns_busy_while_query_is_active() { + let runtime = DuckDbWorkerRuntime::default(); + let stdout = Arc::new(tokio::sync::Mutex::new(tokio::io::stdout())); + *runtime.active_interrupt.lock().unwrap_or_else(|e| e.into_inner()) = + Some(duckdb::Connection::open_in_memory().expect("connection").interrupt_handle()); + + let result = runtime + .handle_request( + DuckDbWorkerRequest::new("req-busy", DuckDbWorkerMethod::ListDatabases, serde_json::json!({})) + .expect("request"), + stdout, + ) + .await; + + let response = result.response.expect("response"); + assert!(!response.ok); + assert_eq!(response.error.expect("error").code, "duckdb_worker_busy"); + } + + #[test] + fn attach_sql_alias_parser_handles_generated_sql() { + assert_eq!( + duckdb_attached_name_from_attach_sql("ATTACH 'D:\\tmp\\sales.duckdb' AS \"sales db\";"), + Some("sales db".to_string()) + ); + assert_eq!(duckdb_attached_name_from_attach_sql("select 'not attach' as value"), None); + } +} diff --git a/crates/dbx-core/src/db/mod.rs b/crates/dbx-core/src/db/mod.rs index 35af96fb7f..1277c0ae65 100644 --- a/crates/dbx-core/src/db/mod.rs +++ b/crates/dbx-core/src/db/mod.rs @@ -1,6 +1,12 @@ pub mod agent_driver; pub mod clickhouse_driver; pub mod duckdb_driver; +#[cfg(feature = "duckdb-bundled")] +pub mod duckdb_worker_process; +#[cfg(feature = "duckdb-bundled")] +pub mod duckdb_worker_protocol; +#[cfg(feature = "duckdb-bundled")] +pub mod duckdb_worker_runtime; pub mod elasticsearch_driver; pub mod elasticsearch_sql; pub mod file_validator; diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index fcbe5fc808..544258a2b7 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -627,7 +627,7 @@ async fn wait_for_duckdb_task_with_interrupt( } #[cfg(feature = "duckdb-bundled")] -fn duckdb_execute_for_database( +pub(crate) fn duckdb_execute_for_database( con: &duckdb::Connection, attached_names: &[String], database: Option<&str>, @@ -1015,10 +1015,34 @@ pub async fn do_execute( } } } + #[cfg(feature = "duckdb-bundled")] + PoolKind::DuckDbWorker(client) => { + let client = client.clone(); + if let Some(ref execution_id) = options.execution_id { + let cancel_client = client.clone(); + state.running_queries.register_interrupt(execution_id, move || { + let cancel_client = cancel_client.clone(); + tokio::spawn(async move { + if let Err(error) = cancel_client.cancel().await { + log::warn!("Failed to cancel DuckDB worker query: {error}"); + } + }); + }); + } + let sql = sql.to_string(); + let database = database.map(str::to_string); + let max_rows = options.max_rows; + drop(connections); + client.execute(database, sql, max_rows, cancel_token, query_timeout).await + } #[cfg(not(feature = "duckdb-bundled"))] PoolKind::DuckDb(_) => { return Err("DuckDB support is not compiled in this build".to_string()); } + #[cfg(not(feature = "duckdb-bundled"))] + PoolKind::DuckDbWorker(_) => { + return Err("DuckDB worker support is not compiled in this build".to_string()); + } PoolKind::Mysql(p, mode) => { let p = p.clone(); let bare = *mode == crate::connection::MysqlMode::Bare; @@ -2045,6 +2069,7 @@ pub async fn execute_statements_in_transaction( PoolKind::MessageQueue | PoolKind::Nacos => TxPath::None, #[cfg(feature = "duckdb-bundled")] PoolKind::DuckDb(_) + | PoolKind::DuckDbWorker(_) | PoolKind::Redis(_) | PoolKind::MongoDb(_) | PoolKind::Elasticsearch(_) @@ -2054,6 +2079,7 @@ pub async fn execute_statements_in_transaction( | PoolKind::ExternalDriver { .. } => TxPath::None, #[cfg(not(feature = "duckdb-bundled"))] PoolKind::DuckDb(_) + | PoolKind::DuckDbWorker(_) | PoolKind::Redis(_) | PoolKind::MongoDb(_) | PoolKind::Elasticsearch(_) diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 4473d22b8e..20b0895c09 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -586,6 +586,12 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result { + let client = client.clone(); + drop(connections); + client.list_databases().await + } _ => Ok(vec![]), } } @@ -786,6 +792,16 @@ async fn list_schemas_once( duckdb_list_schemas_with_attached(&con, database, &duckdb_attached_names) .map(|schemas| filter_visible_schema_names(schemas, visible_schema_filter.as_deref())) } + #[cfg(feature = "duckdb-bundled")] + PoolKind::DuckDbWorker(client) => { + let client = client.clone(); + let database = database.to_string(); + drop(connections); + client + .list_schemas(database) + .await + .map(|schemas| filter_visible_schema_names(schemas, visible_schema_filter.as_deref())) + } _ => Ok(vec![]), } } @@ -1453,6 +1469,16 @@ async fn list_tables_once( return duckdb_query_tables_in_database_with_attached(&con, database, schema, &duckdb_attached_names) .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } + #[cfg(feature = "duckdb-bundled")] + if let Some(client) = extract_pool!(&connections, &pool_key, DuckDbWorker) { + let database = database.to_string(); + let schema = schema.to_string(); + drop(connections); + return client + .list_tables(database, schema) + .await + .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); + } if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { drop(connections); return db::clickhouse_driver::list_tables(&client, clickhouse_metadata_database(database, schema)) @@ -3458,6 +3484,14 @@ pub async fn get_columns_core( &duckdb_attached_names, ); } + #[cfg(feature = "duckdb-bundled")] + if let Some(client) = extract_pool!(&connections, &pool_key, DuckDbWorker) { + let database = database.to_string(); + let schema = schema.to_string(); + let table = table.to_string(); + drop(connections); + return client.list_columns(database, schema, table).await; + } if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { drop(connections); return db::clickhouse_driver::get_columns(&client, clickhouse_metadata_database(database, schema), table) diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index cc33ef7369..a6a1f76f4d 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -119,6 +119,8 @@ pub struct DesktopSettings { #[serde(default)] pub debug_logging_enabled: bool, #[serde(default)] + pub duckdb_worker_process_isolation: bool, + #[serde(default)] pub saved_sql_sync_dir: Option, #[serde(default)] pub driver_store_dir: Option, @@ -142,6 +144,7 @@ impl Default for DesktopSettings { quit_on_close: false, close_action_prompted: false, debug_logging_enabled: false, + duckdb_worker_process_isolation: false, saved_sql_sync_dir: None, driver_store_dir: None, plugin_store_dir: None, @@ -797,6 +800,10 @@ impl Storage { "debug_logging_enabled".to_string(), serde_json::Value::Bool(desktop_settings.debug_logging_enabled), ); + settings.insert( + "duckdb_worker_process_isolation".to_string(), + serde_json::Value::Bool(desktop_settings.duckdb_worker_process_isolation), + ); match desktop_settings.saved_sql_sync_dir.as_ref().filter(|path| !path.trim().is_empty()) { Some(path) => { settings.insert("saved_sql_sync_dir".to_string(), serde_json::Value::String(path.clone())); @@ -857,6 +864,10 @@ impl Storage { .get("debug_logging_enabled") .and_then(|value| value.as_bool()) .unwrap_or_else(|| DesktopSettings::default().debug_logging_enabled), + duckdb_worker_process_isolation: settings + .get("duckdb_worker_process_isolation") + .and_then(|value| value.as_bool()) + .unwrap_or_else(|| DesktopSettings::default().duckdb_worker_process_isolation), saved_sql_sync_dir: settings .get("saved_sql_sync_dir") .and_then(|value| value.as_str()) @@ -2738,6 +2749,7 @@ mod tests { quit_on_close: true, close_action_prompted: false, debug_logging_enabled: true, + duckdb_worker_process_isolation: false, saved_sql_sync_dir: None, driver_store_dir: Some("/tmp/dbx-drivers".to_string()), plugin_store_dir: Some("/tmp/dbx-plugins".to_string()), @@ -2756,6 +2768,7 @@ mod tests { quit_on_close: true, close_action_prompted: false, debug_logging_enabled: true, + duckdb_worker_process_isolation: false, saved_sql_sync_dir: None, driver_store_dir: Some("/tmp/dbx-drivers".to_string()), plugin_store_dir: Some("/tmp/dbx-plugins".to_string()), diff --git a/crates/dbx-core/tests/duckdb_worker_process.rs b/crates/dbx-core/tests/duckdb_worker_process.rs new file mode 100644 index 0000000000..35c0c305e9 --- /dev/null +++ b/crates/dbx-core/tests/duckdb_worker_process.rs @@ -0,0 +1,125 @@ +#![cfg(feature = "duckdb-bundled")] + +use std::path::PathBuf; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use dbx_core::db::duckdb_worker_process::DuckDbWorkerClient; +use dbx_core::query_cancel::{RunningQueries, RunningTaskMetadata}; +use tokio_util::sync::CancellationToken; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_recovers_immediately_after_cancelled_long_query() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-test-host")); + let db_path = temp_duckdb_path(); + let _ = std::fs::remove_file(&db_path); + + let client = + DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new()) + .await + .expect("worker process connects"); + + let token = CancellationToken::new(); + let long_query = client.execute( + None, + "SELECT sum(sin(i::DOUBLE) * cos(i::DOUBLE / 3.0)) FROM range(100000000000) AS t(i)".to_string(), + Some(10), + Some(token.clone()), + Some(Duration::from_secs(30)), + ); + tokio::pin!(long_query); + + tokio::time::sleep(Duration::from_millis(200)).await; + token.cancel(); + + let cancelled = tokio::time::timeout(Duration::from_secs(5), &mut long_query) + .await + .expect("cancelled query should return promptly"); + assert_eq!(cancelled.expect_err("long query should be cancelled"), dbx_core::query::canceled_error()); + + let probe = tokio::time::timeout( + Duration::from_secs(5), + client.execute( + None, + "SELECT 1 AS after_cancel_probe".to_string(), + Some(10), + None, + Some(Duration::from_secs(5)), + ), + ) + .await + .expect("probe query should not hang") + .expect("probe query should succeed"); + + assert_eq!(probe.columns, vec!["after_cancel_probe".to_string()]); + assert_eq!(probe.rows, vec![vec![serde_json::json!(1)]]); + + client.shutdown().await; + let _ = std::fs::remove_file(&db_path); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worker_process_recovers_after_registered_cancel_interrupt() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_duckdb-worker-test-host")); + let db_path = temp_duckdb_path(); + let _ = std::fs::remove_file(&db_path); + + let client = + DuckDbWorkerClient::open_with_executable(executable, db_path.to_string_lossy().to_string(), Vec::new()) + .await + .expect("worker process connects"); + + let running_queries = RunningQueries::default(); + let execution_id = "duckdb-worker-cancel-test"; + let registered = running_queries + .register_task(execution_id.to_string(), RunningTaskMetadata::query("duckdb-conn", "main", None)); + let cancel_client = client.clone(); + running_queries.register_interrupt(execution_id, move || { + let cancel_client = cancel_client.clone(); + tokio::spawn(async move { + let _ = cancel_client.cancel().await; + }); + }); + + let long_query = client.execute( + None, + "SELECT sum(sin(i::DOUBLE) * cos(i::DOUBLE / 3.0)) FROM range(100000000000) AS t(i)".to_string(), + Some(10), + Some(registered.token()), + Some(Duration::from_secs(30)), + ); + tokio::pin!(long_query); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(running_queries.cancel(execution_id)); + + let cancelled = tokio::time::timeout(Duration::from_secs(5), &mut long_query) + .await + .expect("cancelled query should return promptly"); + assert_eq!(cancelled.expect_err("long query should be cancelled"), dbx_core::query::canceled_error()); + drop(registered); + + let probe = tokio::time::timeout( + Duration::from_secs(5), + client.execute( + None, + "SELECT 1 AS after_cancel_probe".to_string(), + Some(10), + None, + Some(Duration::from_secs(5)), + ), + ) + .await + .expect("probe query should not hang") + .expect("probe query should succeed"); + + assert_eq!(probe.columns, vec!["after_cancel_probe".to_string()]); + assert_eq!(probe.rows, vec![vec![serde_json::json!(1)]]); + + client.shutdown().await; + let _ = std::fs::remove_file(&db_path); +} + +fn temp_duckdb_path() -> PathBuf { + let suffix = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos(); + std::env::temp_dir().join(format!("dbx-duckdb-worker-process-{suffix}.duckdb")) +} diff --git a/crates/dbx-core/tests/support/duckdb_worker_test_host.rs b/crates/dbx-core/tests/support/duckdb_worker_test_host.rs new file mode 100644 index 0000000000..6f4e392e87 --- /dev/null +++ b/crates/dbx-core/tests/support/duckdb_worker_test_host.rs @@ -0,0 +1,9 @@ +#![cfg(feature = "duckdb-bundled")] + +fn main() { + let runtime = tokio::runtime::Runtime::new().expect("Failed to create DuckDB worker test runtime"); + if let Err(err) = runtime.block_on(dbx_core::db::duckdb_worker_runtime::run_stdio_worker()) { + eprintln!("DuckDB worker test host failed: {err}"); + std::process::exit(1); + } +} diff --git a/src-tauri/src/commands/app_settings.rs b/src-tauri/src/commands/app_settings.rs index d9d0f80294..de54bfc080 100644 --- a/src-tauri/src/commands/app_settings.rs +++ b/src-tauri/src/commands/app_settings.rs @@ -34,6 +34,7 @@ pub async fn save_desktop_settings( settings: DesktopSettings, ) -> Result<(), String> { state.storage.save_desktop_settings(&settings).await?; + state.apply_duckdb_worker_process_isolation(settings.duckdb_worker_process_isolation).await; apply_debug_log_level(settings.debug_logging_enabled); if let Err(err) = apply_desktop_settings(&app, &settings) { eprintln!("Failed to apply desktop settings: {err}"); diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 9443bf24e8..ce41001adb 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -652,15 +652,17 @@ pub fn run() { ); let state = if let Some(agent_dir) = agent_dir { - Arc::new(AppState::new_with_plugin_and_agent_dir_and_app_version( + AppState::new_with_plugin_and_agent_dir_and_app_version( storage, plugin_dir, agent_dir, env!("CARGO_PKG_VERSION"), - )) + ) } else { - Arc::new(AppState::new_with_plugin_dir_and_app_version(storage, plugin_dir, env!("CARGO_PKG_VERSION"))) + AppState::new_with_plugin_dir_and_app_version(storage, plugin_dir, env!("CARGO_PKG_VERSION")) }; + state.set_duckdb_worker_process_isolation_enabled(desktop_settings.duckdb_worker_process_isolation); + let state = Arc::new(state); app.manage(state.clone()); commands::redis_pubsub_server::start_pubsub_server(state.clone()); app.manage(commands::saved_sql::SavedSqlStorageState { data_dir: data_dir.clone() }); diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs index 1e0016908e..fdfb97e96e 100644 --- a/src-tauri/src/main.rs +++ b/src-tauri/src/main.rs @@ -2,5 +2,15 @@ #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")] fn main() { + #[cfg(feature = "duckdb-bundled")] + if std::env::args().any(|arg| arg == "--duckdb-worker") { + let runtime = tokio::runtime::Runtime::new().expect("Failed to create DuckDB worker runtime"); + if let Err(err) = runtime.block_on(dbx_core::db::duckdb_worker_runtime::run_stdio_worker()) { + eprintln!("DuckDB worker failed: {err}"); + std::process::exit(1); + } + return; + } + dbx_lib::run(); }