diff --git a/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts index a53a1ecfe6..b53252f85b 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts @@ -1021,6 +1021,57 @@ describe("sqlCompletion scoped context classification", () => { expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "recent_orders", columns: ["id", "total"] }), expect.objectContaining({ name: "recent_orders", alias: "ro" })])); }); + it("scopes CTE bodies out of the outer query's referenced tables", () => { + const sql = "WITH cte AS (SELECT id, name FROM orders) SELECT id, na FROM cte"; + const cursor = sql.indexOf(" FROM cte"); + const context = getSqlCompletionContext(sql, cursor); + + expect(context.referencedTables.map((table) => table.name)).toEqual(["cte"]); + }); + + it("suggests bare CTE column names instead of forcing a cte. qualifier (issue #7396)", () => { + const sql = "WITH cte AS (SELECT id, name FROM orders) SELECT id, na FROM cte"; + const cursor = sql.indexOf(" FROM cte"); + const items = buildSqlCompletionItems(sql, cursor, { + tables: [], + columnsByTable: new Map([ + [ + "orders", + [ + { name: "id", table: "orders", dataType: "int" }, + { name: "name", table: "orders", dataType: "text" }, + ], + ], + [ + "cte", + [ + { name: "id", table: "cte" }, + { name: "name", table: "cte" }, + ], + ], + ]), + }); + + expect(items).toEqual(expect.arrayContaining([expect.objectContaining({ label: "name", type: "column" })])); + expect(items.filter((item) => item.type === "column").map((item) => item.label)).not.toContain("orders.name"); + }); + + it("keeps the CTE body's own tables while completing inside that body", () => { + const sql = "WITH cte AS (SELECT id, na FROM orders) SELECT * FROM cte"; + const cursor = sql.indexOf(" FROM orders"); + const context = getSqlCompletionContext(sql, cursor); + + expect(context.referencedTables.map((table) => table.name)).toEqual(expect.arrayContaining(["orders", "cte"])); + }); + + it("keeps the underlying table when a CTE body projects an unresolved star", () => { + const sql = "WITH cte AS (SELECT * FROM orders) SELECT na FROM cte"; + const cursor = sql.indexOf(" FROM cte"); + const context = getSqlCompletionContext(sql, cursor); + + expect(context.referencedTables.map((table) => table.name)).toEqual(expect.arrayContaining(["orders", "cte"])); + }); + it("extracts subquery aliases and projected columns", () => { const sql = "SELECT * FROM (SELECT id, name AS user_name FROM users) sq WHERE sq."; const context = getSqlCompletionContext(sql, sql.length); diff --git a/apps/desktop/src/lib/sql/sqlCompletion.ts b/apps/desktop/src/lib/sql/sqlCompletion.ts index 2980f145e8..e90a9b181f 100644 --- a/apps/desktop/src/lib/sql/sqlCompletion.ts +++ b/apps/desktop/src/lib/sql/sqlCompletion.ts @@ -2210,7 +2210,9 @@ function skipSqlWhitespaceAndComments(sql: string, pos: number): number { export function getSqlCompletionContext(sql: string, cursor: number, options: SqlSemanticBuildOptions = {}): SqlCompletionContext { const statementSpan = sqlCompletionStatementSpan(sql, cursor, options); // Extract the full statement at cursor position for referenced tables - const fullStatement = sql.slice(statementSpan.start, statementSpan.end).trim(); + const rawStatement = sql.slice(statementSpan.start, statementSpan.end); + const fullStatement = rawStatement.trim(); + const cursorInStatement = cursor - statementSpan.start - (rawStatement.length - rawStatement.trimStart().length); // Content before cursor within the current statement const beforeCursor = sql.slice(statementSpan.start, cursor); @@ -2223,10 +2225,10 @@ export function getSqlCompletionContext(sql: string, cursor: number, options: Sq const beforeToken = beforeCursor.slice(0, Math.max(0, bareStart)).trimEnd(); const lastWord = /([A-Za-z_][\w$]*)$/.exec(beforeToken)?.[1]?.toLowerCase() ?? ""; - let referencedTables = extractReferencedTables(fullStatement, options.databaseType); - - // Merge CTE definitions into referenced tables - const cteDefs = extractCteDefinitions(fullStatement); + // CTE bodies are their own scope: resolve them first so the outer query's + // referenced tables are read from a statement with those bodies blanked out. + const cteDefs = scanCteDefinitions(fullStatement); + let referencedTables = extractReferencedTables(maskResolvedCteBodies(fullStatement, cursorInStatement, cteDefs), options.databaseType); for (const cte of cteDefs) { if (!referencedTables.some((rt) => rt.name.toLowerCase() === cte.name.toLowerCase())) { referencedTables.push({ name: cte.name, columns: cte.columns }); @@ -3287,8 +3289,21 @@ function extractSelectColumnNames(sql: string): string[] { return names; } -export function extractCteDefinitions(sql: string): Array<{ name: string; columns: string[] }> { - const ctes: Array<{ name: string; columns: string[] }> = []; +interface ScannedCteDefinition { + name: string; + columns: string[]; + /** Index of the `(` opening the CTE body. */ + bodyStart: number; + /** Index of the `)` closing the CTE body. */ + bodyEnd: number; +} + +/** + * Scans `WITH` definitions, keeping each body's span so callers can tell which + * part of the statement belongs to a CTE rather than to the outer query. + */ +function scanCteDefinitions(sql: string): ScannedCteDefinition[] { + const ctes: ScannedCteDefinition[] = []; let lower = sql.toLowerCase(); const withMatch = /\bwith\b/.exec(lower); if (!withMatch) return ctes; @@ -3347,13 +3362,41 @@ export function extractCteDefinitions(sql: string): Array<{ name: string; column columns = extractSelectColumnNames(body); } - ctes.push({ name: cteName, columns }); + ctes.push({ name: cteName, columns, bodyStart: pos, bodyEnd }); pos = bodyEnd + 1; } return ctes; } +export function extractCteDefinitions(sql: string): Array<{ name: string; columns: string[] }> { + return scanCteDefinitions(sql).map(({ name, columns }) => ({ name, columns })); +} + +/** + * Blanks out CTE bodies whose projected columns are already known, so the outer + * query's referenced tables stay scoped to its own row sources. Without this the + * tables read inside a CTE (`WITH cte AS (SELECT id, name FROM t) SELECT na|`) + * would count as extra outer row sources and force every column suggestion to be + * qualified, hiding the bare `name` candidate the CTE actually provides. + * + * The body holding the cursor is never blanked -- completion inside a CTE body + * still needs that body's own tables -- and neither is a body whose columns could + * not be resolved (e.g. `SELECT *`), where the underlying table remains the only + * source of column names. + */ +function maskResolvedCteBodies(statement: string, cursorOffset: number, ctes: readonly ScannedCteDefinition[]): string { + let masked = statement; + for (const cte of ctes) { + if (cte.columns.length === 0) continue; + if (cursorOffset > cte.bodyStart && cursorOffset <= cte.bodyEnd) continue; + const bodyLength = cte.bodyEnd - cte.bodyStart - 1; + if (bodyLength <= 0) continue; + masked = `${masked.slice(0, cte.bodyStart + 1)}${" ".repeat(bodyLength)}${masked.slice(cte.bodyEnd)}`; + } + return masked; +} + function extractSubqueryReferences(sql: string): SqlCompletionReferencedTable[] { const refs: SqlCompletionReferencedTable[] = []; const pattern = /\b(?:from|join)\s*\(/gi; diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 4b849e27b9..974159829e 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -9,7 +9,7 @@ use rustls::client::verify_server_cert_signed_by_trust_anchor; use rustls::crypto::{verify_tls12_signature, verify_tls13_signature, CryptoProvider}; use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; use rustls::server::ParsedCertificate; -use sqlparser::ast::Statement; +use sqlparser::ast::{SetExpr, Statement}; use sqlparser::dialect::PostgreSqlDialect; use sqlparser::parser::Parser; use std::collections::{BTreeMap, BTreeSet, HashMap}; @@ -5956,27 +5956,60 @@ fn query_result_row_limit(max_rows: Option) -> usize { max_rows.unwrap_or(crate::query::MAX_ROWS).max(1) } +/// Returns whether a parsed DML statement produces a result set, or `None` when +/// the statement is not DML. +fn postgres_dml_statement_returns_rows(statement: &Statement) -> Option { + match statement { + Statement::Insert(insert) => Some(insert.returning.is_some()), + Statement::Update(update) => Some(update.returning.is_some()), + Statement::Delete(delete) => Some(delete.returning.is_some()), + Statement::Merge(merge) => Some(merge.output.is_some()), + _ => None, + } +} + +/// `WITH cte AS (...) UPDATE ...` parses as a query whose body is the wrapped +/// DML statement; returns that statement so its `RETURNING` clause can decide. +fn postgres_query_body_dml(body: &SetExpr) -> Option<&Statement> { + match body { + SetExpr::Insert(statement) + | SetExpr::Update(statement) + | SetExpr::Delete(statement) + | SetExpr::Merge(statement) => Some(statement), + _ => None, + } +} + /// Returns whether PostgreSQL should execute this statement through the row /// retrieval path. DML without `RETURNING` needs `execute` for its command /// tag/affected-row count, while DML with `RETURNING` produces a result set. pub(crate) fn postgres_statement_returns_rows(sql: &str) -> bool { - if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { + if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "TABLE"]) { return true; } + // A leading `WITH` is usually a CTE query, but PostgreSQL also allows a CTE + // to be followed directly by INSERT/UPDATE/DELETE/MERGE, which only returns + // rows with `RETURNING`. The prefix alone cannot decide, so parse it and + // keep the row-returning assumption whenever parsing does not say otherwise. + let leads_with_cte = starts_with_executable_sql_keyword(sql, &["WITH"]); + let Ok(statements) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else { - return false; + return leads_with_cte; }; let [statement] = statements.as_slice() else { - return false; + return leads_with_cte; }; + if let Some(returns_rows) = postgres_dml_statement_returns_rows(statement) { + return returns_rows; + } + match statement { - Statement::Insert(insert) => insert.returning.is_some(), - Statement::Update(update) => update.returning.is_some(), - Statement::Delete(delete) => delete.returning.is_some(), - Statement::Merge(merge) => merge.output.is_some(), - _ => false, + Statement::Query(query) => { + postgres_query_body_dml(&query.body).and_then(postgres_dml_statement_returns_rows).unwrap_or(true) + } + _ => leads_with_cte, } } @@ -8497,6 +8530,47 @@ mod tests { } } + #[test] + fn cte_wrapped_dml_parses_as_query_body() { + let statements = Parser::parse_sql( + &PostgreSqlDialect {}, + "WITH cte AS (SELECT id FROM t) UPDATE t SET x = 1 WHERE id IN (SELECT id FROM cte)", + ) + .expect("CTE-wrapped UPDATE should parse"); + let [Statement::Query(query)] = statements.as_slice() else { + panic!("expected a single Statement::Query, got: {statements:?}"); + }; + assert!(query.with.is_some(), "expected the CTE to be attached to the query"); + assert!( + matches!(postgres_query_body_dml(&query.body), Some(Statement::Update(_))), + "expected an UPDATE body, got: {:?}", + query.body + ); + } + + #[test] + fn postgres_statement_returns_rows_for_cte_wrapped_dml() { + for sql in [ + "WITH cte AS (SELECT id FROM t) SELECT * FROM cte", + "WITH cte AS (SELECT id FROM t) UPDATE t SET x = 1 WHERE id IN (SELECT id FROM cte) RETURNING id", + "WITH cte AS (SELECT id FROM t) DELETE FROM t WHERE id IN (SELECT id FROM cte) RETURNING id", + "WITH cte AS (SELECT id FROM t) INSERT INTO u (id) SELECT id FROM cte RETURNING id", + // Not parseable by sqlparser: keep assuming a CTE query returns rows. + "WITH RECURSIVE cte AS MATERIALIZED (SELECT 1) SELECT * FROM cte", + ] { + assert!(postgres_statement_returns_rows(sql), "expected result rows for: {sql}"); + } + + for sql in [ + "WITH cte AS (SELECT id FROM t) UPDATE t SET x = 1 WHERE id IN (SELECT id FROM cte)", + "WITH cte AS (SELECT id FROM t) DELETE FROM t WHERE id IN (SELECT id FROM cte)", + "WITH cte AS (SELECT id FROM t) INSERT INTO u (id) SELECT id FROM cte", + " with cte as (select id from t)\n update t set x = 1 where id in (select id from cte) ", + ] { + assert!(!postgres_statement_returns_rows(sql), "expected command result for: {sql}"); + } + } + #[test] fn database_list_does_not_collect_storage_usage() { assert!(list_databases_sql().contains("pg_database"));