From 52499f4c9fc4029bf04c7c482922b2ea0bd86ae2 Mon Sep 17 00:00:00 2001 From: amwps290 Date: Wed, 22 Jul 2026 15:08:51 +0800 Subject: [PATCH] feat(postgres): add extension management --- .../objects/InstallExtensionDialog.vue | 162 ++++++++----- .../src/components/sidebar/ConnectionTree.vue | 4 +- .../sidebar/SidebarTreeRuntimeHost.vue | 9 + apps/desktop/src/lib/backend/http.ts | 2 +- apps/desktop/src/lib/backend/tauri.ts | 2 +- .../src/lib/sidebar/sidebarTreeItemLayout.ts | 1 + .../connectionStore.metadataLoading.spec.ts | 2 +- apps/desktop/src/stores/connectionStore.ts | 95 ++++---- crates/dbx-core/src/database_export.rs | 9 +- crates/dbx-core/src/db/postgres.rs | 39 ++-- crates/dbx-core/src/schema.rs | 216 +++++++++++++++++- crates/dbx-web/src/routes/schema.rs | 3 +- src-tauri/src/commands/schema.rs | 4 +- 13 files changed, 426 insertions(+), 122 deletions(-) diff --git a/apps/desktop/src/components/objects/InstallExtensionDialog.vue b/apps/desktop/src/components/objects/InstallExtensionDialog.vue index 702e1c35bf..d0ae49a2b2 100644 --- a/apps/desktop/src/components/objects/InstallExtensionDialog.vue +++ b/apps/desktop/src/components/objects/InstallExtensionDialog.vue @@ -5,6 +5,7 @@ import { useI18n } from "vue-i18n"; import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { ScrollArea } from "@/components/ui/scroll-area"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { buildCreateExtensionSql, buildDropExtensionSql } from "@/lib/database/dbAdminSql"; import * as api from "@/lib/backend/api"; import { useConnectionStore } from "@/stores/connectionStore"; @@ -23,15 +24,34 @@ const props = defineProps<{ const emit = defineEmits<{ close: []; + changed: []; }>(); const open = ref(false); const available = ref([]); const installed = ref([]); +const schemas = ref([]); +const selectedSchema = ref("__default__"); const loading = ref(false); const installing = ref(null); const dropping = ref(null); +function normalizeSchemaOptions(value: string[]): string[] { + return [...new Set(value.map((item) => item.trim()).filter(Boolean))]; +} + +function preferredSchemaOption(options: string[]): string { + if (options.includes("public")) return "public"; + return options[0] ?? "__default__"; +} + +function installedExtensionSummary(ext: ExtensionInfo): string { + const parts = [ext.version]; + if (ext.schema) parts.push(`schema: ${ext.schema}`); + if (ext.comment) parts.push(ext.comment); + return parts.filter(Boolean).join(" - "); +} + function show() { open.value = true; void loadData(); @@ -41,9 +61,21 @@ async function loadData() { if (!props.node.connectionId || !props.node.database) return; loading.value = true; try { - const [avail, inst] = await Promise.all([api.listAvailableExtensions(props.node.connectionId, props.node.database).catch(() => [] as ExtensionInfo[]), api.listExtensions(props.node.connectionId, props.node.database, props.node.schema || "public").catch(() => [] as ExtensionInfo[])]); - available.value = avail; - installed.value = inst; + const [availResult, instResult, schemaResult] = await Promise.allSettled([ + api.listAvailableExtensions(props.node.connectionId, props.node.database), + api.listExtensions(props.node.connectionId, props.node.database, props.node.schema), + api.listSchemaInfos(props.node.connectionId, props.node.database), + ]); + + available.value = availResult.status === "fulfilled" ? availResult.value : []; + installed.value = instResult.status === "fulfilled" ? instResult.value : []; + schemas.value = schemaResult.status === "fulfilled" ? normalizeSchemaOptions(schemaResult.value.map((schema) => schema.name)) : []; + selectedSchema.value = props.node.schema && schemas.value.includes(props.node.schema) ? props.node.schema : selectedSchema.value !== "__default__" && schemas.value.includes(selectedSchema.value) ? selectedSchema.value : preferredSchemaOption(schemas.value); + + const firstError = [availResult, instResult, schemaResult].find((result): result is PromiseRejectedResult => result.status === "rejected")?.reason; + if (firstError) { + throw firstError; + } } catch (e: any) { toast(t("connection.connectFailed", { message: translateBackendError(t, e?.message || String(e)) }), 5000); } finally { @@ -55,17 +87,18 @@ async function installExtension(name: string) { if (!props.node.connectionId || !props.node.database) return; installing.value = name; try { - const sql = buildCreateExtensionSql(name, props.node.schema ?? null); + const schema = selectedSchema.value === "__default__" ? null : selectedSchema.value; + const sql = buildCreateExtensionSql(name, schema); const result = await executeWithProductionSqlGuard({ connection: connectionStore.getConfig(props.node.connectionId), database: props.node.database, sql, source: t("production.sourceExtension"), - execute: () => api.executeQuery(props.node.connectionId!, props.node.database!, sql, props.node.schema ?? undefined), + execute: () => api.executeQuery(props.node.connectionId!, props.node.database!, sql, schema ?? undefined), }); if (!result) return; await loadData(); - emit("close"); + emit("changed"); } catch (e: any) { toast(t("connection.connectFailed", { message: translateBackendError(t, e?.message || String(e)) }), 5000); } finally { @@ -87,6 +120,7 @@ async function dropExtension(name: string) { }); if (!result) return; await loadData(); + emit("changed"); } catch (e: any) { toast(t("connection.connectFailed", { message: translateBackendError(t, e?.message || String(e)) }), 5000); } finally { @@ -108,59 +142,81 @@ defineExpose({ show }); -
- -
-
- - {{ t("extension.available") }} - ({{ available.length }}) +
+
+
+ {{ t("extension.available") }} + {{ available.length }} + / + {{ t("extension.installed") }} + {{ installed.length }} +
+
+
Schema
+
- -
- {{ t("extension.noAvailable") }} -
-
-
-
-
{{ ext.name }}
-
{{ ext.comment || ext.version }}
-
- -
-
-
- -
-
- - {{ t("extension.installed") }} - ({{ installed.length }}) -
- -
- {{ t("extension.noInstalled") }} +
+
+
+ +
{{ t("extension.available") }}
+ {{ available.length }}
-
-
-
-
{{ ext.name }}
-
{{ ext.version }}{{ ext.comment ? ` — ${ext.comment}` : "" }}
-
- + +
+ {{ t("extension.noAvailable") }}
+
+
+
+
{{ ext.name }}
+
{{ ext.comment || ext.version }}
+
+ +
+
+ +
+ +
+
+ +
{{ t("extension.installed") }}
+ {{ installed.length }}
- + +
+ {{ t("extension.noInstalled") }} +
+
+
+
+
{{ ext.name }}
+
{{ installedExtensionSummary(ext) }}
+
+ +
+
+ +
diff --git a/apps/desktop/src/components/sidebar/ConnectionTree.vue b/apps/desktop/src/components/sidebar/ConnectionTree.vue index 29b42687a8..1e6e6e60bd 100644 --- a/apps/desktop/src/components/sidebar/ConnectionTree.vue +++ b/apps/desktop/src/components/sidebar/ConnectionTree.vue @@ -1051,7 +1051,7 @@ async function executeSidebarProcedureSql(sql: string) { } async function refreshSidebarActionTarget() { - const target = sidebarObjectSourceTarget.value?.node || sidebarDdlTarget.value; + const target = sidebarObjectSourceTarget.value?.node || sidebarDdlTarget.value || sidebarInstallExtensionTarget.value; if (!target) return; const currentTarget = findSidebarActionTarget(store.treeNodes, target); if (!currentTarget) return; @@ -1521,7 +1521,7 @@ defineExpose({ focusSearch, createNewGroup, collapseAllTreeNodes }); - +
{{ t("sidebar.noConnections") }}
diff --git a/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue b/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue index be4c70b7d6..581d8d70c1 100644 --- a/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue +++ b/apps/desktop/src/components/sidebar/SidebarTreeRuntimeHost.vue @@ -518,6 +518,13 @@ async function toggle() { return; } + if (node.type === "group-extensions" && connectionStore.isTreeNodeChildrenLoaded(node.id)) { + node.isExpanded = !node.isExpanded; + if (wasExpanded && !connectionStore.sidebarSearchQuery) connectionStore.releaseCollapsedTreeNodeChildren(node.id); + emit("node-toggled", node, wasExpanded); + return; + } + if (node.isExpanded) { node.isExpanded = false; if (!connectionStore.sidebarSearchQuery) connectionStore.releaseCollapsedTreeNodeChildren(node.id); @@ -641,6 +648,8 @@ async function toggle() { await connectionStore.loadIndexes(node.connectionId, node.database, node.tableName, node.schema, node.id, node.catalog); } else if (node.type === "group-fkeys" && node.connectionId && hasTreeNodeDatabaseContext(node) && node.tableName) { await connectionStore.loadForeignKeys(node.connectionId, node.database, node.tableName, node.schema, node.id, node.catalog); + } else if (node.type === "group-extensions" && node.connectionId && hasTreeNodeDatabaseContext(node)) { + await connectionStore.refreshTreeNode(node); } emit("node-toggled", node, wasExpanded); } catch (e: any) { diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 9631ac7e7f..f9d3294129 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -715,7 +715,7 @@ export async function listOwners(connectionId: string, database: string, schema: return get(`/api/schema/owners?${qs({ connection_id: connectionId, database, schema })}`); } -export async function listExtensions(connectionId: string, database: string, schema: string): Promise { +export async function listExtensions(connectionId: string, database: string, schema?: string): Promise { return get(`/api/schema/extensions?${qs({ connection_id: connectionId, database, schema })}`); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index e919f7b645..04d32e23ca 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -1201,7 +1201,7 @@ export async function listOwners(connectionId: string, database: string, schema: return invoke("list_owners", { connectionId, database, schema }); } -export async function listExtensions(connectionId: string, database: string, schema: string): Promise { +export async function listExtensions(connectionId: string, database: string, schema?: string): Promise { return invoke("list_extensions", { connectionId, database, schema }); } diff --git a/apps/desktop/src/lib/sidebar/sidebarTreeItemLayout.ts b/apps/desktop/src/lib/sidebar/sidebarTreeItemLayout.ts index 7fd17a5e03..15294281ac 100644 --- a/apps/desktop/src/lib/sidebar/sidebarTreeItemLayout.ts +++ b/apps/desktop/src/lib/sidebar/sidebarTreeItemLayout.ts @@ -23,6 +23,7 @@ const leafTypes: Set = new Set([ "saved-sql-file", "table-search-control", "load-more", + "extension", ]); const fullWidthLabelTypes: Set = new Set(["table", "view", "materialized_view", "mongo-collection", "mongo-bucket", "vector-collection", "elasticsearch-index"]); diff --git a/apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts index 3c0ed654b6..dc0c1e6b70 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.metadataLoading.spec.ts @@ -262,7 +262,7 @@ describe("connectionStore metadata loading", () => { await store.loadSchemas(connection.id, "app", { force: true }); expect(store.connectionErrors[connection.id]).toBeUndefined(); - expect(store.treeNodes[0]?.children?.[0]?.children?.map((node) => node.label)).toEqual(["public"]); + expect(store.treeNodes[0]?.children?.[0]?.children?.map((node) => node.label)).toEqual(["public", "tree.extensions"]); }); it("clears a failed metadata warning when the driver hint finishes during retry", async () => { diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 6c70a74b09..fdbf123d62 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -1111,6 +1111,18 @@ export const useConnectionStore = defineStore("connection", () => { return sortSidebarNames([...byName.keys()]).map((name) => byName.get(name)!); } + function buildExtensionManagementNode(connectionId: string, database: string): TreeNode { + return { + id: `${connectionId}:${database}:__extensions`, + label: "tree.extensions", + type: "group-extensions", + connectionId, + database, + isExpanded: false, + children: [], + }; + } + function objectGroupCacheKey(node: TreeNode): string { return schemaCacheKey(node.connectionId || "", node.database || "", node.schema || "", node.type, "objects-v5"); } @@ -2840,7 +2852,7 @@ export const useConnectionStore = defineStore("connection", () => { try { await ensureConnected(connectionId); if (useCachedChildren(node, options)) return; - const cacheKey = schemaCacheKey(connectionId, database, "schemas-v2"); + const cacheKey = schemaCacheKey(connectionId, database, "schemas-v3"); if (!options?.force) { const cached = await loadPersistedTreeChildren(node, cacheKey); if (cached.hit) { @@ -2857,7 +2869,7 @@ export const useConnectionStore = defineStore("connection", () => { database, ), ); - const children = schemas + const children: TreeNode[] = schemas .filter((schema) => visibleSchemaNames.has(schema.name)) .map((schema) => { const s = schema.name; @@ -2873,6 +2885,9 @@ export const useConnectionStore = defineStore("connection", () => { children: [], }; }); + if (isPostgresLikeForExtensions(getConfig(connectionId)?.db_type)) { + children.push(buildExtensionManagementNode(connectionId, database)); + } if (isSidebarSearchQueryChanged(options)) return; if (!canApplyTreeMetadataResult(node)) return; setChildren(node, children); @@ -3190,7 +3205,7 @@ export const useConnectionStore = defineStore("connection", () => { await ensureConnected(connectionId); if (useCachedChildren(node, options)) return; const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; - const cacheKey = schemaCacheKey(connectionId, database, schema || "", simpleObjectDisplay ? "objects-simple-v5" : "objects-grouped-v5"); + const cacheKey = schemaCacheKey(connectionId, database, schema || "", simpleObjectDisplay ? "objects-simple-v5" : "objects-grouped-v6"); const searchFilter = activeTreeLoadSearchFilter(options); const isSidebarTableSearch = !!options?.sidebarTableSearchParentId; if (!options?.force && !searchFilter) { @@ -3231,17 +3246,8 @@ export const useConnectionStore = defineStore("connection", () => { schema: effectiveSchema, objectTypes: supportedSidebarObjectTypes(config), }); - if (isPostgresLikeForExtensions(config?.db_type)) { - children.push({ - id: `${nodeId}:__extensions`, - label: "tree.extensions", - type: "group-extensions", - connectionId, - database, - schema: effectiveSchema, - isExpanded: false, - children: [], - }); + if (!schema && isPostgresLikeForExtensions(config?.db_type)) { + children.push(buildExtensionManagementNode(connectionId, database)); } } if (isTreeLoadSearchChanged(searchFilter, options)) return; @@ -3371,35 +3377,6 @@ export const useConnectionStore = defineStore("connection", () => { ); } - async function loadExtensions(connectionId: string, database: string, schema: string) { - const node = findNode(treeNodes.value, `${connectionId}:${database}:${schema}:__extensions`); - if (!node) return; - node.isLoading = true; - try { - await ensureConnected(connectionId); - if (useCachedChildren(node)) return; - const extensions = await withMetadataLoadTimeout(connectionId, api.listExtensions(connectionId, database, schema), "extensions"); - const children: TreeNode[] = extensions.map((ext) => ({ - id: `${node.id}:${ext.name}`, - label: ext.name, - type: "extension" as const, - connectionId, - database, - schema, - comment: ext.comment ?? null, - meta: ext, - isExpanded: false, - })); - setChildren(node, children); - node.isExpanded = true; - } catch (e) { - recordMetadataLoadError(connectionId, e); - throw e; - } finally { - node.isLoading = false; - } - } - async function loadMoreObjectGroupChildren(node: TreeNode) { if (node.type !== "load-more" || !node.loadMore) return; const loadMore = node.loadMore; @@ -3508,6 +3485,36 @@ export const useConnectionStore = defineStore("connection", () => { ); } + async function loadExtensions(connectionId: string, database: string) { + const node = findNode(treeNodes.value, `${connectionId}:${database}:__extensions`); + if (!node) return; + node.isLoading = true; + try { + await ensureConnected(connectionId); + if (useCachedChildren(node)) return; + const extensions = await withMetadataLoadTimeout(connectionId, api.listExtensions(connectionId, database), "extensions"); + const children: TreeNode[] = extensions.map((ext) => ({ + id: `${node.id}:${ext.schema || ""}:${ext.name}`, + label: ext.name, + type: "extension" as const, + connectionId, + database, + schema: ext.schema ?? undefined, + comment: ext.comment ?? null, + meta: ext, + isExpanded: false, + })); + setChildren(node, children); + node.objectCount = children.length; + node.isExpanded = true; + } catch (e) { + recordMetadataLoadError(connectionId, e); + throw e; + } finally { + node.isLoading = false; + } + } + async function loadTableForLocate(target: LocateTableTarget): Promise { const config = getConfig(target.connectionId); if (!config) return false; @@ -3987,7 +3994,7 @@ export const useConnectionStore = defineStore("connection", () => { } else if (node.type === "group-partitions") { node.isExpanded = true; } else if (node.type === "group-extensions" && node.connectionId && hasTreeNodeDatabaseContext(node)) { - await loadExtensions(node.connectionId, node.database || "", node.schema || ""); + await loadExtensions(node.connectionId, node.database || ""); } } diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 187fb5061a..51ba090243 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -1300,8 +1300,13 @@ pub async fn export_database_sql_core( PostgresExtensionMembers::default() }; let postgres_extensions = if request.include_structure && matches!(db_type, DatabaseType::Postgres) { - match crate::schema::list_extensions_core(state, &request.connection_id, &request.database, &request.schema) - .await + match crate::schema::list_extensions_core( + state, + &request.connection_id, + &request.database, + Some(&request.schema), + ) + .await { Ok(extensions) => extensions .into_iter() diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 4594c6a88b..3d5c50d4a5 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -3329,19 +3329,32 @@ pub async fn list_rules(pool: &Pool, schema: &str) -> Result, Stri .collect()) } -pub async fn list_extensions(pool: &Pool, schema: &str) -> Result, String> { +pub async fn list_extensions(pool: &Pool, schema: Option<&str>) -> Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; - let rows = postgres_query_cached( - &client, - "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ - FROM pg_catalog.pg_extension e \ - JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \ - LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \ - WHERE n.nspname = $1 \ - ORDER BY e.extname", - &[&schema], - ) - .await + let rows = if let Some(schema) = schema.filter(|value| !value.is_empty()) { + postgres_query_cached( + &client, + "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ + FROM pg_catalog.pg_extension e \ + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \ + LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \ + WHERE n.nspname = $1 \ + ORDER BY e.extname", + &[&schema], + ) + .await + } else { + postgres_query_cached( + &client, + "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ + FROM pg_catalog.pg_extension e \ + JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \ + LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \ + ORDER BY n.nspname, e.extname", + &[], + ) + .await + } .map_err(|e| e.to_string())?; Ok(rows @@ -3350,7 +3363,7 @@ pub async fn list_extensions(pool: &Pool, schema: &str) -> Result>(2).ok().flatten().filter(|s| !s.is_empty()), - schema: Some(schema.to_string()), + schema: row.try_get::<_, Option>(3).ok().flatten().filter(|s| !s.is_empty()), }) .collect()) } diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 4503d3ec22..50444d63a3 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -3217,6 +3217,57 @@ mod tests { assert!(!is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Mysql))); } + #[test] + fn kingbase_extension_sql_uses_sys_catalog_and_escapes_schema() { + let sql = super::kingbase_list_extensions_sql(Some("app's"), super::KingbaseExtensionCatalog::Sys); + + assert!(sql.contains("FROM sys_catalog.sys_extension e")); + assert!(sql.contains("JOIN sys_catalog.sys_namespace n")); + assert!(sql.contains("LEFT JOIN sys_catalog.sys_description d")); + assert!(sql.contains("WHERE n.nspname = 'app''s'")); + } + + #[test] + fn kingbase_available_extension_sql_supports_pg_catalog_fallback() { + let sql = super::kingbase_list_available_extensions_sql(super::KingbaseExtensionCatalog::Pg); + + assert!(sql.contains("FROM pg_catalog.pg_available_extensions")); + assert!(sql.contains("WHERE installed_version IS NULL")); + } + + #[test] + fn extension_infos_from_query_result_maps_installed_extensions() { + let result = db::QueryResult { + columns: vec![ + "extname".to_string(), + "extversion".to_string(), + "description".to_string(), + "nspname".to_string(), + ], + column_types: Vec::new(), + column_sortables: Vec::new(), + rows: vec![vec![ + serde_json::json!("kdb_utils"), + serde_json::json!("1.0"), + serde_json::json!("KingBase utilities"), + serde_json::json!("public"), + ]], + affected_rows: 0, + execution_time_ms: 0, + truncated: false, + session_id: None, + has_more: false, + }; + + let extensions = super::extension_infos_from_query_result(result, true); + + assert_eq!(extensions.len(), 1); + assert_eq!(extensions[0].name, "kdb_utils"); + assert_eq!(extensions[0].version, "1.0"); + assert_eq!(extensions[0].comment.as_deref(), Some("KingBase utilities")); + assert_eq!(extensions[0].schema.as_deref(), Some("public")); + } + #[test] fn agent_metadata_timeout_defaults_to_sixty_seconds_and_honors_longer_config() { assert_eq!(super::agent_metadata_timeout(None), Some(std::time::Duration::from_secs(60))); @@ -4827,10 +4878,25 @@ pub async fn list_extensions_core( state: &AppState, connection_id: &str, database: &str, - schema: &str, + schema: Option<&str>, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let db_config = connection_config(state, connection_id).await; + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { + let connections = state.connections.read().await; + if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + drop(connections); + return kingbase_agent_list_extensions( + client, + database, + schema, + agent_metadata_timeout(db_config.as_ref()), + ) + .await; + } + } + let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; @@ -4849,6 +4915,20 @@ pub async fn list_available_extensions_core( ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || async { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let db_config = connection_config(state, connection_id).await; + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) { + let connections = state.connections.read().await; + if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + drop(connections); + return kingbase_agent_list_available_extensions( + client, + database, + agent_metadata_timeout(db_config.as_ref()), + ) + .await; + } + } + let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; @@ -4860,6 +4940,140 @@ pub async fn list_available_extensions_core( .await } +#[derive(Clone, Copy)] +enum KingbaseExtensionCatalog { + Sys, + Pg, +} + +impl KingbaseExtensionCatalog { + fn catalog_name(self) -> &'static str { + match self { + Self::Sys => "sys_catalog", + Self::Pg => "pg_catalog", + } + } + + fn prefix(self) -> &'static str { + match self { + Self::Sys => "sys", + Self::Pg => "pg", + } + } +} + +fn kingbase_list_extensions_sql(schema: Option<&str>, catalog: KingbaseExtensionCatalog) -> String { + let catalog_name = catalog.catalog_name(); + let prefix = catalog.prefix(); + let schema_filter = schema + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|schema| format!("WHERE n.nspname = {}", sql_string(schema))) + .unwrap_or_default(); + let order_by = if schema_filter.is_empty() { "n.nspname, e.extname" } else { "e.extname" }; + format!( + "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \ + FROM {catalog_name}.{prefix}_extension e \ + JOIN {catalog_name}.{prefix}_namespace n ON n.oid = e.extnamespace \ + LEFT JOIN {catalog_name}.{prefix}_description d ON d.objoid = e.oid AND d.objsubid = 0 \ + {schema_filter} \ + ORDER BY {order_by}" + ) +} + +fn kingbase_list_available_extensions_sql(catalog: KingbaseExtensionCatalog) -> String { + let catalog_name = catalog.catalog_name(); + let prefix = catalog.prefix(); + format!( + "SELECT name, default_version, comment \ + FROM {catalog_name}.{prefix}_available_extensions \ + WHERE installed_version IS NULL \ + ORDER BY name" + ) +} + +async fn kingbase_agent_query_result( + client: Arc>, + database: &str, + sql: &str, + max_rows: usize, + timeout_duration: Option, +) -> Result { + let params = agent_execute_query_params( + sql, + if database.is_empty() { None } else { Some(database) }, + None, + QueryExecutionOptions { max_rows: Some(max_rows), ..Default::default() }, + ); + let mut client = client.lock().await; + client.execute_query_with_timeout(params, timeout_duration).await +} + +async fn kingbase_agent_query_result_with_catalog_fallback( + client: Arc>, + database: &str, + sys_sql: String, + pg_sql: String, + max_rows: usize, + timeout_duration: Option, +) -> Result { + match kingbase_agent_query_result(client.clone(), database, &sys_sql, max_rows, timeout_duration).await { + Ok(result) => Ok(result), + Err(sys_error) => kingbase_agent_query_result(client, database, &pg_sql, max_rows, timeout_duration) + .await + .map_err(|pg_error| format!("{sys_error}; pg_catalog fallback failed: {pg_error}")), + } +} + +async fn kingbase_agent_list_extensions( + client: Arc>, + database: &str, + schema: Option<&str>, + timeout_duration: Option, +) -> Result, String> { + let result = kingbase_agent_query_result_with_catalog_fallback( + client, + database, + kingbase_list_extensions_sql(schema, KingbaseExtensionCatalog::Sys), + kingbase_list_extensions_sql(schema, KingbaseExtensionCatalog::Pg), + 10_000, + timeout_duration, + ) + .await?; + Ok(extension_infos_from_query_result(result, true)) +} + +async fn kingbase_agent_list_available_extensions( + client: Arc>, + database: &str, + timeout_duration: Option, +) -> Result, String> { + let result = kingbase_agent_query_result_with_catalog_fallback( + client, + database, + kingbase_list_available_extensions_sql(KingbaseExtensionCatalog::Sys), + kingbase_list_available_extensions_sql(KingbaseExtensionCatalog::Pg), + 10_000, + timeout_duration, + ) + .await?; + Ok(extension_infos_from_query_result(result, false)) +} + +fn extension_infos_from_query_result(result: db::QueryResult, include_schema: bool) -> Vec { + result + .rows + .into_iter() + .filter_map(|row| { + let name = query_result_cell_string(&row, 0)?; + let version = query_result_cell_string(&row, 1).unwrap_or_default(); + let comment = query_result_cell_string(&row, 2).filter(|value| !value.trim().is_empty()); + let schema = include_schema.then(|| query_result_cell_string(&row, 3)).flatten(); + Some(db::ExtensionInfo { name, version, comment, schema }) + }) + .collect() +} + pub async fn list_owners_core( state: &AppState, connection_id: &str, diff --git a/crates/dbx-web/src/routes/schema.rs b/crates/dbx-web/src/routes/schema.rs index dbb3068572..5bebf8c86b 100644 --- a/crates/dbx-web/src/routes/schema.rs +++ b/crates/dbx-web/src/routes/schema.rs @@ -471,8 +471,7 @@ pub async fn list_extensions( Query(q): Query, ) -> Result, AppError> { let database = q.database.as_deref().unwrap_or(""); - let schema = q.schema.as_deref().unwrap_or(""); - let result = dbx_core::schema::list_extensions_core(&state.app, &q.connection_id, database, schema) + let result = dbx_core::schema::list_extensions_core(&state.app, &q.connection_id, database, q.schema.as_deref()) .await .map_err(AppError::from)?; Ok(Json(serde_json::to_value(result).map_err(|e| AppError::from(e.to_string()))?)) diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index 213dee96f0..19a5c22428 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -437,9 +437,9 @@ pub async fn list_extensions( state: State<'_, Arc>, connection_id: String, database: String, - schema: String, + schema: Option, ) -> Result, String> { - dbx_core::schema::list_extensions_core(&state, &connection_id, &database, &schema).await + dbx_core::schema::list_extensions_core(&state, &connection_id, &database, schema.as_deref()).await } #[tauri::command]