fix(presto): optimize metadata completion
This commit is contained in:
parent
f5f6bc0b43
commit
95e3cb6571
|
|
@ -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) => {
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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([]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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) }))
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
|
|
|
|||
|
|
@ -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" } };
|
||||
|
|
|
|||
Loading…
Reference in New Issue