From a3f3cadb399b58bedd6f9197261c7af5d80963d4 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 13 Jul 2026 00:26:04 +0800 Subject: [PATCH] fix(postgres): distinguish overloaded routines --- .../components/objects/ObjectSourceDialog.vue | 5 +- .../src/components/sidebar/TreeItem.vue | 8 ++- .../src/lib/__tests__/table/tableTree.spec.ts | 25 ++++++++- apps/desktop/src/lib/backend/http.ts | 4 +- apps/desktop/src/lib/backend/tauri.ts | 4 +- apps/desktop/src/lib/database/dbAdminSql.ts | 1 + apps/desktop/src/lib/table/tableTree.ts | 9 ++- apps/desktop/src/types/database.ts | 2 + crates/dbx-core/src/database_export.rs | 3 + crates/dbx-core/src/db/postgres.rs | 24 ++++---- crates/dbx-core/src/db_admin_sql.rs | 25 ++++++++- crates/dbx-core/src/schema.rs | 56 ++++++++++++++----- crates/dbx-web/src/routes/schema.rs | 16 ++++-- src-tauri/src/commands/schema.rs | 12 +++- 14 files changed, 146 insertions(+), 48 deletions(-) diff --git a/apps/desktop/src/components/objects/ObjectSourceDialog.vue b/apps/desktop/src/components/objects/ObjectSourceDialog.vue index 74b74855a0..8f726937f6 100644 --- a/apps/desktop/src/components/objects/ObjectSourceDialog.vue +++ b/apps/desktop/src/components/objects/ObjectSourceDialog.vue @@ -22,6 +22,7 @@ const props = withDefaults( database: string; schema?: string; name: string; + signature?: string; objectType: ObjectSourceKind; databaseType?: DatabaseType; dialect: "mysql" | "postgres" | "sqlserver"; @@ -58,7 +59,7 @@ const canEdit = computed(() => sourceEditable.value && props.objectType !== "SEQ const title = computed(() => `${editing.value ? t("contextMenu.editView") : t("contextMenu.viewSource")} - ${props.name}`); watch( - () => [props.open, props.connectionId, props.database, props.schema, props.name, props.objectType, props.initialEditing] as const, + () => [props.open, props.connectionId, props.database, props.schema, props.name, props.signature, props.objectType, props.initialEditing] as const, () => { if (props.open) void loadSource(); }, @@ -78,7 +79,7 @@ async function loadSource(nextEditing = props.initialEditing && canEdit.value) { try { if (!props.databaseType) throw new Error("Connection type is unavailable."); const schema = props.schema || props.database; - const result = await api.getObjectSource(props.connectionId, props.database, schema, props.name, props.objectType); + const result = await api.getObjectSource(props.connectionId, props.database, schema, props.name, props.objectType, props.signature); const editableAllowed = result.editable !== false; const editable = await buildEditableObjectSource({ databaseType: props.databaseType, diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index 654c4072f0..3558077442 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -1881,7 +1881,8 @@ function dropObjectSqlOptionsForNode(node: TreeNode): DropObjectSqlOptions | nul databaseType: tableStructureDatabaseTypeForNode(node), objectType: node.type === "view" ? "VIEW" : node.type === "materialized_view" ? "MATERIALIZED_VIEW" : node.type === "procedure" ? "PROCEDURE" : "FUNCTION", schema: node.schema, - name: node.label, + name: node.objectName || node.label, + signature: node.signature, }; } @@ -2402,7 +2403,7 @@ async function confirmRenameObject() { await connectionStore.ensureConnected(node.connectionId); if (supportsSourceBackedRoutineRename(dbType, objectType as any)) { const schema = node.schema || node.database; - const source = await api.getObjectSource(node.connectionId, node.database, schema, node.label, objectType as any); + const source = await api.getObjectSource(node.connectionId, node.database, schema, node.objectName || node.label, objectType as any, node.signature); const statements = await buildRoutineRenameObjectSourceStatements({ databaseType: dbType!, objectType: objectType as any, @@ -5938,7 +5939,8 @@ function treeItemMenuItems(): ContextMenuItem[] { :connection-id="objectSourceTarget.node.connectionId!" :database="objectSourceTarget.node.database!" :schema="objectSourceTarget.node.schema" - :name="objectSourceTarget.node.label" + :name="objectSourceTarget.node.objectName || objectSourceTarget.node.label" + :signature="objectSourceTarget.node.signature" :object-type="objectSourceType" :database-type="objectSourceDatabaseType" :dialect="objectSourceDialect" diff --git a/apps/desktop/src/lib/__tests__/table/tableTree.spec.ts b/apps/desktop/src/lib/__tests__/table/tableTree.spec.ts index 15df48fa7d..8fb947c7cf 100644 --- a/apps/desktop/src/lib/__tests__/table/tableTree.spec.ts +++ b/apps/desktop/src/lib/__tests__/table/tableTree.spec.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; -import { appendTableTreeLoadMoreNode, buildTableTreeNodes, mergeTableTreePageChildren, tablePartitionGroups, withoutTableTreeLoadMoreNodes } from "@/lib/table/tableTree"; -import type { TableInfo, TreeNode } from "@/types/database"; +import { appendTableTreeLoadMoreNode, buildSimpleObjectTreeNodes, buildTableTreeNodes, mergeTableTreePageChildren, tablePartitionGroups, withoutTableTreeLoadMoreNodes } from "@/lib/table/tableTree"; +import type { ObjectInfo, TableInfo, TreeNode } from "@/types/database"; const context = { nodeId: "connection:db", @@ -8,6 +8,27 @@ const context = { database: "db", }; +describe("PostgreSQL overloaded routines", () => { + it("keeps routines with the same name distinct by identity arguments", () => { + const objects: ObjectInfo[] = [ + { name: "calc", object_type: "FUNCTION", schema: "public", signature: "integer" }, + { name: "calc", object_type: "FUNCTION", schema: "public", signature: "integer, integer" }, + { name: "calc", object_type: "FUNCTION", schema: "public", signature: "numeric" }, + ]; + + const nodes = buildSimpleObjectTreeNodes({ ...context, schema: "public", objects }); + + expect(nodes.map((node) => ({ label: node.label, objectName: node.objectName, signature: node.signature }))).toEqual( + expect.arrayContaining([ + { label: "calc(integer)", objectName: "calc", signature: "integer" }, + { label: "calc(integer, integer)", objectName: "calc", signature: "integer, integer" }, + { label: "calc(numeric)", objectName: "calc", signature: "numeric" }, + ]), + ); + expect(new Set(nodes.map((node) => node.id)).size).toBe(3); + }); +}); + describe("TDengine table hierarchy", () => { it("groups child tables under their supertable and keeps ordinary tables flat", () => { const tables: TableInfo[] = [ diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 2253f9702b..cab4f0cfda 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -592,8 +592,8 @@ export async function completionAssistantSearch(request: CompletionAssistantRequ return post("/api/schema/completion-assistant", request); } -export async function getObjectSource(connectionId: string, database: string, schema: string, name: string, objectType: ObjectSourceKind): Promise { - return get(`/api/schema/object-source?${qs({ connection_id: connectionId, database, schema, table: name, object_type: objectType })}`); +export async function getObjectSource(connectionId: string, database: string, schema: string, name: string, objectType: ObjectSourceKind, signature?: string): Promise { + return get(`/api/schema/object-source?${qs({ connection_id: connectionId, database, schema, table: name, object_type: objectType, signature })}`); } export async function getColumns(connectionId: string, database: string, schema: string, table: string, catalog?: string): Promise { diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index ffefecb74d..f8356cf586 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -718,8 +718,8 @@ export async function completionAssistantSearch(request: CompletionAssistantRequ return invoke("completion_assistant_search", { request }); } -export async function getObjectSource(connectionId: string, database: string, schema: string, name: string, objectType: ObjectSourceKind): Promise { - return invoke("get_object_source", { connectionId, database, schema, name, objectType }); +export async function getObjectSource(connectionId: string, database: string, schema: string, name: string, objectType: ObjectSourceKind, signature?: string): Promise { + return invoke("get_object_source", { connectionId, database, schema, name, objectType, signature }); } export async function listSchemas(connectionId: string, database: string, applyVisibleFilter = false): Promise { diff --git a/apps/desktop/src/lib/database/dbAdminSql.ts b/apps/desktop/src/lib/database/dbAdminSql.ts index 791660df2d..f1e1043461 100644 --- a/apps/desktop/src/lib/database/dbAdminSql.ts +++ b/apps/desktop/src/lib/database/dbAdminSql.ts @@ -6,6 +6,7 @@ export interface DropObjectSqlOptions { objectType: DatabaseObjectType; schema?: string | null; name: string; + signature?: string | null; } export interface TableAdminSqlOptions { diff --git a/apps/desktop/src/lib/table/tableTree.ts b/apps/desktop/src/lib/table/tableTree.ts index 540f0a408c..38abbe0dc0 100644 --- a/apps/desktop/src/lib/table/tableTree.ts +++ b/apps/desktop/src/lib/table/tableTree.ts @@ -503,7 +503,8 @@ export function buildSimpleObjectTreeNodes({ nodeId, connectionId, database, sch if (!name) continue; const childSchema = obj.schema ? normalizeDatabaseObjectName(obj.schema) : schema; - const dedupeKey = `${objectType}\0${(childSchema || "").toLowerCase()}\0${name.toLowerCase()}`; + const signature = obj.signature?.trim() || ""; + const dedupeKey = `${objectType}\0${(childSchema || "").toLowerCase()}\0${name.toLowerCase()}\0${signature.toLowerCase()}`; if (seen.has(dedupeKey)) continue; seen.add(dedupeKey); @@ -524,9 +525,11 @@ export function buildSimpleObjectTreeNodes({ nodeId, connectionId, database, sch } else { const simpleNodeType = simpleObjectNodeType(objectType); objectNodes.push({ - id: objectType === "VIEW" || objectType === "MATERIALIZED_VIEW" ? entry.node.id : `${nodeId}:${childSchema ? `${childSchema}:` : ""}${name}:${objectType}`, - label: name, + id: objectType === "VIEW" || objectType === "MATERIALIZED_VIEW" ? entry.node.id : `${nodeId}:${childSchema ? `${childSchema}:` : ""}${name}:${signature}:${objectType}`, + label: signature && (objectType === "FUNCTION" || objectType === "PROCEDURE") ? `${name}(${signature})` : name, type: simpleNodeType, + objectName: name, + signature: signature || undefined, comment: obj.comment, connectionId, database, diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 8453378dc8..324aee318c 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -659,6 +659,8 @@ export interface TreeNode { nacosNamespaceName?: string; schema?: string; tableName?: string; + objectName?: string; + signature?: string; tableType?: string; comment?: string | null; objectCount?: number; diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 061770e14e..3ea9851c8a 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -1359,6 +1359,7 @@ pub async fn export_database_sql_core( &request.schema, view_name, crate::db::ObjectSourceKind::View, + None, ) .await { @@ -1402,6 +1403,7 @@ pub async fn export_database_sql_core( &request.schema, proc_name, crate::db::ObjectSourceKind::Procedure, + None, ) .await { @@ -1448,6 +1450,7 @@ pub async fn export_database_sql_core( &request.schema, func_name, crate::db::ObjectSourceKind::Function, + None, ) .await { diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index c7c9f438a1..6c019eb360 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1808,7 +1808,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prokind = 'p' OR p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -1822,7 +1822,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prokind = 'p' OR p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -1839,7 +1839,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ CASE p.prokind WHEN 'p' THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -1853,7 +1853,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ CASE p.prokind WHEN 'p' THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -1870,7 +1870,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -1884,7 +1884,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ CASE WHEN p.prosp THEN 2 ELSE 3 END AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -1900,7 +1900,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha THEN pg_xact_commit_timestamp(p.xmin)::text END AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ 3 AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -1914,7 +1914,7 @@ fn list_object_routines_sql(include_timestamps: bool, has_proc_prokind: bool, ha NULL::text AS updated_at, \ NULL::text AS parent_schema, \ NULL::text AS parent_name, \ - pg_get_function_arguments(p.oid) AS signature, \ + pg_get_function_identity_arguments(p.oid) AS signature, \ 3 AS sort_order \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ @@ -3899,7 +3899,7 @@ mod tests { assert!(sql.contains("parent_schema")); assert!(sql.contains("parent_name")); assert!(sql.contains("NULL::text AS signature")); - assert!(sql.contains("pg_get_function_arguments(p.oid) AS signature")); + assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); assert!(sql.contains("pc.relkind = 'p'")); assert!(sql.contains("pg_stat_file")); assert!(sql.contains("pg_xact_commit_timestamp")); @@ -3946,7 +3946,7 @@ mod tests { assert!(!sql.contains("p.prosp")); assert!(sql.contains("NOT p.proisagg")); assert!(sql.contains("NOT p.proiswindow")); - assert!(sql.contains("pg_get_function_arguments(p.oid) AS signature")); + assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); assert!(sql.contains("'FUNCTION' AS object_type")); assert!(!sql.contains("'PROCEDURE'")); } @@ -3959,7 +3959,7 @@ mod tests { assert!(sql.contains("CASE WHEN p.prosp THEN 2 ELSE 3 END AS sort_order")); assert!(sql.contains("NOT p.proisagg")); assert!(sql.contains("NOT p.proiswindow")); - assert!(sql.contains("pg_get_function_arguments(p.oid) AS signature")); + assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); } #[test] @@ -3970,7 +3970,7 @@ mod tests { ); assert!(sql.contains("CASE WHEN p.prokind = 'p' OR p.prosp THEN 2 ELSE 3 END AS sort_order")); assert!(sql.contains("p.prokind IN ('p','f') OR p.prosp")); - assert!(sql.contains("pg_get_function_arguments(p.oid) AS signature")); + assert!(sql.contains("pg_get_function_identity_arguments(p.oid) AS signature")); } #[test] diff --git a/crates/dbx-core/src/db_admin_sql.rs b/crates/dbx-core/src/db_admin_sql.rs index 2d187b8126..f0d2249fd7 100644 --- a/crates/dbx-core/src/db_admin_sql.rs +++ b/crates/dbx-core/src/db_admin_sql.rs @@ -78,6 +78,8 @@ pub struct DropObjectSqlOptions { #[serde(default, skip_serializing_if = "Option::is_none")] pub schema: Option, pub name: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub signature: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -311,10 +313,18 @@ pub fn build_create_user_sql(username: &str, password: &str, tablespace: &str) - } pub fn build_drop_object_sql(options: DropObjectSqlOptions) -> String { + let signature = if matches!(options.database_type, Some(DatabaseType::Postgres)) + && matches!(options.object_type, DatabaseObjectType::Function | DatabaseObjectType::Procedure) + { + options.signature.as_deref().map(|value| format!("({value})")).unwrap_or_default() + } else { + String::new() + }; format!( - "DROP {} {};", + "DROP {} {}{};", object_type_keyword(options.object_type), - qualified_name(options.database_type, options.schema.as_deref(), &options.name) + qualified_name(options.database_type, options.schema.as_deref(), &options.name), + signature ) } @@ -1206,9 +1216,20 @@ mod tests { object_type: DatabaseObjectType::Procedure, schema: Some("dbo".to_string()), name: "refresh_cache".to_string(), + signature: None, }), "DROP PROCEDURE [dbo].[refresh_cache];" ); + assert_eq!( + build_drop_object_sql(DropObjectSqlOptions { + database_type: Some(DatabaseType::Postgres), + object_type: DatabaseObjectType::Function, + schema: Some("public".to_string()), + name: "calc".to_string(), + signature: Some("integer, integer".to_string()), + }), + "DROP FUNCTION \"public\".\"calc\"(integer, integer);" + ); assert_eq!( build_drop_database_sql(DatabaseNameSqlOptions { database_type: Some(DatabaseType::Mysql), diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 860bad7c44..de2a821030 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -4526,7 +4526,8 @@ pub async fn get_table_ddl_core( } if matches!(object_type, Some(db::ObjectSourceKind::View)) { let source = - get_object_source_core(state, connection_id, database, schema, table, db::ObjectSourceKind::View).await?; + get_object_source_core(state, connection_id, database, schema, table, db::ObjectSourceKind::View, None) + .await?; let database_type = connection_config(state, connection_id).await.map(|config| config.db_type); return Ok(crate::object_source_sql::build_view_ddl_sql(crate::object_source_sql::BuildViewDdlInput { database_type, @@ -4543,6 +4544,7 @@ pub async fn get_table_ddl_core( schema, table, db::ObjectSourceKind::MaterializedView, + None, ) .await?; return Ok(source.source); @@ -4801,12 +4803,22 @@ pub fn sqlserver_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSo ) } -pub fn postgres_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { - postgres_object_source_sql_inner(schema, name, kind, true) +pub fn postgres_object_source_sql( + schema: &str, + name: &str, + kind: &db::ObjectSourceKind, + signature: Option<&str>, +) -> String { + postgres_object_source_sql_inner(schema, name, kind, signature, true) } -fn postgres_object_source_sql_without_relispopulated(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { - postgres_object_source_sql_inner(schema, name, kind, false) +fn postgres_object_source_sql_without_relispopulated( + schema: &str, + name: &str, + kind: &db::ObjectSourceKind, + signature: Option<&str>, +) -> String { + postgres_object_source_sql_inner(schema, name, kind, signature, false) } fn postgres_function_object_source_sql_without_prokind(schema: &str, name: &str) -> String { @@ -4825,6 +4837,7 @@ fn postgres_object_source_sql_inner( schema: &str, name: &str, kind: &db::ObjectSourceKind, + signature: Option<&str>, include_relispopulated: bool, ) -> String { match kind { @@ -4856,15 +4869,19 @@ fn postgres_object_source_sql_inner( } db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function => { let prokind = if matches!(kind, db::ObjectSourceKind::Procedure) { "p" } else { "f" }; + let signature_filter = signature + .map(|value| format!(" AND pg_get_function_identity_arguments(p.oid) = {}", sql_string(value))) + .unwrap_or_default(); format!( "SELECT pg_get_functiondef(p.oid) \ FROM pg_proc p \ JOIN pg_namespace n ON n.oid = p.pronamespace \ - WHERE n.nspname = {} AND p.proname = {} AND p.prokind = '{}' \ + WHERE n.nspname = {} AND p.proname = {} AND p.prokind = '{}'{} \ ORDER BY p.oid LIMIT 1", sql_string(schema), sql_string(name), - prokind + prokind, + signature_filter ) } db::ObjectSourceKind::Sequence => { @@ -4988,9 +5005,10 @@ pub async fn get_object_source_core( schema: &str, name: &str, object_type: db::ObjectSourceKind, + signature: Option<&str>, ) -> Result { retry_metadata_connection(state, connection_id, Some(database), || { - get_object_source_once(state, connection_id, database, schema, name, object_type.clone()) + get_object_source_once(state, connection_id, database, schema, name, object_type.clone(), signature) }) .await } @@ -5002,6 +5020,7 @@ async fn get_object_source_once( schema: &str, name: &str, object_type: db::ObjectSourceKind, + signature: Option<&str>, ) -> Result { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; let db_config = connection_config(state, connection_id).await; @@ -5069,7 +5088,7 @@ async fn get_object_source_once( // only view db::questdb::questdb_object_source(pool, name).await? } - PoolKind::Postgres(pool) => postgres_object_source(pool, schema, name, &object_type).await?, + PoolKind::Postgres(pool) => postgres_object_source(pool, schema, name, &object_type, signature).await?, PoolKind::Sqlite(pool) => first_string_cell( db::sqlite::execute_query(pool, &sqlite_object_source_sql(name, &object_type)).await?, )?, @@ -5399,15 +5418,16 @@ async fn postgres_object_source( schema: &str, name: &str, object_type: &db::ObjectSourceKind, + signature: Option<&str>, ) -> Result { - let sql = postgres_object_source_sql(schema, name, object_type); + let sql = postgres_object_source_sql(schema, name, object_type, signature); match db::postgres::execute_query(pool, &sql).await.and_then(first_string_cell) { Ok(source) => Ok(source), Err(primary_err) if postgres_missing_relispopulated_error(&primary_err) && matches!(object_type, db::ObjectSourceKind::View | db::ObjectSourceKind::MaterializedView) => { - let fallback_sql = postgres_object_source_sql_without_relispopulated(schema, name, object_type); + let fallback_sql = postgres_object_source_sql_without_relispopulated(schema, name, object_type, signature); db::postgres::execute_query(pool, &fallback_sql) .await .and_then(first_string_cell) @@ -5465,7 +5485,7 @@ mod object_source_tests { #[test] fn builds_postgres_object_source_sql_for_views_and_functions() { - let view_sql = postgres_object_source_sql("public", "active_users", &ObjectSourceKind::View); + let view_sql = postgres_object_source_sql("public", "active_users", &ObjectSourceKind::View, None); assert!(view_sql.contains("CREATE MATERIALIZED VIEW")); assert!(view_sql.contains("CREATE OR REPLACE VIEW")); @@ -5474,9 +5494,14 @@ mod object_source_tests { assert!(view_sql.contains("c.relname = 'active_users'")); assert_eq!( - postgres_object_source_sql("public", "recalc_score", &ObjectSourceKind::Function), + postgres_object_source_sql("public", "recalc_score", &ObjectSourceKind::Function, None), "SELECT pg_get_functiondef(p.oid) FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace WHERE n.nspname = 'public' AND p.proname = 'recalc_score' AND p.prokind = 'f' ORDER BY p.oid LIMIT 1" ); + + assert_eq!( + postgres_object_source_sql("public", "recalc_score", &ObjectSourceKind::Function, Some("integer, integer")), + "SELECT pg_get_functiondef(p.oid) FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace WHERE n.nspname = 'public' AND p.proname = 'recalc_score' AND p.prokind = 'f' AND pg_get_function_identity_arguments(p.oid) = 'integer, integer' ORDER BY p.oid LIMIT 1" + ); } #[test] @@ -5485,6 +5510,7 @@ mod object_source_tests { "public", "active_users", &ObjectSourceKind::MaterializedView, + None, ); assert!(sql.contains("CREATE MATERIALIZED VIEW")); @@ -5504,7 +5530,7 @@ mod object_source_tests { #[test] fn keeps_legacy_materialized_viewdef_when_it_already_contains_create_statement() { - let sql = postgres_object_source_sql("public", "active_users", &ObjectSourceKind::MaterializedView); + let sql = postgres_object_source_sql("public", "active_users", &ObjectSourceKind::MaterializedView, None); assert!( sql.contains( @@ -5531,7 +5557,7 @@ mod object_source_tests { #[test] fn builds_postgres_view_source_sql_without_regclass_cast() { - let sql = postgres_object_source_sql("tenant's schema", "active users", &ObjectSourceKind::View); + let sql = postgres_object_source_sql("tenant's schema", "active users", &ObjectSourceKind::View, None); assert!(!sql.contains("::regclass")); assert!(sql.contains("pg_get_viewdef(c.oid, 0)")); diff --git a/crates/dbx-web/src/routes/schema.rs b/crates/dbx-web/src/routes/schema.rs index 284d5cf682..39c30a34d9 100644 --- a/crates/dbx-web/src/routes/schema.rs +++ b/crates/dbx-web/src/routes/schema.rs @@ -19,6 +19,7 @@ pub struct SchemaQuery { pub limit: Option, pub offset: Option, pub object_type: Option, + pub signature: Option, pub object_types: Option, pub apply_visible_filter: Option, } @@ -258,10 +259,17 @@ pub async fn get_object_source( let schema = q.schema.as_deref().unwrap_or(""); let name = q.table.as_deref().unwrap_or(""); let object_type = q.object_type.ok_or_else(|| AppError("Missing object_type".to_string()))?; - let result = - dbx_core::schema::get_object_source_core(&state.app, &q.connection_id, database, schema, name, object_type) - .await - .map_err(AppError)?; + let result = dbx_core::schema::get_object_source_core( + &state.app, + &q.connection_id, + database, + schema, + name, + object_type, + q.signature.as_deref(), + ) + .await + .map_err(AppError)?; Ok(Json(result)) } diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index b9daa5c174..6b0fbc37d3 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -269,8 +269,18 @@ pub async fn get_object_source( schema: String, name: String, object_type: db::ObjectSourceKind, + signature: Option, ) -> Result { - dbx_core::schema::get_object_source_core(&state, &connection_id, &database, &schema, &name, object_type).await + dbx_core::schema::get_object_source_core( + &state, + &connection_id, + &database, + &schema, + &name, + object_type, + signature.as_deref(), + ) + .await } #[tauri::command]