support column completion for alter table statement (joe-re#125)

Author: joe-re

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

@@ -17,8 +17,7 @@ export function createCandidatesForColumnsOfAnyTable(
         lastToken,
         column.columnName,
         column.description,
-        ICONS.TABLE,
-        'FROM'
+        ICONS.TABLE
       )
     })
     .filter((item) => item.matchesLastToken())

packages/server/src/complete/complete.ts

@@ -9,6 +9,7 @@ import {
   ParseError,
   ExpectedLiteralNode,
   AST,
+  AlterTableStatement,
 } from '@joe-re/sql-parser'
 import log4js from 'log4js'
 import { CompletionItem } from 'vscode-languageserver-types'
@@ -281,6 +282,17 @@ class Completer {
     }
   }
 
+  addCandidatesForParsedAlterTableStatement(ast: AlterTableStatement) {
+    if (ast.command.type === 'alter_table_drop_column') {
+      if (isPosInLocation(ast.command.column.location, this.pos)) {
+        const table = this.schema.tables.find((v) => v.tableName === ast.table)
+        this.addCandidatesForColumnsOfAnyTable(
+          table ? [table] : this.schema.tables
+        )
+      }
+    }
+  }
+
   addCandidatesForParsedSelectQuery(ast: SelectStatement) {
     this.addCandidatesForBasicKeyword()
     if (Array.isArray(ast.columns)) {
@@ -325,6 +337,8 @@ class Completer {
       this.addCandidatesForParsedDeleteStatement(ast)
     } else if (ast.type === 'select') {
       this.addCandidatesForParsedSelectQuery(ast)
+    } else if (ast.type === 'alter_table') {
+      this.addCandidatesForParsedAlterTableStatement(ast)
     } else {
       console.log(`AST type not supported yet: ${ast.type}`)
     }

packages/server/test/complete/alter_table.test.ts

@@ -56,17 +56,72 @@ describe('complete ALTER TABLE statement', () => {
   })
 
   describe('Column name completion', () => {
-    // TODO: Deal with column completion
-    it.skip('completes column name', () => {
-      const sql = 'ALTER TABLE TABLE1 MODIFY C'
-      const result = complete(
-        sql,
-        { line: 0, column: sql.length },
-        SIMPLE_SCHEMA
-      )
-      expect(result.candidates.length).toEqual(2)
-      expect(result.candidates[0].label).toEqual('COLUMN1')
-      expect(result.candidates[0].insertText).toEqual('COLUMN1')
+    describe('ADD command keyword', () => {
+      it("doen't completes column name", () => {
+        const sql = 'ALTER TABLE TABLE1 ADD C'
+        const result = complete(
+          sql,
+          { line: 0, column: sql.length },
+          SIMPLE_SCHEMA
+        )
+        expect(result.candidates.length).toEqual(0)
+      })
+    })
+
+    describe('DROP COLUMN command keyword', () => {
+      it("doen't completes column name", () => {
+        const sql = 'ALTER TABLE TABLE1 DROP COLUMN C'
+        const result = complete(
+          sql,
+          { line: 0, column: sql.length },
+          SIMPLE_SCHEMA
+        )
+        expect(result.candidates.length).toEqual(2)
+        expect(result.candidates[0].label).toEqual('COLUMN1')
+        expect(result.candidates[0].insertText).toEqual('COLUMN1')
+      })
+    })
+
+    describe('MODIFY command keyword', () => {
+      it('completes column name', () => {
+        const sql = 'ALTER TABLE TABLE1 MODIFY C'
+        const result = complete(
+          sql,
+          { line: 0, column: sql.length },
+          SIMPLE_SCHEMA
+        )
+        expect(result.candidates.length).toEqual(2)
+        expect(result.candidates[0].label).toEqual('COLUMN1')
+        expect(result.candidates[0].insertText).toEqual('COLUMN1')
+      })
+    })
+
+    describe('MODIFY COLUMN command keyword', () => {
+      it('completes column name', () => {
+        const sql = 'ALTER TABLE TABLE1 MODIFY COLUMN C'
+        const result = complete(
+          sql,
+          { line: 0, column: sql.length },
+          SIMPLE_SCHEMA
+        )
+        expect(result.candidates.length).toEqual(2)
+        expect(result.candidates[0].label).toEqual('COLUMN1')
+        expect(result.candidates[0].insertText).toEqual('COLUMN1')
+      })
+    })
+
+    describe('ALTER COLUMN command keyword', () => {
+      it('completes column name', () => {
+        const sql = 'ALTER TABLE TABLE1 ALTER COLUMN C'
+        const result = complete(
+          sql,
+          { line: 0, column: sql.length },
+          SIMPLE_SCHEMA
+        )
+        expect(result.candidates.length).toEqual(2)
+        expect(result.candidates[0].label).toEqual('COLUMN1')
+        expect(result.candidates[0].insertText).toEqual('COLUMN1')
+      })
     })
   })
 })