fix(presto): optimize metadata completion

This commit is contained in:
luoianun 2026-07-01 19:05:45 +08:00 committed by GitHub
parent f5f6bc0b43
commit 95e3cb6571
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 528 additions and 37 deletions

View File

@ -156,6 +156,8 @@ const completionTranslations = computed(() => ({
functionDescriptions: Object.fromEntries(SQL_FUNCTION_NAMES.map((name) => [name, t(`editor.completion.functionDescriptions.${name}`)])) as Record<string, string>,
}));
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) => {

View File

@ -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", () => {

View File

@ -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);

View File

@ -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,
};
}

View File

@ -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([]);
});
});

View File

@ -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) }))

View File

@ -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::<Vec<db::TableInfo>>(
@ -1677,7 +1685,11 @@ async fn external_driver_presto_like_tables(
config: &ConnectionConfig,
database: &str,
schema: &str,
filter: Option<&str>,
limit: Option<usize>,
offset: Option<usize>,
) -> Result<Vec<db::TableInfo>, 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<Vec<db::ObjectInfo>, 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<crate::plugins::PluginDriverSession>,
config: &ConnectionConfig,
database: &str,
schema: &str,
table: &str,
) -> Result<Vec<db::ColumnInfo>, 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<usize>,
) -> 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<db::Tab
.collect()
}
fn presto_like_columns_from_query_result(result: &db::QueryResult) -> Vec<db::ColumnInfo> {
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<String> {
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<S
value.as_str().map(ToString::to_string).or_else(|| Some(value.to_string()))
}
fn presto_like_numeric_precision(data_type: &str) -> Option<i32> {
presto_like_type_argument(data_type, &["decimal", "numeric"], 0)
}
fn presto_like_numeric_scale(data_type: &str) -> Option<i32> {
presto_like_type_argument(data_type, &["decimal", "numeric"], 1)
}
fn presto_like_character_maximum_length(data_type: &str) -> Option<i32> {
presto_like_type_argument(data_type, &["char", "varchar"], 0)
}
fn presto_like_type_argument(data_type: &str, type_names: &[&str], index: usize) -> Option<i32> {
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::<i32>().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::<Vec<db::ColumnInfo>>(
"getColumns",

View File

@ -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));

View File

@ -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<String> 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<String> 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<String> calls) {
private static PreparedStatement prestoMetadataStatement(List<String> 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" } };