From da95bd435fd65a9785e23ba6e3de7e27636437d9 Mon Sep 17 00:00:00 2001 From: zipg Date: Wed, 23 Sep 2026 14:43:14 +0800 Subject: [PATCH] fix(updates): auto-continue client upgrade on update-all --- apps/desktop/src/App.vue | 31 +++++--- .../components/config/DriverStoreDialog.vue | 13 +++- .../__tests__/useComponentUpdates.spec.ts | 13 ++-- .../src/composables/useComponentUpdates.ts | 6 +- apps/desktop/src/lib/backend/tauri.ts | 1 + .../componentUpdateOrchestration.spec.ts | 73 ++++++++++++++++++- .../updates/componentUpdateOrchestration.ts | 17 ++++- crates/dbx-core/src/connection/mod.rs | 55 ++++++++++---- crates/dbx-web/src/routes/agents.rs | 30 +++++--- src-tauri/src/commands/agents.rs | 22 ++++-- 10 files changed, 200 insertions(+), 61 deletions(-) diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue index 296de3a0fe..b9f075f463 100644 --- a/apps/desktop/src/App.vue +++ b/apps/desktop/src/App.vue @@ -33,6 +33,7 @@ import type { PluginUpdateBlock } from "@/composables/useComponentUpdates"; import { COMPONENT_UPDATES_CHANGED_EVENT, notifyComponentPluginsUpdated, notifyComponentUpdatesChanged } from "@/lib/updates/componentUpdateEvents"; import { driverStoreUpdateBadgeCount, showMcpUpdateBadge, showToolbarUpdateAction } from "@/lib/updates/updateBadges"; import { + continuePreparedAppUpdate, hasPendingComponentUpdatesAfterAppRestart, markPendingComponentUpdatesAfterAppUpdate, resolveUpdateAllAction, @@ -40,6 +41,7 @@ import { runPendingComponentUpdatePlan, shouldCloseUpdateCenterAfterComponentUpdate, takePendingComponentUpdatesAfterAppRestart, + updateBlockerLabels, type PendingComponentUpdatePlan, } from "@/lib/updates/componentUpdateOrchestration"; import { isUpdatePreviewMockEnabled } from "@/lib/updates/updatePreviewMock"; @@ -1289,7 +1291,9 @@ function reportComponentUpdateResult(result: Awaited 0) notifyComponentPluginsUpdated(); if (updatedComponents.length) toast(t("updates.componentsAutoUpdated", { components: updatedComponents.join(t("updates.componentListSeparator")) })); - if (result.skippedDrivers > 0) toast(t("updates.componentsAutoUpdateSkipped"), 6000); + if (result.blockedDrivers.length) { + toast(t("driverStore.driverUpdateBlocked", { labels: updateBlockerLabels(result.blockedDrivers).join(", ") }), 8000); + } else if (result.skippedDrivers > 0) toast(t("updates.componentsAutoUpdateSkipped"), 6000); const otherFailureCount = result.failed.length - result.blockedPlugins.length; const failureMessages = [result.blockedPlugins.map(pluginUpdateBlockMessage).join("\n"), otherFailureCount > 0 ? t("updates.componentsAutoUpdateFailed", { count: otherFailureCount }) : ""].filter(Boolean); if (failureMessages.length) toast(failureMessages.join("\n"), 8000); @@ -1342,6 +1346,19 @@ function availableComponentUpdateCategories(): ComponentUpdateCategory[] { return categories; } +function continueAppUpdateWithComponents(categories: ComponentUpdateCategory[]) { + return continuePreparedAppUpdate({ + hasComponentUpdates: categories.length > 0, + restartOnly: updateReady.value, + rememberComponentUpdates: () => rememberComponentUpdatesForRestartedApp({ kind: "manual", categories }), + installComponents: async () => { + reportComponentUpdateResult(await componentUpdates.installCategories(categories)); + }, + installDownloadedUpdate, + restartApp, + }); +} + async function updateAllAvailable() { if (updatingAllUpdates.value) return; updatingAllUpdates.value = true; @@ -1359,17 +1376,9 @@ async function updateAllAvailable() { reportComponentUpdateResult(await componentUpdates.installCategories(categories)); return; } - if (action === "defer-components") { - const remembered = await rememberComponentUpdatesForRestartedApp({ kind: "manual", categories }); - if (remembered) { - toast(t("settings.updateRestartHint"), 6000); - return; - } - reportComponentUpdateResult(await componentUpdates.installCategories(categories)); - return; - } + if (action === "install-app") return continueAppUpdateWithComponents(categories); if (action === "download-app") await downloadUpdateInBackground(); - if (updateDownloaded.value || updateReady.value) await rememberComponentUpdatesForRestartedApp({ kind: "manual", categories }); + if (updateDownloaded.value || updateReady.value) await continueAppUpdateWithComponents(categories); } finally { updatingAllUpdates.value = false; } diff --git a/apps/desktop/src/components/config/DriverStoreDialog.vue b/apps/desktop/src/components/config/DriverStoreDialog.vue index 39c28449c4..1e0b429b67 100644 --- a/apps/desktop/src/components/config/DriverStoreDialog.vue +++ b/apps/desktop/src/components/config/DriverStoreDialog.vue @@ -41,6 +41,7 @@ import { runAgentOfflineExportAction } from "@/lib/driverStore/agentOfflineExpor import { DRIVER_CATEGORIES, getCategoryForAgentDriver, assertAgentDriverCategoriesComplete } from "@/lib/connection/driver-category-definitions"; import { hasAnyUpdatableDriverMatching, countInstalledDrivers, countAvailableDrivers, partitionDriversByInstallStatus, upgradeAllDriverTypes, upgradeAllMatchesFullUpdateSet, type DriverInstallStatusFilter } from "@/lib/connection/driverListFilter"; import { notifyComponentUpdatesChanged } from "@/lib/updates/componentUpdateEvents"; +import { updateBlockerLabels } from "@/lib/updates/componentUpdateOrchestration"; const { t } = useI18n(); const { toast } = useToast(); @@ -53,6 +54,10 @@ function backendError(e: unknown): string { return translateBackendError(t, message); } +function driverUpdateBlockedMessage(blockers: Awaited>): string { + return updateBlockerLabels(blockers).join(", "); +} + const props = withDefaults( defineProps<{ updateNotificationsEnabled?: boolean; @@ -504,7 +509,7 @@ async function runDriverInstall(dbType: string) { } const blockers = await api.checkAgentUpdateBlockers([dbType]); if (blockers.length > 0) { - toast(t("driverStore.driverUpdateBlocked", { labels: blockers.map((blocker) => blocker.label).join(", ") })); + toast(t("driverStore.driverUpdateBlocked", { labels: driverUpdateBlockedMessage(blockers) })); return; } const operationId = uuid(); @@ -548,7 +553,7 @@ async function upgradeAll() { if (updatableDbTypes.length === 0) return; const blockers = await api.checkAgentUpdateBlockers(updatableDbTypes); if (blockers.length > 0) { - toast(t("driverStore.driverUpdateBlocked", { labels: blockers.map((blocker) => blocker.label).join(", ") })); + toast(t("driverStore.driverUpdateBlocked", { labels: driverUpdateBlockedMessage(blockers) })); return; } // Do not expose Cancel until the backend operation has registered its token. @@ -623,7 +628,7 @@ async function uninstallDriver(dbType: string) { } const blockers = await api.checkAgentUpdateBlockers([dbType]); if (blockers.length > 0) { - toast(t("driverStore.driverUpdateBlocked", { labels: blockers.map((blocker) => blocker.label).join(", ") })); + toast(t("driverStore.driverUpdateBlocked", { labels: driverUpdateBlockedMessage(blockers) })); return; } await api.uninstallAgent(dbType); @@ -791,7 +796,7 @@ async function importDriverFile(driver: AgentDriverInfo) { } const blockers = await api.checkAgentUpdateBlockers([dbType]); if (blockers.length > 0) { - toast(t("driverStore.driverUpdateBlocked", { labels: blockers.map((blocker) => blocker.label).join(", ") })); + toast(t("driverStore.driverUpdateBlocked", { labels: driverUpdateBlockedMessage(blockers) })); return; } const label = driverLabel(dbType); diff --git a/apps/desktop/src/composables/__tests__/useComponentUpdates.spec.ts b/apps/desktop/src/composables/__tests__/useComponentUpdates.spec.ts index cfa6ce2c46..2f7f12f0c7 100644 --- a/apps/desktop/src/composables/__tests__/useComponentUpdates.spec.ts +++ b/apps/desktop/src/composables/__tests__/useComponentUpdates.spec.ts @@ -131,7 +131,7 @@ describe("useComponentUpdates", () => { expect(mocks.installJdbcPlugin).toHaveBeenCalledOnce(); expect(mocks.installMcpServer).toHaveBeenCalledOnce(); expect(mocks.installMarketplacePlugin).toHaveBeenCalledWith({ repositoryId: "official", pluginId: "example", version: "1.1.0" }); - expect(result).toEqual({ drivers: 1, jdbc: true, mcp: true, plugins: 1, skippedDrivers: 0, failed: [], blockedPlugins: [] }); + expect(result).toEqual({ drivers: 1, jdbc: true, mcp: true, plugins: 1, skippedDrivers: 0, blockedDrivers: [], failed: [], blockedPlugins: [] }); }); it("shares one update operation between automatic and manual callers", async () => { @@ -172,7 +172,7 @@ describe("useComponentUpdates", () => { expect(mocks.installJdbcPlugin).not.toHaveBeenCalled(); expect(mocks.installMcpServer).not.toHaveBeenCalled(); expect(mocks.installMarketplacePlugin).not.toHaveBeenCalled(); - expect(result).toEqual({ drivers: 0, jdbc: false, mcp: false, plugins: 0, skippedDrivers: 0, failed: [], blockedPlugins: [] }); + expect(result).toEqual({ drivers: 0, jdbc: false, mcp: false, plugins: 0, skippedDrivers: 0, blockedDrivers: [], failed: [], blockedPlugins: [] }); }); it("installs every manually selected category when no DBX update exists and automatic updates are disabled", async () => { @@ -190,7 +190,7 @@ describe("useComponentUpdates", () => { expect(mocks.installJdbcPlugin).toHaveBeenCalledOnce(); expect(mocks.installMcpServer).toHaveBeenCalledOnce(); expect(mocks.installMarketplacePlugin).toHaveBeenCalledOnce(); - expect(result).toEqual({ drivers: 1, jdbc: true, mcp: true, plugins: 1, skippedDrivers: 0, failed: [], blockedPlugins: [] }); + expect(result).toEqual({ drivers: 1, jdbc: true, mcp: true, plugins: 1, skippedDrivers: 0, blockedDrivers: [], failed: [], blockedPlugins: [] }); }); it("resumes a persisted manual update-all plan after restart while automatic updates are disabled", async () => { @@ -213,7 +213,7 @@ describe("useComponentUpdates", () => { expect(mocks.installJdbcPlugin).toHaveBeenCalledOnce(); expect(mocks.installMcpServer).toHaveBeenCalledOnce(); expect(mocks.installMarketplacePlugin).toHaveBeenCalledOnce(); - expect(result).toEqual({ drivers: 1, jdbc: true, mcp: true, plugins: 1, skippedDrivers: 0, failed: [], blockedPlugins: [] }); + expect(result).toEqual({ drivers: 1, jdbc: true, mcp: true, plugins: 1, skippedDrivers: 0, blockedDrivers: [], failed: [], blockedPlugins: [] }); }); it("keeps the toolbar update action hidden while a persisted restart plan is still installing", async () => { @@ -247,7 +247,7 @@ describe("useComponentUpdates", () => { }); it("skips blocked agent updates while allowing other component updates", async () => { - mocks.checkAgentUpdateBlockers.mockResolvedValue([{ db_type: "mysql", label: "MySQL" }]); + mocks.checkAgentUpdateBlockers.mockResolvedValue([{ db_type: "mysql", label: "MySQL", connections: ["生产 MySQL", "报表 MySQL"] }]); const updates = useComponentUpdates({ isDesktop: true }); const result = await updates.autoUpdateEnabledComponents(); @@ -257,6 +257,7 @@ describe("useComponentUpdates", () => { expect(mocks.installMcpServer).toHaveBeenCalledOnce(); expect(mocks.installMarketplacePlugin).toHaveBeenCalledOnce(); expect(result.skippedDrivers).toBe(1); + expect(result.blockedDrivers).toEqual([{ db_type: "mysql", label: "MySQL", connections: ["生产 MySQL", "报表 MySQL"] }]); }); it("does not install from stale state when update detection fails", async () => { @@ -302,7 +303,7 @@ describe("useComponentUpdates", () => { const result = await updates.installCategory("jdbc"); expect(mocks.installJdbcPlugin).toHaveBeenCalledOnce(); - expect(result).toEqual({ drivers: 0, jdbc: true, mcp: false, plugins: 0, skippedDrivers: 0, failed: [], blockedPlugins: [] }); + expect(result).toEqual({ drivers: 0, jdbc: true, mcp: false, plugins: 0, skippedDrivers: 0, blockedDrivers: [], failed: [], blockedPlugins: [] }); }); it("reports the blocking connections when a K8S plugin update cannot start", async () => { diff --git a/apps/desktop/src/composables/useComponentUpdates.ts b/apps/desktop/src/composables/useComponentUpdates.ts index 6121278eba..c3841582ef 100644 --- a/apps/desktop/src/composables/useComponentUpdates.ts +++ b/apps/desktop/src/composables/useComponentUpdates.ts @@ -6,7 +6,7 @@ import { mcpUpdateAvailability } from "@/lib/mcp/mcpUpdateStatus"; import { useSettingsStore } from "@/stores/settingsStore"; import { isUpdatePreviewMockEnabled, previewDriverUpdates, previewJdbcUpdate, previewMcpUpdate, previewPluginUpdates } from "@/lib/updates/updatePreviewMock"; import type { ComponentUpdateCategory } from "@/lib/updates/componentUpdateOrchestration"; -import type { AgentDriverInfo, McpServerStatus } from "@/lib/backend/tauri"; +import type { AgentDriverInfo, AgentUpdateBlocker, McpServerStatus } from "@/lib/backend/tauri"; import type { JdbcPluginStatus } from "@/types/database"; export type { ComponentUpdateCategory } from "@/lib/updates/componentUpdateOrchestration"; @@ -19,6 +19,7 @@ export interface ComponentUpdateResult { mcp: boolean; plugins: number; skippedDrivers: number; + blockedDrivers: AgentUpdateBlocker[]; failed: string[]; blockedPlugins: PluginUpdateBlock[]; } @@ -28,7 +29,7 @@ interface ComponentUpdateRefreshOptions { } function emptyResult(): ComponentUpdateResult { - return { drivers: 0, jdbc: false, mcp: false, plugins: 0, skippedDrivers: 0, failed: [], blockedPlugins: [] }; + return { drivers: 0, jdbc: false, mcp: false, plugins: 0, skippedDrivers: 0, blockedDrivers: [], failed: [], blockedPlugins: [] }; } function pluginUpdateBlock(pluginName: string, message: string): PluginUpdateBlock | null { @@ -139,6 +140,7 @@ export function useComponentUpdates(options: { isDesktop: boolean }) { const blockers = await api.checkAgentUpdateBlockers(updatable); if (blockers.length) { result.skippedDrivers = blockers.length; + result.blockedDrivers = blockers; return; } const upgraded = await api.upgradeAllAgents(); diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 6ee68240de..36130eb8df 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -204,6 +204,7 @@ export interface UpgradeAllAgentDriversResult { export interface AgentUpdateBlocker { db_type: string; label: string; + connections: string[]; } export type AgentOfflineArtifactKind = "jar" | "native"; diff --git a/apps/desktop/src/lib/updates/__tests__/componentUpdateOrchestration.spec.ts b/apps/desktop/src/lib/updates/__tests__/componentUpdateOrchestration.spec.ts index e74bd84f42..1077611e54 100644 --- a/apps/desktop/src/lib/updates/__tests__/componentUpdateOrchestration.spec.ts +++ b/apps/desktop/src/lib/updates/__tests__/componentUpdateOrchestration.spec.ts @@ -3,6 +3,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { clearPendingComponentUpdatesAfterAppUpdate, + continuePreparedAppUpdate, hasPendingComponentUpdatesAfterAppRestart, markPendingComponentUpdatesAfterAppUpdate, resolveUpdateAllAction, @@ -10,6 +11,7 @@ import { runPendingComponentUpdatePlan, shouldCloseUpdateCenterAfterComponentUpdate, takePendingComponentUpdatesAfterAppRestart, + updateBlockerLabels, } from "@/lib/updates/componentUpdateOrchestration"; describe("component update orchestration", () => { @@ -17,6 +19,16 @@ describe("component update orchestration", () => { localStorage.clear(); }); + it("lists blocked connection names and falls back to the driver label", () => { + expect( + updateBlockerLabels([ + { label: "MySQL", connections: ["生产 MySQL", "报表 MySQL"] }, + { label: "Oracle", connections: [] }, + { label: "DuckDB", connections: ["生产 MySQL"] }, + ]), + ).toEqual(["生产 MySQL", "报表 MySQL", "Oracle"]); + }); + it("keeps pending component updates across a failed or ordinary startup", () => { expect(markPendingComponentUpdatesAfterAppUpdate("0.6.16", "0.6.17")).toBe(true); @@ -162,19 +174,74 @@ describe("component update orchestration", () => { expect(takePendingComponentUpdatesAfterAppRestart("0.6.17")?.plan).toEqual({ kind: "auto" }); }); - it("downloads DBX before deferring component updates", () => { + it("downloads DBX before installing it with component updates", () => { expect(resolveUpdateAllAction({ hasAppUpdate: true, appUpdateCanInstall: true, appUpdatePrepared: false, hasComponentUpdates: true })).toBe("download-app"); - expect(resolveUpdateAllAction({ hasAppUpdate: true, appUpdateCanInstall: true, appUpdatePrepared: true, hasComponentUpdates: true })).toBe("defer-components"); + expect(resolveUpdateAllAction({ hasAppUpdate: true, appUpdateCanInstall: true, appUpdatePrepared: true, hasComponentUpdates: true })).toBe("install-app"); }); it("persists the selected component categories when the DBX package is already prepared", () => { const action = resolveUpdateAllAction({ hasAppUpdate: true, appUpdateCanInstall: true, appUpdatePrepared: true, hasComponentUpdates: true }); - expect(action).toBe("defer-components"); + expect(action).toBe("install-app"); expect(markPendingComponentUpdatesAfterAppUpdate("0.6.16", "0.6.17", { kind: "manual", categories: ["drivers", "mcp"] })).toBe(true); expect(takePendingComponentUpdatesAfterAppRestart("0.6.17")?.plan).toEqual({ kind: "manual", categories: ["drivers", "mcp"] }); }); + it("continues from update all into the prepared app install", async () => { + const rememberComponentUpdates = vi.fn(async () => true); + const installComponents = vi.fn(async () => {}); + const installDownloadedUpdate = vi.fn(async () => {}); + const restartApp = vi.fn(async () => {}); + + await continuePreparedAppUpdate({ + hasComponentUpdates: true, + restartOnly: false, + rememberComponentUpdates, + installComponents, + installDownloadedUpdate, + restartApp, + }); + + expect(rememberComponentUpdates).toHaveBeenCalledOnce(); + expect(installDownloadedUpdate).toHaveBeenCalledOnce(); + expect(restartApp).not.toHaveBeenCalled(); + expect(installComponents).not.toHaveBeenCalled(); + }); + + it("continues from update all into restart when the app update is already staged", async () => { + const restartApp = vi.fn(async () => {}); + + await continuePreparedAppUpdate({ + hasComponentUpdates: true, + restartOnly: true, + rememberComponentUpdates: async () => true, + installComponents: async () => {}, + installDownloadedUpdate: async () => {}, + restartApp, + }); + + expect(restartApp).toHaveBeenCalledOnce(); + }); + + it("falls back to component updates without restarting when the restart plan cannot be persisted", async () => { + const installComponents = vi.fn(async () => {}); + const installDownloadedUpdate = vi.fn(async () => {}); + const restartApp = vi.fn(async () => {}); + + await continuePreparedAppUpdate({ + hasComponentUpdates: true, + restartOnly: false, + rememberComponentUpdates: async () => false, + installComponents, + installDownloadedUpdate, + restartApp, + }); + + expect(installComponents).toHaveBeenCalledOnce(); + expect(installDownloadedUpdate).not.toHaveBeenCalled(); + expect(restartApp).not.toHaveBeenCalled(); + }); + it("updates components immediately when no DBX update exists or the package cannot be installed", () => { expect(resolveUpdateAllAction({ hasAppUpdate: false, appUpdateCanInstall: true, appUpdatePrepared: false, hasComponentUpdates: true })).toBe("update-components"); expect(resolveUpdateAllAction({ hasAppUpdate: true, appUpdateCanInstall: false, appUpdatePrepared: false, hasComponentUpdates: true })).toBe("update-components"); diff --git a/apps/desktop/src/lib/updates/componentUpdateOrchestration.ts b/apps/desktop/src/lib/updates/componentUpdateOrchestration.ts index 0bcef1cfd7..40171065b3 100644 --- a/apps/desktop/src/lib/updates/componentUpdateOrchestration.ts +++ b/apps/desktop/src/lib/updates/componentUpdateOrchestration.ts @@ -14,7 +14,11 @@ export interface PendingComponentUpdates { plan: PendingComponentUpdatePlan; } -export type UpdateAllAction = "none" | "download-app" | "defer-components" | "update-components"; +export type UpdateAllAction = "none" | "download-app" | "install-app" | "update-components"; + +export function updateBlockerLabels(blockers: Array<{ label: string; connections?: string[] }>): string[] { + return [...new Set(blockers.flatMap((blocker) => (blocker.connections?.length ? blocker.connections : [blocker.label])))]; +} function normalizeVersion(version: string): string { return version.trim().replace(/^[vV]/, ""); @@ -112,7 +116,16 @@ export function takePendingComponentUpdatesAfterAppRestart(currentVersion: strin export function resolveUpdateAllAction(options: { hasAppUpdate: boolean; appUpdateCanInstall: boolean; appUpdatePrepared: boolean; hasComponentUpdates: boolean }): UpdateAllAction { if (!options.hasAppUpdate || !options.appUpdateCanInstall) return options.hasComponentUpdates ? "update-components" : "none"; - return options.appUpdatePrepared ? "defer-components" : "download-app"; + return options.appUpdatePrepared ? "install-app" : "download-app"; +} + +export async function continuePreparedAppUpdate(options: { hasComponentUpdates: boolean; restartOnly: boolean; rememberComponentUpdates: () => Promise; installComponents: () => Promise; installDownloadedUpdate: () => Promise; restartApp: () => Promise }): Promise { + if (options.hasComponentUpdates && !(await options.rememberComponentUpdates())) { + await options.installComponents(); + return; + } + if (options.restartOnly) await options.restartApp(); + else await options.installDownloadedUpdate(); } export function shouldCloseUpdateCenterAfterComponentUpdate(options: { failedCount: number; skippedDriverCount: number; hasAppUpdate: boolean; remainingComponentUpdateCount: number }): boolean { diff --git a/crates/dbx-core/src/connection/mod.rs b/crates/dbx-core/src/connection/mod.rs index ceaf1863d5..7ba65b7889 100644 --- a/crates/dbx-core/src/connection/mod.rs +++ b/crates/dbx-core/src/connection/mod.rs @@ -4904,15 +4904,17 @@ impl AppState { Ok(closed) } - pub async fn active_agent_connection_driver_keys(&self) -> HashSet { + pub async fn active_agent_connection_driver_connections(&self) -> HashMap> { let configs = self.configs.read().await; let connections = self.connection_pools_snapshot().await; - let mut keys = HashSet::new(); + let mut connections_by_key: HashMap> = HashMap::new(); for (pool_key, pool) in connections.iter() { #[cfg(feature = "duckdb-sidecar")] if matches!(pool, PoolKind::DuckDbWorker(_)) { - keys.insert("duckdb".to_string()); + if let Some(config) = config_for_pool_key(pool_key, &configs) { + connections_by_key.entry("duckdb".to_string()).or_default().push(config.name.clone()); + } continue; } if !matches!(pool, PoolKind::Agent(_)) { @@ -4925,25 +4927,33 @@ impl AppState { &config.db_type, config.driver_profile.as_deref(), ) { - keys.insert(agent_key.to_string()); + connections_by_key.entry(agent_key.to_string()).or_default().push(config.name.clone()); } } - keys + for names in connections_by_key.values_mut() { + names.sort(); + names.dedup(); + } + connections_by_key } - pub async fn prepare_agent_driver_updates(&self, driver_keys: &[String]) -> HashSet { + pub async fn active_agent_connection_driver_keys(&self) -> HashSet { + self.active_agent_connection_driver_connections().await.into_keys().collect() + } + + pub async fn prepare_agent_driver_updates(&self, driver_keys: &[String]) -> HashMap> { let candidates = driver_keys.iter().cloned().collect::>(); if candidates.is_empty() { - return HashSet::new(); + return HashMap::new(); } let blockers = self - .active_agent_connection_driver_keys() + .active_agent_connection_driver_connections() .await .into_iter() - .filter(|key| candidates.contains(key)) - .collect::>(); + .filter(|(key, _)| candidates.contains(key)) + .collect::>(); if !blockers.is_empty() { return blockers; } @@ -4953,7 +4963,11 @@ impl AppState { } // A connection may have started while idle runtimes were stopping. - self.active_agent_connection_driver_keys().await.into_iter().filter(|key| candidates.contains(key)).collect() + self.active_agent_connection_driver_connections() + .await + .into_iter() + .filter(|(key, _)| candidates.contains(key)) + .collect() } pub async fn connection_identifier_quote( @@ -7842,8 +7856,13 @@ mod tests { let (state, dir) = test_app_state().await; let mut config = mysql_config(None); config.id = "dameng-conn".to_string(); + config.name = "达梦生产".to_string(); config.db_type = DatabaseType::Dameng; + let mut replica_config = config.clone(); + replica_config.id = "dameng-replica".to_string(); + replica_config.name = "达梦报表".to_string(); state.configs.write().await.insert(config.id.clone(), config); + state.configs.write().await.insert(replica_config.id.clone(), replica_config); state .agent_manager .daemons @@ -7851,14 +7870,19 @@ mod tests { .await .insert("oracle".to_string(), crate::db::agent_driver::AgentDriverClient::test_stub()); state.connections.write().await.insert("dameng-conn".to_string(), agent_pool_stub()); + state.connections.write().await.insert("dameng-replica".to_string(), agent_pool_stub()); assert_eq!( - state.active_agent_connection_driver_keys().await, - std::collections::HashSet::from(["dameng".to_string()]) + state.active_agent_connection_driver_connections().await, + std::collections::HashMap::from([( + "dameng".to_string(), + vec!["达梦报表".to_string(), "达梦生产".to_string()] + )]) ); state.connections.write().await.remove("dameng-conn"); - assert!(state.active_agent_connection_driver_keys().await.is_empty()); + state.connections.write().await.remove("dameng-replica"); + assert!(state.active_agent_connection_driver_connections().await.is_empty()); let _ = std::fs::remove_dir_all(dir); } @@ -7886,6 +7910,7 @@ mod tests { let (state, dir) = test_app_state().await; let mut config = mysql_config(None); config.id = "dameng-conn".to_string(); + config.name = "达梦生产".to_string(); config.db_type = DatabaseType::Dameng; state.configs.write().await.insert(config.id.clone(), config); state.connections.write().await.insert("dameng-conn".to_string(), agent_pool_stub()); @@ -7898,7 +7923,7 @@ mod tests { let blockers = state.prepare_agent_driver_updates(&["dameng".to_string()]).await; - assert_eq!(blockers, std::collections::HashSet::from(["dameng".to_string()])); + assert_eq!(blockers, std::collections::HashMap::from([("dameng".to_string(), vec!["达梦生产".to_string()])])); assert_eq!(state.agent_manager.active_daemon_keys().await, vec!["dameng".to_string()]); let _ = std::fs::remove_dir_all(dir); diff --git a/crates/dbx-web/src/routes/agents.rs b/crates/dbx-web/src/routes/agents.rs index 15002e6d4e..75788ee688 100644 --- a/crates/dbx-web/src/routes/agents.rs +++ b/crates/dbx-web/src/routes/agents.rs @@ -53,6 +53,7 @@ pub struct AgentUpdateBlockersRequest { pub struct AgentUpdateBlocker { pub db_type: String, pub label: String, + pub connections: Vec, } #[derive(Deserialize)] @@ -407,32 +408,37 @@ async fn ensure_no_agent_update_blockers( state: &dbx_core::connection::AppState, db_types: &[String], ) -> Result<(), String> { - let blockers = update_blockers_from_keys(state.prepare_agent_driver_updates(db_types).await, db_types); + let blockers = update_blockers_from_connections(state.prepare_agent_driver_updates(db_types).await, db_types); if blockers.is_empty() { return Ok(()); } - let labels = blockers.into_iter().map(|blocker| blocker.label).collect::>().join(", "); + let labels = blockers + .into_iter() + .flat_map(|blocker| if blocker.connections.is_empty() { vec![blocker.label] } else { blocker.connections }) + .collect::>() + .join(", "); Err(format!("Close these database connections before updating drivers: {labels}")) } async fn agent_update_blockers(state: &dbx_core::connection::AppState, db_types: &[String]) -> Vec { - update_blockers_from_keys(state.active_agent_connection_driver_keys().await, db_types) + update_blockers_from_connections(state.active_agent_connection_driver_connections().await, db_types) } -fn update_blockers_from_keys( - active_keys: std::collections::HashSet, +fn update_blockers_from_connections( + active_connections: std::collections::HashMap>, db_types: &[String], ) -> Vec { let candidate_keys: std::collections::HashSet<&str> = db_types.iter().map(String::as_str).collect(); if candidate_keys.is_empty() { return Vec::new(); } - let mut blockers = active_keys + let mut blockers = active_connections .into_iter() - .filter(|key| candidate_keys.contains(key.as_str())) - .map(|db_type| AgentUpdateBlocker { + .filter(|(key, _)| candidate_keys.contains(key.as_str())) + .map(|(db_type, connections)| AgentUpdateBlocker { label: dbx_core::agent_catalog::label_for_key(&db_type).unwrap_or(&db_type).to_string(), db_type, + connections, }) .collect::>(); blockers.sort_by(|left, right| left.label.cmp(&right.label)); @@ -472,12 +478,16 @@ mod tests { #[test] fn update_blockers_include_active_duckdb_connections() { - let active_keys = ["duckdb".to_string(), "oracle".to_string()].into_iter().collect(); - let blockers = update_blockers_from_keys(active_keys, &["duckdb".to_string()]); + let active_connections = std::collections::HashMap::from([ + ("duckdb".to_string(), vec!["本地分析".to_string()]), + ("oracle".to_string(), vec!["生产 Oracle".to_string()]), + ]); + let blockers = update_blockers_from_connections(active_connections, &["duckdb".to_string()]); assert_eq!(blockers.len(), 1); assert_eq!(blockers[0].db_type, "duckdb"); assert_eq!(blockers[0].label, "DuckDB"); + assert_eq!(blockers[0].connections, vec!["本地分析"]); } #[test] diff --git a/src-tauri/src/commands/agents.rs b/src-tauri/src/commands/agents.rs index 51973d5d36..4b2a9ed088 100644 --- a/src-tauri/src/commands/agents.rs +++ b/src-tauri/src/commands/agents.rs @@ -24,6 +24,7 @@ use dbx_core::DownloadSource; pub struct AgentUpdateBlocker { pub db_type: String, pub label: String, + pub connections: Vec, } #[tauri::command] @@ -314,11 +315,15 @@ fn emit_agent_progress(app: &tauri::AppHandle, operation_id: &str, event: AgentP } async fn ensure_no_agent_update_blockers(state: &AppState, db_types: &[String]) -> Result<(), String> { - let blockers = update_blockers_from_keys(state.prepare_agent_driver_updates(db_types).await, db_types); + let blockers = update_blockers_from_connections(state.prepare_agent_driver_updates(db_types).await, db_types); if blockers.is_empty() { return Ok(()); } - let labels = blockers.into_iter().map(|blocker| blocker.label).collect::>().join(", "); + let labels = blockers + .into_iter() + .flat_map(|blocker| if blocker.connections.is_empty() { vec![blocker.label] } else { blocker.connections }) + .collect::>() + .join(", "); Err(format!("Close these database connections before updating drivers: {labels}")) } @@ -336,23 +341,24 @@ async fn ensure_no_offline_import_blockers(state: &AppState, plan: &OfflineImpor } async fn agent_update_blockers(state: &AppState, db_types: &[String]) -> Vec { - update_blockers_from_keys(state.active_agent_connection_driver_keys().await, db_types) + update_blockers_from_connections(state.active_agent_connection_driver_connections().await, db_types) } -fn update_blockers_from_keys( - active_keys: std::collections::HashSet, +fn update_blockers_from_connections( + active_connections: std::collections::HashMap>, db_types: &[String], ) -> Vec { let candidate_keys: std::collections::HashSet<&str> = db_types.iter().map(String::as_str).collect(); if candidate_keys.is_empty() { return Vec::new(); } - let mut blockers = active_keys + let mut blockers = active_connections .into_iter() - .filter(|key| candidate_keys.contains(key.as_str())) - .map(|db_type| AgentUpdateBlocker { + .filter(|(key, _)| candidate_keys.contains(key.as_str())) + .map(|(db_type, connections)| AgentUpdateBlocker { label: dbx_core::agent_catalog::label_for_key(&db_type).unwrap_or(&db_type).to_string(), db_type, + connections, }) .collect::>(); blockers.sort_by(|left, right| left.label.cmp(&right.label));