Add types for database and catalog and use them

Author: the21st

packages/server/src/complete/CompletionItemUtils.ts

@@ -2,11 +2,13 @@ import { CompletionItem, CompletionItemKind } from 'vscode-languageserver-types'
 import { DbFunction } from '../database_libs/AbstractClient'
 
 export const ICONS = {
-  KEYWORD: CompletionItemKind.Text,
-  COLUMN: CompletionItemKind.Interface,
-  TABLE: CompletionItemKind.Field,
-  FUNCTION: CompletionItemKind.Property,
-  ALIAS: CompletionItemKind.Variable,
+  KEYWORD: CompletionItemKind.Keyword,
+  COLUMN: CompletionItemKind.Field,
+  TABLE: CompletionItemKind.Constant,
+  DATABASE: CompletionItemKind.Enum,
+  CATALOG: CompletionItemKind.Folder,
+  FUNCTION: CompletionItemKind.Event,
+  ALIAS: CompletionItemKind.Constant,
   UTILITY: CompletionItemKind.Event,
 }
 

packages/server/src/complete/Identifier.ts

@@ -1,13 +1,5 @@
 import { CompletionItem, CompletionItemKind } from 'vscode-languageserver-types'
-
-export const ICONS = {
-  KEYWORD: CompletionItemKind.Text,
-  COLUMN: CompletionItemKind.Interface,
-  TABLE: CompletionItemKind.Field,
-  FUNCTION: CompletionItemKind.Property,
-  ALIAS: CompletionItemKind.Variable,
-  UTILITY: CompletionItemKind.Event,
-}
+import { ICONS } from './CompletionItemUtils'
 
 type OnClause = 'FROM' | 'ALTER TABLE' | 'OTHERS'
 export class Identifier {
@@ -44,18 +36,37 @@ export class Identifier {
   toCompletionItem(): CompletionItem {
     const idx = this.lastToken.lastIndexOf('.')
     const label = this.identifier.substring(idx + 1)
-    let kindName: string
-    if (this.kind === ICONS.TABLE) {
+    if (
+      this.kind === ICONS.TABLE ||
+      this.kind === ICONS.DATABASE ||
+      this.kind === ICONS.CATALOG
+    ) {
       let tableName = label
       const i = tableName.lastIndexOf('.')
       if (i > 0) {
         tableName = label.substring(i + 1)
       }
-      kindName = 'table'
-    } else {
-      kindName = 'column'
     }
 
+    const kindName = (() => {
+      switch (this.kind) {
+        case ICONS.TABLE:
+          return 'table'
+        case ICONS.DATABASE:
+          return 'schema'
+        case ICONS.CATALOG:
+          return 'database'
+        case ICONS.FUNCTION:
+          return 'function'
+        case ICONS.ALIAS:
+          return 'table'
+        case ICONS.COLUMN:
+          return 'column'
+        default:
+          return 'column'
+      }
+    })()
+
     const item: CompletionItem = {
       label: label,
       detail: `${kindName} ${this.detail}`,

packages/server/src/complete/candidates/createTableCandidates.ts

@@ -36,35 +36,51 @@ export function createCatalogDatabaseAndTableCandidates(
     }
     const qualificationLevelNeeded = qualificationNeeded - qualificationLevel
     switch (qualificationLevelNeeded) {
-      case 0:
-        return [getFullyQualifiedTableName(table)]
-      case 1:
-        if (table.catalog && table.database) {
-          return [table.catalog + '.' + table.database]
-        }
-        if (table.database) {
-          return [table.database]
+      case 0: {
+        const tableIdentifier = new Identifier(
+          lastToken,
+          getFullyQualifiedTableName(table),
+          '',
+          ICONS.TABLE,
+          onFromClause ? 'FROM' : 'OTHERS'
+        )
+        return [tableIdentifier]
+      }
+      case 1: {
+        const qualifiedDatabaseName =
+          table.catalog && table.database
+            ? table.catalog + '.' + table.database
+            : table.database
+
+        if (qualifiedDatabaseName !== null) {
+          const databaseIdentifier = new Identifier(
+            lastToken,
+            qualifiedDatabaseName,
+            '',
+            ICONS.DATABASE,
+            onFromClause ? 'FROM' : 'OTHERS'
+          )
+          return [databaseIdentifier]
         }
         break
+      }
       case 2:
         if (table.catalog) {
-          return [table.catalog]
+          const catalogIdentifier = new Identifier(
+            lastToken,
+            table.catalog,
+            '',
+            ICONS.CATALOG,
+            onFromClause ? 'FROM' : 'OTHERS'
+          )
+          return [catalogIdentifier]
         }
         break
     }
     return []
   })
 
   return qualifiedEntities
-    .map((databaseEntity) => {
-      return new Identifier(
-        lastToken,
-        databaseEntity,
-        '',
-        ICONS.TABLE,
-        onFromClause ? 'FROM' : 'OTHERS'
-      )
-    })
     .filter((item) => item.matchesLastToken())
     .map((item) => item.toCompletionItem())
 }