diff --git a/src/autocomplete/content-assist.ts b/src/autocomplete/content-assist.ts index cb1a569..e32ae7b 100644 --- a/src/autocomplete/content-assist.ts +++ b/src/autocomplete/content-assist.ts @@ -3,7 +3,10 @@ import { parser, parse as parseRaw } from "../parser/parser" import { visitor } from "../parser/visitor" import { QuestDBLexer } from "../parser/lexer" import type { Statement } from "../parser/ast" -import { IDENTIFIER_KEYWORD_TOKENS } from "./token-classification" +import { + IDENTIFIER_KEYWORD_TOKENS, + EXPRESSION_OPERATORS, +} from "./token-classification" // ============================================================================= // Constants @@ -85,6 +88,8 @@ export interface ContentAssistResult { * context. Used by the provider to boost tables containing all these columns. */ referencedColumns: Set + /** Whether the cursor is inside a WHERE clause expression */ + isConditionContext: boolean } // ============================================================================= @@ -683,6 +688,7 @@ interface ComputeResult { nextTokenTypes: TokenType[] suggestColumns: boolean suggestTables: boolean + isConditionContext: boolean } /** @@ -767,7 +773,21 @@ function computeSuggestions(tokens: IToken[]): ComputeResult { } } - return { nextTokenTypes: result, suggestColumns, suggestTables } + // Check if an expression operator's ruleStack includes "whereClause". + // Must check operators specifically — Chevrotain explores ahead into + // not-yet-started WHERE paths even from JOIN ON positions. + const isConditionContext = effectiveSuggestions.some( + (s) => + EXPRESSION_OPERATORS.has(s.nextTokenType.name) && + s.ruleStack.includes("whereClause"), + ) + + return { + nextTokenTypes: result, + suggestColumns, + suggestTables, + isConditionContext, + } } /** @@ -906,6 +926,7 @@ export function getContentAssist( suggestColumns: false, suggestTables: false, referencedColumns: new Set(), + isConditionContext: false, } } } @@ -934,11 +955,13 @@ export function getContentAssist( let nextTokenTypes: TokenType[] = [] let suggestColumns = false let suggestTables = false + let isConditionContext = false try { const computed = computeSuggestions(tokensForAssist) nextTokenTypes = computed.nextTokenTypes suggestColumns = computed.suggestColumns suggestTables = computed.suggestTables + isConditionContext = computed.isConditionContext } catch (e) { // If content assist fails, return empty suggestions // This can happen with malformed input @@ -1025,6 +1048,7 @@ export function getContentAssist( suggestColumns, suggestTables, referencedColumns, + isConditionContext, } } diff --git a/src/autocomplete/provider.ts b/src/autocomplete/provider.ts index b78a89d..555373b 100644 --- a/src/autocomplete/provider.ts +++ b/src/autocomplete/provider.ts @@ -21,10 +21,36 @@ import type { IToken } from "chevrotain" import { getContentAssist } from "./content-assist" import { buildSuggestions } from "./suggestion-builder" -import { shouldSkipToken } from "./token-classification" +import { + shouldSkipToken, + SKIP_TOKENS, + PUNCTUATION_TOKENS, + EXPRESSION_OPERATORS, + tokenNameToKeyword, +} from "./token-classification" import type { AutocompleteProvider, SchemaInfo, Suggestion } from "./types" import { SuggestionKind, SuggestionPriority } from "./types" +const EXPRESSION_OPERATOR_LABELS = new Set( + [...EXPRESSION_OPERATORS].map(tokenNameToKeyword), +) + +function isSchemaColumn( + image: string, + tablesInScope: { table: string; alias?: string }[], + schema: SchemaInfo, +): boolean { + const nameLower = image.toLowerCase() + for (const ref of tablesInScope) { + const cols = schema.columns[ref.table.toLowerCase()] + if (cols?.some((c) => c.name.toLowerCase() === nameLower)) return true + } + for (const cols of Object.values(schema.columns)) { + if (cols.some((c) => c.name.toLowerCase() === nameLower)) return true + } + return false +} + const TABLE_NAME_TOKENS = new Set([ "From", "Join", @@ -163,6 +189,7 @@ export function createAutocompleteProvider( suggestColumns, suggestTables, referencedColumns, + isConditionContext, } = getContentAssist(query, cursorOffset) // Merge CTE columns into the schema so getColumnsInScope() can find them @@ -226,6 +253,42 @@ export function createAutocompleteProvider( if (suggestTables) { rankTableSuggestions(suggestions, referencedColumns, columnIndex) } + + // In WHERE after a column, boost operators over clause keywords. + if (isConditionContext) { + const hasExpressionOperators = nextTokenTypes.some((t) => + EXPRESSION_OPERATORS.has(t.name), + ) + if (hasExpressionOperators) { + const effectiveTokens = + isMidWord && tokensBefore.length > 0 + ? tokensBefore.slice(0, -1) + : tokensBefore + const lastToken = effectiveTokens[effectiveTokens.length - 1] + const lastTokenName = lastToken?.tokenType?.name + + if ( + lastTokenName && + !SKIP_TOKENS.has(lastTokenName) && + !PUNCTUATION_TOKENS.has(lastTokenName) && + isSchemaColumn( + lastToken.image, + effectiveTablesInScope, + effectiveSchema, + ) + ) { + for (const s of suggestions) { + if (s.kind !== SuggestionKind.Keyword) continue + if (s.label === "IN") { + s.priority = SuggestionPriority.High + } else if (!EXPRESSION_OPERATOR_LABELS.has(s.label)) { + s.priority = SuggestionPriority.MediumLow + } + } + } + } + } + return suggestions } diff --git a/src/autocomplete/suggestion-builder.ts b/src/autocomplete/suggestion-builder.ts index 1441783..a2b52ea 100644 --- a/src/autocomplete/suggestion-builder.ts +++ b/src/autocomplete/suggestion-builder.ts @@ -17,7 +17,6 @@ import { import { SKIP_TOKENS, PUNCTUATION_TOKENS, - EXPRESSION_OPERATORS, tokenNameToKeyword, } from "./token-classification" import { functions } from "../grammar/index" @@ -170,16 +169,13 @@ export function buildSuggestions( // All parser keyword tokens are keywords (not functions). // Functions are suggested separately in the functions loop below. const kind = SuggestionKind.Keyword - const priority = EXPRESSION_OPERATORS.has(name) - ? SuggestionPriority.MediumLow - : SuggestionPriority.Medium suggestions.push({ label: keyword, kind, insertText: keyword, filterText: keyword.toLowerCase(), - priority, + priority: SuggestionPriority.Medium, }) } diff --git a/src/autocomplete/token-classification.ts b/src/autocomplete/token-classification.ts index c70d7c7..b5d9558 100644 --- a/src/autocomplete/token-classification.ts +++ b/src/autocomplete/token-classification.ts @@ -33,9 +33,7 @@ export const IDENTIFIER_TOKENS = new Set([ export const IDENTIFIER_KEYWORD_TOKENS = IDENTIFIER_KEYWORD_NAMES /** - * Expression-continuation operators that are valid after any expression but - * should be deprioritized so clause-level keywords (ASC, DESC, LIMIT, etc.) - * appear first in the suggestion list. + * Expression-level operators and keywords (as opposed to clause-level keywords). */ export const EXPRESSION_OPERATORS = new Set([ "And", @@ -47,18 +45,11 @@ export const EXPRESSION_OPERATORS = new Set([ "Like", "Ilike", "Within", - // Subquery/set operators "All", "Any", "Some", - // Expression-start keywords that continue an expression context "Case", "Cast", - // Query connectors — valid after any complete query but should not - // overshadow clause-level keywords the user is more likely typing. - "Union", - "Except", - "Intersect", ]) /** diff --git a/tests/autocomplete.test.ts b/tests/autocomplete.test.ts index 5ed07ce..ed3cc94 100644 --- a/tests/autocomplete.test.ts +++ b/tests/autocomplete.test.ts @@ -3139,3 +3139,206 @@ describe("CTE autocomplete", () => { }) }) }) + +describe("condition-aware operator priority", () => { + function getPriority( + sql: string, + label: string, + offset?: number, + ): SuggestionPriority | undefined { + const suggestions = provider.getSuggestions(sql, offset ?? sql.length) + return suggestions.find((s) => s.label === label)?.priority + } + + describe("WHERE column | (condition context, column detected)", () => { + const sql = "SELECT * FROM trades WHERE amount " + + it("boosts IN to High priority", () => { + expect(getPriority(sql, "IN")).toBe(SuggestionPriority.High) + }) + + it("keeps other expression operators at Medium", () => { + expect(getPriority(sql, "AND")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "BETWEEN")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "IS")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "NOT")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "OR")).toBe(SuggestionPriority.Medium) + }) + + it("demotes clause keywords to MediumLow", () => { + expect(getPriority(sql, "GROUP")).toBe(SuggestionPriority.MediumLow) + expect(getPriority(sql, "ORDER")).toBe(SuggestionPriority.MediumLow) + expect(getPriority(sql, "LIMIT")).toBe(SuggestionPriority.MediumLow) + }) + }) + + describe("SELECT column | (not a condition context)", () => { + const sql = "SELECT amount FROM trades" + + it("does not adjust priorities outside WHERE", () => { + // cursor at position 14 (after "amount ") + expect(getPriority(sql, "AND", 14)).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "IN", 14)).toBe(SuggestionPriority.Medium) + }) + }) + + describe("WHERE | (no column before cursor)", () => { + const sql = "SELECT * FROM trades WHERE " + + it("does not adjust priorities when no column precedes cursor", () => { + // After WHERE, parser expects an expression to start — operators + // like IN are not among suggestions at this position. + // Columns should be at their default High priority. + const suggestions = provider.getSuggestions(sql, sql.length) + const columns = suggestions.filter( + (s) => s.kind === SuggestionKind.Column, + ) + expect(columns.length).toBeGreaterThan(0) + for (const c of columns) { + expect(c.priority).toBe(SuggestionPriority.High) + } + }) + }) + + describe("WHERE col > 5 | (last token is literal, not column)", () => { + const sql = "SELECT * FROM trades WHERE amount > 5 " + + it("does not adjust priorities when last token is a literal", () => { + expect(getPriority(sql, "AND")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "GROUP")).toBe(SuggestionPriority.Medium) + }) + }) + + describe("WHERE col > 5 AND col2 | (second condition column)", () => { + const sql = "SELECT * FROM trades WHERE amount > 5 AND price " + + it("boosts IN for the second condition column", () => { + expect(getPriority(sql, "IN")).toBe(SuggestionPriority.High) + }) + + it("demotes clause keywords for the second condition column", () => { + expect(getPriority(sql, "GROUP")).toBe(SuggestionPriority.MediumLow) + expect(getPriority(sql, "ORDER")).toBe(SuggestionPriority.MediumLow) + }) + }) + + describe("WHERE col OR col2 | (OR continuation)", () => { + const sql = "SELECT * FROM trades WHERE amount > 5 OR price " + + it("boosts IN after OR", () => { + expect(getPriority(sql, "IN")).toBe(SuggestionPriority.High) + }) + + it("keeps expression operators at Medium after OR", () => { + expect(getPriority(sql, "AND")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "BETWEEN")).toBe(SuggestionPriority.Medium) + }) + + it("demotes clause keywords after OR", () => { + expect(getPriority(sql, "GROUP")).toBe(SuggestionPriority.MediumLow) + expect(getPriority(sql, "ORDER")).toBe(SuggestionPriority.MediumLow) + }) + }) + + describe("WHERE (col) | (parenthesized expression)", () => { + const sql = "SELECT * FROM trades WHERE (amount " + + it("boosts IN inside parenthesized condition", () => { + expect(getPriority(sql, "IN")).toBe(SuggestionPriority.High) + }) + + it("keeps expression operators at Medium", () => { + expect(getPriority(sql, "AND")).toBe(SuggestionPriority.Medium) + }) + }) + + describe("WHERE col = 'x' | (last token is string literal)", () => { + const sql = "SELECT * FROM trades WHERE symbol = 'BTC' " + + it("does not adjust priorities when last token is a string literal", () => { + expect(getPriority(sql, "AND")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "GROUP")).toBe(SuggestionPriority.Medium) + }) + }) + + describe("WHERE NOT col | (NOT before column)", () => { + const sql = "SELECT * FROM trades WHERE NOT amount " + + it("boosts IN after NOT column", () => { + expect(getPriority(sql, "IN")).toBe(SuggestionPriority.High) + }) + + it("demotes clause keywords after NOT column", () => { + expect(getPriority(sql, "GROUP")).toBe(SuggestionPriority.MediumLow) + }) + }) + + describe("set operations are not treated as expression operators", () => { + const sql = "SELECT * FROM trades WHERE amount " + + it("demotes UNION as clause keyword in condition context", () => { + const priority = getPriority(sql, "UNION") + if (priority !== undefined) { + expect(priority).toBe(SuggestionPriority.MediumLow) + } + }) + + it("demotes EXCEPT as clause keyword in condition context", () => { + const priority = getPriority(sql, "EXCEPT") + if (priority !== undefined) { + expect(priority).toBe(SuggestionPriority.MediumLow) + } + }) + + it("demotes INTERSECT as clause keyword in condition context", () => { + const priority = getPriority(sql, "INTERSECT") + if (priority !== undefined) { + expect(priority).toBe(SuggestionPriority.MediumLow) + } + }) + }) + + describe("mid-word typing in WHERE context", () => { + it("adjusts priorities when typing operator prefix after column", () => { + const sql = "SELECT * FROM trades WHERE amount b" + // mid-word "b" could match BETWEEN — still in condition context + expect(getPriority(sql, "BETWEEN")).toBe(SuggestionPriority.Medium) + }) + + it("adjusts priorities when typing 'i' after column (matches IN)", () => { + const sql = "SELECT * FROM trades WHERE amount i" + const priority = getPriority(sql, "IN") + if (priority !== undefined) { + expect(priority).toBe(SuggestionPriority.High) + } + }) + }) + + describe("ORDER BY col | (not a condition context)", () => { + const sql = "SELECT * FROM trades WHERE amount > 5 ORDER BY price " + + it("does not adjust priorities in ORDER BY", () => { + // After ORDER BY column, no condition-aware adjustment + const andPriority = getPriority(sql, "AND") + const limitPriority = getPriority(sql, "LIMIT") + if (andPriority !== undefined && limitPriority !== undefined) { + // Both should be at default Medium (no condition boost) + expect(andPriority).toBe(SuggestionPriority.Medium) + expect(limitPriority).toBe(SuggestionPriority.Medium) + } + }) + }) + + describe("UPDATE WHERE col | (condition context in UPDATE)", () => { + const sql = "UPDATE trades SET amount = 1 WHERE symbol " + + it("boosts IN in UPDATE WHERE clause", () => { + expect(getPriority(sql, "IN")).toBe(SuggestionPriority.High) + }) + + it("keeps expression operators at Medium", () => { + expect(getPriority(sql, "AND")).toBe(SuggestionPriority.Medium) + expect(getPriority(sql, "LIKE")).toBe(SuggestionPriority.Medium) + }) + }) +}) diff --git a/tests/content-assist.test.ts b/tests/content-assist.test.ts index 041e74a..74c0cac 100644 --- a/tests/content-assist.test.ts +++ b/tests/content-assist.test.ts @@ -923,4 +923,203 @@ describe("Content Assist", () => { expect(tokens).toContain("As") }) }) + + describe("isConditionContext", () => { + function isCondition(sql: string, offset?: number): boolean { + const result = getContentAssist(sql, offset ?? sql.length) + return result.isConditionContext + } + + describe("WHERE clause positions", () => { + it("true after WHERE column", () => { + expect(isCondition("SELECT * FROM t WHERE col ")).toBe(true) + }) + + it("true after WHERE column with complex query", () => { + expect(isCondition("SELECT a, b FROM trades WHERE amount ")).toBe(true) + }) + + it("true after WHERE col > 5 AND col2", () => { + expect(isCondition("SELECT * FROM t WHERE a > 5 AND b ")).toBe(true) + }) + + it("true after WHERE col BETWEEN 1 AND col2", () => { + expect(isCondition("SELECT * FROM t WHERE a BETWEEN 1 AND b ")).toBe( + true, + ) + }) + + it("true after WHERE with nested conditions", () => { + expect(isCondition("SELECT * FROM t WHERE (a > 1 AND b ")).toBe(true) + }) + + it("true after WHERE col > 5", () => { + // Even after a literal — we're still in WHERE context + expect(isCondition("SELECT * FROM t WHERE col > 5 ")).toBe(true) + }) + + it("true at WHERE start (before any column)", () => { + expect(isCondition("SELECT * FROM t WHERE ")).toBe(true) + }) + + it("true after WHERE with OR", () => { + expect(isCondition("SELECT * FROM t WHERE a = 1 OR b ")).toBe(true) + }) + + it("true after WHERE NOT", () => { + expect(isCondition("SELECT * FROM t WHERE NOT a ")).toBe(true) + }) + }) + + describe("non-WHERE positions", () => { + it("false after SELECT column", () => { + expect(isCondition("SELECT col ")).toBe(false) + }) + + it("false after SELECT column with FROM", () => { + expect(isCondition("SELECT col FROM t", 11)).toBe(false) + }) + + it("false after FROM", () => { + expect(isCondition("SELECT * FROM ")).toBe(false) + }) + + it("false at start of query", () => { + expect(isCondition("")).toBe(false) + }) + + it("false after ORDER BY", () => { + expect(isCondition("SELECT * FROM t ORDER BY ")).toBe(false) + }) + + it("false after GROUP BY", () => { + expect(isCondition("SELECT * FROM t GROUP BY ")).toBe(false) + }) + + it("false after HAVING column", () => { + // HAVING is semantically similar to WHERE, but ruleStack won't + // include "whereClause" — it's a separate "havingClause" rule. + expect( + isCondition("SELECT a, count() FROM t GROUP BY a HAVING count() "), + ).toBe(false) + }) + + it("false after INSERT INTO", () => { + expect(isCondition("INSERT INTO ")).toBe(false) + }) + + it("false after CREATE TABLE", () => { + expect(isCondition("CREATE TABLE ")).toBe(false) + }) + + it("false after LIMIT", () => { + expect(isCondition("SELECT * FROM t LIMIT ")).toBe(false) + }) + }) + + describe("subquery WHERE positions", () => { + it("true in subquery WHERE", () => { + expect(isCondition("SELECT * FROM (SELECT * FROM t WHERE col ")).toBe( + true, + ) + }) + + it("true in WHERE with subquery completed", () => { + expect( + isCondition( + "SELECT * FROM t WHERE col IN (SELECT id FROM s) AND col2 ", + ), + ).toBe(true) + }) + }) + + describe("UPDATE and DELETE WHERE positions", () => { + it("true after UPDATE SET ... WHERE column", () => { + expect(isCondition("UPDATE t SET a = 1 WHERE b ")).toBe(true) + }) + }) + + describe("JOIN ON positions (should NOT be condition context)", () => { + it("false after JOIN ON column", () => { + expect(isCondition("SELECT * FROM a JOIN b ON a.id ")).toBe(false) + }) + + it("false after LEFT JOIN ON column", () => { + expect(isCondition("SELECT * FROM a LEFT JOIN b ON a.id ")).toBe(false) + }) + + it("false after CROSS JOIN ON column", () => { + expect(isCondition("SELECT * FROM a CROSS JOIN b ON a.id ")).toBe(false) + }) + + it("false after JOIN ON col = col AND col2", () => { + expect( + isCondition("SELECT * FROM a JOIN b ON a.id = b.id AND a.name "), + ).toBe(false) + }) + }) + + describe("CASE WHEN positions (should NOT be condition context)", () => { + it("false inside CASE WHEN in SELECT", () => { + expect(isCondition("SELECT CASE WHEN col ")).toBe(false) + }) + + it("false after CASE WHEN col = val THEN ... WHEN col2", () => { + expect(isCondition("SELECT CASE WHEN a = 1 THEN 'x' WHEN b ")).toBe( + false, + ) + }) + + it("false inside CASE WHEN ELSE in SELECT", () => { + expect( + isCondition("SELECT CASE WHEN a > 1 THEN a ELSE b FROM t", 38), + ).toBe(false) + }) + }) + + describe("SELECT expression positions (should NOT be condition context)", () => { + it("false after SELECT with arithmetic", () => { + expect(isCondition("SELECT a + b FROM t", 13)).toBe(false) + }) + + it("false after SELECT function call", () => { + expect(isCondition("SELECT count(a) FROM t", 16)).toBe(false) + }) + + it("false after SELECT DISTINCT column", () => { + expect(isCondition("SELECT DISTINCT col ")).toBe(false) + }) + }) + + describe("other clause positions (should NOT be condition context)", () => { + it("false in SAMPLE BY clause", () => { + expect(isCondition("SELECT * FROM t SAMPLE BY ")).toBe(false) + }) + + it("false in LATEST ON clause", () => { + expect(isCondition("SELECT * FROM t LATEST ON ")).toBe(false) + }) + + it("false after LIMIT expression", () => { + expect(isCondition("SELECT * FROM t WHERE a > 1 LIMIT 10 ")).toBe(false) + }) + + it("false after ORDER BY column with trailing position", () => { + expect(isCondition("SELECT * FROM t WHERE a > 1 ORDER BY col ")).toBe( + false, + ) + }) + + it("false in INSERT VALUES", () => { + expect(isCondition("INSERT INTO t VALUES (1, ")).toBe(false) + }) + }) + + describe("CASE WHEN inside WHERE (should be condition context)", () => { + it("true for CASE WHEN inside WHERE", () => { + // CASE WHEN nested inside WHERE — ruleStack should include whereClause + expect(isCondition("SELECT * FROM t WHERE CASE WHEN a ")).toBe(true) + }) + }) + }) })