Merge pull request #5 from deepnote/ss/better-empty-select-candidates

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

packages/server/src/complete/complete.ts

@@ -244,13 +244,29 @@ class Completer {
     this.addCandidatesForSelectStar(fromNodes, schemaAndSubqueries)
     const expectedLiteralNodes =
       e.expected?.filter(
-        (v): v is ExpectedLiteralNode => v.type === 'literal'
+        (v): v is ExpectedLiteralNode =>
+          v.type === 'literal' && hasAtLeastTwoLetters(v.text)
       ) || []
     this.addCandidatesForExpectedLiterals(expectedLiteralNodes)
     this.addCandidatesForFunctions()
-    this.addCandidatesForScopedColumns(fromNodes, schemaAndSubqueries)
+
+    const { addedSome: addedSomeScopedColumnCandidates } =
+      this.addCandidatesForScopedColumns(fromNodes, schemaAndSubqueries)
+    if (!addedSomeScopedColumnCandidates) {
+      this.addCandidatesForUnscopedColumns(fromNodes, schemaAndSubqueries)
+    }
+
     this.addCandidatesForAliases(fromNodes)
-    this.addCandidatesForTables(schemaAndSubqueries, true)
+
+    const fromNodesContainingCursor = fromNodes.filter((tableNode) =>
+      isPosInLocation(tableNode.location, this.pos)
+    )
+    const isCursorInsideFromClause = fromNodesContainingCursor.length > 0
+    if (isCursorInsideFromClause) {
+      // only add table candidates if the cursor is inside a FROM clause or JOIN clause, etc.
+      this.addCandidatesForTables(schemaAndSubqueries, true)
+    }
+
     if (logger.isDebugEnabled())
       logger.debug(
         `candidates for error returns: ${JSON.stringify(this.candidates)}`
@@ -378,14 +394,20 @@ class Completer {
     console.timeEnd('addCandidatesForSelectStar')
   }
 
-  addCandidatesForScopedColumns(fromNodes: FromTableNode[], tables: Table[]) {
+  addCandidatesForScopedColumns(
+    fromNodes: FromTableNode[],
+    tables: Table[]
+  ): { addedSome: boolean } {
     console.time('addCandidatesForScopedColumns')
+    let addedSome = false
     createCandidatesForScopedColumns(fromNodes, tables, this.lastToken).forEach(
       (v) => {
+        addedSome = true
         this.addCandidate(v)
       }
     )
     console.timeEnd('addCandidatesForScopedColumns')
+    return { addedSome }
   }
 
   addCandidatesForUnscopedColumns(fromNodes: FromTableNode[], tables: Table[]) {
@@ -419,3 +441,7 @@ export function complete(
   console.timeEnd('complete')
   return { candidates: candidates, error: completer.error }
 }
+
+function hasAtLeastTwoLetters(value: string): boolean {
+  return /[a-zA-Z].*[a-zA-Z]/.test(value)
+}

packages/server/test/complete.test.ts

@@ -174,14 +174,24 @@ describe('on blank space', () => {
 
   test('complete inside SELECT', () => {
     const result = complete('SELECT ', { line: 0, column: 7 }, SIMPLE_SCHEMA)
-    expect(result.candidates.length).toEqual(12) // TODO whare are they?
     const expected = [
       expect.objectContaining({ label: 'array_concat()' }),
       expect.objectContaining({ label: 'array_contains()' }),
     ]
     expect(result.candidates).toEqual(expect.arrayContaining(expected))
   })
 
+  test('complete after SELECT  FROM schema2.table2', () => {
+    const result = complete(
+      'SELECT  FROM schema2.table2',
+      { line: 0, column: 8 },
+      SIMPLE_NESTED_SCHEMA
+    )
+    expect(result.candidates).not.toContainEqual(
+      expect.objectContaining({ label: 'TABLE1' })
+    )
+  })
+
   test('complete function inside WHERE select star', () => {
     const result = complete(
       'SELECT * FROM tab1 WHERE arr',
@@ -450,7 +460,6 @@ describe('Fully qualified table names', () => {
       { line: 0, column: 31 },
       SIMPLE_NESTED_SCHEMA
     )
-    expect(result.candidates.length).toEqual(1)
     const expected = [expect.objectContaining({ label: 'table3' })]
     expect(result.candidates).toEqual(expect.arrayContaining(expected))
   })
@@ -880,18 +889,18 @@ test('complete aliased column inside function', () => {
   expect(result.candidates[0].label).toEqual('department_id')
 })
 
-test('complete column inside function', () => {
-  const sql = `SELECT TO_CHAR(empl, 'MM/DD/YYYY') FROM employees x`
+test('complete table inside function', () => {
+  const sql = `SELECT TO_CHAR(empl, 'MM/DD/YYYY') FROM employees`
   const result = complete(sql, { line: 0, column: 19 }, COMPLEX_SCHEMA)
-  expect(result.candidates.length).toEqual(1)
-  expect(result.candidates[0].label).toEqual('employees')
+  const expected = [expect.objectContaining({ label: 'employees' })]
+  expect(result.candidates).toEqual(expect.arrayContaining(expected))
 })
 
 test('complete an alias inside function', () => {
   const sql = `SELECT TO_CHAR(an_ali, 'MM/DD/YYYY') FROM employees an_alias`
   const result = complete(sql, { line: 0, column: 21 }, COMPLEX_SCHEMA)
-  expect(result.candidates.length).toEqual(1)
-  expect(result.candidates[0].label).toEqual('an_alias')
+  const expected = [expect.objectContaining({ label: 'an_alias' })]
+  expect(result.candidates).toEqual(expect.arrayContaining(expected))
 })
 
 describe('From clause subquery', () => {