UNPKG

flink-sql-language-server

Version:

A LSP-based language server for Apache Flink SQL

486 lines (484 loc) 18.5 kB
"use strict"; Object.defineProperty(exports, "__esModule", { value: true }); exports.CompletionVisitor = void 0; const cursor_1 = require("../cursor"); const FlinkSQLParser_1 = require("../lib/FlinkSQLParser"); const utils_1 = require("../utils"); const flinksql_relation_visitor_1 = require("./flinksql-relation.visitor"); const getRelationName = (relationPrimaryString) => { const splits = relationPrimaryString.split(flinksql_relation_visitor_1.RELATION_SEPARATOR); return splits[splits.length - 1]; }; function availableColumns(relation) { const columns = []; if (relation instanceof flinksql_relation_visitor_1.QueryRelation) { relation.relations.forEach((rel, name) => { const relationName = name !== rel.id ? name : undefined; rel.columns.forEach(col => { if (col.label !== flinksql_relation_visitor_1.STAR_LABEL && !col.label.includes(cursor_1.defaultCursorPlaceholder)) { columns.push({ relation: relationName !== undefined ? { category: 'table', label: relationName } : undefined, name: col.label, desc: col.data !== undefined ? col.data : undefined }); } }); }); } else { if (relation.sinkTableRelation !== undefined) { relation.sinkTableRelation.columns.forEach(col => { if (col.label !== flinksql_relation_visitor_1.STAR_LABEL && !col.label.includes(cursor_1.defaultCursorPlaceholder)) { columns.push({ relation: { category: 'table', label: relation.sinkTableRelation.tablePrimary.relationName }, name: col.label, desc: col.data !== undefined ? col.data : undefined }); } }); } } return columns; } function getTableMetadataFnFactory(getColumns) { return (relationPrimary) => { if (relationPrimary.type !== 'table') { return undefined; } const columns = getColumns({ table: relationPrimary.relationName, catalog: relationPrimary.catalogName, database: relationPrimary.databaseName }); return { columns: columns.map(c => { return { id: `${(0, flinksql_relation_visitor_1.relationPrimaryToString)(relationPrimary)}_${c.label}`, label: c.label, data: c }; }) }; }; } const SQL_SNIPPETS = [ { label: 'BEGIN STATEMENT SET', filterText: 'BEGIN', template: `BEGIN STATEMENT SET ; $0` }, { label: 'CREATE TEMPORARY TABLE ?', filterText: 'CREATE', template: `CREATE TEMPORARY TABLE \${1:__tableName} ( \${2:__columnDef} \${3:__columnType} -- ... ) COMMENT \${4:'__comment'} WITH ( \${5:'__config_key'} = \${6:'__config_value'} -- ... ) ;` }, { label: 'CREATE TEMPORARY VIEW ?', filterText: 'CREATE', template: `CREATE TEMPORARY VIEW \${1:__viewName} AS SELECT \${2:__projectItems} FROM \${3:__tableReferences} ;` }, { label: 'CREATE MATERIALIZED TABLE ?', filterText: 'CREATE', template: `CREATE MATERIALIZED TABLE \${1:__tableName} FRESHNESS = INTERVAL '\${2:__positiveInteger}' \${3:__timeUnit} AS SELECT \${4:__projectItems} FROM \${5:__tableReferences} ;` }, { label: 'CREATE MODEL ?', filterText: 'CREATE', template: `CREATE MODEL \${1:__modelName} INPUT (\${2:__inputColumns}) OUTPUT (\${3:__outputColumns}) WITH ( \${4:'__config_key'} = \${5:'__config_value'} -- ... ) ;` }, { label: 'INSERT INTO ?', filterText: 'INSERT', template: `INSERT INTO \${1:__tableName} SELECT \${2:__projectItems} FROM \${3:__tableReferences} ;` }, { label: 'SELECT ? FROM ?', filterText: 'SELECT', template: `SELECT \${1:__projectItems} FROM \${2:__tableReferences}` } ]; class CompletionVisitor extends flinksql_relation_visitor_1.FlinkSQLRelationVisitor { constructor(cursor, getColumns) { super(getColumns && getTableMetadataFnFactory(getColumns)); this.cursor = cursor; this.getColumns = getColumns; this.hasCompletions = false; this.completions = { type: 'other', snippets: [] }; } defaultResult() { return; } aggregateResult() { return; } onColumnReference() { return; } updateQueryCompletionItems(relation) { if (!this.hasCompletions && this.caretScope !== undefined) { this.hasCompletions = true; switch (this.caretScope.type) { case 'select-column': case 'spec-column': { const columns = availableColumns(relation); this.completions = { type: 'column', columns, relations: [...new Set(columns.filter(c => c.relation !== undefined).map(c => c.relation))], snippets: this.completions.snippets }; break; } case 'scoped-column': { let newCompletions = []; const relationName = this.caretScope.relation; const relationFound = this.relationDDLColumns.get(relationName); if (relationFound) { newCompletions = relationFound.map(c => ({ relation: { category: 'table', label: getRelationName(relationName) }, name: c.label })); } else { const cteFound = relation.findCTE(getRelationName(relationName)); if (cteFound) { newCompletions = cteFound.columns.map(c => ({ relation: { category: 'CTE', label: getRelationName(relationName) }, name: c.label })); } else { newCompletions = relation .findLocalRelation(relationName) ?.columns.filter(c => c.label !== flinksql_relation_visitor_1.STAR_LABEL && !c.label.includes(cursor_1.defaultCursorPlaceholder)) .map(c => ({ relation: { category: 'table', label: getRelationName(relationName) }, name: c.label })) ?? []; } } this.completions = { type: 'column', relations: [], columns: newCompletions, snippets: this.completions.snippets }; break; } case 'relation': { const ctes = relation.getCTENames().map(label => ({ label, category: 'CTE' })); const tableDDLNames = [...this.tableDDLs.keys()].map(label => ({ label: getRelationName(label), category: 'table' })); const viewDDLNames = [...this.viewDDLs.keys()].map(label => ({ label: getRelationName(label), category: 'view' })); this.completions = { type: 'relation', relations: [...ctes, ...tableDDLNames, ...viewDDLNames], snippets: this.completions.snippets, incompleteReference: this.caretScope.prefix.length > 0 ? this.caretScope.prefix : undefined }; break; } default: break; } } } handleRelationContext(ctx) { const multipartRelationName = ctx.text.split('.').map(t => (0, utils_1.sanitizeText)(t)); const lastPart = multipartRelationName[multipartRelationName.length - 1]; const category = ctx instanceof FlinkSQLParser_1.TableContext ? 'table' : 'view'; this.relationInStatementContext = this.relationPrimaryFromMultipart(multipartRelationName, category); if (!this.hasCompletions && this.cursor.isIn(lastPart)) { this.hasCompletions = true; this.caretScope = { type: 'relation', prefix: multipartRelationName.slice(0, -1) }; this.completions = { type: 'relation', relations: [...(category === 'table' ? this.tableDDLs : this.viewDDLs).keys()].map(label => ({ label: getRelationName(label), category })), snippets: this.completions.snippets, incompleteReference: this.caretScope.prefix.length > 0 ? this.caretScope.prefix : undefined }; } } handleScopedColumn() { if (!this.hasCompletions && this.relationInStatementContext) { this.hasCompletions = true; this.caretScope = { type: 'scoped-column', relation: (0, flinksql_relation_visitor_1.relationPrimaryToString)(this.relationInStatementContext) }; let newCompletions = []; const relationName = this.relationInStatementContext.relationName; const relationDDL = this.relationDDLColumns.get(relationName); if (relationDDL) { newCompletions = relationDDL.map(c => ({ relation: { category: 'table', label: relationName }, name: c.label })); } else if (this.getColumns) { newCompletions = this.getColumns({ catalog: this.relationInStatementContext.catalogName, database: this.relationInStatementContext.databaseName, table: this.relationInStatementContext.relationName }).map(c => ({ relation: { category: 'table', label: relationName }, name: c.label })); } this.completions = { type: 'column', relations: [], columns: newCompletions, snippets: this.completions.snippets }; } } handleDatabaseContext(ctx) { const multipartDatabaseName = ctx.text.split('.').map(t => (0, utils_1.sanitizeText)(t)); const lastPart = multipartDatabaseName[multipartDatabaseName.length - 1]; if (this.cursor.isIn(lastPart)) { this.hasCompletions = true; const incompleteReference = multipartDatabaseName.length > 1 ? multipartDatabaseName[0] : undefined; this.completions = { type: 'database', incompleteReference, snippets: [] }; } } getSuggestions() { return this.completions; } onRelation(relation) { if (relation instanceof flinksql_relation_visitor_1.TableRelation || relation instanceof flinksql_relation_visitor_1.InsertRelation) { return; } this.updateQueryCompletionItems(relation); } visitErrorNode(node) { super.visitErrorNode(node); if (this.cursor.isIn(node.text)) { if (node.parent instanceof FlinkSQLParser_1.StatementContext || node.parent instanceof FlinkSQLParser_1.ProgramContext) { this.completions.snippets.push(...SQL_SNIPPETS.filter(s => s.filterText.toUpperCase().indexOf(this.cursor.revert(node.text).toUpperCase()) === 0)); this.caretScope = { type: 'other' }; } } } visitStatement(ctx) { this.relationInStatementContext = undefined; super.visitStatement(ctx); } visitUse(ctx) { super.visitUse(ctx); if (ctx.CATALOG() && ctx.catalog()) { if (this.cursor.isIn(ctx.catalog().text)) { this.hasCompletions = true; this.completions = { type: 'catalog', snippets: [] }; } } else if (ctx.database()) { this.handleDatabaseContext(ctx.database()); } } visitAnalyze(ctx) { if (ctx.TABLE().text.toUpperCase() === 'TABLE') { this.handleRelationContext(ctx.table()); } if (ctx.partition()) { const cursorCol = ctx .partition() .columnAssignments() .columnAssignment() .find(assign => this.cursor.isIn(assign.columnName().text)); if (cursorCol !== undefined) { return this.handleScopedColumn(); } } if (ctx.columnsWithoutParenthesis()) { const cursorCol = ctx .columnsWithoutParenthesis() .columnName() .find(col => this.cursor.isIn(col.text)); if (cursorCol !== undefined) { return this.handleScopedColumn(); } } } visitCreateTable(ctx) { super.visitCreateTable(ctx); if (ctx.likeTable()) { super.visit(ctx.likeTable()); } } visitLikeTable(ctx) { this.handleRelationContext(ctx.table()); super.visitChildren(ctx); } visitAsDatabase(ctx) { this.handleDatabaseContext(ctx.database()); super.visitChildren(ctx); } visitDropDatabase(ctx) { this.handleDatabaseContext(ctx.database()); } visitDropTable(ctx) { this.handleRelationContext(ctx.table()); super.visitChildren(ctx); } visitAutoOptimizeTable(ctx) { this.handleRelationContext(ctx.table()); super.visitChildren(ctx); } visitDescribeCatalog(ctx) { if (this.cursor.isIn(ctx.catalog().text)) { this.hasCompletions = true; this.completions = { type: 'catalog', snippets: [] }; } } visitDescribeDatabase(ctx) { this.handleDatabaseContext(ctx.database()); } visitDescribeTable(ctx) { this.handleRelationContext(ctx.table()); super.visitChildren(ctx); } visitAlterDatabase(ctx) { this.handleDatabaseContext(ctx.database()); super.visitChildren(ctx); } visitAlterTable(ctx) { this.handleRelationContext(ctx.table()); super.visitChildren(ctx); } visitAlterView(ctx) { this.handleRelationContext(ctx.view()); super.visitChildren(ctx); } visitInsert(ctx) { super.visitInsert(ctx); if (ctx.INTO()?.text.toUpperCase() === 'INTO' || ctx.OVERWRITE()?.text.toUpperCase() === 'OVERWRITE') { this.handleRelationContext(ctx.table()); } if (this.relationInStatementContext) { if (ctx.partition()) { const cursorCol = ctx .partition() .columnAssignments() .columnAssignment() .find(assign => this.cursor.isIn(assign.columnName().text)); if (cursorCol !== undefined) { return this.handleScopedColumn(); } } if (ctx.columns()) { const cursorCol = ctx .columns() .columnName() .find(col => this.cursor.isIn(col.text)); if (cursorCol !== undefined) { return this.handleScopedColumn(); } } } } visitQueryPrimary(ctx) { super.visitQueryPrimary(ctx); if (this.currentQueryRelation !== undefined) { this.updateQueryCompletionItems(this.currentQueryRelation); } } visitTableReference(ctx) { if (ctx.tablePrimary().tableSource()?.table() !== undefined) { const multipartTableName = ctx .tablePrimary() .tableSource() .table() .text.split('.') .map(t => (0, utils_1.sanitizeText)(t)); const lastPart = multipartTableName[multipartTableName.length - 1]; if (this.cursor.isIn(lastPart)) { this.caretScope = { type: 'relation', prefix: multipartTableName.slice(0, -1) }; } } super.visitTableReference(ctx); } visitColumnReference(ctx) { super.visitColumnReference(ctx); if (this.currentQueryRelation !== undefined && this.cursor.isIn(ctx.text)) { if (this.currentQueryRelation.currentClause === 'select') { this.caretScope = { type: 'select-column' }; } else { this.caretScope = { type: 'spec-column' }; } } } visitDereference(ctx) { super.visitDereference(ctx); if (this.currentQueryRelation !== undefined && this.cursor.isIn(ctx.columnName().text)) { this.caretScope = { type: 'scoped-column', relation: (0, utils_1.sanitizeText)(ctx.table().tableName().text) }; } } visitJoinCondition(ctx) { super.visitJoinCondition(ctx); if (ctx.columns()) { this.handleColumnsContext(ctx.columns()); } } visitPartitionedBy(ctx) { super.visitChildren(ctx); if (ctx.columns()) { this.handleColumnsContext(ctx.columns()); } } handleColumnsContext(ctx) { if (this.currentQueryRelation !== undefined) { const cursorCol = ctx?.columnName().find(col => this.cursor.isIn(col.text)); if (cursorCol !== undefined) { if (this.currentQueryRelation.currentClause === 'select') { this.caretScope = { type: 'select-column' }; } else { this.caretScope = { type: 'spec-column' }; } } } } } exports.CompletionVisitor = CompletionVisitor;