boost expression operators over clause keywords in WHERE context

Author: emrberk

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<string>
+  /** 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,
   }
 }
 

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
       }
 

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

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",
 ])
 
 /**