diff --git a/.changeset/identifier-and-semantic-floors.md b/.changeset/identifier-and-semantic-floors.md new file mode 100644 index 0000000..6758cbf --- /dev/null +++ b/.changeset/identifier-and-semantic-floors.md @@ -0,0 +1,39 @@ +--- +'@asksql/core': patch +--- + +Quote the identifiers a database would not read back as itself. A mixed-case Postgres schema failed +every query, because an unquoted name folds to lower case and resolves to nothing; Oracle folds the +other way, and MySQL on Linux compares table names case-sensitively. Table and column names are now +quoted from the catalog before the query is validated. A name that is already correct is left +untouched, so MySQL output is unchanged, and a name spelled two ways across the catalog is skipped +rather than guessed. A table named like a parser keyword, such as `order` or `Nulls`, also works now: +the bare form could not be parsed at all and the question failed after three attempts. + +Reserved words now come from each database itself, through `pg_get_keywords()`, +`information_schema.KEYWORDS`, `V$RESERVED_WORDS` and `duckdb_keywords()`, rather than one shared list +that applied MySQL's rules to Postgres and missed most of MySQL's own: MySQL reserves 262 words where +the shared list had about a hundred. Regenerate with `node tools/generate-sql-keywords.mjs`. + +Quoting knows where a word is syntax rather than a name, so `CAST(x AS DATE)` and +`EXTRACT(MONTH FROM d)` are left alone, and a CTE is still recognised once its name is quoted. + +When a database rejects a name, the corrected query is derived from the catalog rather than from a +second model call, and the table repair names the closest match it already knew. + +Tell the model which database and schema it is connected to. Without that it wrote +`table_schema = 'your_database_name'` against `information_schema` and returned nothing at all, which +reads as an empty database rather than an error. Structure questions also get a correct catalog query +for the engine to build on. + +Say what is actually wrong when a statement will not parse. An apostrophe inside a value, as in +`'O'Brien'`, only produced "could not parse", so the model returned the same statement until it ran +out of attempts; it is now told to double the quote. + +Reject `AVG(SUM(x))` before it reaches the database. Nested aggregates are invalid everywhere, so the +query is repaired rather than run and failed. + +Fix two false alarms. A `UNION ALL` of per-table counts was blocked as a hallucinated column, because +each branch's columns were judged against every branch's tables; per-table row counts now work. Index +columns arrived already quoted from introspection and were quoted a second time, so every prompt +carried `"""ColumnName"""`. diff --git a/docs/screenshots/web-05-delete-refused.png b/docs/screenshots/web-05-delete-refused.png deleted file mode 100644 index edde23e..0000000 Binary files a/docs/screenshots/web-05-delete-refused.png and /dev/null differ diff --git a/packages/browser-extension/STORE-CERTIFICATION-NOTES.md b/packages/browser-extension/STORE-CERTIFICATION-NOTES.md new file mode 100644 index 0000000..9583452 --- /dev/null +++ b/packages/browser-extension/STORE-CERTIFICATION-NOTES.md @@ -0,0 +1,50 @@ +# Notes for Certification + +Paste the section below into **Submission Options > Notes for Certification** when resubmitting to the +Microsoft Edge Add-ons store. It answers policy 1.3.1 (Product is Testable), which the 08/13/2026 +review flagged. Product ID: 248cd48a-7dbe-4cfd-8ec0-df1e07231acd + +--- + +Product ID: 248cd48a-7dbe-4cfd-8ec0-df1e07231acd + +**Why no test account credentials are provided** + +AskSQL has no accounts, no sign-in, and no server of our own. Nothing is hosted by us, so there is no +credential we could issue. The extension stores its settings locally and talks only to two things the +user chooses: their own data files, and their own AI model provider. + +Because of that, testing needs no credentials from us. It needs a model provider and a data file, and +both can be supplied at no cost in a few minutes. + +**Fastest way to test, with no API key and no account (about 5 minutes)** + +1. Install Ollama from https://ollama.com (free, no account required) and run: + `ollama pull qwen2.5-coder:7b` +2. Start Ollama with `OLLAMA_ORIGINS=* ollama serve` so it serves on http://127.0.0.1:11434. + The variable matters: fetching the model list works without it, but asking a question fails with + 403, because Ollama rejects the extension's origin on POST requests. +3. Open the extension's Options page, choose provider **Ollama**, click **Fetch models**, pick the + model, and click **Test provider**. It should report success. +4. Add a connection: click **Add connection**, choose **Data files**, and select any CSV or Excel + file. Any small spreadsheet works; no database server is needed. +5. Open the side panel and ask a question about the file, for example "how many rows are there?" or + "show me the first 10 records". + +**Alternative, if you prefer a hosted provider** + +Any OpenAI, Anthropic or Groq API key works. Enter it in the Options page under the matching provider +and follow steps 3 to 5 above. We cannot include one of our keys in this submission, because the key +would be visible to anyone who reads the listing and would be billed to us. + +**What the extension sends where** + +Questions and database schema go only to the provider the user configures, over a connection they +control. Data files are read in the browser and never uploaded to us. The extension has no analytics +and no backend. Generated SQL is read-only and is checked before it runs, so a query cannot modify +the user's data. + +**If anything blocks the review** + +Please include the Product ID in any reply and we will respond quickly with whatever else is helpful, +including a recorded walkthrough if that is easier than running it locally. diff --git a/packages/core/README.md b/packages/core/README.md index 0e8a5b6..fb5424f 100644 --- a/packages/core/README.md +++ b/packages/core/README.md @@ -97,6 +97,15 @@ Beyond ask -> approve -> run, all optional: column (a common small-model slip), it is handed the real column list and re-asked, so the fix happens before the database ever sees the query. The schema is also auto-shrunk and retried once on context overflow. +- **Identifier quoting** - names a database would not read back as themselves are quoted from the + catalog before the query is validated, following each engine's own rule: Postgres folds unquoted + names down, Oracle folds them up, MySQL on Linux compares table names case-sensitively, and every + engine has reserved words. A mixed-case schema therefore works without the model having to + remember quotes, and names that are already correct are left alone. If a database still rejects a + name, the corrected query comes from the catalog rather than a second model call. +- **Semantic floors** - a query that would be rejected or would answer the wrong question is + repaired before it runs: an aggregate beside a bare column with no `GROUP BY`, an aggregate nested + inside another (`AVG(SUM(x))`), and a one-to-many join that inflates a `SUM`. - **Follow-up context** - prior turns are threaded into the prompt so "now break that down by month" works. - **Query history** - `config.history` records every attempt (status, duration), backed by an diff --git a/packages/core/src/catalog.ts b/packages/core/src/catalog.ts index 7ee68dd..992ab64 100644 --- a/packages/core/src/catalog.ts +++ b/packages/core/src/catalog.ts @@ -5,6 +5,7 @@ */ import type { PrunerSettings, SchemaCatalog, TableInfo } from './types.js'; +import { reservedWordsFor } from './sql-keywords.js'; import { VALUE_SAMPLE_MAX_DISTINCT } from './types.js'; import { dialectFor } from './dialects.js'; @@ -35,33 +36,22 @@ function sanitizeComment(comment: string | null | undefined): string | null { /** A name that can be written without quotes; anything else is rendered quoted. */ const PLAIN_IDENTIFIER_RE = /^[A-Za-z_][A-Za-z0-9_]*$/u; -/** Words an engine will not accept as a bare identifier; not exhaustive across every dialect. */ -const RESERVED_WORDS: ReadonlySet = new Set( - ( - 'select from where group by order having limit offset union all distinct join inner outer left ' + - 'right full cross natural on using as into insert update delete set values create drop alter ' + - 'table column view index key primary foreign unique constraint references default check null ' + - 'not and or in is like between case when then else end exists any some cast collate with ' + - 'recursive returning window over partition range rows current session system user grant revoke ' + - 'to begin commit rollback transaction lock database schema trigger procedure function ' + - 'desc asc date time timestamp interval level size type comment position language' - ).split(' '), -); - /** * True when the engine would not read the bare name back as itself; an unquoted identifier folds * case - PostgreSQL to lower, Oracle to upper. */ -function needsQuoting(name: string, engine: string): boolean { +export function needsQuoting(name: string, engine: string): boolean { if (!PLAIN_IDENTIFIER_RE.test(name)) return true; - if (RESERVED_WORDS.has(name.toLowerCase())) return true; + if (reservedWordsFor(engine).has(name.toLowerCase())) return true; if (engine === 'oracle') return name !== name.toUpperCase(); // MySQL, SQLite and DuckDB match identifiers case-insensitively, so folding cannot lose a name. if (engine === 'mysql' || engine === 'sqlite' || engine === 'duckdb') return false; return name !== name.toLowerCase(); } -function promptIdentifier(name: string, quote: string, engine: string): string { +function promptIdentifier(raw: string, quote: string, engine: string): string { + // Index columns arrive already quoted from introspection; quoting twice escapes the quotes into the name. + const name = raw.length > 1 && raw.startsWith(quote) && raw.endsWith(quote) ? raw.slice(1, -1) : raw; if (!needsQuoting(name, engine)) return name; // Doubling is how every supported engine escapes its own quote character inside an identifier. return `${quote}${name.split(quote).join(quote + quote)}${quote}`; diff --git a/packages/core/src/engine.ts b/packages/core/src/engine.ts index 45947b7..3d3a514 100644 --- a/packages/core/src/engine.ts +++ b/packages/core/src/engine.ts @@ -5,11 +5,19 @@ * The guard runs on every SQL string before execution; no DB session is held open across an LLM call. */ -import { joinGraph, pruneCatalog } from './catalog.js'; +import { joinGraph, needsQuoting, pruneCatalog } from './catalog.js'; +import { + correctTableCase, + foldingFor, + looksLikeUnknownTable, + hasUnterminatedLiteral, + quoteCatalogIdentifiers, + withoutLiteralsAndComments, +} from './identifier-case.js'; import { AskSqlError } from './errors.js'; import { extractImpossible, extractSql } from './extract.js'; import { guardSql, resolveGuardPolicy } from './guard.js'; -import { fanOutAggregate, ungroupedAggregate } from './semantics.js'; +import { fanOutAggregate, nestedAggregate, ungroupedAggregate } from './semantics.js'; import { historyId, MemoryHistoryStore } from './history.js'; import { callModel } from './llm.js'; import { @@ -517,6 +525,22 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { emit({ type: 'stage', stage: 'catalog' }, opts); const fullCatalog = await getCatalog(conn); + // Names the engine would not read back as themselves: folded case, reserved words, symbols. + // A name spelled two ways across the catalog is skipped: rewriting "status" to "Status" would + // ask one table for another table's column. + const allNames = fullCatalog.tables.flatMap((t) => [t.name, ...t.columns.map((c) => c.name)]); + const spellings = new Map>(); + for (const n of allNames) { + const key = n.toLowerCase(); + const set = spellings.get(key) ?? new Set(); + set.add(n); + spellings.set(key, set); + } + const quotableNames = allNames.filter( + (n) => needsQuoting(n, conn.engine) && spellings.get(n.toLowerCase())?.size === 1, + ); + // Only a table may be quoted before a dot; a schema qualifier that matched a column name broke it. + const quotableTables = fullCatalog.tables.map((t) => t.name).filter((n) => quotableNames.includes(n)); emit({ type: 'stage', stage: 'prune' }, opts); let pruned = pruneCatalog(fullCatalog, q, config.pruner); @@ -538,6 +562,9 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { fewShots, glossary: config.glossary, rerunPrevious: isRerunPreviousRequest(q), + database: conn.database, + schemas: fullCatalog.schemas, + catalogHint: isMetadataQuestion(q) ? catalogQueryHint(conn.dialect.engine) : undefined, }); const usageTotal: { input: number; output: number } = { input: 0, output: 0 }; @@ -581,6 +608,8 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { maxRows: policy.maxRows, context: opts.context, glossary: config.glossary, + database: conn.database, + schemas: fullCatalog.schemas, }); attempt -= 1; // does not consume a repair attempt continue; @@ -655,7 +684,13 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { lastSql = extraction.sql; emit({ type: 'stage', stage: 'guard' }, opts); - const verdict = guardSql({ sql: extraction.sql, dialect: conn.dialect, policy }); + // Quote first: a folding engine resolves a bare name elsewhere, and the parser cannot read a + // bare table named like a keyword. Falls back untouched if quoting makes it unparseable. + const normalised = quoteCatalogIdentifiers(extraction.sql, quotableNames, conn.dialect.quoteChar, quotableTables); + const normalisedVerdict = normalised ? guardSql({ sql: normalised, dialect: conn.dialect, policy }) : null; + const verdict = normalisedVerdict?.allowed + ? normalisedVerdict + : guardSql({ sql: extraction.sql, dialect: conn.dialect, policy }); if (!verdict.allowed) { if (attempt >= MAX_REPAIRS) { await recordHistory({ @@ -673,10 +708,14 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { detail: `ruleId=${verdict.ruleId ?? 'unknown'} after ${attempt + 1} attempts`, }); } + // "could not parse" alone leaves the model repeating the same statement; name the real cause. + const quoteHint = hasUnterminatedLiteral(extraction.sql, conn.dialect.quoteChar === '`') + ? " A text value contains an apostrophe that is not escaped: write it doubled, as 'O''Brien'." + : ''; userPrompt = buildRepairUser({ question: q, failedSql: extraction.sql, - failure: `The SQL validator rejected it: ${verdict.reason ?? verdict.ruleId ?? 'not allowed'}. Produce a single read-only SELECT.`, + failure: `The SQL validator rejected it: ${verdict.reason ?? verdict.ruleId ?? 'not allowed'}.${quoteHint} Produce a single read-only SELECT.`, schemaText, dialect: conn.dialect, }); @@ -711,10 +750,14 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { retryable: false, }); } + // The column repair already names the real columns; give the table repair the same head start. + const nearest = closestTableName(unknownTable, fullCatalog); userPrompt = buildRepairUser({ question: q, failedSql: verdict.sql, - failure: `Table "${unknownTable}" does not exist in the schema. Use only tables from the block.`, + failure: + `Table "${unknownTable}" does not exist in the schema.${nearest ? ` Did you mean "${nearest}"?` : ''} ` + + 'Use only tables from the block.', schemaText, dialect: conn.dialect, }); @@ -736,6 +779,21 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { continue; } + // Semantic floor: AVG(SUM(x)) and friends. Every engine rejects it, so repair before executing. + const nested = nestedAggregate(verdict.sql, conn.dialect.grammar); + if (nested && attempt < MAX_REPAIRS) { + userPrompt = buildRepairUser({ + question: q, + failedSql: verdict.sql, + failure: + `${nested}() contains another aggregate, which no SQL engine allows. ` + + 'Aggregate once over the rows, or aggregate the inner result in a subquery or CTE and then aggregate that.', + schemaText, + dialect: conn.dialect, + }); + continue; + } + // Semantic floor: a one-to-many join multiplies the rows a SUM sees, so the total is inflated. const fanOut = fanOutAggregate(verdict.sql, conn.dialect.grammar, fullCatalog); if (fanOut && attempt < MAX_REPAIRS) { @@ -779,6 +837,7 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { } emit({ type: 'stage', stage: 'done' }, opts); + const folding = foldingFor(conn.engine); const finalSql = verdict.sql; const explanation = extraction.explanation; const usage: LlmUsage = { inputTokens: usageTotal.input, outputTokens: usageTotal.output }; @@ -800,7 +859,8 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { // Never after a cancel: a repair would fire a fresh provider request the user just declined to wait for. const wasCancelled = (execOpts?.signal?.aborted ?? false) || (opts.signal?.aborted ?? false); if (AskSqlError.is(err) && err.code === 'DB_QUERY_ERROR' && !wasCancelled) { - const suggestion = await tryRepairAfterDbError(err); + // A wrong-cased table is repairable from the catalog alone, so try that before the model. + const suggestion = caseFixFor(err) ?? (await tryRepairAfterDbError(err)); if (suggestion) (err as DbErrorWithSuggestion).suggestedSql = suggestion; } throw err; @@ -808,6 +868,20 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { }, }; + function caseFixFor(dbErr: AskSqlError): string | null { + const message = dbErr.detail ?? dbErr.userMessage ?? ''; + if (!looksLikeUnknownTable(message)) return null; + const fixed = correctTableCase( + finalSql, + fullCatalog.tables.map((t) => t.name), + conn.dialect.quoteChar, + folding, + ); + if (!fixed) return null; + const v = guardSql({ sql: fixed, dialect: conn.dialect, policy }); + return v.allowed ? v.sql : null; + } + async function tryRepairAfterDbError(dbErr: AskSqlError): Promise { try { const repairPrompt = buildRepairUser({ @@ -1020,6 +1094,19 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { try { const catalog = await getCatalog(conn).catch(() => null); if (!catalog) return null; + // A wrong-cased table is repairable from the catalog alone, the same as ask().run() does. + if (looksLikeUnknownTable(opts.errorDetail ?? '')) { + const cased = correctTableCase( + bad, + catalog.tables.map((t) => t.name), + conn.dialect.quoteChar, + foldingFor(conn.engine), + ); + if (cased) { + const casedVerdict = guardSql({ sql: cased, dialect: conn.dialect, policy }); + if (casedVerdict.allowed) return casedVerdict.sql; + } + } const schemaText = pruneCatalog(catalog, question, config.pruner).schemaText; const repaired = await callModel({ model: config.model, @@ -1037,7 +1124,11 @@ export function createAskSql(config: AskSqlConfig): AskSqlEngine { const ex = extractSql(repaired.text); if (!ex) return null; const v = guardSql({ sql: ex.sql, dialect: conn.dialect, policy }); - return v.allowed && v.sql !== bad ? v.sql : null; + if (!v.allowed || v.sql === bad) return null; + // The same hallucination floors ask() enforces; a fix naming a missing table is not a fix. + if (firstUnknownTable(v.sql, catalog, conn.dialect.grammar, v.tables)) return null; + if (firstUnknownColumn(v.sql, catalog, conn.dialect.grammar)) return null; + return v.sql; } catch { return null; // best-effort; the original error stands } @@ -1078,7 +1169,9 @@ function collectCteNames(sql: string): ReadonlySet { const names = new Set(); if (!/\bwith\b/iu.test(sql)) return names; // Scans the whole statement; over-collecting CTE names only makes the floor more lenient. - for (const m of sql.matchAll(/([A-Za-z_][A-Za-z0-9_]*)\s+as\s*\(/giu)) { + // The quote characters matter: normalisation may have quoted a CTE named like a catalog column, + // and a model can quote one itself. An unrecognised CTE reads as a hallucinated table. + for (const m of sql.matchAll(/["`[]?([A-Za-z_][A-Za-z0-9_]*)["`\]]?\s+as\s*\(/giu)) { names.add(m[1]!.toLowerCase()); } return names; @@ -1123,7 +1216,10 @@ export function firstUnknownColumn(sql: string, catalog: SchemaCatalog, grammar: // The query's base tables that we know; unqualified columns are judged only when all are known. const queryTables: string[] = []; - let attributable = !/\(\s*select\b/iu.test(sql); + // A set operation has one column list per branch, and the parser reports them merged, so a column + // from one branch would be judged against another branch's tables. Not attributable, like a subquery. + const code = withoutLiteralsAndComments(sql); + let attributable = !/\(\s*select\b/iu.test(code) && !/\b(union|intersect|except)\b/iu.test(code); try { for (const t of tableParser.tableList(sql, { database: grammar })) { let name = (t.split('::')[2] ?? '').toLowerCase(); diff --git a/packages/core/src/identifier-case.ts b/packages/core/src/identifier-case.ts new file mode 100644 index 0000000..3f83bd9 --- /dev/null +++ b/packages/core/src/identifier-case.ts @@ -0,0 +1,288 @@ +/** + * Identifiers a database would not read back as itself: a name spelled in the wrong case, a + * mixed-case name left unquoted on an engine that folds, a reserved word, a symbol. MySQL on Linux + * compares table names case-sensitively; Postgres folds unquoted names down and Oracle folds up. + * Quoting from the catalog covers all of it, and the catalog can repair it without a model call. + */ + +import { reservedWordsFor } from './sql-keywords.js'; + +// The union across engines: treating a word as syntax is the conservative side of this decision. +const ANY_RESERVED = reservedWordsFor('*'); + +/** Only these lead-ins put an identifier in table position, which keeps a same-named column alone. */ +const TABLE_POSITION = + /\b(from|join|update|into)(\s+)([`"[]?)([A-Za-z_][\w$]*)[`"\]]?(\s*\.\s*([`"[]?)([A-Za-z_][\w$]*)[`"\]]?)?/gi; + +const CLOSING: Record = { '`': '`', '"': '"', '[': ']' }; + +function quoted(name: string, quoteChar: string): string { + return `${quoteChar}${name}${CLOSING[quoteChar] ?? quoteChar}`; +} + +/** Where a literal or comment ends, so identifiers are never rewritten inside one. */ +function skipTo(sql: string, i: number, doubleQuoteIsLiteral: boolean, backslashEscapes = false): number { + const ch = sql[i]; + const next = sql[i + 1]; + // $$body$$ and $tag$body$tag$ are literals in Postgres and DuckDB, and may contain anything. + if (ch === '$') { + const open = /^\$[A-Za-z_][\w]*\$|^\$\$/.exec(sql.slice(i)); + if (open) { + const close = sql.indexOf(open[0], i + open[0].length); + return close === -1 ? sql.length : close + open[0].length; + } + } + if (ch === '-' && next === '-') { + const nl = sql.indexOf('\n', i); + return nl === -1 ? sql.length : nl; + } + if (ch === '/' && next === '*') { + const close = sql.indexOf('*/', i + 2); + return close === -1 ? sql.length : close + 2; + } + if (ch === "'" || (ch === '"' && doubleQuoteIsLiteral)) { + let j = i + 1; + while (j < sql.length) { + if (backslashEscapes && sql[j] === '\\') j += 2; + else if (sql[j] === ch) { + // A doubled quote is an escaped one, so the literal continues past it. + if (sql[j + 1] === ch) j += 2; + else return j + 1; + } else j++; + } + return sql.length; + } + return -1; +} + +/** How an engine resolves an unquoted identifier: Postgres lower-cases it, Oracle upper-cases it. */ +export type Folding = 'lower' | 'upper' | 'none'; + +export function foldingFor(engine: string): Folding { + if (engine === 'postgres') return 'lower'; + if (engine === 'oracle') return 'upper'; + return 'none'; +} + +function folded(name: string, folding: Folding): string { + if (folding === 'lower') return name.toLowerCase(); + if (folding === 'upper') return name.toUpperCase(); + return name; +} + +/** + * Rewrites table references the database will not resolve to the catalog's table. That covers a name + * spelled in the wrong case, and on a folding engine a mixed-case name left unquoted, which resolves + * to something else entirely. Returns null when nothing changes. + */ +export function correctTableCase( + sql: string, + tableNames: readonly string[], + quoteChar: string, + folding: Folding = 'none', +): string | null { + const byLower = new Map(); + for (const name of tableNames) { + const lower = name.toLowerCase(); + // An ambiguous fold has no single right answer, so leave those names untouched. + if (byLower.has(lower) && byLower.get(lower) !== name) byLower.set(lower, ''); + else if (!byLower.has(lower)) byLower.set(lower, name); + } + + let changed = false; + const fixCode = (code: string): string => + code.replace( + TABLE_POSITION, + ( + whole: string, + keyword: string, + gap: string, + open: string, + first: string, + qualified: string | undefined, + _open2: string | undefined, + second: string | undefined, + offset: number, + ) => { + const target = second ?? first; + const canonical = byLower.get(target.toLowerCase()); + if (!canonical) return whole; + // A third part means what matched is a qualifier: prod.sales.orders names orders, not sales. + if (/^\s*\./.test(sql.slice(offset + whole.length))) return whole; + // An unquoted name is resolved folded, so what matters is what the database will look up. + const wasQuoted = (second === undefined ? open : (_open2 ?? '')) !== ''; + const resolvesTo = wasQuoted ? target : folded(target, folding); + if (resolvesTo === canonical) return whole; + changed = true; + const fixed = quoted(canonical, quoteChar); + if (second === undefined) return `${keyword}${gap}${fixed}`; + return `${keyword}${gap}${open ? quoted(first, open) : first}.${fixed}`; + }, + ); + + // A double quote is a string in MySQL but an identifier in Postgres, so the dialect decides. + const doubleQuoteIsLiteral = quoteChar !== '"'; + // MySQL is the only engine here that escapes with a backslash, and the backtick identifies it. + const backslashEscapes = quoteChar === '`'; + let out = ''; + let start = 0; + let i = 0; + while (i < sql.length) { + const end = skipTo(sql, i, doubleQuoteIsLiteral, backslashEscapes); + if (end >= 0) { + out += fixCode(sql.slice(start, i)) + sql.slice(i, end); + start = end; + i = end; + } else i++; + } + out += fixCode(sql.slice(start)); + return changed ? out : null; +} + +/** A bare identifier, and whatever follows it, so a function call can be told from a column. */ +const BARE_IDENTIFIER = /([A-Za-z_][\w$]*)(\s*[.(]?)/g; + +/** + * A reserved word is only treated as a name where it cannot be syntax: after FROM/JOIN/UPDATE/INTO, + * or qualified by a dot. Accepting AS or "(" quoted the type in CAST(x AS DATE) and the field in + * EXTRACT(MONTH FROM d), both of which are valid SQL that quoting breaks. + */ +const NAME_POSITION = /(?:\bfrom|\bjoin|\bupdate|\binto|\.)\s*$/i; + +/** The first argument of these is a keyword, not a name: EXTRACT(MONTH FROM d), TRIM(BOTH x FROM s). */ +const KEYWORD_ARGUMENT = /\b(?:extract|trim|position|overlay|substring)\s*\(\s*$/i; + +/** + * Quotes every table and column the engine would otherwise fold away. The schema text already shows + * these names quoted and models still drop the quotes, so the query is normalised before it runs + * rather than left to fail. Returns null when nothing needed quoting. + */ +export function quoteCatalogIdentifiers( + sql: string, + names: readonly string[], + quoteChar: string, + tableNames: readonly string[] = names, +): string | null { + const tables = new Set(tableNames.map((n) => n.toLowerCase())); + const byLower = new Map(); + for (const name of names) { + const lower = name.toLowerCase(); + if (byLower.has(lower) && byLower.get(lower) !== name) byLower.set(lower, ''); + else if (!byLower.has(lower)) byLower.set(lower, name); + } + if (byLower.size === 0) return null; + + let changed = false; + const fixCode = (code: string, chunkStart: number): string => + code.replace(BARE_IDENTIFIER, (whole, token: string, tail: string, offset: number) => { + if (tail.trimStart().startsWith('(')) return whole; // a function call, not an identifier + const canonical = byLower.get(token.toLowerCase()); + if (!canonical) return whole; + const before = code.slice(0, offset); + if (KEYWORD_ARGUMENT.test(before)) return whole; + // TIMESTAMP '2024-01-01' and DATE '...' are typed literals: the word is syntax, not a name. + // The literal is its own segment, so this reads the statement rather than the chunk. + if (/^\s*'/.test(sql.slice(chunkStart + offset + token.length))) return whole; + // A token before a dot qualifies what follows; quoting a schema name breaks a working query. + if (tail.trimStart().startsWith('.') && !tables.has(token.toLowerCase())) return whole; + // Rewriting a keyword blindly turns ORDER BY into "order" BY, so one must announce itself. + if (ANY_RESERVED.has(token.toLowerCase()) && !NAME_POSITION.test(before)) return whole; + changed = true; + return `${quoted(canonical, quoteChar)}${tail}`; + }); + + const doubleQuoteIsLiteral = quoteChar !== '"'; + const backslashEscapes = quoteChar === '`'; + let out = ''; + let start = 0; + let i = 0; + while (i < sql.length) { + // An already-quoted identifier is opaque: re-quoting it would double the quote characters. + if (sql[i] === quoteChar) { + const close = sql.indexOf(CLOSING[quoteChar] ?? quoteChar, i + 1); + const end = close === -1 ? sql.length : close + 1; + out += fixCode(sql.slice(start, i), start) + sql.slice(i, end); + start = end; + i = end; + continue; + } + const end = skipTo(sql, i, doubleQuoteIsLiteral, backslashEscapes); + if (end >= 0) { + out += fixCode(sql.slice(start, i), start) + sql.slice(i, end); + start = end; + i = end; + } else i++; + } + out += fixCode(sql.slice(start), start); + return changed ? out : null; +} + +/** + * The statement with string literals and comments blanked out, so a keyword search cannot match a + * value like 'except this' or a table named in a comment. Length and offsets are preserved. + */ +export function withoutLiteralsAndComments(sql: string): string { + let out = ''; + let start = 0; + let i = 0; + while (i < sql.length) { + // Quoted identifiers are code, not literals, so only ' and comments are blanked here. + const end = skipTo(sql, i, false); + if (end >= 0) { + out += sql.slice(start, i) + ' '.repeat(end - i); + start = end; + i = end; + } else i++; + } + return out + sql.slice(start); +} + +/** + * True when a text value opens and never closes, which is what an unescaped apostrophe looks like: + * 'O'Brien' reads as the value 'O', then Brien, then a literal running to the end of the statement. + * The parser only reports "could not parse", so naming the real cause is what makes the repair land. + */ +export function hasUnterminatedLiteral(sql: string, backslashEscapes = false): boolean { + let open = false; + let i = 0; + while (i < sql.length) { + const ch = sql[i]; + if (open) { + if (backslashEscapes && ch === '\\') i += 2; + else if (ch === "'" && sql[i + 1] === "'") + i += 2; // an escaped apostrophe, the value continues + else if (ch === "'") { + open = false; + i++; + } else i++; + continue; + } + // A dollar-quoted body is a literal that needs no escaping, so an apostrophe inside is fine. + if (ch === '$') { + const dollar = skipTo(sql, i, false); + if (dollar > i) { + i = dollar; + continue; + } + } + if (ch === '-' && sql[i + 1] === '-') { + const nl = sql.indexOf('\n', i); + i = nl === -1 ? sql.length : nl; + } else if (ch === '/' && sql[i + 1] === '*') { + const close = sql.indexOf('*/', i + 2); + i = close === -1 ? sql.length : close + 2; + } else if (ch === "'") { + open = true; + i++; + } else i++; + } + return open; +} + +/** Matches the unknown-table wording of every engine AskSQL supports. */ +const UNKNOWN_TABLE = + /\b(doesn't exist|does not exist|not found|unknown table|invalid object name|undefined table|no such table|table or view does not exist)\b/i; + +export function looksLikeUnknownTable(message: string): boolean { + return UNKNOWN_TABLE.test(message); +} diff --git a/packages/core/src/prompt.ts b/packages/core/src/prompt.ts index cbe3a83..bb6eae5 100644 --- a/packages/core/src/prompt.ts +++ b/packages/core/src/prompt.ts @@ -19,6 +19,11 @@ export interface SqlPromptInput { readonly glossary?: readonly { term: string; definition: string }[]; /** The question asks to re-run the previous query rather than for a new one. */ readonly rerunPrevious?: boolean; + /** Named so a question about system catalogs does not invent a placeholder database name. */ + readonly database?: string; + readonly schemas?: readonly string[]; + /** A correct catalog query for this engine, offered when the question is about structure. */ + readonly catalogHint?: string; } export function buildSqlSystem(dialect: DialectInfo, maxRows: number, prompts?: PromptSettings): string { @@ -53,6 +58,25 @@ export function buildSqlUser(input: SqlPromptInput): string { const parts: string[] = []; parts.push('', input.schemaText, ''); + // Without these, a question about information_schema gets a guessed name like 'your_database_name'. + const where: string[] = []; + if (input.database) where.push(`database/catalog is "${input.database}"`); + if (input.schemas && input.schemas.length > 0) where.push(`schema is "${input.schemas[0]}"`); + if (where.length > 0) { + parts.push( + '', + `You are connected to: the ${where.join(', the ')}. Use these exact names when a query filters on system catalogs such as information_schema; never write a placeholder.`, + ); + } + + // System-catalog column names are not in the schema block, so a structure question otherwise guesses them. + if (input.catalogHint) { + parts.push( + '', + `This question is about the database's structure. Build on this correct query for this engine: ${input.catalogHint}`, + ); + } + if (input.glossary && input.glossary.length > 0) { parts.push('', 'Business glossary (use these definitions when the question uses these terms):'); for (const g of input.glossary.slice(0, 40)) parts.push(`- ${g.term}: ${g.definition}`); diff --git a/packages/core/src/semantics.ts b/packages/core/src/semantics.ts index a88957f..15a514b 100644 --- a/packages/core/src/semantics.ts +++ b/packages/core/src/semantics.ts @@ -213,3 +213,50 @@ export function fanOutAggregate(sql: string, grammar: string, catalog: FanOutCat } return null; } + +/** + * An aggregate nested inside another aggregate, like AVG(x + SUM(y)). Every engine rejects it, so + * catching it before execution turns a database error into a repair. Returns the outer function name. + */ +export function nestedAggregate(sql: string, grammar: string): string | null { + let ast: unknown; + try { + ast = parser.parse(sql, { database: grammar }).ast; + } catch { + return null; // the guard already parsed it; never double-block here + } + + let outer: string | null = null; + const walk = (node: unknown, insideAggregate: string | null): void => { + if (outer) return; + if (Array.isArray(node)) { + for (const item of node) walk(item, insideAggregate); + return; + } + if (!isNode(node)) return; + const isAgg = isBareAggregate(node); + if (isAgg && insideAggregate) { + outer = insideAggregate; + return; + } + const within = isAgg ? aggregateName(node) : insideAggregate; + for (const key of Object.keys(node)) { + // A subquery has its own scope, so an aggregate inside one is not nested in the outer call. + if (key === 'ast' || key === 'from') continue; + walk(node[key], within); + } + }; + walk(ast, null); + return outer; +} + +function aggregateName(node: Node): string { + const name = node['name']; + const text = + typeof name === 'string' + ? name + : isNode(name) && Array.isArray(name['name']) + ? String((name['name'][0] as Node)?.['value'] ?? '') + : ''; + return text.toUpperCase(); +} diff --git a/packages/core/src/sql-keywords.ts b/packages/core/src/sql-keywords.ts new file mode 100644 index 0000000..95d02d3 --- /dev/null +++ b/packages/core/src/sql-keywords.ts @@ -0,0 +1,65 @@ +/** + * Reserved words per engine, read from each database itself rather than guessed: + * pg_get_keywords(), information_schema.KEYWORDS, V$RESERVED_WORDS, duckdb_keywords(). + * SQLite publishes a fixed list and has no catalog to query. + * + * Regenerate with: node tools/generate-sql-keywords.mjs + * One word list plus a bit per engine, which costs a few hundred bytes instead of repeating words. + */ + +const WORDS = + 'abort|accessible|action|add|after|all|alter|always|analyse|analyze|and|any|array|as|asc|asensitive|asymmetric|' + + 'attach|authorization|autoincrement|before|begin|between|bigint|binary|blob|both|by|call|cascade|case|cast|chan' + + 'ge|char|character|check|cluster|collate|collation|column|commit|compress|concurrently|condition|conflict|conne' + + 'ct|constraint|continue|convert|create|cross|cube|cume_dist|current|current_catalog|current_date|current_role|c' + + 'urrent_schema|current_time|current_timestamp|current_user|cursor|database|databases|date|day_hour|day_microsec' + + 'ond|day_minute|day_second|dec|decimal|declare|default|deferrable|deferred|delayed|delete|dense_rank|desc|descr' + + 'ibe|detach|deterministic|distinct|distinctrow|div|do|double|drop|dual|each|else|elseif|empty|enclosed|end|esca' + + 'pe|escaped|except|exclude|exclusive|exists|exit|explain|fail|false|fetch|filter|first|first_value|float|float4' + + '|float8|following|for|force|foreign|freeze|from|full|fulltext|function|generated|get|glob|grant|group|grouping' + + '|groups|having|high_priority|hour_microsecond|hour_minute|hour_second|identified|if|ignore|ilike|immediate|in|' + + 'index|indexed|infile|initially|inner|inout|insensitive|insert|instead|int|int1|int2|int3|int4|int8|integer|int' + + 'ersect|interval|into|io_after_gtids|io_before_gtids|is|isnull|iterate|join|json_table|key|keys|kill|lag|lambda' + + '|last|last_value|lateral|lead|leading|leave|left|like|limit|linear|lines|load|localtime|localtimestamp|lock|lo' + + 'ng|longblob|longtext|loop|low_priority|match|materialized|maxvalue|mediumblob|mediumint|mediumtext|middleint|m' + + 'inus|minute_microsecond|minute_second|mod|mode|modifies|natural|no|no_write_to_binlog|nocompress|not|nothing|n' + + 'otnull|nowait|nth_value|ntile|null|nulls|number|numeric|of|offset|on|only|optimize|optimizer_costs|option|opti' + + 'onally|or|order|others|out|outer|outfile|over|overlaps|partition|pctfree|percent_rank|pivot|pivot_longer|pivot' + + '_wider|placing|plan|pragma|preceding|precision|primary|prior|procedure|public|purge|qualify|query|raise|range|' + + 'rank|raw|read|read_write|reads|real|recursive|references|regexp|reindex|release|rename|repeat|replace|require|' + + 'resignal|resource|restrict|return|returning|revoke|right|rlike|rollback|row|row_number|rows|savepoint|schema|s' + + 'chemas|second_microsecond|select|sensitive|separator|session_user|set|share|show|signal|similar|size|smallint|' + + 'some|spatial|specific|sql|sql_big_result|sql_calc_found_rows|sql_small_result|sqlexception|sqlstate|sqlwarning' + + '|ssl|start|starting|stored|straight_join|summarize|symmetric|synonym|system|system_user|table|tablesample|temp' + + '|temporary|terminated|then|ties|tinyblob|tinyint|tinytext|to|trailing|transaction|trigger|true|unbounded|undo|' + + 'union|unique|unlock|unpivot|unsigned|update|usage|use|user|using|utc_date|utc_time|utc_timestamp|vacuum|values' + + '|varbinary|varchar|varchar2|varcharacter|variadic|varying|verbose|view|virtual|when|where|while|window|with|wi' + + 'thout|write|xor|year_month|zerofill'; + +const WORD_LIST = WORDS.split('|'); + +/** One bit per word in WORD_LIST, base64 encoded. */ +const BY_ENGINE: Record = { + sqlite: '/WZ66KhRpkwAV6XG3gxrqsHeDCgrBAfAAJhjDq4CFxz4RbURAQAAaKM1Iga8Aw==', + postgres: 'IH8FxehExh8AQyREAgN6MAHFACgLUMcAAIgiHCaBEAAQAAWQkACAHGEyMICyAQ==', + mysql: 'aubQf6/In/z++d4/c/Ou99+s9//9+f9/36046/cKWLO/3+5+Tf87mn177135PQ==', + oracle: 'YGxACBoiAgBBUYQEGiAiMCEMBCwBAAIDIMKkigYEoEAAIQIQYwAECSExAjQkAQ==', + duckdb: 'IH8BxKhAAgAAwyREAgMqIAFEACgAUgQAAIAgHAbwEAIQAAEQhADACGGyIICwAQ==', +}; + +const cache = new Map>(); + +/** The engine's reserved words; an unknown engine falls back to the union, which is the safe side. */ +export function reservedWordsFor(engine: string): ReadonlySet { + const key = engine.toLowerCase(); + let set = cache.get(key); + if (!set) { + const packed = BY_ENGINE[key]; + if (packed) { + const bytes = atob(packed); + set = new Set(WORD_LIST.filter((_, i) => (bytes.charCodeAt(i >> 3) >> (i & 7)) & 1)); + } else set = new Set(WORD_LIST); + cache.set(key, set); + } + return set; +} diff --git a/packages/core/test/identifier-case.test.ts b/packages/core/test/identifier-case.test.ts new file mode 100644 index 0000000..29cbcc1 --- /dev/null +++ b/packages/core/test/identifier-case.test.ts @@ -0,0 +1,435 @@ +import { describe, expect, it } from 'vitest'; +import { pruneCatalog } from '../src/catalog.js'; +import { buildSqlUser } from '../src/prompt.js'; +import { POSTGRES_DIALECT } from '../src/dialects.js'; +import { + correctTableCase, + hasUnterminatedLiteral, + looksLikeUnknownTable, + quoteCatalogIdentifiers, +} from '../src/identifier-case.js'; + +const TABLES = ['Customers', 'OrderItems']; + +describe('correctTableCase', () => { + it('corrects a lower-cased table name', () => { + expect(correctTableCase('SELECT * FROM customers', TABLES, '`')).toBe('SELECT * FROM `Customers`'); + }); + + it('corrects an upper-cased table name after JOIN', () => { + const sql = 'SELECT * FROM Customers c JOIN ORDERITEMS o ON c.id = o.id'; + expect(correctTableCase(sql, TABLES, '`')).toBe( + 'SELECT * FROM Customers c JOIN `OrderItems` o ON c.id = o.id', + ); + }); + + it('leaves an alias after the table alone', () => { + expect(correctTableCase('SELECT * FROM orderitems oi', TABLES, '`')).toBe( + 'SELECT * FROM `OrderItems` oi', + ); + }); + + it('returns null when every name already matches', () => { + expect(correctTableCase('SELECT * FROM Customers', TABLES, '`')).toBeNull(); + }); + + it('quotes with the dialect character', () => { + expect(correctTableCase('SELECT * FROM customers', TABLES, '"')).toBe('SELECT * FROM "Customers"'); + }); + + it('corrects a name that was already quoted in the wrong case', () => { + expect(correctTableCase('SELECT * FROM `customers`', TABLES, '`')).toBe('SELECT * FROM `Customers`'); + }); + + it('keeps a schema prefix and corrects only the table', () => { + expect(correctTableCase('SELECT * FROM shop.orderitems', TABLES, '`')).toBe( + 'SELECT * FROM shop.`OrderItems`', + ); + }); + + /** A column sharing a table's name must not be rewritten: it is not in table position. */ + it('leaves a same-named column alone', () => { + expect(correctTableCase('SELECT customers FROM Customers', TABLES, '`')).toBeNull(); + }); + + /** Rewriting inside a literal would change the query's meaning, not just its spelling. */ + it('leaves a string literal alone', () => { + expect(correctTableCase("SELECT * FROM Customers WHERE note = 'from customers'", TABLES, '`')).toBeNull(); + }); + + /** Two tables differing only by case have no single right answer. */ + it('leaves an ambiguous fold untouched', () => { + expect(correctTableCase('SELECT * FROM orders', ['Orders', 'ORDERS'], '`')).toBeNull(); + }); + + it('corrects an UPDATE target', () => { + expect(correctTableCase('UPDATE customers SET x = 1', TABLES, '`')).toBe('UPDATE `Customers` SET x = 1'); + }); + + it('leaves an unknown table alone', () => { + expect(correctTableCase('SELECT * FROM invoices', TABLES, '`')).toBeNull(); + }); +}); + +describe('looksLikeUnknownTable', () => { + it.each([ + "Table 'asksql_test.customers' doesn't exist", + 'relation "customers" does not exist', + 'no such table: customers', + 'ORA-00942: table or view does not exist', + 'Invalid object name customers.', + ])('recognises %s', (message) => { + expect(looksLikeUnknownTable(message)).toBe(true); + }); + + it('does not fire on an unrelated failure', () => { + expect(looksLikeUnknownTable('Unknown column x in field list')).toBe(false); + }); +}); + +describe('folding engines', () => { + /** Postgres folds an unquoted name to lower case, so a mixed-case table needs quoting even when spelled right. */ + it('quotes a correctly spelled mixed-case table on Postgres', () => { + expect(correctTableCase('SELECT * FROM Customers', TABLES, '"', 'lower')).toBe('SELECT * FROM "Customers"'); + }); + + it('leaves an all-lower-case table alone on Postgres', () => { + expect(correctTableCase('SELECT * FROM orders', ['orders'], '"', 'lower')).toBeNull(); + }); + + /** Unquoted Orders folds to orders, which is the catalog table, so nothing needs changing. */ + it('leaves a name the fold already resolves alone on Postgres', () => { + expect(correctTableCase('SELECT * FROM Orders', ['orders'], '"', 'lower')).toBeNull(); + }); + + it('leaves an already quoted mixed-case name alone on Postgres', () => { + expect(correctTableCase('SELECT * FROM "Customers"', TABLES, '"', 'lower')).toBeNull(); + }); + + it('quotes a mixed-case table on Oracle, which folds upper', () => { + expect(correctTableCase('SELECT * FROM MixedCase', ['MixedCase'], '"', 'upper')).toBe('SELECT * FROM "MixedCase"'); + }); + + it('leaves an upper-case Oracle table alone', () => { + expect(correctTableCase('SELECT * FROM employees', ['EMPLOYEES'], '"', 'upper')).toBeNull(); + }); +}); + +describe('quoteCatalogIdentifiers', () => { + const NAMES = ['Customers', 'OrderItems', 'FirstName', 'Country', 'CustomerId']; + const q = (sql: string) => quoteCatalogIdentifiers(sql, NAMES, '"'); + + it('quotes both the table and the columns', () => { + expect(q("SELECT FirstName FROM Customers WHERE Country = 'UK'")).toBe( + `SELECT "FirstName" FROM "Customers" WHERE "Country" = 'UK'`, + ); + }); + + it('quotes a qualified column', () => { + expect(q('SELECT c.CustomerId FROM Customers c')).toBe('SELECT c."CustomerId" FROM "Customers" c'); + }); + + /** Doubling the quotes would make the identifier unreadable. */ + it('leaves an already quoted identifier alone', () => { + expect(q('SELECT "FirstName" FROM "Customers"')).toBeNull(); + }); + + /** A reserved word used as a function must not become an identifier. */ + it('leaves a function call alone', () => { + expect(quoteCatalogIdentifiers('SELECT COUNT(*) FROM Customers', ['Customers', 'count'], '"')).toBe( + 'SELECT COUNT(*) FROM "Customers"', + ); + }); + + it('leaves a string literal alone', () => { + expect(q("SELECT 1 WHERE x = 'FirstName'")).toBeNull(); + }); + + it('returns null when nothing needs quoting', () => { + expect(quoteCatalogIdentifiers('SELECT id FROM orders', [], '"')).toBeNull(); + }); + + /** A reserved word is quoted too, not only a folded name. */ + /** A table called "order" once turned ORDER BY into "order" BY, which the guard then rejected. */ + it('does not rewrite a keyword that is not naming the table', () => { + expect( + quoteCatalogIdentifiers('SELECT x FROM Customers ORDER BY x DESC', ['Customers', 'order'], '"'), + ).toBe('SELECT x FROM "Customers" ORDER BY x DESC'); + }); + + it('still quotes GROUP BY and other keyword-adjacent columns', () => { + expect(quoteCatalogIdentifiers('SELECT Country FROM t GROUP BY Country', ['Country'], '"')).toBe( + 'SELECT "Country" FROM t GROUP BY "Country"', + ); + }); + + /** A table called Nulls broke the parser: NULLS is a keyword, so the bare name would not parse. */ + it('quotes a table named like a parser keyword', () => { + expect(quoteCatalogIdentifiers('SELECT Val FROM Nulls', ['Nulls', 'Val'], '"')).toBe( + 'SELECT "Val" FROM "Nulls"', + ); + }); + + it('quotes a table whose name is a reserved word', () => { + expect(quoteCatalogIdentifiers('SELECT * FROM order', ['order'], '"')).toBe('SELECT * FROM "order"'); + }); +}); + +describe('schema text quoting', () => { + /** Index columns arrive already quoted from introspection, and were being quoted a second time. */ + it('does not double-quote a name that is already quoted', () => { + const text = pruneCatalog( + { + engine: 'postgres', + schemas: ['public'], + tables: [ + { + name: 'Customers', + kind: 'table', + columns: [{ name: 'CustomerId', dbType: 'integer', nullable: false }], + primaryKey: ['CustomerId'], + foreignKeys: [], + uniques: [], + checks: [], + indexes: [{ name: 'Customers_pkey', columns: ['"CustomerId"'], unique: true }], + source: 'db', + }, + ], + enums: [], + sequences: [], + triggers: [], + routines: [], + warnings: [], + fetchedAt: 'now', + } as never, + 'customers', + ).schemaText; + + expect(text).not.toContain('"""'); + expect(text).toContain('"CustomerId"'); + }); +}); + +describe('connection identity in the prompt', () => { + /** Without the real name a model writes table_schema = 'your_database_name', which silently returns nothing. */ + it('names the database and schema so system-catalog filters are real', () => { + const text = buildSqlUser({ + question: 'what views exist?', + schemaText: 'TABLE film', + dialect: POSTGRES_DIALECT, + maxRows: 100, + database: 'sakila', + schemas: ['public'], + }); + + expect(text).toContain('"sakila"'); + expect(text).toContain('"public"'); + expect(text).toMatch(/never write a placeholder/i); + }); + + it('says nothing when the connection does not report a database', () => { + const text = buildSqlUser({ question: 'q', schemaText: 's', dialect: POSTGRES_DIALECT, maxRows: 10 }); + expect(text).not.toMatch(/You are connected to/); + }); +}); + +describe('catalog hint for structure questions', () => { + /** System-catalog columns are not in the schema block, so the model used to guess them. */ + it('offers a correct catalog query when one is given', () => { + const text = buildSqlUser({ + question: 'what tables exist?', + schemaText: 'TABLE film', + dialect: POSTGRES_DIALECT, + maxRows: 100, + catalogHint: 'SELECT table_name FROM information_schema.tables', + }); + + expect(text).toContain('SELECT table_name FROM information_schema.tables'); + }); + + it('says nothing about structure for an ordinary data question', () => { + const text = buildSqlUser({ question: 'how many films?', schemaText: 's', dialect: POSTGRES_DIALECT, maxRows: 10 }); + expect(text).not.toMatch(/about the database's structure/); + }); +}); + +describe('escaped quotes inside literals', () => { + const NAMES = ['Notes', 'Body', 'Customers', 'Author']; + + /** A doubled quote is SQL's escaped apostrophe; treating it as the close rewrote text inside the value. */ + it('does not rewrite identifiers inside a literal containing an escaped quote', () => { + const sql = "SELECT * FROM Notes WHERE Body = 'it''s about Customers' AND Author = 'x'"; + + expect(quoteCatalogIdentifiers(sql, NAMES, '"')).toBe( + 'SELECT * FROM "Notes" WHERE "Body" = \'it\'\'s about Customers\' AND "Author" = \'x\'', + ); + }); + + it('keeps scanning as code after the literal really ends', () => { + const sql = "SELECT Body FROM Notes WHERE Author = 'o''brien'"; + expect(quoteCatalogIdentifiers(sql, NAMES, '"')).toBe( + 'SELECT "Body" FROM "Notes" WHERE "Author" = \'o\'\'brien\'', + ); + }); + + /** correctTableCase shares the scanner, so it has the same hazard. */ + it('leaves a table name inside an escaped literal alone when repairing case', () => { + const sql = "SELECT * FROM notes WHERE Body = 'it''s notes'"; + expect(correctTableCase(sql, ['Notes'], '"', 'lower')).toBe( + 'SELECT * FROM "Notes" WHERE Body = \'it\'\'s notes\'', + ); + }); +}); + +describe('keywords that are syntax, not names', () => { + const NAMES = ['Date', 'Status', 'Orders', 'Month', 'Amount', 'order']; + const q = (sql: string) => quoteCatalogIdentifiers(sql, NAMES, '"'); + + /** Quoting the type turns CAST(x AS DATE) into a reference to a column that does not exist. */ + it('leaves the type in a CAST alone', () => { + expect(q('SELECT CAST(x AS Date) FROM t')).toBeNull(); + }); + + /** EXTRACT's first argument is a field keyword; the source after FROM is a real column. */ + it('quotes the source of an EXTRACT but not the field', () => { + expect(q('SELECT EXTRACT(month FROM Date) FROM t')).toBe('SELECT EXTRACT(month FROM "Date") FROM t'); + }); + + it('leaves the leading keyword of a TRIM alone', () => { + expect(q("SELECT TRIM(both 'x' FROM Status) FROM t")).toBe('SELECT TRIM(both \'x\' FROM "Status") FROM t'); + }); + + it('still quotes a reserved word in table position', () => { + expect(q('SELECT * FROM order')).toBe('SELECT * FROM "order"'); + }); + + it('still quotes a reserved word qualified by a dot', () => { + expect(q('SELECT t.order FROM t')).toBe('SELECT t."order" FROM t'); + }); + + it('still quotes an ordinary column inside a function call', () => { + expect(q('SELECT COUNT(Amount) FROM Orders')).toBe('SELECT COUNT("Amount") FROM "Orders"'); + }); +}); + +describe('unterminated text values', () => { + /** 'O'Brien' is what a model writes when it forgets to double the apostrophe. */ + it('spots an unescaped apostrophe', () => { + expect(hasUnterminatedLiteral("SELECT Note FROM Person WHERE Name = 'O'Brien'")).toBe(true); + }); + + it('accepts a correctly doubled apostrophe', () => { + expect(hasUnterminatedLiteral("SELECT Note FROM Person WHERE Name = 'O''Brien'")).toBe(false); + }); + + it.each([ + "SELECT * FROM t WHERE a = 'x' AND b = 'y'", + 'SELECT * FROM t', + "SELECT * FROM t -- it's fine", + "SELECT * FROM t /* it's fine */", + ])('accepts %s', (sql) => { + expect(hasUnterminatedLiteral(sql)).toBe(false); + }); + + it('spots a value that runs off the end', () => { + expect(hasUnterminatedLiteral("SELECT * FROM t WHERE a = 'oops")).toBe(true); + }); +}); + +describe('literals and qualifiers the rewriter must not touch', () => { + const NAMES = ['Users', 'Notes', 'Sales']; + + /** Postgres and DuckDB dollar-quote bodies, which may contain anything at all. */ + it('leaves a dollar-quoted body alone', () => { + expect(quoteCatalogIdentifiers('SELECT * FROM Users WHERE Notes = $$from users now$$', NAMES, '"')).toBe( + 'SELECT * FROM "Users" WHERE "Notes" = $$from users now$$', + ); + }); + + it('leaves a tagged dollar-quoted body alone', () => { + expect(quoteCatalogIdentifiers('SELECT * FROM Users WHERE Notes = $x$from users$x$', NAMES, '"')).toBe( + 'SELECT * FROM "Users" WHERE "Notes" = $x$from users$x$', + ); + }); + + /** Quoting a schema qualifier turns a working query into "schema does not exist". */ + it('does not quote a qualifier that is not a table', () => { + expect(quoteCatalogIdentifiers('SELECT * FROM sales.orders', ['Sales'], '"', [])).toBeNull(); + }); + + it('still quotes a qualifier that is a real table', () => { + expect( + quoteCatalogIdentifiers('SELECT Customers.FirstName FROM Customers', ['Customers', 'FirstName'], '"', [ + 'Customers', + ]), + ).toBe('SELECT "Customers"."FirstName" FROM "Customers"'); + }); + + /** In prod.sales.orders the table is orders; sales is a qualifier. */ + it('does not recase the middle of a three-part name', () => { + expect(correctTableCase('SELECT * FROM prod.sales.orders', ['Sales'], '"', 'lower')).toBeNull(); + }); +}); + +describe('dialect-specific literal rules', () => { + const BS = String.fromCharCode(92); + + /** Backslash escapes a quote in MySQL only; Postgres, Oracle, SQLite and DuckDB read it literally. */ + it('treats a backslash as an escape only for the backtick dialect', () => { + expect(hasUnterminatedLiteral(`SELECT * FROM t WHERE n = 'O${BS}'Brien'`, true)).toBe(false); + expect(hasUnterminatedLiteral(`SELECT * FROM t WHERE n = 'O${BS}'Brien'`)).toBe(true); + }); + + it('still spots a genuinely unescaped apostrophe', () => { + expect(hasUnterminatedLiteral("SELECT * FROM t WHERE n = 'O'Brien'")).toBe(true); + }); + + /** A dollar-quoted body needs no escaping, so an apostrophe inside it is not a defect. */ + it('does not flag an apostrophe inside a dollar-quoted body', () => { + expect(hasUnterminatedLiteral("SELECT * FROM t WHERE n = $$don't$$")).toBe(false); + }); + + /** TIMESTAMP '2024-01-01' is one typed literal; quoting the word makes it a missing column. */ + it.each([ + ["SELECT * FROM t WHERE created > TIMESTAMP '2024-01-01'", 'Timestamp'], + ["SELECT * FROM t WHERE d > DATE '2024-01-01'", 'Date'], + ])('leaves a typed literal alone: %s', (sql, name) => { + expect(quoteCatalogIdentifiers(sql, [name], '"')).toBeNull(); + }); + + it('still quotes the same word used as a real column', () => { + expect(quoteCatalogIdentifiers('SELECT Timestamp FROM t', ['Timestamp'], '"')).toBe('SELECT "Timestamp" FROM t'); + }); +}); + +describe('names split across literal boundaries', () => { + const NAMES = ['Customers', 'FirstName', 'Country', 'City', 'Timestamp']; + + /** + * Literals split the statement into segments, and the typed-literal check reads the whole + * statement rather than one segment. If it read the segment, every name sitting at a segment + * boundary would silently stop being quoted. + */ + it('quotes names before, between and after several literals', () => { + const sql = "SELECT FirstName FROM Customers WHERE Country = 'UK' AND City = 'York' AND FirstName <> 'x'"; + + expect(quoteCatalogIdentifiers(sql, NAMES, '"')).toBe( + 'SELECT "FirstName" FROM "Customers" WHERE "Country" = \'UK\' AND "City" = \'York\' AND "FirstName" <> \'x\'', + ); + }); + + it('quotes a name that ends the statement, with a literal earlier', () => { + expect(quoteCatalogIdentifiers("SELECT * FROM Customers WHERE Country = 'UK' ORDER BY City", NAMES, '"')).toBe( + 'SELECT * FROM "Customers" WHERE "Country" = \'UK\' ORDER BY "City"', + ); + }); + + /** The one word that must not be quoted is the one a literal directly follows. */ + it('skips only the typed literal, quoting every other name in the same statement', () => { + const sql = "SELECT FirstName FROM Customers WHERE created > TIMESTAMP '2024-01-01'"; + + expect(quoteCatalogIdentifiers(sql, NAMES, '"')).toBe( + 'SELECT "FirstName" FROM "Customers" WHERE created > TIMESTAMP \'2024-01-01\'', + ); + }); +}); diff --git a/packages/core/test/nested-aggregate.test.ts b/packages/core/test/nested-aggregate.test.ts new file mode 100644 index 0000000..886cc0e --- /dev/null +++ b/packages/core/test/nested-aggregate.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from 'vitest'; +import { nestedAggregate } from '../src/semantics.js'; + +describe('nestedAggregate', () => { + /** Reported by a real question: AVG over a SUM is rejected by every engine. */ + it('flags an aggregate inside another aggregate', () => { + const sql = + 'SELECT c.country, AVG(o.freight + SUM(od.unit_price)) AS v FROM orders o ' + + 'JOIN customers c ON o.customer_id = c.customer_id JOIN order_details od ON o.order_id = od.order_id GROUP BY c.country'; + expect(nestedAggregate(sql, 'Postgresql')).toBe('AVG'); + }); + + it('allows aggregates side by side', () => { + expect(nestedAggregate('SELECT SUM(a), AVG(b) FROM t', 'Postgresql')).toBeNull(); + }); + + it('allows an aggregate over an expression', () => { + expect(nestedAggregate('SELECT SUM(a * b + 1) FROM t', 'Postgresql')).toBeNull(); + }); + + /** A subquery has its own scope, so its aggregate is not nested in the outer call. */ + it('allows an aggregate inside a subquery argument', () => { + expect(nestedAggregate('SELECT SUM((SELECT COUNT(*) FROM u WHERE u.id = t.id)) FROM t', 'Postgresql')).toBeNull(); + }); + + it('returns null for unparsable sql rather than blocking', () => { + expect(nestedAggregate('NOT SQL AT ALL', 'Postgresql')).toBeNull(); + }); + + it('flags nesting in HAVING', () => { + expect(nestedAggregate('SELECT a FROM t GROUP BY a HAVING SUM(COUNT(b)) > 1', 'Postgresql')).toBe('SUM'); + }); + + it('flags nesting in ORDER BY', () => { + expect(nestedAggregate('SELECT a FROM t GROUP BY a ORDER BY AVG(SUM(b))', 'Postgresql')).toBe('AVG'); + }); +}); diff --git a/packages/core/test/sql-keywords.test.ts b/packages/core/test/sql-keywords.test.ts new file mode 100644 index 0000000..ba78e04 --- /dev/null +++ b/packages/core/test/sql-keywords.test.ts @@ -0,0 +1,35 @@ +import { describe, expect, it } from 'vitest'; +import { reservedWordsFor } from '../src/sql-keywords.js'; + +/** Read from each engine's own catalog, so the lists differ; a shared guess is what they replaced. */ +describe('reservedWordsFor', () => { + it('gives each engine its own list', () => { + expect(reservedWordsFor('mysql').size).toBeGreaterThan(reservedWordsFor('postgres').size); + }); + + it.each([ + ['postgres', 'select'], + ['mysql', 'select'], + ['oracle', 'select'], + ['sqlite', 'select'], + ['duckdb', 'select'], + ])('%s reserves %s', (engine, word) => { + expect(reservedWordsFor(engine).has(word)).toBe(true); + }); + + /** MySQL reserves it, Postgres does not: proof the lists are not one shared set. */ + it('separates a word that only some engines reserve', () => { + expect(reservedWordsFor('mysql').has('rank')).toBe(true); + expect(reservedWordsFor('postgres').has('rank')).toBe(false); + }); + + /** An unknown engine gets the union, so a name is over-quoted rather than left broken. */ + it('falls back to the union for an unknown engine', () => { + const union = reservedWordsFor('not-an-engine'); + expect(union.size).toBeGreaterThan(reservedWordsFor('mysql').size); + }); + + it('is case-insensitive on the engine name', () => { + expect(reservedWordsFor('PostgreSQL'.toLowerCase().slice(0, 8)).size).toBe(reservedWordsFor('postgres').size); + }); +}); diff --git a/packages/core/test/table-case-repair.test.ts b/packages/core/test/table-case-repair.test.ts new file mode 100644 index 0000000..7c4247d --- /dev/null +++ b/packages/core/test/table-case-repair.test.ts @@ -0,0 +1,219 @@ +/** + * The engine wiring for the catalog-driven case repair: a wrong-cased table must come back with a + * corrected query, and must not spend a model round trip getting there. + */ +import { describe, expect, it, vi } from 'vitest'; +import { createAskSql, firstUnknownColumn, firstUnknownTable } from '../src/engine.js'; +import { AskSqlError } from '../src/errors.js'; +import { MYSQL_DIALECT, POSTGRES_DIALECT } from '../src/dialects.js'; +import type { Connector, CustomModel, ResultSet, SchemaCatalog } from '../src/types.js'; + +const CATALOG: SchemaCatalog = { + engine: 'mysql', + schemas: ['public'], + tables: [ + { + name: 'Users', + kind: 'table', + columns: [ + { name: 'id', dbType: 'bigint', nullable: false }, + { name: 'name', dbType: 'text', nullable: false }, + ], + primaryKey: ['id'], + foreignKeys: [], + uniques: [], + checks: [], + indexes: [], + source: 'db', + }, + ], + enums: [], + sequences: [], + triggers: [], + routines: [], + warnings: [], + fetchedAt: 'now', +}; + +function connThatRejects(detail: string, execute?: () => Promise): Connector { + return { + // MySQL folds nothing, so generation leaves the case alone and the repair path is what runs. + engine: 'mysql', + dialect: MYSQL_DIALECT, + capabilities: { + supportsCancel: true, + supportsExplain: true, + supportsSchemas: true, + readOnlySession: true, + supportsMatViews: true, + supportsTriggers: true, + supportsRoutines: true, + }, + id: 'db', + name: 'DB', + async connect() {}, + async close() {}, + async introspect() { + return CATALOG; + }, + execute: + execute ?? + (async () => { + throw new AskSqlError('DB_QUERY_ERROR', { userMessage: 'query failed', detail }); + }), + } as unknown as Connector; +} + +const modelSaying = + (sql: string): CustomModel => + async () => + `\`\`\`sql\n${sql}\n\`\`\``; + +describe('wrong-cased table repair', () => { + it('suggests the catalog spelling when the database rejects the case', async () => { + const conn = connThatRejects('relation "users" does not exist'); + const engine = createAskSql({ connectors: [conn], model: modelSaying('SELECT * FROM users') }); + + const answer = await engine.ask('list the users'); + const err = await answer.run().catch((e: unknown) => e); + + expect(AskSqlError.is(err)).toBe(true); + expect((err as { suggestedSql?: string }).suggestedSql).toContain('`Users`'); + }); + + it('derives the fix from the catalog rather than a second model call', async () => { + const model = vi.fn(async () => '```sql\nSELECT * FROM users\n```'); + const engine = createAskSql({ + connectors: [connThatRejects('relation "users" does not exist')], + model: model as unknown as CustomModel, + }); + + await engine + .ask('list the users') + .then((a) => a.run()) + .catch(() => undefined); + + expect(model).toHaveBeenCalledTimes(1); // the ask itself, with no repair round trip + }); + + /** A failure the catalog cannot explain still belongs to the model repair. */ + it('leaves an unrelated database error to the model', async () => { + const model = vi.fn(async () => '```sql\nSELECT * FROM `Users`\n```'); + const engine = createAskSql({ + connectors: [connThatRejects('column "nope" does not exist')], + model: model as unknown as CustomModel, + }); + + await engine + .ask('list the users') + .then((a) => a.run()) + .catch(() => undefined); + + expect(model.mock.calls.length).toBeGreaterThan(1); // ask, then the repair attempt + }); + + it('does not suggest anything when the query already matches the catalog', async () => { + let calls = 0; + const conn = connThatRejects('', async () => { + calls++; + throw new AskSqlError('DB_QUERY_ERROR', { userMessage: 'boom', detail: 'deadlock detected' }); + }); + const engine = createAskSql({ connectors: [conn], model: modelSaying('SELECT * FROM `Users`') }); + + const answer = await engine.ask('list the users'); + const err = await answer.run().catch((e: unknown) => e); + + expect(calls).toBe(1); + expect((err as { suggestedSql?: string }).suggestedSql).toBeUndefined(); + }); + + /** On a folding engine the query is corrected before it runs, so nothing has to fail first. */ + it('quotes a folded name at generation time on Postgres', async () => { + const pg = { + ...connThatRejects('unused'), + engine: 'postgres', + dialect: POSTGRES_DIALECT, + async execute() { + return { columns: [], rows: [], rowCount: 0, truncated: false, durationMs: 1, warnings: [] }; + }, + } as unknown as Connector; + const engine = createAskSql({ connectors: [pg], model: modelSaying('SELECT name FROM users') }); + + const answer = await engine.ask('list the users'); + + expect(answer.sql).toContain('"Users"'); + }); +}); + +describe('unknown-column floor on set operations', () => { + const catalog = { + engine: 'sqlite', + schemas: [], + tables: [ + { name: 'Album', kind: 'table', columns: [{ name: 'AlbumId', dbType: 'int', nullable: false }], primaryKey: [], foreignKeys: [], uniques: [], checks: [], indexes: [], source: 'db' }, + { name: 'Artist', kind: 'table', columns: [{ name: 'ArtistId', dbType: 'int', nullable: false }], primaryKey: [], foreignKeys: [], uniques: [], checks: [], indexes: [], source: 'db' }, + ], + enums: [], sequences: [], triggers: [], routines: [], warnings: [], fetchedAt: 'now', + } as unknown as SchemaCatalog; + + /** Per-table row counts are a normal DBA question, and this blocked them outright. */ + it('does not flag a UNION ALL of per-table counts', () => { + const sql = + "SELECT 'Album' AS TableName, COUNT(*) AS RowCount FROM Album " + + "UNION ALL SELECT 'Artist' AS TableName, COUNT(*) AS RowCount FROM Artist"; + + expect(firstUnknownColumn(sql, catalog, 'sqlite')).toBeNull(); + }); + + it('still flags a hallucinated column on a plain select', () => { + expect(firstUnknownColumn('SELECT Nope FROM Album', catalog, 'sqlite')).not.toBeNull(); + }); +}); + +describe('set-operation detection ignores literals', () => { + const catalog = { + engine: 'postgres', + schemas: [], + tables: [ + { name: 'notes', kind: 'table', columns: [{ name: 'body', dbType: 'text', nullable: false }], primaryKey: [], foreignKeys: [], uniques: [], checks: [], indexes: [], source: 'db' }, + ], + enums: [], sequences: [], triggers: [], routines: [], warnings: [], fetchedAt: 'now', + } as unknown as SchemaCatalog; + + /** A value containing "except" once disabled the column floor for an ordinary query. */ + it('still flags a hallucinated column when a literal contains a set-operation word', () => { + const sql = "SELECT nope FROM notes WHERE body = 'except this'"; + expect(firstUnknownColumn(sql, catalog, 'Postgresql')).not.toBeNull(); + }); + + it('still skips attribution for a real set operation', () => { + const sql = "SELECT nope FROM notes UNION ALL SELECT body FROM notes"; + expect(firstUnknownColumn(sql, catalog, 'Postgresql')).toBeNull(); + }); +}); + +describe('catalog-driven guards', () => { + const base = { + engine: 'postgres', + schemas: [], + enums: [], sequences: [], triggers: [], routines: [], warnings: [], fetchedAt: 'now', + }; + const table = (name: string, cols: string[]) => ({ + name, kind: 'table', + columns: cols.map((c) => ({ name: c, dbType: 'text', nullable: true })), + primaryKey: [], foreignKeys: [], uniques: [], checks: [], indexes: [], source: 'db', + }); + + /** A quoted CTE was read as a hallucinated table, rejecting a valid query. */ + it('recognises a CTE whose name is quoted', () => { + const catalog = { ...base, tables: [table('Orders', ['Amount'])] } as unknown as SchemaCatalog; + const sql = 'WITH "Amount" AS (SELECT "Amount" FROM "Orders") SELECT * FROM "Amount"'; + + expect(firstUnknownTable(sql, catalog, 'Postgresql')).toBeNull(); + }); + + it('still reports a genuinely unknown table', () => { + const catalog = { ...base, tables: [table('Orders', ['Amount'])] } as unknown as SchemaCatalog; + expect(firstUnknownTable('SELECT * FROM nosuchtable', catalog, 'Postgresql')).toBe('nosuchtable'); + }); +}); diff --git a/packages/jetbrains/CHANGELOG.md b/packages/jetbrains/CHANGELOG.md index 5ab214d..fa4d318 100644 --- a/packages/jetbrains/CHANGELOG.md +++ b/packages/jetbrains/CHANGELOG.md @@ -3,7 +3,27 @@ All notable changes to the AskSQL JetBrains plugin are documented here. Format loosely follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). -## [Unreleased] +## [0.5.2] - 2026-08-14 + +### Fixed +- A mixed-case Postgres schema failed every query. An unquoted name folds to lower case and resolves + to nothing; Oracle folds the other way, and MySQL on Linux compares table names case-sensitively. + Table and column names are now quoted from the catalog before the query is validated, and correct + names are left untouched. +- Reserved words now come from each database itself rather than one shared list that applied MySQL's + rules to Postgres and missed most of MySQL's own. +- Quoting knows where a word is syntax rather than a name, so `CAST(x AS DATE)` and + `EXTRACT(MONTH FROM d)` are left alone. +- A table named like a parser keyword, such as `order` or `Nulls`, could not be parsed in its bare + form, so the question failed after three attempts. +- An apostrophe inside a value, as in `'O'Brien'`, produced only "could not parse", so the model + returned the same statement until it ran out of attempts. It is now told to double the quote. +- Questions about structure no longer invent a database name. Without being told which database it + is connected to, the model wrote `table_schema = 'your_database_name'` and returned nothing, which + reads as an empty database rather than an error. +- `AVG(SUM(x))` is repaired instead of run. Nested aggregates are invalid in every engine. +- Per-table row counts work again. A `UNION ALL` across tables was blocked as a hallucinated column, + because each branch's columns were judged against every branch's tables. ## [0.5.1] - 2026-08-12 diff --git a/packages/jetbrains/gradle.properties b/packages/jetbrains/gradle.properties index 93a7ad8..68f9e15 100644 --- a/packages/jetbrains/gradle.properties +++ b/packages/jetbrains/gradle.properties @@ -4,7 +4,7 @@ pluginGroup = com.rahulmahadik.asksql pluginName = AskSQL -pluginVersion = 0.5.1 +pluginVersion = 0.5.2 # IntelliJ Platform target used to COMPILE and RUN the sandbox. Broad # compatibility is governed by pluginSinceBuild/pluginUntilBuild in diff --git a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/CatalogPruner.kt b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/CatalogPruner.kt index a90c3f9..59a2b98 100644 --- a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/CatalogPruner.kt +++ b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/CatalogPruner.kt @@ -68,24 +68,14 @@ object CatalogPruner { */ private fun quoteCharFor(engine: EngineKind): Char = if (engine == EngineKind.MYSQL) '`' else '"' - /** Words an engine will not accept as a bare identifier; the ones that turn up as real column names, not every dialect. */ - private val RESERVED_WORDS = ( - "select from where group by order having limit offset union all distinct join inner outer left " + - "right full cross natural on using as into insert update delete set values create drop alter " + - "table column view index key primary foreign unique constraint references default check null " + - "not and or in is like between case when then else end exists any some cast collate with " + - "recursive returning window over partition range rows current session system user grant revoke " + - "to begin commit rollback transaction lock database schema trigger procedure function " + - "desc asc date time timestamp interval level size type comment position language" - ).split(" ").toSet() - + /** * True when the engine would not read the bare name back as itself: an unquoted identifier * folds case - PostgreSQL to lower, Oracle to upper. */ - private fun needsQuoting(name: String, engine: EngineKind): Boolean { + internal fun needsQuoting(name: String, engine: EngineKind): Boolean { if (!PLAIN_IDENTIFIER_RE.matches(name)) return true - if (name.lowercase() in RESERVED_WORDS) return true + if (name.lowercase() in SqlKeywords.reservedWordsFor(engine.name)) return true return when (engine) { EngineKind.ORACLE -> name != name.uppercase() // MySQL, SQLite and DuckDB match identifiers case-insensitively. diff --git a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/EnginePipeline.kt b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/EnginePipeline.kt index 51b1586..7bc9763 100644 --- a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/EnginePipeline.kt +++ b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/EnginePipeline.kt @@ -335,6 +335,16 @@ class EnginePipeline( onEvent?.onEvent(EngineEvent.StageEvent(Stage.CATALOG)) val fullCatalog = catalog(descriptor, password) + // Names the engine would not read back as themselves: folded case, reserved words, symbols. + // A name spelled two ways across the catalog is skipped: rewriting "status" to "Status" would + // ask one table for another table's column. + val allNames = fullCatalog.tables.flatMap { t -> listOf(t.name) + t.columns.map { it.name } } + val spellings = allNames.groupBy { it.lowercase() } + val quotableNames = allNames.filter { + CatalogPruner.needsQuoting(it, descriptor.engine) && spellings[it.lowercase()]?.distinct()?.size == 1 + } + // Only a table may be quoted before a dot; a schema qualifier that matched a column name broke it. + val quotableTables = fullCatalog.tables.map { it.name }.filter { it in quotableNames } onEvent?.onEvent(EngineEvent.StageEvent(Stage.PRUNE)) val initialPrunerSettings = CatalogPruner.PrunerSettings(maxSchemaTokens = maxSchemaTokens) @@ -351,6 +361,9 @@ class EnginePipeline( glossary = glossary, context = context, rerunPrevious = isRerunPreviousRequest(q), + database = descriptor.database, + schemas = fullCatalog.schemas, + catalogHint = if (isMetadataQuestion(q)) catalogQueryHint(descriptor.engine) else null, ) var lastSql = "" @@ -382,7 +395,10 @@ class EnginePipeline( ) pruned = tighter schemaText = tighter.schemaText - userPrompt = Prompts.buildSqlUser(question = q, schemaText = schemaText, context = context) + userPrompt = Prompts.buildSqlUser( + question = q, schemaText = schemaText, context = context, + database = descriptor.database, schemas = fullCatalog.schemas, + ) continue } throw e @@ -449,7 +465,13 @@ class EnginePipeline( lastSql = extraction.sql onEvent?.onEvent(EngineEvent.StageEvent(Stage.GUARD)) - val verdict = SqlGuard.guard(extraction.sql, dialect, policy) + // Quote first: a folding engine resolves a bare name elsewhere, and the parser cannot read a + // bare table named like a keyword. Falls back untouched if quoting makes it unparseable. + val normalised = + IdentifierCase.quoteCatalogIdentifiers(extraction.sql, quotableNames, dialect.quoteChar, quotableTables) + val normalisedVerdict = normalised?.let { SqlGuard.guard(it, dialect, policy) } + val verdict = if (normalisedVerdict?.allowed == true) normalisedVerdict + else SqlGuard.guard(extraction.sql, dialect, policy) if (!verdict.allowed) { if (attempt >= MAX_REPAIRS) { history.add(auditEntry(descriptor.id, q, extraction.sql, HistoryStatus.BLOCKED, verdict.ruleId)) @@ -459,9 +481,13 @@ class EnginePipeline( detail = "ruleId=${verdict.ruleId} after ${attempt + 1} attempts", ) } + // "could not parse" alone leaves the model repeating the same statement; name the real cause. + val quoteHint = if (IdentifierCase.hasUnterminatedLiteral(extraction.sql, dialect.quoteChar == '`')) { + " A text value contains an apostrophe that is not escaped: write it doubled, as 'O''Brien'." + } else "" userPrompt = Prompts.buildRepairUser( question = q, failedSql = extraction.sql, - failure = "The SQL validator rejected it: ${verdict.reason ?: verdict.ruleId ?: "not allowed"}. Produce a single read-only SELECT.", + failure = "The SQL validator rejected it: ${verdict.reason ?: verdict.ruleId ?: "not allowed"}.$quoteHint Produce a single read-only SELECT.", schemaText = schemaText, dialect = dialect, ) attempt++ @@ -492,9 +518,12 @@ class EnginePipeline( retryable = false, ) } + // The column repair already names the real columns; give the table repair the same head start. + val nearest = SchemaFuzzyMatch.closestTableName(unknownTable, fullCatalog) + val didYouMean = if (nearest != null) " Did you mean \"$nearest\"?" else "" userPrompt = Prompts.buildRepairUser( question = q, failedSql = verdict.sql, - failure = "Table \"$unknownTable\" does not exist in the schema. Use only tables from the block.", + failure = "Table \"$unknownTable\" does not exist in the schema.$didYouMean Use only tables from the block.", schemaText = schemaText, dialect = dialect, ) attempt++ @@ -516,6 +545,19 @@ class EnginePipeline( continue } + // Semantic floor: AVG(SUM(x)) and friends. Every engine rejects it, so repair before executing. + val nested = Semantics.nestedAggregate(verdict.sql) + if (nested != null && attempt < MAX_REPAIRS) { + userPrompt = Prompts.buildRepairUser( + question = q, failedSql = verdict.sql, + failure = "$nested() contains another aggregate, which no SQL engine allows. Aggregate once " + + "over the rows, or aggregate the inner result in a subquery or CTE and then aggregate that.", + schemaText = schemaText, dialect = dialect, + ) + attempt++ + continue + } + val unknownColumn = HallucinationChecks.firstUnknownColumn(verdict.sql, fullCatalog) if (unknownColumn != null) { if (attempt >= MAX_REPAIRS) { @@ -613,6 +655,17 @@ class EnginePipeline( return try { val dialect = Dialects.of(descriptor.engine) val catalog = catalog(descriptor, password) + // A wrong-cased table is repairable from the catalog alone, so try that before the model. + if (errorDetail != null && IdentifierCase.looksLikeUnknownTable(errorDetail)) { + val cased = IdentifierCase.correctTableCase( + bad, catalog.tables.map { it.name }, dialect.quoteChar, + IdentifierCase.foldingFor(descriptor.engine.name), + ) + if (cased != null) { + val casedVerdict = SqlGuard.guard(cased, dialect, policy) + if (casedVerdict.allowed) return casedVerdict.sql + } + } val schemaText = CatalogPruner.pruneCatalog(catalog, q).schemaText val repairPrompt = Prompts.buildRepairUser( question = q, failedSql = bad, diff --git a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/HallucinationChecks.kt b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/HallucinationChecks.kt index 1993f47..4f00893 100644 --- a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/HallucinationChecks.kt +++ b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/HallucinationChecks.kt @@ -19,9 +19,15 @@ object HallucinationChecks { data class UnknownColumn(val table: String, val column: String, val available: List) private val SYSTEM_SCHEMAS = setOf("information_schema", "pg_catalog", "mysql", "performance_schema", "sys") - private val CTE_NAME = Regex("""([A-Za-z_][A-Za-z0-9_]*)\s+as\s*\(""", RegexOption.IGNORE_CASE) + /** + * The quote characters matter: normalisation may have quoted a CTE named like a catalog column, + * and a model can quote one itself. An unrecognised CTE reads as a hallucinated table. + */ + private val CTE_NAME = + Regex("""["`\[]?([A-Za-z_][A-Za-z0-9_]*)["`\]]?\s+as\s*\(""", RegexOption.IGNORE_CASE) private val HAS_WITH = Regex("""\bwith\b""", RegexOption.IGNORE_CASE) private val SELECT_ALIAS_RE = Regex("""\bas\s+["'`]?([A-Za-z_][A-Za-z0-9_]*)["'`]?""", RegexOption.IGNORE_CASE) + private val SET_OPERATION_RE = Regex("""\b(union|intersect|except)\b""", RegexOption.IGNORE_CASE) private val SUBQUERY_OPEN_RE = Regex("""\(\s*select\b""", RegexOption.IGNORE_CASE) /** Scans the whole statement; over-collecting CTE names only makes the floor more lenient. */ @@ -85,8 +91,14 @@ object HallucinationChecks { val aliases = SELECT_ALIAS_RE.findAll(sql).map { it.groupValues[1].lowercase() }.toSet() val tableAliases = collectTableAliases(statement) - val hasSubquery = SUBQUERY_OPEN_RE.containsMatchIn(sql) - var attributable = !hasSubquery + // Both probes read blanked text: a literal or comment must not disable the floor. + val code = withoutLiterals(sql) + val hasSubquery = SUBQUERY_OPEN_RE.containsMatchIn(code) + // A set operation has one column list per branch, and the parser reports them merged, so a column + // from one branch would be judged against another branch's tables. Not attributable, like a subquery. + // Blank literals first: a value like 'except this' would otherwise disable the floor entirely. + val hasSetOperation = SET_OPERATION_RE.containsMatchIn(code) + var attributable = !hasSubquery && !hasSetOperation val queryTables = mutableListOf() val tableNames = try { @@ -198,4 +210,37 @@ object HallucinationChecks { else -> Unit } } + + /** The statement with string literals and line comments blanked out; offsets are preserved. */ + private fun withoutLiterals(sql: String): String { + val out = StringBuilder(sql) + var i = 0 + while (i < sql.length) { + when { + sql[i] == '\'' -> { + var j = i + 1 + while (j < sql.length) { + if (sql[j] == '\'' && j + 1 < sql.length && sql[j + 1] == '\'') j += 2 + else if (sql[j] == '\'') break + else j++ + } + for (k in i..minOf(j, sql.length - 1)) out[k] = ' ' + i = j + 1 + } + sql[i] == '-' && i + 1 < sql.length && sql[i + 1] == '-' -> { + var j = i + while (j < sql.length && sql[j] != '\n') { out[j] = ' '; j++ } + i = j + } + sql[i] == '/' && i + 1 < sql.length && sql[i + 1] == '*' -> { + val close = sql.indexOf("*/", i + 2) + val end = if (close == -1) sql.length else close + 2 + for (k in i until end) out[k] = ' ' + i = end + } + else -> i++ + } + } + return out.toString() + } } diff --git a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/IdentifierCase.kt b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/IdentifierCase.kt new file mode 100644 index 0000000..c0b0f0b --- /dev/null +++ b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/IdentifierCase.kt @@ -0,0 +1,259 @@ +package com.rahulmahadik.asksql.ide.engine + +/** + * A table named in the wrong case is the one database rejection a catalog can repair on its own, + * with no model round trip. MySQL on Linux compares table names case-sensitively, and Postgres and + * Oracle fold unquoted names, so the corrected name is quoted to survive either rule. + */ +object IdentifierCase { + + /** The union across engines: treating a word as syntax is the conservative side of this decision. */ + private val ANY_RESERVED = SqlKeywords.reservedWordsFor("*") + + /** Only these lead-ins put an identifier in table position, which keeps a same-named column alone. */ + private val TABLE_POSITION = + Regex("""\b(from|join|update|into)(\s+)([`"\[]?)([A-Za-z_][\w$]*)[`"\]]?(\s*\.\s*([`"\[]?)([A-Za-z_][\w$]*)[`"\]]?)?""", RegexOption.IGNORE_CASE) + + private val UNKNOWN_TABLE = + Regex("""\b(doesn't exist|does not exist|not found|unknown table|invalid object name|undefined table|no such table|table or view does not exist)\b""", RegexOption.IGNORE_CASE) + + fun looksLikeUnknownTable(message: String): Boolean = UNKNOWN_TABLE.containsMatchIn(message) + + private fun quoted(name: String, quoteChar: Char): String { + val close = if (quoteChar == '[') ']' else quoteChar + return "$quoteChar$name$close" + } + + /** Where a literal or comment ends, or -1 when the position starts neither. */ + /** A dollar-quoted body is a literal in Postgres and DuckDB, and may contain anything. */ + private val DOLLAR_OPEN = Regex("""\$[A-Za-z_]\w*\$|\$\$""") + + private fun skipTo(sql: String, i: Int, doubleQuoteIsLiteral: Boolean, backslashEscapes: Boolean = false): Int { + val ch = sql[i] + val next = if (i + 1 < sql.length) sql[i + 1] else ' ' + if (ch == '$') { + val open = DOLLAR_OPEN.matchAt(sql, i) + if (open != null) { + val close = sql.indexOf(open.value, i + open.value.length) + return if (close == -1) sql.length else close + open.value.length + } + } + if (ch == '-' && next == '-') { + val nl = sql.indexOf('\n', i) + return if (nl == -1) sql.length else nl + } + if (ch == '/' && next == '*') { + val close = sql.indexOf("*/", i + 2) + return if (close == -1) sql.length else close + 2 + } + if (ch == '\'' || (ch == '"' && doubleQuoteIsLiteral)) { + var j = i + 1 + while (j < sql.length) { + when { + backslashEscapes && sql[j] == '\\' -> j += 2 + sql[j] == ch -> if (j + 1 < sql.length && sql[j + 1] == ch) j += 2 else return j + 1 + else -> j++ + } + } + return sql.length + } + return -1 + } + + /** How an engine resolves an unquoted identifier: Postgres lower-cases it, Oracle upper-cases it. */ + enum class Folding { LOWER, UPPER, NONE } + + fun foldingFor(engine: String): Folding = when (engine.lowercase()) { + "postgres" -> Folding.LOWER + "oracle" -> Folding.UPPER + else -> Folding.NONE + } + + private fun folded(name: String, folding: Folding): String = when (folding) { + Folding.LOWER -> name.lowercase() + Folding.UPPER -> name.uppercase() + Folding.NONE -> name + } + + /** Returns the rewritten SQL, or null when no table reference needed correcting. */ + fun correctTableCase( + sql: String, + tableNames: List, + quoteChar: Char, + folding: Folding = Folding.NONE, + ): String? { + val byLower = HashMap() + for (name in tableNames) { + val lower = name.lowercase() + // An ambiguous fold has no single right answer, so leave those names untouched. + if (byLower.containsKey(lower) && byLower[lower] != name) byLower[lower] = "" else byLower.putIfAbsent(lower, name) + } + + var changed = false + fun fixCode(code: String): String = TABLE_POSITION.replace(code) { m -> + val keyword = m.groupValues[1] + val gap = m.groupValues[2] + val open = m.groupValues[3] + val first = m.groupValues[4] + val second = m.groupValues[7] + val target = second.ifEmpty { first } + val open2 = m.groupValues[6] + val canonical = byLower[target.lowercase()] + // A third part means what matched is a qualifier: prod.sales.orders names orders, not sales. + val restStartsWithDot = code.substring(m.range.last + 1).trimStart().startsWith(".") + // An unquoted name is resolved folded, so what matters is what the database will look up. + val wasQuoted = (if (second.isEmpty()) open else open2).isNotEmpty() + val resolvesTo = if (wasQuoted) target else folded(target, folding) + if (canonical.isNullOrEmpty() || restStartsWithDot || resolvesTo == canonical) { + m.value + } else { + changed = true + val fixed = quoted(canonical, quoteChar) + if (second.isEmpty()) "$keyword$gap$fixed" + else "$keyword$gap${if (open.isNotEmpty()) quoted(first, open[0]) else first}.$fixed" + } + } + + // A double quote is a string in MySQL but an identifier in Postgres, so the dialect decides. + val doubleQuoteIsLiteral = quoteChar != '"' + // MySQL is the only engine here that escapes with a backslash, and the backtick identifies it. + val backslashEscapes = quoteChar == '`' + val out = StringBuilder() + var start = 0 + var i = 0 + while (i < sql.length) { + val end = skipTo(sql, i, doubleQuoteIsLiteral, backslashEscapes) + if (end >= 0) { + out.append(fixCode(sql.substring(start, i))).append(sql, i, end) + start = end + i = end + } else i++ + } + out.append(fixCode(sql.substring(start))) + return if (changed) out.toString() else null + } + + /** A bare identifier, and whatever follows it, so a function call can be told from a column. */ + private val BARE_IDENTIFIER = Regex("""([A-Za-z_][\w$]*)(\s*[.(]?)""") + + /** + * A reserved word is only treated as a name where it cannot be syntax: after FROM/JOIN/UPDATE/INTO, + * or qualified by a dot. Accepting AS or "(" quoted the type in CAST(x AS DATE) and the field in + * EXTRACT(MONTH FROM d), both of which are valid SQL that quoting breaks. + */ + private val NAME_POSITION = Regex("""(?:\bfrom|\bjoin|\bupdate|\binto|\.)\s*$""", RegexOption.IGNORE_CASE) + + /** The first argument of these is a keyword, not a name: EXTRACT(MONTH FROM d), TRIM(BOTH x FROM s). */ + private val KEYWORD_ARGUMENT = + Regex("""\b(?:extract|trim|position|overlay|substring)\s*\(\s*$""", RegexOption.IGNORE_CASE) + + /** + * Quotes every table and column the engine would not read back as itself. The schema text already + * shows these names quoted and models still drop the quotes, so the query is normalised before it + * runs rather than left to fail. Returns null when nothing needed quoting. + */ + fun quoteCatalogIdentifiers( + sql: String, + names: List, + quoteChar: Char, + tableNames: List = names, + ): String? { + val tables = tableNames.map { it.lowercase() }.toSet() + val byLower = HashMap() + for (name in names) { + val lower = name.lowercase() + if (byLower.containsKey(lower) && byLower[lower] != name) byLower[lower] = "" else byLower.putIfAbsent(lower, name) + } + if (byLower.isEmpty()) return null + + var changed = false + fun fixCode(code: String, chunkStart: Int): String = BARE_IDENTIFIER.replace(code) { m -> + val token = m.groupValues[1] + val tail = m.groupValues[2] + val canonical = byLower[token.lowercase()] + val before = code.substring(0, m.range.first) + // TIMESTAMP '2024-01-01' is one typed literal; the literal is its own segment, so this + // reads the statement rather than the chunk. + val typedLiteral = sql.drop(chunkStart + m.range.first + token.length).trimStart().startsWith("'") + // Rewriting a keyword blindly turns ORDER BY into "order" BY, so one must announce itself. + val keywordOutOfPlace = token.lowercase() in ANY_RESERVED && + !NAME_POSITION.containsMatchIn(before) + // A token before a dot qualifies what follows; quoting a schema name breaks a working query. + val qualifierNotATable = tail.trimStart().startsWith(".") && token.lowercase() !in tables + if (tail.trimStart().startsWith("(") || canonical.isNullOrEmpty() || keywordOutOfPlace || + KEYWORD_ARGUMENT.containsMatchIn(before) || qualifierNotATable || typedLiteral + ) { + m.value + } else { + changed = true + "${quoted(canonical, quoteChar)}$tail" + } + } + + val doubleQuoteIsLiteral = quoteChar != '"' + val backslashEscapes = quoteChar == '`' + val out = StringBuilder() + var start = 0 + var i = 0 + while (i < sql.length) { + // An already-quoted identifier is opaque: re-quoting it would double the quote characters. + if (sql[i] == quoteChar) { + val closeChar = if (quoteChar == '[') ']' else quoteChar + val close = sql.indexOf(closeChar, i + 1) + val end = if (close == -1) sql.length else close + 1 + out.append(fixCode(sql.substring(start, i), start)).append(sql, i, end) + start = end + i = end + continue + } + val end = skipTo(sql, i, doubleQuoteIsLiteral, backslashEscapes) + if (end >= 0) { + out.append(fixCode(sql.substring(start, i), start)).append(sql, i, end) + start = end + i = end + } else i++ + } + out.append(fixCode(sql.substring(start), start)) + return if (changed) out.toString() else null + } + + /** + * True when a text value opens and never closes, which is what an unescaped apostrophe looks like: + * 'O'Brien' reads as the value 'O', then Brien, then a literal running to the end of the statement. + * The parser only reports "could not parse", so naming the real cause is what makes the repair land. + */ + fun hasUnterminatedLiteral(sql: String, backslashEscapes: Boolean = false): Boolean { + var open = false + var i = 0 + while (i < sql.length) { + val ch = sql[i] + if (open) { + when { + backslashEscapes && ch == '\\' -> i += 2 + ch == '\'' && i + 1 < sql.length && sql[i + 1] == '\'' -> i += 2 + ch == '\'' -> { open = false; i++ } + else -> i++ + } + continue + } + // A dollar-quoted body needs no escaping, so an apostrophe inside it is not a defect. + if (ch == '$') { + val dollar = skipTo(sql, i, false) + if (dollar > i) { i = dollar; continue } + } + when { + ch == '-' && i + 1 < sql.length && sql[i + 1] == '-' -> { + val nl = sql.indexOf('\n', i) + i = if (nl == -1) sql.length else nl + } + ch == '/' && i + 1 < sql.length && sql[i + 1] == '*' -> { + val close = sql.indexOf("*/", i + 2) + i = if (close == -1) sql.length else close + 2 + } + ch == '\'' -> { open = true; i++ } + else -> i++ + } + } + return open + } +} diff --git a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Prompts.kt b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Prompts.kt index 8ec304e..0dde33f 100644 --- a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Prompts.kt +++ b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Prompts.kt @@ -46,9 +46,30 @@ object Prompts { fewShots: List = emptyList(), context: List = emptyList(), rerunPrevious: Boolean = false, + /** Named so a question about system catalogs does not invent a placeholder database name. */ + database: String? = null, + schemas: List = emptyList(), + /** A correct catalog query for this engine, offered when the question is about structure. */ + catalogHint: String? = null, ): String { val parts = mutableListOf("", schemaText, "") + // Without these, a question about information_schema gets a guessed name like 'your_database_name'. + val where = mutableListOf() + if (!database.isNullOrBlank()) where += "database/catalog is \"$database\"" + if (schemas.isNotEmpty()) where += "schema is \"${schemas.first()}\"" + if (where.isNotEmpty()) { + parts += "" + parts += "You are connected to: the ${where.joinToString(", the ")}. Use these exact names when a " + + "query filters on system catalogs such as information_schema; never write a placeholder." + } + + // System-catalog column names are not in the schema block, so a structure question otherwise guesses them. + if (!catalogHint.isNullOrBlank()) { + parts += "" + parts += "This question is about the database's structure. Build on this correct query for this engine: $catalogHint" + } + if (glossary.isNotEmpty()) { parts += "" parts += "Business glossary (use these definitions when the question uses these terms):" diff --git a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Semantics.kt b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Semantics.kt index 1bad6a6..d5e9924 100644 --- a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Semantics.kt +++ b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/Semantics.kt @@ -120,4 +120,43 @@ object Semantics { } return null } + + /** + * An aggregate nested inside another aggregate, like AVG(x + SUM(y)). Every engine rejects it, so + * catching it before execution turns a database error into a repair. Returns the outer function name. + */ + fun nestedAggregate(sql: String): String? { + val statement = try { + CCJSqlParserUtil.parse(sql) + } catch (e: Exception) { + return null // the guard already fails closed on unparsable SQL + } + val select = statement as? Select ?: return null + for (body in plainSelects(select)) { + // Every clause an aggregate can appear in, matching the TypeScript walk over the statement. + val clauses = (body.selectItems ?: emptyList()).map { it.expression } + + listOfNotNull(body.having, body.where) + + (body.orderByElements ?: emptyList()).map { it.expression } + + (body.groupBy?.groupByExpressionList?.toList() ?: emptyList()) + for (clause in clauses) { + val outer = findNested(clause, null) + if (outer != null) return outer + } + } + return null + } + + private fun findNested(node: Expression?, insideAggregate: String?): String? { + if (node == null) return null + // A subquery has its own scope, so an aggregate inside one is not nested in the outer call. + if (node is Select || node is ParenthesedSelect) return null + val isAgg = node is Function && isBareAggregate(node) + if (isAgg && insideAggregate != null) return insideAggregate + val within = if (isAgg) (node as Function).name?.uppercase() else insideAggregate + for (child in childrenOf(node)) { + val found = findNested(child, within) + if (found != null) return found + } + return null + } } diff --git a/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/SqlKeywords.kt b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/SqlKeywords.kt new file mode 100644 index 0000000..4cae0af --- /dev/null +++ b/packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/SqlKeywords.kt @@ -0,0 +1,66 @@ +package com.rahulmahadik.asksql.ide.engine + +/** + * Reserved words per engine, generated alongside packages/core/src/sql-keywords.ts so the plugin and + * the npm engine cannot drift. Regenerate both with: node tools/generate-sql-keywords.mjs + */ +object SqlKeywords { + + private val WORDS = ( + "abort|accessible|action|add|after|all|alter|always|analyse|analyze|and|any|array|as|asc|asensitive|a" + + "symmetric|attach|authorization|autoincrement|before|begin|between|bigint|binary|blob|both|by|call|ca" + + "scade|case|cast|change|char|character|check|cluster|collate|collation|column|commit|compress|concurr" + + "ently|condition|conflict|connect|constraint|continue|convert|create|cross|cube|cume_dist|current|cur" + + "rent_catalog|current_date|current_role|current_schema|current_time|current_timestamp|current_user|cu" + + "rsor|database|databases|date|day_hour|day_microsecond|day_minute|day_second|dec|decimal|declare|defa" + + "ult|deferrable|deferred|delayed|delete|dense_rank|desc|describe|detach|deterministic|distinct|distin" + + "ctrow|div|do|double|drop|dual|each|else|elseif|empty|enclosed|end|escape|escaped|except|exclude|excl" + + "usive|exists|exit|explain|fail|false|fetch|filter|first|first_value|float|float4|float8|following|fo" + + "r|force|foreign|freeze|from|full|fulltext|function|generated|get|glob|grant|group|grouping|groups|ha" + + "ving|high_priority|hour_microsecond|hour_minute|hour_second|identified|if|ignore|ilike|immediate|in|" + + "index|indexed|infile|initially|inner|inout|insensitive|insert|instead|int|int1|int2|int3|int4|int8|i" + + "nteger|intersect|interval|into|io_after_gtids|io_before_gtids|is|isnull|iterate|join|json_table|key|" + + "keys|kill|lag|lambda|last|last_value|lateral|lead|leading|leave|left|like|limit|linear|lines|load|lo" + + "caltime|localtimestamp|lock|long|longblob|longtext|loop|low_priority|match|materialized|maxvalue|med" + + "iumblob|mediumint|mediumtext|middleint|minus|minute_microsecond|minute_second|mod|mode|modifies|natu" + + "ral|no|no_write_to_binlog|nocompress|not|nothing|notnull|nowait|nth_value|ntile|null|nulls|number|nu" + + "meric|of|offset|on|only|optimize|optimizer_costs|option|optionally|or|order|others|out|outer|outfile" + + "|over|overlaps|partition|pctfree|percent_rank|pivot|pivot_longer|pivot_wider|placing|plan|pragma|pre" + + "ceding|precision|primary|prior|procedure|public|purge|qualify|query|raise|range|rank|raw|read|read_w" + + "rite|reads|real|recursive|references|regexp|reindex|release|rename|repeat|replace|require|resignal|r" + + "esource|restrict|return|returning|revoke|right|rlike|rollback|row|row_number|rows|savepoint|schema|s" + + "chemas|second_microsecond|select|sensitive|separator|session_user|set|share|show|signal|similar|size" + + "|smallint|some|spatial|specific|sql|sql_big_result|sql_calc_found_rows|sql_small_result|sqlexception" + + "|sqlstate|sqlwarning|ssl|start|starting|stored|straight_join|summarize|symmetric|synonym|system|syst" + + "em_user|table|tablesample|temp|temporary|terminated|then|ties|tinyblob|tinyint|tinytext|to|trailing|" + + "transaction|trigger|true|unbounded|undo|union|unique|unlock|unpivot|unsigned|update|usage|use|user|u" + + "sing|utc_date|utc_time|utc_timestamp|vacuum|values|varbinary|varchar|varchar2|varcharacter|variadic|" + + "varying|verbose|view|virtual|when|where|while|window|with|without|write|xor|year_month|zerofill" + ).split("|") + + /** One bit per word in WORDS, base64 encoded. */ + private val BY_ENGINE = mapOf( + "sqlite" to "/WZ66KhRpkwAV6XG3gxrqsHeDCgrBAfAAJhjDq4CFxz4RbURAQAAaKM1Iga8Aw==", + "postgres" to "IH8FxehExh8AQyREAgN6MAHFACgLUMcAAIgiHCaBEAAQAAWQkACAHGEyMICyAQ==", + "mysql" to "aubQf6/In/z++d4/c/Ou99+s9//9+f9/36046/cKWLO/3+5+Tf87mn177135PQ==", + "oracle" to "YGxACBoiAgBBUYQEGiAiMCEMBCwBAAIDIMKkigYEoEAAIQIQYwAECSExAjQkAQ==", + "duckdb" to "IH8BxKhAAgAAwyREAgMqIAFEACgAUgQAAIAgHAbwEAIQAAEQhADACGGyIICwAQ==", + ) + + private val cache = HashMap>() + + /** The engine's reserved words; an unknown engine falls back to the union, which is the safe side. */ + fun reservedWordsFor(engine: String): Set { + val key = engine.lowercase() + cache[key]?.let { return it } + val packed = BY_ENGINE[key] + val set = if (packed != null) { + val bytes = java.util.Base64.getDecoder().decode(packed) + WORDS.filterIndexed { i, _ -> (bytes[i shr 3].toInt() shr (i and 7)) and 1 == 1 }.toSet() + } else { + WORDS.toSet() + } + cache[key] = set + return set + } +} diff --git a/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/IdentifierCaseTest.kt b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/IdentifierCaseTest.kt new file mode 100644 index 0000000..0a5c8bb --- /dev/null +++ b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/IdentifierCaseTest.kt @@ -0,0 +1,325 @@ +package com.rahulmahadik.asksql.ide.engine + +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertNull +import org.junit.Assert.assertTrue +import org.junit.Test + +/** Mirrors packages/core/test/identifier-case.test.ts; the two must agree. */ +class IdentifierCaseTest { + + private val tables = listOf("Customers", "OrderItems") + + private fun fix(sql: String, quote: Char = '`') = IdentifierCase.correctTableCase(sql, tables, quote) + + @Test fun `corrects a lower-cased table name`() { + assertEquals("SELECT * FROM `Customers`", fix("SELECT * FROM customers")) + } + + @Test fun `corrects an upper-cased table name after JOIN`() { + assertEquals( + "SELECT * FROM Customers c JOIN `OrderItems` o ON c.id = o.id", + fix("SELECT * FROM Customers c JOIN ORDERITEMS o ON c.id = o.id"), + ) + } + + @Test fun `leaves an alias after the table alone`() { + assertEquals("SELECT * FROM `OrderItems` oi", fix("SELECT * FROM orderitems oi")) + } + + @Test fun `returns null when every name already matches`() { + assertNull(fix("SELECT * FROM Customers")) + } + + @Test fun `quotes with the dialect character`() { + assertEquals("""SELECT * FROM "Customers"""", fix("SELECT * FROM customers", '"')) + } + + @Test fun `corrects a name that was already quoted in the wrong case`() { + assertEquals("SELECT * FROM `Customers`", fix("SELECT * FROM `customers`")) + } + + @Test fun `keeps a schema prefix and corrects only the table`() { + assertEquals("SELECT * FROM shop.`OrderItems`", fix("SELECT * FROM shop.orderitems")) + } + + /** A column sharing a table's name must not be rewritten: it is not in table position. */ + @Test fun `leaves a same-named column alone`() { + assertNull(fix("SELECT customers FROM Customers")) + } + + /** Rewriting inside a literal would change the query's meaning, not just its spelling. */ + @Test fun `leaves a string literal alone`() { + assertNull(fix("SELECT * FROM Customers WHERE note = 'from customers'")) + } + + @Test fun `leaves a comment alone`() { + assertNull(fix("SELECT * FROM Customers -- from customers")) + } + + /** Two tables differing only by case have no single right answer. */ + @Test fun `leaves an ambiguous fold untouched`() { + assertNull(IdentifierCase.correctTableCase("SELECT * FROM orders", listOf("Orders", "ORDERS"), '`')) + } + + @Test fun `corrects an UPDATE target`() { + assertEquals("UPDATE `Customers` SET x = 1", fix("UPDATE customers SET x = 1")) + } + + @Test fun `leaves an unknown table alone`() { + assertNull(fix("SELECT * FROM invoices")) + } + + @Test fun `recognises every engine's unknown-table wording`() { + for (message in listOf( + "Table 'asksql_test.customers' doesn't exist", + """relation "customers" does not exist""", + "no such table: customers", + "ORA-00942: table or view does not exist", + "Invalid object name customers.", + )) { + assertTrue(message, IdentifierCase.looksLikeUnknownTable(message)) + } + assertFalse(IdentifierCase.looksLikeUnknownTable("Unknown column x in field list")) + } + + /** Postgres folds an unquoted name to lower case, so a mixed-case table needs quoting even when spelled right. */ + @Test fun `quotes a correctly spelled mixed-case table on Postgres`() { + assertEquals( + """SELECT * FROM "Customers"""", + IdentifierCase.correctTableCase("SELECT * FROM Customers", tables, '"', IdentifierCase.Folding.LOWER), + ) + } + + @Test fun `leaves a name the fold already resolves alone on Postgres`() { + assertNull(IdentifierCase.correctTableCase("SELECT * FROM Orders", listOf("orders"), '"', IdentifierCase.Folding.LOWER)) + } + + @Test fun `leaves an already quoted mixed-case name alone on Postgres`() { + assertNull(IdentifierCase.correctTableCase("""SELECT * FROM "Customers"""", tables, '"', IdentifierCase.Folding.LOWER)) + } + + @Test fun `quotes a mixed-case table on Oracle, which folds upper`() { + assertEquals( + """SELECT * FROM "MixedCase"""", + IdentifierCase.correctTableCase("SELECT * FROM MixedCase", listOf("MixedCase"), '"', IdentifierCase.Folding.UPPER), + ) + } + + @Test fun `leaves an upper-case Oracle table alone`() { + assertNull(IdentifierCase.correctTableCase("SELECT * FROM employees", listOf("EMPLOYEES"), '"', IdentifierCase.Folding.UPPER)) + } + + private val quotable = listOf("Customers", "OrderItems", "FirstName", "Country", "CustomerId") + private fun q(sql: String) = IdentifierCase.quoteCatalogIdentifiers(sql, quotable, '"') + + @Test fun `quotes both the table and the columns`() { + assertEquals( + """SELECT "FirstName" FROM "Customers" WHERE "Country" = 'UK'""", + q("""SELECT FirstName FROM Customers WHERE Country = 'UK'"""), + ) + } + + @Test fun `quotes a qualified column`() { + assertEquals("""SELECT c."CustomerId" FROM "Customers" c""", q("SELECT c.CustomerId FROM Customers c")) + } + + /** Doubling the quotes would make the identifier unreadable. */ + @Test fun `leaves an already quoted identifier alone when quoting`() { + assertNull(q("""SELECT "FirstName" FROM "Customers"""")) + } + + /** A reserved word used as a function must not become an identifier. */ + @Test fun `leaves a function call alone`() { + assertEquals( + """SELECT COUNT(*) FROM "Customers"""", + IdentifierCase.quoteCatalogIdentifiers("SELECT COUNT(*) FROM Customers", listOf("Customers", "count"), '"'), + ) + } + + @Test fun `leaves a string literal alone when quoting`() { + assertNull(q("""SELECT 1 WHERE x = 'FirstName'""")) + } + + /** A table called "order" once turned ORDER BY into "order" BY, which the guard then rejected. */ + @Test fun `does not rewrite a keyword that is not naming the table`() { + assertEquals( + """SELECT x FROM "Customers" ORDER BY x DESC""", + IdentifierCase.quoteCatalogIdentifiers("SELECT x FROM Customers ORDER BY x DESC", listOf("Customers", "order"), '"'), + ) + } + + @Test fun `still quotes GROUP BY and other keyword-adjacent columns`() { + assertEquals( + """SELECT "Country" FROM t GROUP BY "Country"""", + IdentifierCase.quoteCatalogIdentifiers("SELECT Country FROM t GROUP BY Country", listOf("Country"), '"'), + ) + } + + /** A table called Nulls broke the parser: NULLS is a keyword, so the bare name would not parse. */ + @Test fun `quotes a table named like a parser keyword`() { + assertEquals( + """SELECT "Val" FROM "Nulls"""", + IdentifierCase.quoteCatalogIdentifiers("SELECT Val FROM Nulls", listOf("Nulls", "Val"), '"'), + ) + } + + @Test fun `quotes a table whose name is a reserved word`() { + assertEquals("""SELECT * FROM "order"""", IdentifierCase.quoteCatalogIdentifiers("SELECT * FROM order", listOf("order"), '"')) + } + + @Test fun `returns null when nothing needs quoting`() { + assertNull(IdentifierCase.quoteCatalogIdentifiers("SELECT id FROM orders", emptyList(), '"')) + } + + private val syntaxNames = listOf("Date", "Status", "Orders", "Month", "Amount", "order") + private fun qs(sql: String) = IdentifierCase.quoteCatalogIdentifiers(sql, syntaxNames, '"') + + /** Quoting the type turns CAST(x AS DATE) into a reference to a column that does not exist. */ + @Test fun `leaves the type in a CAST alone`() { + assertNull(qs("SELECT CAST(x AS Date) FROM t")) + } + + /** EXTRACT's first argument is a field keyword; the source after FROM is a real column. */ + @Test fun `quotes the source of an EXTRACT but not the field`() { + assertEquals("""SELECT EXTRACT(month FROM "Date") FROM t""", qs("SELECT EXTRACT(month FROM Date) FROM t")) + } + + @Test fun `leaves the leading keyword of a TRIM alone`() { + assertEquals("""SELECT TRIM(both 'x' FROM "Status") FROM t""", qs("SELECT TRIM(both 'x' FROM Status) FROM t")) + } + + @Test fun `still quotes a reserved word in table position`() { + assertEquals("""SELECT * FROM "order"""", qs("SELECT * FROM order")) + } + + @Test fun `still quotes a reserved word qualified by a dot`() { + assertEquals("""SELECT t."order" FROM t""", qs("SELECT t.order FROM t")) + } + + @Test fun `still quotes an ordinary column inside a function call`() { + assertEquals("""SELECT COUNT("Amount") FROM "Orders"""", qs("SELECT COUNT(Amount) FROM Orders")) + } + + /** A doubled quote is SQL's escaped apostrophe; treating it as the close rewrote text inside the value. */ + @Test fun `does not rewrite identifiers inside a literal containing an escaped quote`() { + assertEquals( + """SELECT * FROM "Orders" WHERE note = 'it''s about Status'""", + IdentifierCase.quoteCatalogIdentifiers( + "SELECT * FROM Orders WHERE note = 'it''s about Status'", listOf("Orders"), '"', + ), + ) + } + + /** 'O'Brien' is what a model writes when it forgets to double the apostrophe. */ + @Test fun `spots an unescaped apostrophe`() { + assertTrue(IdentifierCase.hasUnterminatedLiteral("SELECT Note FROM Person WHERE Name = 'O'Brien'")) + } + + @Test fun `accepts a correctly doubled apostrophe`() { + assertFalse(IdentifierCase.hasUnterminatedLiteral("SELECT Note FROM Person WHERE Name = 'O''Brien'")) + } + + @Test fun `accepts ordinary statements`() { + for (sql in listOf( + "SELECT * FROM t WHERE a = 'x' AND b = 'y'", + "SELECT * FROM t", + "SELECT * FROM t -- it's fine", + "SELECT * FROM t /* it's fine */", + )) { + assertFalse(sql, IdentifierCase.hasUnterminatedLiteral(sql)) + } + } + + @Test fun `spots a value that runs off the end`() { + assertTrue(IdentifierCase.hasUnterminatedLiteral("SELECT * FROM t WHERE a = 'oops")) + } + + private val litNames = listOf("Users", "Notes", "Sales") + + /** Postgres and DuckDB dollar-quote bodies, which may contain anything at all. */ + @Test fun `leaves a dollar-quoted body alone`() { + val d = "\u0024\u0024" // a dollar-quote delimiter, written as escapes so Kotlin sees no template + assertEquals( + "SELECT * FROM \"Users\" WHERE \"Notes\" = ${d}from users now$d", + IdentifierCase.quoteCatalogIdentifiers("SELECT * FROM Users WHERE Notes = ${d}from users now$d", litNames, '"'), + ) + } + + /** Quoting a schema qualifier turns a working query into "schema does not exist". */ + @Test fun `does not quote a qualifier that is not a table`() { + assertNull(IdentifierCase.quoteCatalogIdentifiers("SELECT * FROM sales.orders", listOf("Sales"), '"', emptyList())) + } + + @Test fun `still quotes a qualifier that is a real table`() { + assertEquals( + """SELECT "Customers"."FirstName" FROM "Customers"""", + IdentifierCase.quoteCatalogIdentifiers( + "SELECT Customers.FirstName FROM Customers", listOf("Customers", "FirstName"), '"', listOf("Customers"), + ), + ) + } + + /** In prod.sales.orders the table is orders; sales is a qualifier. */ + @Test fun `does not recase the middle of a three-part name`() { + assertNull(IdentifierCase.correctTableCase("SELECT * FROM prod.sales.orders", listOf("Sales"), '"', IdentifierCase.Folding.LOWER)) + } + + private val BS = "\u005c" // a single backslash + + /** Backslash escapes a quote in MySQL only; the other engines read it literally. */ + @Test fun `treats a backslash as an escape only for the backtick dialect`() { + val sql = "SELECT * FROM t WHERE n = 'O${BS}'Brien'" + assertFalse(IdentifierCase.hasUnterminatedLiteral(sql, true)) + assertTrue(IdentifierCase.hasUnterminatedLiteral(sql)) + } + + /** A dollar-quoted body needs no escaping, so an apostrophe inside it is not a defect. */ + @Test fun `does not flag an apostrophe inside a dollar-quoted body`() { + val d = "\u0024\u0024" + assertFalse(IdentifierCase.hasUnterminatedLiteral("SELECT * FROM t WHERE n = ${d}don't$d")) + } + + /** TIMESTAMP '2024-01-01' is one typed literal; quoting the word makes it a missing column. */ + @Test fun `leaves a typed literal alone`() { + assertNull( + IdentifierCase.quoteCatalogIdentifiers( + "SELECT * FROM t WHERE created > TIMESTAMP '2024-01-01'", listOf("Timestamp"), '"', + ), + ) + } + + @Test fun `still quotes the same word used as a real column`() { + assertEquals( + "SELECT \"Timestamp\" FROM t", + IdentifierCase.quoteCatalogIdentifiers("SELECT Timestamp FROM t", listOf("Timestamp"), '"'), + ) + } + + private val boundaryNames = listOf("Customers", "FirstName", "Country", "City", "Timestamp") + + /** + * Literals split the statement into segments, and the typed-literal check reads the whole + * statement rather than one segment. If it read the segment, every name sitting at a segment + * boundary would silently stop being quoted. + */ + @Test fun `quotes names before, between and after several literals`() { + assertEquals( + "SELECT \"FirstName\" FROM \"Customers\" WHERE \"Country\" = 'UK' AND \"City\" = 'York'", + IdentifierCase.quoteCatalogIdentifiers( + "SELECT FirstName FROM Customers WHERE Country = 'UK' AND City = 'York'", boundaryNames, '"', + ), + ) + } + + /** The one word that must not be quoted is the one a literal directly follows. */ + @Test fun `skips only the typed literal, quoting every other name`() { + assertEquals( + "SELECT \"FirstName\" FROM \"Customers\" WHERE created > TIMESTAMP '2024-01-01'", + IdentifierCase.quoteCatalogIdentifiers( + "SELECT FirstName FROM Customers WHERE created > TIMESTAMP '2024-01-01'", boundaryNames, '"', + ), + ) + } +} diff --git a/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/NestedAggregateTest.kt b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/NestedAggregateTest.kt new file mode 100644 index 0000000..2f5c288 --- /dev/null +++ b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/NestedAggregateTest.kt @@ -0,0 +1,42 @@ +package com.rahulmahadik.asksql.ide.engine + +import org.junit.Assert.assertEquals +import org.junit.Assert.assertNull +import org.junit.Test + +/** Mirrors packages/core/test/nested-aggregate.test.ts; the two must agree. */ +class NestedAggregateTest { + + /** Reported by a real question: AVG over a SUM is rejected by every engine. */ + @Test fun `flags an aggregate inside another aggregate`() { + val sql = "SELECT c.country, AVG(o.freight + SUM(od.unit_price)) AS v FROM orders o " + + "JOIN customers c ON o.customer_id = c.customer_id JOIN order_details od ON o.order_id = od.order_id " + + "GROUP BY c.country" + assertEquals("AVG", Semantics.nestedAggregate(sql)) + } + + @Test fun `allows aggregates side by side`() { + assertNull(Semantics.nestedAggregate("SELECT SUM(a), AVG(b) FROM t")) + } + + @Test fun `allows an aggregate over an expression`() { + assertNull(Semantics.nestedAggregate("SELECT SUM(a * b + 1) FROM t")) + } + + @Test fun `returns null for unparsable sql rather than blocking`() { + assertNull(Semantics.nestedAggregate("NOT SQL AT ALL")) + } + + /** A subquery has its own scope, so its aggregate is not nested in the outer call. */ + @Test fun `allows an aggregate inside a subquery argument`() { + assertNull(Semantics.nestedAggregate("SELECT SUM((SELECT COUNT(*) FROM u WHERE u.id = t.id)) FROM t")) + } + + @Test fun `flags nesting in HAVING`() { + assertEquals("SUM", Semantics.nestedAggregate("SELECT a FROM t GROUP BY a HAVING SUM(COUNT(b)) > 1")) + } + + @Test fun `flags nesting in ORDER BY`() { + assertEquals("AVG", Semantics.nestedAggregate("SELECT a FROM t GROUP BY a ORDER BY AVG(SUM(b))")) + } +} diff --git a/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/PromptIdentityTest.kt b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/PromptIdentityTest.kt new file mode 100644 index 0000000..5528275 --- /dev/null +++ b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/PromptIdentityTest.kt @@ -0,0 +1,39 @@ +package com.rahulmahadik.asksql.ide.engine + +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Test + +/** Mirrors the core prompt test: a missing database name became 'your_database_name' in real queries. */ +class PromptIdentityTest { + + @Test fun `names the database and schema so system-catalog filters are real`() { + val text = Prompts.buildSqlUser( + question = "what views exist?", schemaText = "TABLE film", + database = "sakila", schemas = listOf("public"), + ) + + assertTrue(text, text.contains("\"sakila\"")) + assertTrue(text, text.contains("\"public\"")) + assertTrue(text, text.contains("never write a placeholder")) + } + + @Test fun `says nothing when the connection does not report a database`() { + val text = Prompts.buildSqlUser(question = "q", schemaText = "s") + assertFalse(text, text.contains("You are connected to")) + } + + /** System-catalog columns are not in the schema block, so the model used to guess them. */ + @Test fun `offers a correct catalog query for a structure question`() { + val text = Prompts.buildSqlUser( + question = "what tables exist?", schemaText = "TABLE film", + catalogHint = "SELECT name FROM sqlite_master", + ) + + assertTrue(text, text.contains("SELECT name FROM sqlite_master")) + } + + @Test fun `offers no catalog query for an ordinary data question`() { + assertFalse(Prompts.buildSqlUser(question = "how many films?", schemaText = "s").contains("structure")) + } +} diff --git a/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/SetOperationColumnTest.kt b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/SetOperationColumnTest.kt new file mode 100644 index 0000000..39854d4 --- /dev/null +++ b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/SetOperationColumnTest.kt @@ -0,0 +1,40 @@ +package com.rahulmahadik.asksql.ide.engine + +import com.rahulmahadik.asksql.ide.model.ColumnInfo +import com.rahulmahadik.asksql.ide.model.EngineKind +import com.rahulmahadik.asksql.ide.model.SchemaCatalog +import com.rahulmahadik.asksql.ide.model.TableInfo +import com.rahulmahadik.asksql.ide.model.TableKind +import org.junit.Assert.assertNotNull +import org.junit.Assert.assertNull +import org.junit.Test + +/** Mirrors the core test: per-table row counts were blocked outright by the column floor. */ +class SetOperationColumnTest { + + private fun table(name: String, column: String) = TableInfo( + name = name, kind = TableKind.TABLE, + columns = listOf(ColumnInfo(name = column, dbType = "int", nullable = false)), + ) + + private val catalog = SchemaCatalog( + engine = EngineKind.SQLITE, + tables = listOf(table("Album", "AlbumId"), table("Artist", "ArtistId")), + ) + + @Test fun `does not flag a UNION ALL of per-table counts`() { + val sql = "SELECT 'Album' AS TableName, COUNT(*) AS RowCount FROM Album " + + "UNION ALL SELECT 'Artist' AS TableName, COUNT(*) AS RowCount FROM Artist" + + assertNull(HallucinationChecks.firstUnknownColumn(sql, catalog)) + } + + @Test fun `still flags a hallucinated column on a plain select`() { + assertNotNull(HallucinationChecks.firstUnknownColumn("SELECT Nope FROM Album", catalog)) + } + + /** A value containing "except" once disabled the floor entirely for an ordinary query. */ + @Test fun `still flags a hallucinated column when a literal contains a set-operation word`() { + assertNotNull(HallucinationChecks.firstUnknownColumn("SELECT Nope FROM Album WHERE AlbumId = 'except this'", catalog)) + } +} diff --git a/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/SqlKeywordsTest.kt b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/SqlKeywordsTest.kt new file mode 100644 index 0000000..86233d5 --- /dev/null +++ b/packages/jetbrains/src/test/kotlin/com/rahulmahadik/asksql/ide/engine/SqlKeywordsTest.kt @@ -0,0 +1,36 @@ +package com.rahulmahadik.asksql.ide.engine + +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Test + +/** Mirrors packages/core/test/sql-keywords.test.ts; both are generated from the same source. */ +class SqlKeywordsTest { + + @Test fun `gives each engine its own list`() { + assertTrue(SqlKeywords.reservedWordsFor("mysql").size > SqlKeywords.reservedWordsFor("postgres").size) + } + + @Test fun `every engine reserves select`() { + for (engine in listOf("postgres", "mysql", "oracle", "sqlite", "duckdb")) { + assertTrue(engine, SqlKeywords.reservedWordsFor(engine).contains("select")) + } + } + + /** MySQL reserves it, Postgres does not: proof the lists are not one shared set. */ + @Test fun `separates a word that only some engines reserve`() { + assertTrue(SqlKeywords.reservedWordsFor("mysql").contains("rank")) + assertFalse(SqlKeywords.reservedWordsFor("postgres").contains("rank")) + } + + /** An unknown engine gets the union, so a name is over-quoted rather than left broken. */ + @Test fun `falls back to the union for an unknown engine`() { + assertTrue(SqlKeywords.reservedWordsFor("not-an-engine").size > SqlKeywords.reservedWordsFor("mysql").size) + } + + @Test fun `matches the core list sizes`() { + assertEquals(101, SqlKeywords.reservedWordsFor("postgres").size) + assertEquals(262, SqlKeywords.reservedWordsFor("mysql").size) + } +} diff --git a/packages/vscode/CHANGELOG.md b/packages/vscode/CHANGELOG.md index b08ad30..221f275 100644 --- a/packages/vscode/CHANGELOG.md +++ b/packages/vscode/CHANGELOG.md @@ -4,6 +4,29 @@ All notable changes to the AskSQL VS Code extension are documented here. The for [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.7.1] - 2026-08-14 + +### Fixed +- A mixed-case Postgres schema failed every query. An unquoted name folds to lower case and resolves + to nothing; Oracle folds the other way, and MySQL on Linux compares table names case-sensitively. + Table and column names are now quoted from the catalog before the query is validated, and correct + names are left untouched. +- Reserved words now come from each database itself rather than one shared list that applied MySQL's + rules to Postgres and missed most of MySQL's own. +- Quoting knows where a word is syntax rather than a name, so `CAST(x AS DATE)` and + `EXTRACT(MONTH FROM d)` are left alone. +- A table named like a parser keyword, such as `order` or `Nulls`, could not be parsed in its bare + form, so the question failed after three attempts. +- An apostrophe inside a value, as in `'O'Brien'`, produced only "could not parse", so the model + returned the same statement until it ran out of attempts. It is now told to double the quote. +- Questions about structure no longer invent a database name. Without being told which database it + is connected to, the model wrote `table_schema = 'your_database_name'` and returned nothing, which + reads as an empty database rather than an error. +- `AVG(SUM(x))` is repaired instead of run. Nested aggregates are invalid in every engine. +- Per-table row counts work again. A `UNION ALL` across tables was blocked as a hallucinated column, + because each branch's columns were judged against every branch's tables. +- Index columns are no longer quoted twice in the schema sent to the model. + ## [0.7.0] - 2026-08-09 ### Fixed diff --git a/packages/vscode/package.json b/packages/vscode/package.json index 78b631b..ddd9bb1 100644 --- a/packages/vscode/package.json +++ b/packages/vscode/package.json @@ -3,7 +3,7 @@ "private": true, "displayName": "AskSQL", "description": "AI database chat: ask in plain language, review the query, get answers. Read-only by design, bring your own model.", - "version": "0.7.0", + "version": "0.7.1", "publisher": "RahulMahadik", "license": "Apache-2.0", "pricing": "Free", diff --git a/tests/bundle-size.test.ts b/tests/bundle-size.test.ts index 7ac4d4e..7de9f2f 100644 --- a/tests/bundle-size.test.ts +++ b/tests/bundle-size.test.ts @@ -26,8 +26,12 @@ const BUDGETS: Record = { // stops measuring. History: 45->51 scope guard + Mongo, 51->54 grounding vocabularies, // 54->60 question routing + the ungrouped-aggregate lint, 60->63 routing precision + // identifier quoting, 63->65 routing words disambiguated from identifiers ("the archive table", - // "the best selling products", "the prompts table"). - core: 65, + // "the best selling products", "the prompts table"), 65->66 identifier normalisation plus the + // connection identity and catalog hint that keep system-catalog queries from being guessed, + // 67->69 each engine's own reserved words read from its catalog (MySQL reserves 262, the shared + // guess had ~100), stored as one word list plus a bit per engine, 69->70 dollar-quoted literals, + // qualifier handling and the per-dialect backslash rule. + core: 70, // 20 -> 23: copy controls, streamed-token progress, cell tooltips, export feedback, result-grid copy. react: 23, // 12 -> 13: the CSRF/Host gate every adapter inherits, client-path confinement for @@ -64,7 +68,9 @@ describe('bundle-size budgets (gzipped, own code)', () => { const core = gzippedKb('core'); const react = gzippedKb('react'); if (core === null || react === null) return; - // Own code only (React is a peer): must stay well under the 85 KB budget. - expect(core + react).toBeLessThan(85); + // Own code only (React is a peer). 85 -> 91: identifier normalisation, which is what stops a + // mixed-case Postgres schema from failing every query, the prompt's connection identity, and the + // per-engine reserved-word lists, and the literal rules each dialect actually follows. + expect(core + react).toBeLessThan(92); }); }); diff --git a/tools/generate-sql-keywords.mjs b/tools/generate-sql-keywords.mjs new file mode 100644 index 0000000..f509315 --- /dev/null +++ b/tools/generate-sql-keywords.mjs @@ -0,0 +1,206 @@ +/** + * Regenerates packages/core/src/sql-keywords.ts from the databases themselves, so the reserved-word + * lists are the engines' own rather than a hand-maintained guess that drifts. + * + * Needs the engines reachable; see docs for the local containers. Usage: + * node tools/generate-sql-keywords.mjs + */ +import { writeFileSync } from 'node:fs'; +import { execFileSync } from 'node:child_process'; + +const sh = (cmd, args) => execFileSync(cmd, args, { encoding: 'utf8', maxBuffer: 1 << 24 }); + +/** Each engine's own catalog of reserved words. SQLite has none, so its published list is inlined. */ +const SOURCES = { + postgres: () => + sh('docker', [ + 'exec', + 'asksql-pg', + 'psql', + '-U', + 'postgres', + '-tAc', + "SELECT word FROM pg_get_keywords() WHERE catcode IN ('R','T') ORDER BY word", + ]), + mysql: () => + sh('docker', [ + 'exec', + 'asksql-mysql', + 'mysql', + '-uroot', + '-N', + '-e', + 'SELECT LOWER(WORD) FROM information_schema.KEYWORDS WHERE RESERVED=1 ORDER BY WORD', + ]), + oracle: () => + sh('docker', [ + 'exec', + 'asksql-oracle', + 'bash', + '-lc', + `echo "SET PAGESIZE 0 FEEDBACK OFF +SELECT LOWER(keyword) FROM V\\$RESERVED_WORDS WHERE reserved='Y' ORDER BY keyword; +EXIT;" | sqlplus -s system/oracle@localhost:1521/FREEPDB1`, + ]), + duckdb: async () => { + const { DuckDBInstance } = await import('../packages/duckdb/node_modules/@duckdb/node-api/lib/index.js'); + const instance = await DuckDBInstance.create(':memory:'); + const conn = await instance.connect(); + const result = await conn.runAndReadAll( + "SELECT keyword_name FROM duckdb_keywords() WHERE keyword_category = 'reserved' ORDER BY 1", + ); + return result + .getRows() + .map((r) => String(r[0])) + .join(','); + }, +}; + +const SQLITE = `abort action add after all alter always analyze and as asc attach autoincrement before begin +between by cascade case cast check collate column commit conflict constraint create cross current +current_date current_time current_timestamp database default deferrable deferred delete desc detach +distinct do drop each else end escape except exclude exclusive exists explain fail filter first +following for foreign from full generated glob group groups having if ignore immediate in index +indexed initially inner insert instead intersect into is isnull join key last left like limit match +materialized natural no not nothing notnull null nulls of offset on or order others outer over +partition plan pragma preceding primary query raise range recursive references regexp reindex release +rename replace restrict returning right rollback row rows savepoint select set table temp temporary +then ties to transaction trigger unbounded union unique update using vacuum values view virtual when +where window with without`; + +const clean = (text) => + [ + ...new Set( + text + .split(/[\s,]+/) + .map((w) => w.trim().toLowerCase()) + .filter((w) => /^[a-z_][a-z0-9_]*$/.test(w)), + ), + ].sort(); + +const engines = { sqlite: clean(SQLITE) }; +for (const [name, read] of Object.entries(SOURCES)) { + try { + const words = clean(await read()); + if (words.length > 20) engines[name] = words; + else console.warn(`[skip] ${name}: only ${words.length} words, is the container running?`); + } catch (err) { + console.warn(`[skip] ${name}: ${String(err).split('\n')[0]}`); + } +} + +// Emitting a file that quietly lost an engine would downgrade it to the union without anyone noticing. +const EXPECTED = ['postgres', 'mysql', 'oracle', 'duckdb', 'sqlite']; +const missing = EXPECTED.filter((e) => !engines[e]); +if (missing.length > 0) { + console.error(`refusing to write: no keywords read for ${missing.join(', ')}. Start the databases and retry.`); + process.exit(1); +} + +const all = [...new Set(Object.values(engines).flat())].sort(); +const index = new Map(all.map((w, i) => [w, i])); +const chunks = all.join('|').match(/.{1,110}/g) ?? []; + +/** One bit per word in the shared list, so a set costs 64 characters instead of a few hundred. */ +const bitmap = (words) => { + const bytes = new Uint8Array(Math.ceil(all.length / 8)); + for (const w of words) { + const i = index.get(w); + bytes[i >> 3] |= 1 << (i & 7); + } + return Buffer.from(bytes).toString('base64'); +}; + +const HEADER = `/** + * Reserved words per engine, read from each database itself rather than guessed: + * pg_get_keywords(), information_schema.KEYWORDS, V$RESERVED_WORDS, duckdb_keywords(). + * SQLite publishes a fixed list and has no catalog to query. + * + * Regenerate with: node tools/generate-sql-keywords.mjs + * One word list plus a bit per engine, which costs a few hundred bytes instead of repeating words. + */`; + +writeFileSync( + 'packages/core/src/sql-keywords.ts', + `${HEADER} + +const WORDS = +${chunks.map((c, i) => ` '${c}'${i === chunks.length - 1 ? ';' : ' +'}`).join('\n')} + +const WORD_LIST = WORDS.split('|'); + +/** One bit per word in WORD_LIST, base64 encoded. */ +const BY_ENGINE: Record = { +${Object.entries(engines) + .map(([e, w]) => ` ${e}: '${bitmap(w)}',`) + .join('\n')} +}; + +const cache = new Map>(); + +/** The engine's reserved words; an unknown engine falls back to the union, which is the safe side. */ +export function reservedWordsFor(engine: string): ReadonlySet { + const key = engine.toLowerCase(); + let set = cache.get(key); + if (!set) { + const packed = BY_ENGINE[key]; + if (packed) { + const bytes = atob(packed); + set = new Set(WORD_LIST.filter((_, i) => (bytes.charCodeAt(i >> 3) >> (i & 7)) & 1)); + } else set = new Set(WORD_LIST); + cache.set(key, set); + } + return set; +} +`, +); + +// The plugin mirrors the same data; generating both here is what stops the two from drifting. +const ktChunks = all.join('|').match(/.{1,100}/g) ?? []; +writeFileSync( + 'packages/jetbrains/src/main/kotlin/com/rahulmahadik/asksql/ide/engine/SqlKeywords.kt', + `package com.rahulmahadik.asksql.ide.engine + +/** + * Reserved words per engine, generated alongside packages/core/src/sql-keywords.ts so the plugin and + * the npm engine cannot drift. Regenerate both with: node tools/generate-sql-keywords.mjs + */ +object SqlKeywords { + + private val WORDS = ( +${ktChunks.map((c, i) => ` "${c}"${i === ktChunks.length - 1 ? '' : ' +'}`).join('\n')} + ).split("|") + + /** One bit per word in WORDS, base64 encoded. */ + private val BY_ENGINE = mapOf( +${Object.entries(engines) + .map(([e, w]) => ` "${e}" to "${bitmap(w)}",`) + .join('\n')} + ) + + private val cache = HashMap>() + + /** The engine's reserved words; an unknown engine falls back to the union, which is the safe side. */ + fun reservedWordsFor(engine: String): Set { + val key = engine.lowercase() + cache[key]?.let { return it } + val packed = BY_ENGINE[key] + val set = if (packed != null) { + val bytes = java.util.Base64.getDecoder().decode(packed) + WORDS.filterIndexed { i, _ -> (bytes[i shr 3].toInt() shr (i and 7)) and 1 == 1 }.toSet() + } else { + WORDS.toSet() + } + cache[key] = set + return set + } +} +`, +); + +console.log( + 'engines:', + Object.fromEntries(Object.entries(engines).map(([e, w]) => [e, w.length])), + 'unique:', + all.length, +);