dt-sql-parser
Version:
SQL Parsers for BigData, built with antlr4
333 lines (332 loc) • 13.4 kB
JavaScript
var __rest = (this && this.__rest) || function (s, e) {
var t = {};
for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
t[p] = s[p];
if (s != null && typeof Object.getOwnPropertySymbols === "function")
for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
t[p[i]] = s[p[i]];
}
return t;
};
import { isToken, Token } from 'antlr4ng';
import { SimpleStack } from './simpleStack';
import { ctxToText, isWordRange, tokenToWord, } from './textAndWord';
import { EntityContextType } from './types';
/**
* TODO: more stmt type should be supported.
*/
export var StmtContextType;
(function (StmtContextType) {
/** A self-contained and complete statement */
StmtContextType["COMMON_STMT"] = "commonStmt";
StmtContextType["CREATE_CATALOG_STMT"] = "createCatalogStmt";
StmtContextType["CREATE_DATABASE_STMT"] = "createDatabaseStmt";
StmtContextType["CREATE_TABLE_STMT"] = "createTableStmt";
StmtContextType["CREATE_VIEW_STMT"] = "createViewStmt";
StmtContextType["SELECT_STMT"] = "selectStmt";
StmtContextType["INSERT_STMT"] = "insertStmt";
StmtContextType["CREATE_FUNCTION_STMT"] = "createFunctionStmt";
StmtContextType["ALTER_TABLE_STMT"] = "alterTableStmt";
})(StmtContextType || (StmtContextType = {}));
export function toStmtContext(ctx, type, input, rootStmt, parentStmt, isContainCaret) {
const text = ctxToText(ctx, input);
if (!text)
return null;
const { text: _ } = text, position = __rest(text, ["text"]);
return {
stmtContextType: type,
position,
rootStmt: rootStmt !== null && rootStmt !== void 0 ? rootStmt : null,
parentStmt: parentStmt !== null && parentStmt !== void 0 ? parentStmt : null,
isContainCaret,
};
}
/**
* entity's attribute
* @key comment: entity's comment attribute
* @key colType: column entity's type attribute
* @key alias: entity's alias attribute
* */
export var AttrName;
(function (AttrName) {
AttrName["comment"] = "_comment";
AttrName["colType"] = "_colType";
AttrName["alias"] = "_alias";
})(AttrName || (AttrName = {}));
export function isCommonEntityContext(entity) {
if (!entity)
return false;
return 'relatedEntities' in entity && !('arguments' in entity);
}
export function isFuncEntityContext(entity) {
if (!entity)
return false;
return 'arguments' in entity;
}
export function isColumnEntityContext(entity) {
if (!entity)
return false;
return AttrName.colType in entity;
}
export function toEntityContext(ctx, type, input, belongStmt, attrInfo) {
var _a, _b;
const word = ctxToText(ctx, input);
if (!word)
return null;
const { text, startLine, endLine } = word, rest = __rest(word, ["text", "startLine", "endLine"]);
const position = Object.assign(Object.assign({}, rest), { line: startLine });
let entityInfo = {
entityContextType: type,
text,
position,
belongStmt,
[AttrName.comment]: null,
};
switch (entityInfo.entityContextType) {
case EntityContextType.FUNCTION:
case EntityContextType.FUNCTION_CREATE: {
entityInfo.relatedEntities = null;
entityInfo.arguments = null;
break;
}
case EntityContextType.COLUMN:
case EntityContextType.COLUMN_CREATE:
entityInfo[AttrName.colType] = null;
break;
default:
entityInfo.relatedEntities = null;
break;
}
if (attrInfo) {
for (let k = 0; k < attrInfo.length; k++) {
const attributeName = (_a = attrInfo[k]) === null || _a === void 0 ? void 0 : _a.attrName;
const attrToken = findAttribute(ctx, attributeName, (_b = attrInfo[k]) === null || _b === void 0 ? void 0 : _b.endContextList);
if (attrToken) {
const attrVal = isToken(attrToken)
? tokenToWord(attrToken, input)
: ctxToText(attrToken, input);
if (attrVal) {
if (isWordRange(attrVal)) {
entityInfo = Object.assign(entityInfo, {
[attributeName]: attrVal,
});
}
else {
const { startIndex, endIndex, startColumn, endColumn, text } = attrVal, rest = __rest(attrVal, ["startIndex", "endIndex", "startColumn", "endColumn", "text"]);
entityInfo = Object.assign(entityInfo, {
[attributeName]: {
startIndex,
endIndex,
startColumn,
endColumn,
text,
line: rest === null || rest === void 0 ? void 0 : rest.startLine,
},
});
}
}
}
}
}
return entityInfo;
}
export function findAttribute(ctx, keyName, endContextNameList) {
var _a, _b;
const parent = (ctx === null || ctx === void 0 ? void 0 : ctx.parent) || null;
let attrVal = null;
if (parent === null || parent === void 0 ? void 0 : parent[keyName]) {
attrVal = (parent === null || parent === void 0 ? void 0 : parent[keyName]) || null;
return attrVal;
}
else {
if (((_a = parent === null || parent === void 0 ? void 0 : parent.constructor) === null || _a === void 0 ? void 0 : _a.name) && !endContextNameList.includes((_b = parent === null || parent === void 0 ? void 0 : parent.constructor) === null || _b === void 0 ? void 0 : _b.name)) {
attrVal = findAttribute(parent, keyName, endContextNameList);
}
if (!attrVal) {
if (parent === null || parent === void 0 ? void 0 : parent.children) {
attrVal = findAttributeChildren(parent, keyName);
}
}
}
return attrVal;
}
function findAttributeChildren(ctx, keyName) {
const visitChildren = (ctx === null || ctx === void 0 ? void 0 : ctx.children) || [];
let attrVal = null;
if (visitChildren.length) {
for (let i = 0; i < visitChildren.length; i++) {
const child = visitChildren[i] || null;
if (child === null || child === void 0 ? void 0 : child[keyName]) {
attrVal = (child === null || child === void 0 ? void 0 : child[keyName]) || null;
return attrVal;
}
else {
attrVal = findAttributeChildren(child, keyName);
}
}
}
return attrVal;
}
/**
* @todo: Handle alias, includes column alias, table alias, query as alias and so on.
* @todo: [may be need] Combine the entities in each clause.
*/
export class EntityCollector {
constructor(input, allTokens, caretTokenIndex) {
this._input = input;
this._allTokens = allTokens || [];
this._caretTokenIndex = caretTokenIndex !== null && caretTokenIndex !== void 0 ? caretTokenIndex : -1;
this._entitiesSet = new Set();
this._stmtStack = new SimpleStack();
this._entityStack = new SimpleStack();
this._rootStmt = null;
}
visitTerminal() { }
visitErrorNode() { }
enterEveryRule() { }
exitEveryRule() { }
getRootStmt() {
return this._rootStmt;
}
getEntities() {
return Array.from(this._entitiesSet);
}
enterProgram() {
this._entitiesSet.clear();
this._stmtStack.clear();
this._entityStack.clear();
this._rootStmt = null;
}
/**
* The antlr4 will ignore hidden tokens, if we type whitespace at the end of a statement,
* the whitespace token will not as stop token, so we consider the whitespace token as a part of the nonhidden token in front of it
*/
getPrevNonHiddenTokenIndex(caretTokenIndex) {
if (this._allTokens[caretTokenIndex].channel !== Token.HIDDEN_CHANNEL)
return caretTokenIndex;
for (let i = caretTokenIndex - 1; i >= 0; i--) {
const token = this._allTokens[i];
if (token.channel !== Token.HIDDEN_CHANNEL) {
// If prev nonhidden token is ';', the current token does not belong to any statement.
return token.text === ';' ? Infinity : token.tokenIndex;
}
}
return Infinity;
}
pushStmt(ctx, type) {
var _a;
let isContainCaret;
if (this._caretTokenIndex >= 0) {
isContainCaret =
!!ctx.start &&
!!ctx.stop &&
ctx.start.tokenIndex <= this._caretTokenIndex &&
ctx.stop.tokenIndex >= this.getPrevNonHiddenTokenIndex(this._caretTokenIndex);
}
const stmtContext = toStmtContext(ctx, type, this._input, this._rootStmt, this._stmtStack.peek(), isContainCaret);
if (stmtContext) {
if (this._stmtStack.isEmpty() ||
((_a = this._stmtStack.peek()) === null || _a === void 0 ? void 0 : _a.stmtContextType) === StmtContextType.COMMON_STMT) {
this._rootStmt = stmtContext;
}
this._stmtStack.push(stmtContext);
}
return stmtContext;
}
popStmt() {
const stmtContext = this._stmtStack.pop();
if (stmtContext && this._rootStmt === stmtContext) {
this._rootStmt = this._stmtStack.peek();
if (!this._entityStack.isEmpty()) {
this.combineEntitiesAndAdd(stmtContext);
}
}
return stmtContext;
}
pushEntity(ctx, type, attrInfo) {
const entityContext = toEntityContext(ctx, type, this._input, this._stmtStack.peek(), attrInfo);
if (entityContext) {
if (this._stmtStack.isEmpty()) {
this._entitiesSet.add(entityContext);
}
else {
// If is inside a statement
this._entityStack.push(entityContext);
}
}
return entityContext;
}
/**
* Combine entities that inside a single statement.
* e.g. combine tableName and column if they are inside a same createTableStatement.
* Then add combined entities into result.
*/
combineEntitiesAndAdd(stmtContext) {
const entitiesInsideStmt = [];
while (!this._entityStack.isEmpty() &&
(this._entityStack.peek().belongStmt === stmtContext ||
this._entityStack.peek().belongStmt.rootStmt === stmtContext)) {
entitiesInsideStmt.unshift(this._entityStack.pop());
}
const combinedEntities = this.combineRootStmtEntities(stmtContext, entitiesInsideStmt);
while (combinedEntities.length) {
const entity = combinedEntities.shift();
entity && this._entitiesSet.add(entity);
}
}
/**
* Combined all entities under a rootStmt.
*/
combineRootStmtEntities(stmtContext, entitiesInsideStmt) {
if (stmtContext.stmtContextType === StmtContextType.CREATE_VIEW_STMT ||
stmtContext.stmtContextType === StmtContextType.CREATE_TABLE_STMT) {
return this.combineCreateTableOrViewStmtEntities(stmtContext, entitiesInsideStmt);
}
return entitiesInsideStmt;
}
combineCreateTableOrViewStmtEntities(stmtContext, entitiesInsideStmt) {
const columns = [];
const relatedEntities = [];
let mainEntity = null;
const finalEntities = entitiesInsideStmt.reduce((result, entity) => {
if (entity.belongStmt !== stmtContext) {
if (entity.entityContextType !== EntityContextType.COLUMN &&
entity.entityContextType !== EntityContextType.COLUMN_CREATE) {
relatedEntities.push(entity);
result.push(entity);
}
return result;
}
if (entity.entityContextType === EntityContextType.COLUMN_CREATE) {
columns.push(entity);
}
else if (entity.entityContextType === EntityContextType.TABLE_CREATE ||
entity.entityContextType === EntityContextType.VIEW_CREATE) {
mainEntity = entity;
result.push(entity);
return result;
}
else if (entity.entityContextType !== EntityContextType.COLUMN) {
relatedEntities.push(entity);
result.push(entity);
}
return result;
}, []);
if (mainEntity && columns.length) {
if (isCommonEntityContext(mainEntity)) {
mainEntity = Object.assign(mainEntity, {
columns,
});
}
}
if (mainEntity && relatedEntities.length) {
if (isCommonEntityContext(mainEntity) || isFuncEntityContext(mainEntity)) {
mainEntity = Object.assign(mainEntity, {
relatedEntities,
});
}
}
return finalEntities;
}
}