From 95e3cb6571f4a3516a3831b04cf8a851a5541d5e Mon Sep 17 00:00:00 2001 From: luoianun Date: Wed, 1 Jul 2026 19:05:45 +0800 Subject: [PATCH] fix(presto): optimize metadata completion --- .../src/components/editor/QueryEditor.vue | 20 +- .../__tests__/prestoSqlBuiltinDriver.spec.ts | 2 +- .../__tests__/sqlCompletion.context.spec.ts | 11 + .../desktop/src/lib/prestoSqlBuiltinDriver.ts | 4 +- .../connectionStore.completion.spec.ts | 29 +++ apps/desktop/src/stores/connectionStore.ts | 4 +- crates/dbx-core/src/schema.rs | 233 ++++++++++++++++-- .../main/java/app/dbx/jdbc/DbxJdbcPlugin.java | 134 +++++++++- .../java/app/dbx/jdbc/DbxJdbcPluginTest.java | 128 +++++++++- 9 files changed, 528 insertions(+), 37 deletions(-) diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index 54cf98d82..79bc93b3a 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -156,6 +156,8 @@ const completionTranslations = computed(() => ({ functionDescriptions: Object.fromEntries(SQL_FUNCTION_NAMES.map((name) => [name, t(`editor.completion.functionDescriptions.${name}`)])) as Record, })); const MAX_COMPLETION_TABLES = 200; +const PRESTO_ON_DEMAND_TABLE_COMPLETION_MIN_PREFIX = 2; +const PRESTO_ON_DEMAND_TABLE_COMPLETION_LIMIT = 20; const MAX_JOIN_FK_PREFETCH_TABLES = 24; const MAX_SEMANTIC_DIAGNOSTIC_COLUMN_TABLES = 4; const liveFontSize = ref(settingsStore.editorSettings.fontSize); @@ -717,6 +719,12 @@ function usesOnDemandOnlyCompletionColumns(): boolean { return usesOnDemandOnlyEditorColumnMetadata(props.databaseType); } +function allowsOnDemandQualifiedTableCompletion(prefix: string): boolean { + if (!usesLocalOnlyCompletionMetadata()) return false; + if (props.databaseType !== "prestosql" && props.databaseType !== "trino") return false; + return prefix.trim().length >= PRESTO_ON_DEMAND_TABLE_COMPLETION_MIN_PREFIX; +} + function completionMetadataTarget(table: { name: string; schema?: string | null }): { database: string; schema?: string } | null { if (props.database == null) return null; if (supportsDatabaseQualifierCompletion() && table.schema) { @@ -1814,9 +1822,12 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext // If qualifier didn't match any table names, try it as a schema name let qualifierIsSchema = false; if (completionContext.qualifier && !qualifierDatabase && !isReferencedTableQualifier(completionContext) && tables.length === 0 && (completionContext.suggestTables || completionContext.exclusiveColumnSuggestions)) { - const schemaTables = localOnlyMetadata - ? connectionStore.lookupLocalCompletionTables(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, completionContext.qualifier) - : await listCompletionTablesWithLatencyBudget(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, completionContext.qualifier); + let schemaTables = connectionStore.lookupLocalCompletionTables(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, completionContext.qualifier); + if (!localOnlyMetadata) { + schemaTables = await listCompletionTablesWithLatencyBudget(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, completionContext.qualifier); + } else if (schemaTables.length === 0 && allowsOnDemandQualifiedTableCompletion(completionContext.prefix)) { + schemaTables = await listCompletionTablesWithLatencyBudget(props.connectionId!, props.database!, completionContext.prefix, PRESTO_ON_DEMAND_TABLE_COMPLETION_LIMIT, completionContext.qualifier); + } if (schemaTables.length > 0) { tables = schemaTables; qualifierIsSchema = true; @@ -1868,7 +1879,8 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext } } - const shouldFetchColumnsForCompletion = !onDemandOnlyColumns || completionContext.suggestColumns || completionContext.exclusiveColumnSuggestions || !!completionContext.insertTable; + const isTableNameCompletionContext = completionContext.suggestTables || completionContext.exclusiveTableSuggestions; + const shouldFetchColumnsForCompletion = !onDemandOnlyColumns || ((completionContext.suggestColumns || completionContext.exclusiveColumnSuggestions) && !isTableNameCompletionContext) || !!completionContext.insertTable; if (shouldFetchColumnsForCompletion) { await Promise.all( refs.map(async (refTable) => { diff --git a/apps/desktop/src/lib/__tests__/prestoSqlBuiltinDriver.spec.ts b/apps/desktop/src/lib/__tests__/prestoSqlBuiltinDriver.spec.ts index 2561499f7..14f8b3234 100644 --- a/apps/desktop/src/lib/__tests__/prestoSqlBuiltinDriver.spec.ts +++ b/apps/desktop/src/lib/__tests__/prestoSqlBuiltinDriver.spec.ts @@ -35,7 +35,7 @@ describe("prestoSqlBuiltinDriver", () => { expect(row.version).toBe("350"); expect(row.installed).toBe(false); expect(row.installed_version).toBeNull(); - expect(row.jre).toBe(""); + expect(row.jre).toBe("21"); }); it("marks PrestoSQL installed when its Maven coordinate is present", () => { diff --git a/apps/desktop/src/lib/__tests__/sqlCompletion.context.spec.ts b/apps/desktop/src/lib/__tests__/sqlCompletion.context.spec.ts index 6c7a57a7a..01f61f56f 100644 --- a/apps/desktop/src/lib/__tests__/sqlCompletion.context.spec.ts +++ b/apps/desktop/src/lib/__tests__/sqlCompletion.context.spec.ts @@ -180,6 +180,17 @@ describe("sqlCompletion scoped context classification", () => { expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ schema: "dbo", name: "Users", alias: "u" }), expect.objectContaining({ name: "Orders", alias: "o" })])); }); + it("treats schema-qualified table prefixes in FROM as table completion input", () => { + const sql = "SELECT * FROM dws_game_sdk_base.di"; + const context = getSqlCompletionContext(sql, sql.length); + + expect(context.qualifier).toBe("dws_game_sdk_base"); + expect(context.prefix).toBe("di"); + expect(context.suggestTables).toBe(true); + expect(context.exclusiveTableSuggestions).toBe(true); + expect(context.suggestColumns).toBe(true); + }); + it("exposes CTEs as table-like referenced tables", () => { const sql = "WITH recent_orders(id, total) AS (SELECT id, total FROM orders) SELECT * FROM recent_orders ro WHERE ro."; const context = getSqlCompletionContext(sql, sql.length); diff --git a/apps/desktop/src/lib/prestoSqlBuiltinDriver.ts b/apps/desktop/src/lib/prestoSqlBuiltinDriver.ts index 094a881f2..01c811f30 100644 --- a/apps/desktop/src/lib/prestoSqlBuiltinDriver.ts +++ b/apps/desktop/src/lib/prestoSqlBuiltinDriver.ts @@ -24,8 +24,8 @@ export function prestoSqlBuiltinDriverRow(bundles: JdbcMavenBundleInfo[]): Agent installed: Boolean(installedBundle), installed_version: installedBundle ? PRESTOSQL_JDBC_DRIVER_VERSION : null, update_available: false, - requires_java_runtime: false, - jre: "", + requires_java_runtime: true, + jre: "21", jre_installed: true, }; } diff --git a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts index 5200f7d1f..b43ce25f5 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts @@ -83,4 +83,33 @@ describe("connectionStore completion assistant", () => { expect(listTables).toHaveBeenCalledWith("pg-1", "app", "public", "acc", 20); expect(tables).toEqual([{ name: "accounts", schema: "public", type: "table" }]); }); + + it("keeps schema-qualified local table completion scoped to the selected schema", async () => { + const completionAssistantSearch = vi.fn().mockRejectedValue(new Error("assistant unavailable")); + const listTables = vi.fn(async (_connectionId: string, _database: string, schema: string, filter: string) => { + if (schema === "dim_game_base" && filter === "dim") { + return [{ name: "dim_game", table_type: "BASE TABLE", comment: null }]; + } + return []; + }); + + vi.doMock("@/lib/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + completionAssistantSearch, + listSchemas: vi.fn().mockResolvedValue(["dim_game_base", "dws_game_sdk_base"]), + listTables, + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection()]; + store.connectedIds.add("pg-1"); + + const dimTables = await store.listCompletionTables("pg-1", "app", "dim", 20, "dim_game_base"); + const dwsTables = store.lookupLocalCompletionTables("pg-1", "app", "d", 20, "dws_game_sdk_base"); + + expect(dimTables).toEqual([{ name: "dim_game", schema: "dim_game_base", type: "table" }]); + expect(dwsTables).toEqual([]); + }); }); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 4cb06bb22..a57f3981b 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -2966,7 +2966,7 @@ export const useConnectionStore = defineStore("connection", () => { function lookupLocalCompletionTables(connectionId: string, database: string, filter = "", limit?: number, schema?: string): SqlCompletionTable[] { const allScopes = [...completionTableIndex.entries()].filter(([key]) => key.startsWith(`${connectionId}:${database}:`)).map(([, entry]) => entry); const preferred = schema ? completionTableIndex.get(completionScopeKey(connectionId, database, schema)) : undefined; - const scopes = preferred ? [preferred, ...allScopes.filter((entry) => entry !== preferred)] : allScopes; + const scopes = schema ? (preferred ? [preferred] : []) : allScopes; const treeTables = completionTablesFromTree(treeNodes.value, connectionId, database, schema); const ranked = scopes .flatMap((entry) => entry?.tables ?? []) @@ -2980,7 +2980,7 @@ export const useConnectionStore = defineStore("connection", () => { function lookupLocalCompletionObjects(connectionId: string, database: string, filter = "", limit?: number, schema?: string): SqlCompletionObject[] { const allScopes = [...completionObjectIndex.entries()].filter(([key]) => key.startsWith(`${connectionId}:${database}:`)).map(([, entry]) => entry); const preferred = schema ? completionObjectIndex.get(completionScopeKey(connectionId, database, schema)) : undefined; - const scopes = preferred ? [preferred, ...allScopes.filter((entry) => entry !== preferred)] : allScopes; + const scopes = schema ? (preferred ? [preferred] : []) : allScopes; const ranked = scopes .flatMap((entry) => entry?.objects ?? []) .map((object) => ({ object, score: objectMatchScore(object, filter, schema) })) diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index ef4fa6d53..e154668b0 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -1414,9 +1414,17 @@ async fn list_tables_once( let session = session.clone(); drop(connections); if uses_presto_like_information_schema_tables(&config.db_type) { - return external_driver_presto_like_tables(session, config.as_ref(), database, schema) - .await - .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); + return external_driver_presto_like_tables( + session, + config.as_ref(), + database, + schema, + filter, + limit, + offset, + ) + .await + .map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)); } return session .invoke_with_timeout::>( @@ -1677,7 +1685,11 @@ async fn external_driver_presto_like_tables( config: &ConnectionConfig, database: &str, schema: &str, + filter: Option<&str>, + limit: Option, + offset: Option, ) -> Result, String> { + let query_limit = limit.map(|limit| limit.saturating_add(offset.unwrap_or(0)).max(1)).unwrap_or(100000); let result: db::QueryResult = session .invoke_with_timeout( "executeQuery", @@ -1685,8 +1697,8 @@ async fn external_driver_presto_like_tables( "connection": config, "database": database, "schema": schema, - "sql": presto_like_information_schema_tables_sql(database, schema), - "maxRows": 100000, + "sql": presto_like_information_schema_tables_sql(database, schema, filter, Some(query_limit)), + "maxRows": query_limit, "fetchSize": 1000, "timeoutSecs": 60 }), @@ -1702,7 +1714,7 @@ async fn external_driver_presto_like_objects( database: &str, schema: &str, ) -> Result, String> { - let tables = external_driver_presto_like_tables(session, config, database, schema).await?; + let tables = external_driver_presto_like_tables(session, config, database, schema, None, None, None).await?; Ok(tables .into_iter() .map(|table| db::ObjectInfo { @@ -1719,18 +1731,72 @@ async fn external_driver_presto_like_objects( .collect()) } -fn presto_like_information_schema_tables_sql(database: &str, schema: &str) -> String { +async fn external_driver_presto_like_columns( + session: Arc, + config: &ConnectionConfig, + database: &str, + schema: &str, + table: &str, +) -> Result, String> { + let result: db::QueryResult = session + .invoke( + "executeQuery", + serde_json::json!({ + "connection": config, + "database": database, + "schema": schema, + "sql": presto_like_information_schema_columns_sql(database, schema, table), + "maxRows": 10000, + "fetchSize": 1000, + "timeoutSecs": 60 + }), + ) + .await?; + Ok(presto_like_columns_from_query_result(&result)) +} + +fn presto_like_information_schema_tables_sql( + database: &str, + schema: &str, + filter: Option<&str>, + limit: Option, +) -> String { let source = if database.trim().is_empty() { "information_schema.tables".to_string() } else { format!("{}.information_schema.tables", quote_presto_like_identifier(database)) }; - format!( + let mut sql = format!( "SELECT table_name, CASE table_type WHEN 'BASE TABLE' THEN 'TABLE' ELSE table_type END AS table_type \ FROM {source} \ - WHERE table_schema = {} AND table_type IN ('BASE TABLE', 'VIEW') \ - ORDER BY table_type, table_name", + WHERE table_schema = {} AND table_type IN ('BASE TABLE', 'VIEW')", sql_string_literal(schema) + ); + if let Some(filter) = filter.map(str::trim).filter(|value| !value.is_empty()) { + sql.push_str(" AND lower(table_name) LIKE "); + sql.push_str(&sql_string_literal(&format!("{}%", escape_presto_like_pattern(&filter.to_lowercase())))); + sql.push_str(" ESCAPE '\\'"); + } + sql.push_str(" ORDER BY table_type, table_name"); + if let Some(limit) = limit { + sql.push_str(&format!(" LIMIT {}", limit.max(1))); + } + sql +} + +fn presto_like_information_schema_columns_sql(database: &str, schema: &str, table: &str) -> String { + let source = if database.trim().is_empty() { + "information_schema.columns".to_string() + } else { + format!("{}.information_schema.columns", quote_presto_like_identifier(database)) + }; + format!( + "SELECT column_name, data_type, is_nullable, column_default, comment \ + FROM {source} \ + WHERE table_schema = {} AND table_name = {} \ + ORDER BY ordinal_position", + sql_string_literal(schema), + sql_string_literal(table) ) } @@ -1756,6 +1822,35 @@ fn presto_like_tables_from_query_result(result: &db::QueryResult) -> Vec Vec { + result + .rows + .iter() + .filter_map(|row| { + let name = query_result_cell_string(row, 0)?; + if name.trim().is_empty() { + return None; + } + let data_type = query_result_cell_string(row, 1).unwrap_or_default(); + Some(db::ColumnInfo { + name, + // Presto/Trino do not expose precision/length columns in information_schema.columns. + data_type: data_type.clone(), + is_nullable: query_result_cell_string(row, 2) + .map(|value| value.eq_ignore_ascii_case("YES")) + .unwrap_or(true), + column_default: query_result_cell_string(row, 3), + is_primary_key: false, + extra: None, + comment: query_result_cell_string(row, 4), + numeric_precision: presto_like_numeric_precision(&data_type), + numeric_scale: presto_like_numeric_scale(&data_type), + character_maximum_length: presto_like_character_maximum_length(&data_type), + }) + }) + .collect() +} + fn query_result_cell_string(row: &[serde_json::Value], index: usize) -> Option { let value = row.get(index)?; if value.is_null() { @@ -1764,6 +1859,29 @@ fn query_result_cell_string(row: &[serde_json::Value], index: usize) -> Option Option { + presto_like_type_argument(data_type, &["decimal", "numeric"], 0) +} + +fn presto_like_numeric_scale(data_type: &str) -> Option { + presto_like_type_argument(data_type, &["decimal", "numeric"], 1) +} + +fn presto_like_character_maximum_length(data_type: &str) -> Option { + presto_like_type_argument(data_type, &["char", "varchar"], 0) +} + +fn presto_like_type_argument(data_type: &str, type_names: &[&str], index: usize) -> Option { + let value = data_type.trim(); + let open = value.find('(')?; + let close = value[open + 1..].find(')')? + open + 1; + let name = value[..open].trim().to_ascii_lowercase(); + if !type_names.iter().any(|type_name| *type_name == name) { + return None; + } + value[open + 1..close].split(',').nth(index)?.trim().parse::().ok() +} + fn normalize_information_schema_table_type(table_type: &str) -> String { match table_type.trim().to_ascii_uppercase().replace(' ', "_").as_str() { "BASE_TABLE" => "TABLE".to_string(), @@ -1785,6 +1903,10 @@ fn sql_string_literal(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } +fn escape_presto_like_pattern(value: &str) -> String { + value.replace('\\', "\\\\").replace('%', "\\%").replace('_', "\\_") +} + #[cfg(test)] mod tests { use super::db; @@ -1796,8 +1918,8 @@ mod tests { oracle_object_statistics_from_query_result, oracle_object_statistics_rows_only_sql, oracle_object_statistics_sql, oracle_object_statistics_user_segments_sql, oracle_table_comment_from_query_result, oracle_table_comment_sql, oracle_table_comments_from_query_result, - oracle_table_comments_sql, presto_like_information_schema_tables_sql, presto_like_tables_from_query_result, - visible_schema_filter, + oracle_table_comments_sql, presto_like_columns_from_query_result, presto_like_information_schema_columns_sql, + presto_like_information_schema_tables_sql, presto_like_tables_from_query_result, visible_schema_filter, }; #[cfg(feature = "duckdb-bundled")] use super::{ @@ -2038,7 +2160,7 @@ mod tests { #[test] fn presto_like_information_schema_sql_uses_catalog_and_schema_without_system_jdbc() { - let sql = presto_like_information_schema_tables_sql("hive", "sales_analytics"); + let sql = presto_like_information_schema_tables_sql("hive", "sales_analytics", None, None); assert_eq!( sql, @@ -2049,12 +2171,40 @@ mod tests { #[test] fn presto_like_information_schema_sql_escapes_identifiers_and_literals() { - let sql = presto_like_information_schema_tables_sql("hi\"ve", "sales'analytics"); + let sql = presto_like_information_schema_tables_sql("hi\"ve", "sales'analytics", None, None); assert!(sql.contains("\"hi\"\"ve\".information_schema.tables")); assert!(sql.contains("table_schema = 'sales''analytics'")); } + #[test] + fn presto_like_information_schema_sql_pushes_table_filter_and_limit() { + let sql = presto_like_information_schema_tables_sql("hive", "sales_analytics", Some("Daily_%\\"), Some(20)); + + assert!(sql.contains("AND lower(table_name) LIKE 'daily\\_\\%\\\\%' ESCAPE '\\'")); + assert!(sql.ends_with("ORDER BY table_type, table_name LIMIT 20")); + } + + #[test] + fn presto_like_information_schema_columns_sql_uses_catalog_information_schema() { + let sql = presto_like_information_schema_columns_sql("hive", "sales_analytics", "daily_revenue"); + + assert_eq!( + sql, + "SELECT column_name, data_type, is_nullable, column_default, comment FROM \"hive\".information_schema.columns WHERE table_schema = 'sales_analytics' AND table_name = 'daily_revenue' ORDER BY ordinal_position" + ); + assert!(!sql.contains("system.jdbc.columns")); + } + + #[test] + fn presto_like_information_schema_columns_sql_escapes_identifiers_and_literals() { + let sql = presto_like_information_schema_columns_sql("hi\"ve", "sales'analytics", "daily'revenue"); + + assert!(sql.contains("\"hi\"\"ve\".information_schema.columns")); + assert!(sql.contains("table_schema = 'sales''analytics'")); + assert!(sql.contains("table_name = 'daily''revenue'")); + } + #[test] fn presto_like_tables_from_query_result_normalizes_base_table_type() { let result = super::db::QueryResult { @@ -2081,6 +2231,58 @@ mod tests { assert_eq!(normalize_information_schema_table_type("MATERIALIZED VIEW"), "MATERIALIZED_VIEW"); } + #[test] + fn presto_like_columns_from_query_result_maps_column_metadata() { + let result = super::db::QueryResult { + columns: vec![ + "column_name".to_string(), + "data_type".to_string(), + "is_nullable".to_string(), + "column_default".to_string(), + "comment".to_string(), + ], + column_types: vec![], + column_sortables: vec![], + rows: vec![ + vec![ + serde_json::json!("amount"), + serde_json::json!("decimal(12,2)"), + serde_json::json!("NO"), + serde_json::Value::Null, + serde_json::json!("daily amount"), + ], + vec![ + serde_json::json!("code"), + serde_json::json!("varchar(64)"), + serde_json::json!("YES"), + serde_json::Value::Null, + serde_json::Value::Null, + ], + ], + affected_rows: 0, + execution_time_ms: 1, + truncated: false, + session_id: None, + has_more: false, + }; + + let columns = presto_like_columns_from_query_result(&result); + + assert_eq!(columns[0].name, "amount"); + assert_eq!(columns[0].data_type, "decimal(12,2)"); + assert!(!columns[0].is_nullable); + assert_eq!(columns[0].comment.as_deref(), Some("daily amount")); + assert_eq!(columns[0].numeric_precision, Some(12)); + assert_eq!(columns[0].numeric_scale, Some(2)); + assert_eq!(columns[0].character_maximum_length, None); + assert!(!columns[0].is_primary_key); + assert_eq!(columns[1].name, "code"); + assert!(columns[1].is_nullable); + assert_eq!(columns[1].numeric_precision, None); + assert_eq!(columns[1].numeric_scale, None); + assert_eq!(columns[1].character_maximum_length, Some(64)); + } + #[cfg(feature = "duckdb-bundled")] #[test] fn duckdb_list_databases_includes_attached_database() { @@ -3167,6 +3369,9 @@ pub async fn get_columns_core( let config = config.clone(); let session = session.clone(); drop(connections); + if uses_presto_like_information_schema_tables(&config.db_type) { + return external_driver_presto_like_columns(session, config.as_ref(), database, schema, table).await; + } let columns = session .invoke_with_timeout::>( "getColumns", diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 5112395ea..d4eb29f64 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -256,7 +256,13 @@ public final class DbxJdbcPlugin { case "closeQuerySession", "close_query_session" -> closeQuerySessionResult(requireText(params, "sessionId")); case "listDatabases" -> listDatabases(connection); case "listSchemas" -> listSchemas(connection, optionalText(params, "database")); - case "listTables" -> listTables(connection, optionalText(params, "database"), optionalText(params, "schema")); + case "listTables" -> listTables( + connection, + optionalText(params, "database"), + optionalText(params, "schema"), + optionalText(params, "filter"), + nonNegativeInt(params, "limit", 0) + ); case "listObjects", "list_objects" -> listObjects( connection, optionalText(params, "database"), @@ -361,11 +367,29 @@ public final class DbxJdbcPlugin { private static void applyConnectTimeout(JsonNode connection, Properties properties) { int connectTimeoutSecs = positiveInt(connection, "connect_timeout_secs", 30); DriverManager.setLoginTimeout(connectTimeoutSecs); + if (isPrestoOrTrinoConnection(connection)) { + return; + } String value = Integer.toString(connectTimeoutSecs); properties.putIfAbsent("loginTimeout", value); properties.putIfAbsent("connectTimeout", value); } + private static boolean isPrestoOrTrinoConnection(JsonNode connection) { + String url = jdbcUrl(connection); + if (urlMatchesPrefix(url, "jdbc:presto:") || urlMatchesPrefix(url, "jdbc:trino:")) { + return true; + } + String driverClass = optionalText(connection, "jdbc_driver_class"); + if (driverClass == null) { + return false; + } + String normalized = driverClass.toLowerCase(Locale.ROOT); + return normalized.equals("io.prestosql.jdbc.prestodriver") || + normalized.equals("com.facebook.presto.jdbc.prestodriver") || + normalized.equals("io.trino.jdbc.trinodriver"); + } + private static void applyOracleProperties(JsonNode connection, Properties properties) { properties.putIfAbsent("remarksReporting", "false"); properties.putIfAbsent("restrictGetTables", "true"); @@ -1087,7 +1111,7 @@ public final class DbxJdbcPlugin { return result; } - private static JsonNode listTables(JsonNode connection, String database, String schema) throws SQLException { + private static JsonNode listTables(JsonNode connection, String database, String schema, String filter, int limit) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); Connection conn = openConnection(connection); JdbcDriverQuirks quirks = driverQuirks(connection); @@ -1095,7 +1119,7 @@ public final class DbxJdbcPlugin { return oracleListTables(conn, oracleEffectiveSchema(conn, schema)); } if (usePrestoInformationSchemaTables(connection)) { - return prestoListTables(conn, database, schema); + return prestoListTables(conn, database, schema, filter, limit); } DatabaseMetaData meta = conn.getMetaData(); String[] types = jdbcTableTypes(meta); @@ -1200,6 +1224,9 @@ public final class DbxJdbcPlugin { if (isKingbaseUrl(optionalText(connection, "connection_string"))) { return kingbaseGetColumns(conn, schema, table); } + if (usePrestoInformationSchemaTables(connection)) { + return prestoGetColumns(conn, database, schema, table); + } DatabaseMetaData meta = conn.getMetaData(); JdbcDriverQuirks quirks = driverQuirks(connection); String catalog = metadataCatalog(database, quirks); @@ -1366,10 +1393,13 @@ public final class DbxJdbcPlugin { return urlMatchesPrefix(url, "jdbc:presto:") || urlMatchesPrefix(url, "jdbc:trino:"); } - private static JsonNode prestoListTables(Connection conn, String database, String schema) throws SQLException { + private static JsonNode prestoListTables(Connection conn, String database, String schema, String filter, int limit) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); - try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database))) { + try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database, filter, limit))) { ps.setString(1, schema); + if (emptyToNull(filter) != null) { + ps.setString(2, escapeLikePattern(filter.toLowerCase(Locale.ROOT)) + "%"); + } try (ResultSet rs = ps.executeQuery()) { while (rs.next()) { ObjectNode item = MAPPER.createObjectNode(); @@ -1385,7 +1415,7 @@ public final class DbxJdbcPlugin { private static JsonNode prestoListObjects(Connection conn, String database, String schema) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); - try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database))) { + try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaTablesSql(database, null, 0))) { ps.setString(1, schema); try (ResultSet rs = ps.executeQuery()) { while (rs.next()) { @@ -1401,13 +1431,95 @@ public final class DbxJdbcPlugin { return result; } - static String prestoInformationSchemaTablesSql(String database) { + private static JsonNode prestoGetColumns(Connection conn, String database, String schema, String table) throws SQLException { + ArrayNode result = MAPPER.createArrayNode(); + try (PreparedStatement ps = conn.prepareStatement(prestoInformationSchemaColumnsSql(database))) { + ps.setString(1, schema); + ps.setString(2, table); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + String dataType = rs.getString(2); + ObjectNode item = columnNode(result, rs.getString(1)); + item.put("data_type", dataType); + item.put("is_nullable", !"NO".equalsIgnoreCase(rs.getString(3))); + putNullablePreferValue(item, "column_default", rs.getString(4)); + item.put("is_primary_key", false); + item.putNull("extra"); + putNullablePreferValue(item, "comment", rs.getString(5)); + // Presto/Trino information_schema.columns does not expose precision/length fields. + putNullableInt(item, "numeric_precision", prestoNumericPrecision(dataType)); + putNullableInt(item, "numeric_scale", prestoNumericScale(dataType)); + putNullableInt(item, "character_maximum_length", prestoCharacterMaximumLength(dataType)); + } + } + } + return result; + } + + static String prestoInformationSchemaTablesSql(String database, String filter, int limit) { String source = emptyToNull(database) == null ? "information_schema.tables" : quoteAnsiIdentifier(database) + ".information_schema.tables"; - return "SELECT table_name, table_type FROM " + source + + StringBuilder sql = new StringBuilder("SELECT table_name, table_type FROM " + source + " WHERE table_schema = ? AND table_type IN ('BASE TABLE', 'VIEW')" + - " ORDER BY table_type, table_name"; + (emptyToNull(filter) == null ? "" : " AND lower(table_name) LIKE ? ESCAPE '\\'") + + " ORDER BY table_type, table_name"); + if (limit > 0) { + sql.append(" LIMIT ").append(limit); + } + return sql.toString(); + } + + static String prestoInformationSchemaColumnsSql(String database) { + String source = emptyToNull(database) == null + ? "information_schema.columns" + : quoteAnsiIdentifier(database) + ".information_schema.columns"; + return "SELECT column_name, data_type, is_nullable, column_default, comment FROM " + source + + " WHERE table_schema = ? AND table_name = ?" + + " ORDER BY ordinal_position"; + } + + private static Integer prestoNumericPrecision(String dataType) { + return prestoTypeArgument(dataType, 0, "decimal", "numeric"); + } + + private static Integer prestoNumericScale(String dataType) { + return prestoTypeArgument(dataType, 1, "decimal", "numeric"); + } + + private static Integer prestoCharacterMaximumLength(String dataType) { + return prestoTypeArgument(dataType, 0, "char", "varchar"); + } + + private static Integer prestoTypeArgument(String dataType, int argumentIndex, String... typeNames) { + if (dataType == null) { + return null; + } + int open = dataType.indexOf('('); + int close = open < 0 ? -1 : dataType.indexOf(')', open + 1); + if (open <= 0 || close <= open) { + return null; + } + String name = dataType.substring(0, open).trim().toLowerCase(Locale.ROOT); + boolean matches = false; + for (String typeName : typeNames) { + if (typeName.equals(name)) { + matches = true; + break; + } + } + if (!matches) { + return null; + } + String[] arguments = dataType.substring(open + 1, close).split(","); + if (argumentIndex >= arguments.length) { + return null; + } + try { + return Integer.valueOf(arguments[argumentIndex].trim()); + } catch (NumberFormatException e) { + return null; + } } static String normalizeInformationSchemaTableType(String tableType) { @@ -2129,6 +2241,10 @@ public final class DbxJdbcPlugin { return value == null || value.isBlank() ? null : value; } + private static String escapeLikePattern(String value) { + return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_"); + } + private static Path expandHome(String path) { if (path.equals("~") || path.startsWith("~/")) { return Path.of(System.getProperty("user.home") + path.substring(1)); diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index 8dcf29fc3..44413701f 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -270,6 +270,25 @@ final class DbxJdbcPluginTest { assertEquals("45", properties.getProperty("connectTimeout")); } + @Test + void prestoConnectTimeoutDoesNotSetUnsupportedDriverProperties() throws Exception { + Method method = DbxJdbcPlugin.class.getDeclaredMethod("applyConnectTimeout", JsonNode.class, Properties.class); + method.setAccessible(true); + Properties properties = new Properties(); + JsonNode connection = MAPPER.readTree(""" + { + "connection_string": "jdbc:presto://presto.example.test:8080/hive", + "jdbc_driver_class": "io.prestosql.jdbc.PrestoDriver", + "connect_timeout_secs": 45 + } + """); + + method.invoke(null, connection, properties); + + assertFalse(properties.containsKey("loginTimeout")); + assertFalse(properties.containsKey("connectTimeout")); + } + @Test void jdbcUrlAppendsConnectionUrlParams() throws Exception { JsonNode connection = MAPPER.readTree(""" @@ -774,6 +793,77 @@ final class DbxJdbcPluginTest { } } + @Test + void prestoListTablesPushesFilterAndLimitToInformationSchema() throws Exception { + List calls = new ArrayList<>(); + Driver driver = new PrestoMetadataDriver(calls); + DriverManager.registerDriver(driver); + try { + JsonNode response = request("listTables", """ + { + "connection": { + "connection_string": "jdbc:presto://presto.example.test:8080/hive", + "connect_timeout_secs": 30 + }, + "database": "hive", + "schema": "sales_analytics", + "filter": "Daily_%", + "limit": 20 + } + """); + + assertFalse(response.has("error"), response.toString()); + assertEquals( + List.of( + "prepare:SELECT table_name, table_type FROM \"hive\".information_schema.tables WHERE table_schema = ? AND table_type IN ('BASE TABLE', 'VIEW') AND lower(table_name) LIKE ? ESCAPE '\\' ORDER BY table_type, table_name LIMIT 20", + "setString:1:sales_analytics", + "setString:2:daily\\_\\%%", + "executeQuery" + ), + calls + ); + } finally { + DriverManager.deregisterDriver(driver); + } + } + + @Test + void prestoGetColumnsUsesInformationSchemaInsteadOfJdbcMetadata() throws Exception { + List calls = new ArrayList<>(); + Driver driver = new PrestoMetadataDriver(calls); + DriverManager.registerDriver(driver); + try { + JsonNode response = request("getColumns", """ + { + "connection": { + "connection_string": "jdbc:presto://presto.example.test:8080/hive", + "connect_timeout_secs": 30 + }, + "database": "hive", + "schema": "sales_analytics", + "table": "daily_revenue" + } + """); + + assertFalse(response.has("error"), response.toString()); + assertEquals("amount", response.path("result").path(0).path("name").asText()); + assertEquals("decimal(12,2)", response.path("result").path(0).path("data_type").asText()); + assertEquals(12, response.path("result").path(0).path("numeric_precision").asInt()); + assertEquals(2, response.path("result").path(0).path("numeric_scale").asInt()); + assertEquals( + List.of( + "prepare:SELECT column_name, data_type, is_nullable, column_default, comment FROM \"hive\".information_schema.columns WHERE table_schema = ? AND table_name = ? ORDER BY ordinal_position", + "setString:1:sales_analytics", + "setString:2:daily_revenue", + "executeQuery" + ), + calls + ); + } finally { + DriverManager.deregisterDriver(driver); + } + } + @Test void oracleMetadataObjectTypeAcceptsPackageBodyAliases() throws Exception { Method method = DbxJdbcPlugin.class.getDeclaredMethod("oracleMetadataObjectType", String.class); @@ -1202,10 +1292,11 @@ final class DbxJdbcPluginTest { new Class[] { Connection.class }, (proxy, method, args) -> switch (method.getName()) { case "prepareStatement" -> { - calls.add("prepare:" + args[0]); - yield prestoMetadataStatement(calls); + String sql = String.valueOf(args[0]); + calls.add("prepare:" + sql); + yield prestoMetadataStatement(calls, sql); } - case "getMetaData" -> throw new SQLException("DatabaseMetaData should not be used for Presto listTables"); + case "getMetaData" -> throw new SQLException("DatabaseMetaData should not be used for Presto metadata"); case "isClosed" -> false; case "close" -> null; default -> defaultValue(method.getReturnType()); @@ -1213,7 +1304,7 @@ final class DbxJdbcPluginTest { ); } - private static PreparedStatement prestoMetadataStatement(List calls) { + private static PreparedStatement prestoMetadataStatement(List calls, String sql) { return (PreparedStatement) Proxy.newProxyInstance( DbxJdbcPluginTest.class.getClassLoader(), new Class[] { PreparedStatement.class }, @@ -1224,7 +1315,7 @@ final class DbxJdbcPluginTest { } case "executeQuery" -> { calls.add("executeQuery"); - yield prestoMetadataResultSet(); + yield sql.contains("information_schema.columns") ? prestoColumnMetadataResultSet() : prestoMetadataResultSet(); } case "close" -> null; default -> defaultValue(method.getReturnType()); @@ -1232,6 +1323,33 @@ final class DbxJdbcPluginTest { ); } + private static ResultSet prestoColumnMetadataResultSet() { + String[] labels = { "column_name", "data_type", "is_nullable", "column_default", "comment" }; + Object[][] rows = { { "amount", "decimal(12,2)", "NO", null, "daily amount" } }; + return (ResultSet) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { ResultSet.class }, + new java.lang.reflect.InvocationHandler() { + private int index = -1; + + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return switch (method.getName()) { + case "next" -> ++index < rows.length; + case "getMetaData" -> resultSetMeta(labels); + case "getString" -> { + Object value = rows[index][((Integer) args[0]) - 1]; + yield value == null ? null : value.toString(); + } + case "getObject" -> rows[index][((Integer) args[0]) - 1]; + case "close" -> null; + default -> defaultValue(method.getReturnType()); + }; + } + } + ); + } + private static ResultSet prestoMetadataResultSet() { String[] labels = { "table_name", "table_type" }; String[][] rows = { { "daily_revenue", "BASE TABLE" }, { "revenue_view", "VIEW" } };