diff --git a/src/parser/common/entityCollector.ts b/src/parser/common/entityCollector.ts index 60bbbf97..a7579bcc 100644 --- a/src/parser/common/entityCollector.ts +++ b/src/parser/common/entityCollector.ts @@ -17,6 +17,7 @@ export enum StmtContextType { SELECT_STMT = 'selectStmt', INSERT_STMT = 'insertStmt', CREATE_FUNCTION_STMT = 'createFunctionStmt', + ALTER_TABLE_STMT = 'alterTableStmt', } export interface StmtContext { diff --git a/src/parser/postgresql/postgreEntityCollector.ts b/src/parser/postgresql/postgreEntityCollector.ts index 47eaf237..8bf5308e 100644 --- a/src/parser/postgresql/postgreEntityCollector.ts +++ b/src/parser/postgresql/postgreEntityCollector.ts @@ -1,4 +1,5 @@ import type { + AltertablestmtContext, ColumnCreateTableContext, ColumnNameCreateContext, CreateDatabaseContext, @@ -145,4 +146,10 @@ export class PostgreSqlEntityCollector extends EntityCollector implements Postgr exitCreatefunctionstmt(ctx: CreatefunctionstmtContext) { this.popStmt(); } + enterAltertablestmt(ctx: AltertablestmtContext) { + this.pushStmt(ctx, StmtContextType.ALTER_TABLE_STMT); + } + exitAltertablestmt(ctx: AltertablestmtContext) { + this.popStmt(); + } } diff --git a/test/parser/postgresql/suggestion/fixtures/suggestionWithEntity.sql b/test/parser/postgresql/suggestion/fixtures/suggestionWithEntity.sql index 6bbbf123..8824296e 100644 --- a/test/parser/postgresql/suggestion/fixtures/suggestionWithEntity.sql +++ b/test/parser/postgresql/suggestion/fixtures/suggestionWithEntity.sql @@ -8,4 +8,6 @@ INSERT INTO insert_tb SELECT id, age, FROM from_tb; CREATE TABLE sorted_census_data AS SELECT FROM unsorted_census_data; -CREATE TABLE sorted_census_data AS SELECT id, age, FROM unsorted_census_data; \ No newline at end of file +CREATE TABLE sorted_census_data AS SELECT id, age, FROM unsorted_census_data; + +ALTER TABLE my_table DROP a_column; \ No newline at end of file diff --git a/test/parser/postgresql/suggestion/suggestionWithEntity.test.ts b/test/parser/postgresql/suggestion/suggestionWithEntity.test.ts index f89bc016..ab840ba6 100644 --- a/test/parser/postgresql/suggestion/suggestionWithEntity.test.ts +++ b/test/parser/postgresql/suggestion/suggestionWithEntity.test.ts @@ -153,4 +153,25 @@ describe('PostgreSql Syntax Suggestion with collect entity', () => { expect(entities[1].entityContextType).toBe(EntityContextType.TABLE); expect(entities[1].belongStmt.isContainCaret).toBeTruthy(); }); + + test('alter table drop column', () => { + const pos: CaretPosition = { + lineNumber: 13, + column: 35, + }; + const sql = commentOtherLine(syntaxSql, pos.lineNumber); + + const syntaxes = postgre.getSuggestionAtCaretPosition(sql, pos)?.syntax; + const suggestion = syntaxes?.find( + (syn) => syn.syntaxContextType === EntityContextType.COLUMN + ); + expect(suggestion).not.toBeUndefined(); + expect(suggestion?.wordRanges.map((token) => token.text)).toEqual(['a_column']); + + const entities = postgre.getAllEntities(sql, pos); + expect(entities.length).toBe(1); + expect(entities[0].text).toBe('my_table'); + expect(entities[0].entityContextType).toBe(EntityContextType.TABLE); + expect(entities[0].belongStmt?.isContainCaret).toBeTruthy(); + }); });