rawsql-ts
Version:
[beta]High-performance SQL parser and AST analyzer written in TypeScript. Provides fast parsing and advanced transformation capabilities.
396 lines • 15.6 kB
JavaScript
import { CommonTable, FetchClause, ForClause, FromClause, GroupByClause, HavingClause, JoinClause, JoinOnClause, JoinUsingClause, LimitClause, OffsetClause, OrderByClause, OrderByItem, ParenSource, SelectClause, SelectItem, SourceExpression, SubQuerySource, TableSource, WhereClause, WindowFrameClause, WithClause } from "../models/Clause";
import { BinarySelectQuery, SimpleSelectQuery, ValuesQuery } from "../models/SelectQuery";
import { ArrayExpression, BetweenExpression, BinaryExpression, CaseExpression, CaseKeyValuePair, CastExpression, FunctionCall, InlineQuery, ParenExpression, SwitchCaseArgument, TupleExpression, UnaryExpression, RawString } from "../models/ValueComponent";
import { CTECollector } from "./CTECollector";
/**
* A visitor that collects all table source names from a SQL query structure.
*
* When selectableOnly is true (default behavior):
* - Includes only table sources from FROM and JOIN clauses
* - Excludes inline queries, subqueries, and CTEs
*
* When selectableOnly is false:
* - Scans all parts of the query including WITH clauses, subqueries, etc.
* - Collects all table sources from the entire query
* - Excludes tables that are managed by CTEs
*
* For UNION-like queries, it scans both the left and right parts.
*/
export class TableSourceCollector {
constructor(selectableOnly = true) {
this.tableSources = [];
this.visitedNodes = new Set();
this.tableNameMap = new Map();
this.cteNames = new Set();
this.isRootVisit = true;
this.selectableOnly = selectableOnly;
this.handlers = new Map();
// Setup handlers for query components
this.handlers.set(SimpleSelectQuery.kind, (expr) => this.visitSimpleSelectQuery(expr));
this.handlers.set(BinarySelectQuery.kind, (expr) => this.visitBinarySelectQuery(expr));
this.handlers.set(ValuesQuery.kind, (expr) => this.visitValuesQuery(expr));
// WITH clause and common tables
this.handlers.set(WithClause.kind, (expr) => this.visitWithClause(expr));
this.handlers.set(CommonTable.kind, (expr) => this.visitCommonTable(expr));
// Handlers for FROM and JOIN components
this.handlers.set(FromClause.kind, (expr) => this.visitFromClause(expr));
this.handlers.set(JoinClause.kind, (expr) => this.visitJoinClause(expr));
this.handlers.set(JoinOnClause.kind, (expr) => this.visitJoinOnClause(expr));
this.handlers.set(JoinUsingClause.kind, (expr) => this.visitJoinUsingClause(expr));
// Source components
this.handlers.set(SourceExpression.kind, (expr) => this.visitSourceExpression(expr));
this.handlers.set(TableSource.kind, (expr) => this.visitTableSource(expr));
this.handlers.set(ParenSource.kind, (expr) => this.visitParenSource(expr));
this.handlers.set(SubQuerySource.kind, (expr) => this.visitSubQuerySource(expr));
this.handlers.set(InlineQuery.kind, (expr) => this.visitInlineQuery(expr));
// Only register these handlers when not in selectableOnly mode
if (!selectableOnly) {
// Additional clause handlers for full scanning
this.handlers.set(WhereClause.kind, (expr) => this.visitWhereClause(expr));
this.handlers.set(GroupByClause.kind, (expr) => this.visitGroupByClause(expr));
this.handlers.set(HavingClause.kind, (expr) => this.visitHavingClause(expr));
this.handlers.set(OrderByClause.kind, (expr) => this.visitOrderByClause(expr));
this.handlers.set(WindowFrameClause.kind, (expr) => this.visitWindowFrameClause(expr));
this.handlers.set(LimitClause.kind, (expr) => this.visitLimitClause(expr));
this.handlers.set(OffsetClause.kind, (expr) => this.visitOffsetClause(expr));
this.handlers.set(FetchClause.kind, (expr) => this.visitFetchClause(expr));
this.handlers.set(ForClause.kind, (expr) => this.visitForClause(expr));
this.handlers.set(OrderByItem.kind, (expr) => this.visitOrderByItem(expr));
this.handlers.set(SelectClause.kind, (expr) => this.visitSelectClause(expr));
this.handlers.set(SelectItem.kind, (expr) => this.visitSelectItem(expr));
// Value components that might contain table references
this.handlers.set(ParenExpression.kind, (expr) => this.visitParenExpression(expr));
this.handlers.set(BinaryExpression.kind, (expr) => this.visitBinaryExpression(expr));
this.handlers.set(UnaryExpression.kind, (expr) => this.visitUnaryExpression(expr));
this.handlers.set(CaseExpression.kind, (expr) => this.visitCaseExpression(expr));
this.handlers.set(CaseKeyValuePair.kind, (expr) => this.visitCaseKeyValuePair(expr));
this.handlers.set(SwitchCaseArgument.kind, (expr) => this.visitSwitchCaseArgument(expr));
this.handlers.set(BetweenExpression.kind, (expr) => this.visitBetweenExpression(expr));
this.handlers.set(FunctionCall.kind, (expr) => this.visitFunctionCall(expr));
this.handlers.set(ArrayExpression.kind, (expr) => this.visitArrayExpression(expr));
this.handlers.set(TupleExpression.kind, (expr) => this.visitTupleExpression(expr));
this.handlers.set(CastExpression.kind, (expr) => this.visitCastExpression(expr));
}
}
/**
* Gets all collected table sources
*/
getTableSources() {
return this.tableSources;
}
/**
* Reset the collection of table sources
*/
reset() {
this.tableSources = [];
this.tableNameMap.clear();
this.visitedNodes.clear();
this.cteNames.clear();
}
/**
* Gets a unique identifier for a table source
*/
getTableIdentifier(source) {
// Use QualifiedName for identifier (dot-joined string)
if (source.qualifiedName.namespaces && source.qualifiedName.namespaces.length > 0) {
return source.qualifiedName.namespaces.map(ns => ns.name).join('.') + '.' + (source.qualifiedName.name instanceof RawString ? source.qualifiedName.name.value : source.qualifiedName.name.name);
}
else {
return source.qualifiedName.name instanceof RawString ? source.qualifiedName.name.value : source.qualifiedName.name.name;
}
}
collect(query) {
// Visit the SQL component to collect table sources
this.visit(query);
return this.getTableSources();
}
/**
* Main entry point for the visitor pattern.
* Implements the shallow visit pattern to distinguish between root and recursive visits.
*/
visit(arg) {
// If not a root visit, just visit the node and return
if (!this.isRootVisit) {
this.visitNode(arg);
return;
}
// If this is a root visit, we need to reset the state
this.reset();
this.isRootVisit = false;
try {
// When in full scan mode, collect CTEs first to exclude them from table sources
if (!this.selectableOnly) {
this.collectCTEs(arg);
}
this.visitNode(arg);
}
finally {
// Regardless of success or failure, reset the root visit flag
this.isRootVisit = true;
}
}
/**
* Internal visit method used for all nodes.
* This separates the visit flag management from the actual node visitation logic.
*/
visitNode(arg) {
// Skip if we've already visited this node to prevent infinite recursion
if (this.visitedNodes.has(arg)) {
return;
}
// Mark as visited
this.visitedNodes.add(arg);
const handler = this.handlers.get(arg.getKind());
if (handler) {
handler(arg);
return;
}
// If no handler found, that's ok - we only care about specific components
}
/**
* Collects all CTE names to exclude them from real table sources
*/
collectCTEs(query) {
// Use CommonTableCollector to get all CTEs
const cteCollector = new CTECollector();
cteCollector.visit(query);
const commonTables = cteCollector.getCommonTables();
// Add CTE names to the set
for (const cte of commonTables) {
// aliasExpression.table is TableSource, so use .table getter (IdentifierString)
this.cteNames.add(cte.aliasExpression.table.name);
}
}
visitSimpleSelectQuery(query) {
// Process the FROM and JOIN clauses
if (query.fromClause) {
query.fromClause.accept(this);
}
// If in full scan mode, visit all other clauses too
if (!this.selectableOnly) {
if (query.withClause) {
query.withClause.accept(this);
}
if (query.whereClause) {
query.whereClause.accept(this);
}
if (query.groupByClause) {
query.groupByClause.accept(this);
}
if (query.havingClause) {
query.havingClause.accept(this);
}
if (query.orderByClause) {
query.orderByClause.accept(this);
}
if (query.windowClause) {
for (const win of query.windowClause.windows) {
win.accept(this);
}
}
if (query.limitClause) {
query.limitClause.accept(this);
}
if (query.offsetClause) {
query.offsetClause.accept(this);
}
if (query.fetchClause) {
query.fetchClause.accept(this);
}
if (query.forClause) {
query.forClause.accept(this);
}
query.selectClause.accept(this);
}
}
visitBinarySelectQuery(query) {
// For UNION-like queries, visit both sides
query.left.accept(this);
query.right.accept(this);
}
visitValuesQuery(query) {
if (!this.selectableOnly) {
// VALUES queries might contain subqueries in tuple expressions
for (const tuple of query.tuples) {
tuple.accept(this);
}
}
}
visitWithClause(withClause) {
if (!this.selectableOnly) {
// Visit each CommonTable
for (const table of withClause.tables) {
table.accept(this);
}
}
}
visitCommonTable(commonTable) {
if (!this.selectableOnly) {
// Process the query within the common table
commonTable.query.accept(this);
}
}
visitFromClause(fromClause) {
// Check the main source in FROM clause
fromClause.source.accept(this);
// Check all JOIN clauses
if (fromClause.joins) {
for (const join of fromClause.joins) {
join.accept(this);
}
}
}
visitSourceExpression(source) {
// Process the actual data source, ignoring aliases
source.datasource.accept(this);
}
visitTableSource(source) {
// Get the table identifier for uniqueness check
const identifier = this.getTableIdentifier(source);
// Check if this is a table managed by a CTE
if (!this.tableNameMap.has(identifier) && !this.isCTETable(source.table.name)) {
this.tableNameMap.set(identifier, true);
this.tableSources.push(source);
}
}
/**
* Checks if a table name is a CTE name
*/
isCTETable(tableName) {
return this.cteNames.has(tableName);
}
visitParenSource(source) {
// For parenthesized sources, visit the inner source
source.source.accept(this);
}
visitSubQuerySource(subQuery) {
if (!this.selectableOnly) {
// In full scan mode, we also check subqueries
subQuery.query.accept(this);
}
// In selectableOnly mode, we don't collect sources from subqueries
}
visitInlineQuery(inlineQuery) {
if (!this.selectableOnly) {
// In full scan mode, visit inline queries too
inlineQuery.selectQuery.accept(this);
}
}
visitJoinClause(joinClause) {
// Visit the source being joined
joinClause.source.accept(this);
// If full scanning, also visit the join condition
if (!this.selectableOnly && joinClause.condition) {
joinClause.condition.accept(this);
}
}
visitJoinOnClause(joinOn) {
if (!this.selectableOnly) {
// In full scan mode, check ON condition for table references
joinOn.condition.accept(this);
}
}
visitJoinUsingClause(joinUsing) {
if (!this.selectableOnly) {
// In full scan mode, check USING condition for table references
joinUsing.condition.accept(this);
}
}
// Additional visitor methods only used in full scan mode
visitWhereClause(whereClause) {
whereClause.condition.accept(this);
}
visitGroupByClause(clause) {
for (const item of clause.grouping) {
item.accept(this);
}
}
visitHavingClause(clause) {
clause.condition.accept(this);
}
visitOrderByClause(clause) {
for (const item of clause.order) {
item.accept(this);
}
}
visitWindowFrameClause(clause) {
clause.expression.accept(this);
}
visitLimitClause(clause) {
clause.value.accept(this);
}
visitOffsetClause(clause) {
clause.value.accept(this);
}
visitFetchClause(clause) {
clause.expression.accept(this);
}
visitForClause(clause) {
// FOR clause doesn't contain table sources
}
visitOrderByItem(item) {
item.value.accept(this);
}
visitSelectClause(clause) {
for (const item of clause.items) {
item.accept(this);
}
}
visitSelectItem(item) {
item.value.accept(this);
}
visitParenExpression(expr) {
expr.expression.accept(this);
}
visitBinaryExpression(expr) {
expr.left.accept(this);
expr.right.accept(this);
}
visitUnaryExpression(expr) {
expr.expression.accept(this);
}
visitCaseExpression(expr) {
if (expr.condition) {
expr.condition.accept(this);
}
expr.switchCase.accept(this);
}
visitSwitchCaseArgument(switchCase) {
for (const caseItem of switchCase.cases) {
caseItem.accept(this);
}
if (switchCase.elseValue) {
switchCase.elseValue.accept(this);
}
}
visitCaseKeyValuePair(pair) {
pair.key.accept(this);
pair.value.accept(this);
}
visitBetweenExpression(expr) {
expr.expression.accept(this);
expr.lower.accept(this);
expr.upper.accept(this);
}
visitFunctionCall(func) {
if (func.argument) {
func.argument.accept(this);
}
if (func.over) {
func.over.accept(this);
}
}
visitArrayExpression(expr) {
expr.expression.accept(this);
}
visitTupleExpression(expr) {
for (const value of expr.values) {
value.accept(this);
}
}
visitCastExpression(expr) {
expr.input.accept(this);
expr.castType.accept(this);
}
}
//# sourceMappingURL=TableSourceCollector.js.map