From 5f16cc29e6488e0846a1d1b1b830497fd1f9abd7 Mon Sep 17 00:00:00 2001
From: Marijn Haverbeke <marijn@haverbeke.berlin>
Date: Tue, 1 Oct 2024 09:12:42 +0200
Subject: [PATCH] Add keywordCompletion option

FEATURE: The new `keywordCompletion` option can be used to define what kind of completions
are generated for keywords.

Issue https://github.com/codemirror/lang-sql/pull/17
---
 src/complete.ts       | 14 ++++++++------
 src/sql.ts            | 22 +++++++++++++---------
 test/test-complete.ts | 20 +++++++++++++++++---
 3 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/src/complete.ts b/src/complete.ts
index 743f415..43e1662 100644
--- a/src/complete.ts
+++ b/src/complete.ts
@@ -204,12 +204,14 @@ export function completeFromSchema(schema: SQLNamespace,
   }
 }
 
-export function completeKeywords(keywords: {[name: string]: number}, upperCase: boolean) {
-  let completions =  Object.keys(keywords).map(keyword => ({
-    label: upperCase ? keyword.toUpperCase() : keyword,
-    type: keywords[keyword] == Type ? "type" : keywords[keyword] == Keyword ? "keyword" : "variable",
-    boost: -1
-  }))
+function completionType(tokenType: number) {
+  return tokenType == Type ? "type" : tokenType == Keyword ? "keyword" : "variable"
+}
+
+export function completeKeywords(keywords: {[name: string]: number}, upperCase: boolean,
+                                 build: (name: string, type: string) => Completion) {
+  let completions =  Object.keys(keywords)
+    .map(keyword => build(upperCase ? keyword.toUpperCase() : keyword, completionType(keywords[keyword])))
   return ifNotIn(["QuotedIdentifier", "SpecialVar", "String", "LineComment", "BlockComment", "."],
                  completeFromList(completions))
 }
diff --git a/src/sql.ts b/src/sql.ts
index 376223f..b486be9 100644
--- a/src/sql.ts
+++ b/src/sql.ts
@@ -149,18 +149,17 @@ export interface SQLConfig {
   defaultSchema?: string,
   /// When set to true, keyword completions will be upper-case.
   upperCaseKeywords?: boolean
+  /// Can be used to customize the completions generated for keywords.
+  keywordCompletion?: (label: string, type: string) => Completion
 }
 
+function defaultKeyword(label: string, type: string) { return {label, type, boost: -1} }
+
 /// Returns a completion source that provides keyword completion for
 /// the given SQL dialect.
-export function keywordCompletionSource(dialect: SQLDialect, upperCase = false): CompletionSource {
-  return completeKeywords(dialect.dialect.words, upperCase)
-}
-
-function keywordCompletion(dialect: SQLDialect, upperCase = false): Extension {
-  return dialect.language.data.of({
-    autocomplete: keywordCompletionSource(dialect, upperCase)
-  })
+export function keywordCompletionSource(dialect: SQLDialect, upperCase = false,
+                                        build?: (label: string, type: string) => Completion): CompletionSource {
+  return completeKeywords(dialect.dialect.words, upperCase, build || defaultKeyword)
 }
 
 /// Returns a completion sources that provides schema-based completion
@@ -183,7 +182,12 @@ function schemaCompletion(config: SQLConfig): Extension {
 /// extensions.
 export function sql(config: SQLConfig = {}) {
   let lang = config.dialect || StandardSQL
-  return new LanguageSupport(lang.language, [schemaCompletion(config), keywordCompletion(lang, !!config.upperCaseKeywords)])
+  return new LanguageSupport(lang.language, [
+    schemaCompletion(config),
+    lang.language.data.of({
+      autocomplete: keywordCompletionSource(lang, config.upperCaseKeywords, config.keywordCompletion)
+    })
+  ])
 }
 
 /// The standard SQL dialect.
diff --git a/test/test-complete.ts b/test/test-complete.ts
index ef3e87d..4a21397 100644
--- a/test/test-complete.ts
+++ b/test/test-complete.ts
@@ -1,16 +1,18 @@
 import {EditorState} from "@codemirror/state"
 import {CompletionContext, CompletionResult, CompletionSource} from "@codemirror/autocomplete"
-import {schemaCompletionSource, PostgreSQL, MySQL, SQLConfig, SQLDialect} from "@codemirror/lang-sql"
+import {schemaCompletionSource, keywordCompletionSource, PostgreSQL, MySQL, SQLConfig, SQLDialect} from "@codemirror/lang-sql"
 import ist from "ist"
 
-function get(doc: string, conf: SQLConfig & {explicit?: boolean} = {}) {
+function get(doc: string, conf: SQLConfig & {explicit?: boolean, keywords?: boolean} = {}) {
   let cur = doc.indexOf("|"), dialect = conf.dialect || PostgreSQL
   doc = doc.slice(0, cur) + doc.slice(cur + 1)
   let state = EditorState.create({
     doc,
     selection: {anchor: cur},
     extensions: [dialect, dialect.language.data.of({
-      autocomplete: schemaCompletionSource(Object.assign({dialect}, conf))
+      autocomplete: conf.keywords
+        ? keywordCompletionSource(dialect, conf.upperCaseKeywords, conf.keywordCompletion)
+        : schemaCompletionSource(Object.assign({dialect}, conf))
     })]
   })
   let result = state.languageDataAt<CompletionSource>("autocomplete", cur)[0](new CompletionContext(state, cur, !!conf.explicit))
@@ -213,4 +215,16 @@ describe("SQL completion", () => {
     ist(get("select foo.|", s)!.options[0].type, "constant")
     ist(get("select foo.|", s)!.options.length, 1)
   })
+
+  it("can complete keywords", () => {
+    ist(get("s|", {keywords: true})!.options.some(c => c.label == "select"))
+  })
+
+  it("can complete upper-case keywords", () => {
+    ist(get("s|", {keywords: true, upperCaseKeywords: true})!.options.some(c => c.label == "SELECT"))
+  })
+
+  it("can transform keyword completions", () => {
+    ist(get("s|", {keywords: true, keywordCompletion: l => ({label: l, type: "x"})})!.options.every(c => c.type == "x"))
+  })
 })