rawsql-ts
Version:
High-performance SQL parser and AST analyzer written in TypeScript. Provides fast parsing and advanced transformation capabilities.
308 lines • 12.7 kB
JavaScript
import { FromClause, ParenSource, SelectClause, SourceExpression, SubQuerySource, TableSource } from "../models/Clause";
import { SimpleSelectQuery } from "../models/SelectQuery";
import { InsertQuery } from "../models/InsertQuery";
import { UpdateQuery } from "../models/UpdateQuery";
import { DeleteQuery } from "../models/DeleteQuery";
import { MergeQuery } from "../models/MergeQuery";
import { ColumnReference } from "../models/ValueComponent";
import { CTECollector } from "./CTECollector";
/**
* A visitor that collects all SelectItem instances from a SQL query structure.
* This visitor scans through select clauses and collects all the SelectItem objects.
* It can also resolve wildcard selectors (table.* or *) using a provided table column resolver.
*/
export class SelectValueCollector {
constructor(tableColumnResolver = null, initialCommonTables = null, preserveDuplicateSelectItems = false) {
this.selectValues = [];
this.visitedNodes = new Set();
this.isRootVisit = true;
this.tableColumnResolver = tableColumnResolver !== null && tableColumnResolver !== void 0 ? tableColumnResolver : null;
this.commonTableCollector = new CTECollector();
this.commonTables = [];
this.initialCommonTables = initialCommonTables;
this.preserveDuplicateSelectItems = preserveDuplicateSelectItems;
this.handlers = new Map();
this.handlers.set(SimpleSelectQuery.kind, (expr) => this.visitSimpleSelectQuery(expr));
this.handlers.set(SelectClause.kind, (expr) => this.visitSelectClause(expr));
this.handlers.set(SourceExpression.kind, (expr) => this.visitSourceExpression(expr));
this.handlers.set(FromClause.kind, (expr) => this.visitFromClause(expr));
}
/**
* Get all collected SelectItems as an array of objects with name and value properties
* @returns An array of objects with name (string) and value (ValueComponent) properties
*/
getValues() {
return this.selectValues;
}
/**
* Reset the collection of SelectItems
*/
reset() {
this.selectValues = [];
this.visitedNodes.clear();
if (this.initialCommonTables) {
this.commonTables = this.initialCommonTables;
}
else {
this.commonTables = [];
}
}
collect(arg) {
// Visit the component and return the collected select items
this.visit(arg);
const items = this.getValues();
this.reset(); // Reset after collection
return items;
}
/**
* 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 {
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;
}
}
/**
* Process a SimpleSelectQuery to collect data and store the current context
*/
visitSimpleSelectQuery(query) {
if (this.commonTables.length === 0 && this.initialCommonTables === null) {
this.commonTables = this.commonTableCollector.collect(query);
}
if (query.selectClause) {
query.selectClause.accept(this);
}
// no wildcard
const wildcards = this.selectValues.filter(item => item.name === '*');
if (wildcards.length === 0) {
return;
}
const expandedValues = [];
for (const item of this.selectValues) {
if (item.name !== '*' || !(item.value instanceof ColumnReference)) {
expandedValues.push(item);
continue;
}
expandedValues.push(...this.expandWildcardValue(item.value, query.fromClause));
}
this.selectValues = expandedValues;
}
processFromClause(clause, joinCascade) {
for (const item of this.collectFromClauseValues(clause, joinCascade)) {
this.addSelectValue(item.name, item.value);
}
return;
}
processJoinClause(clause) {
for (const item of this.collectJoinClauseValues(clause)) {
this.addSelectValue(item.name, item.value);
}
}
processSourceExpression(sourceName, source) {
for (const item of this.collectSourceExpressionValues(sourceName, source)) {
this.addSelectValue(item.name, item.value);
}
}
expandWildcardValue(value, fromClause) {
if (!fromClause) {
return [];
}
if (value.namespaces === null) {
return this.collectFromClauseValues(fromClause, true);
}
const sourceName = value.getNamespace();
if (fromClause.getSourceAliasName() === sourceName) {
return this.collectFromClauseValues(fromClause, false);
}
if (!fromClause.joins) {
return [];
}
const join = fromClause.joins.find(item => this.getJoinSourceName(item) === sourceName);
return join ? this.collectJoinClauseValues(join) : [];
}
collectFromClauseValues(clause, joinCascade) {
const values = this.collectSourceExpressionValues(clause.getSourceAliasName(), clause.source);
if (clause.joins && joinCascade) {
for (const join of clause.joins) {
values.push(...this.collectJoinClauseValues(join));
}
}
return values;
}
collectJoinClauseValues(clause) {
return this.collectSourceExpressionValues(this.getJoinSourceName(clause), clause.source);
}
getJoinSourceName(clause) {
return clause.source.getAliasName();
}
collectSourceExpressionValues(sourceName, source) {
// check common table
const tableSourceName = source.datasource instanceof TableSource ? source.datasource.getSourceName() : null;
const commonTable = tableSourceName
? this.commonTables.find(item => item.aliasExpression.table.name === tableSourceName)
: null;
if (commonTable) {
// Exclude this CTE from consideration to prevent self-reference
const innerCommonTables = this.commonTables.filter(item => item.aliasExpression.table.name !== commonTable.aliasExpression.table.name);
const innerSelected = this.collectValuesFromCteQuery(commonTable.query, innerCommonTables);
return innerSelected.map(item => ({
name: item.name,
value: new ColumnReference(sourceName ? [sourceName] : null, item.name)
}));
}
const innerCollector = new SelectValueCollector(this.tableColumnResolver, this.commonTables, true);
const innerSelected = innerCollector.collect(source);
return innerSelected.map(item => ({
name: item.name,
value: new ColumnReference(sourceName ? [sourceName] : null, item.name)
}));
}
visitSelectClause(clause) {
for (const item of clause.items) {
this.processSelectItem(item);
}
}
processSelectItem(item) {
if (item.identifier) {
this.addSelectValueFromSelectItem(item.identifier.name, item.value);
}
else if (item.value instanceof ColumnReference) { // Handle column reference
// columnName can be '*'
const columnName = item.value.column.name;
if (columnName === '*') {
// Force add without checking duplicates
this.selectValues.push({ name: columnName, value: item.value });
}
else {
// Add with duplicate checking
this.addSelectValueFromSelectItem(columnName, item.value);
}
}
}
visitSourceExpression(source) {
// Column aliases have the highest priority if present
// For physical tables, use external function to get column names
// For subqueries, instantiate a new collector and get column names from the subquery
// For parenthesized expressions, treat them the same as subqueries
if (source.aliasExpression && source.aliasExpression.columns) {
const sourceName = source.getAliasName();
source.aliasExpression.columns.forEach(column => {
this.addSelectValueFromSelectItem(column.name, new ColumnReference(sourceName ? [sourceName] : null, column.name));
});
return;
}
else if (source.datasource instanceof TableSource) {
if (this.tableColumnResolver) {
const sourceName = source.datasource.getSourceName();
this.tableColumnResolver(sourceName).forEach(column => {
this.addSelectValueFromSelectItem(column, new ColumnReference([sourceName], column));
});
}
return;
}
else if (source.datasource instanceof SubQuerySource) {
const sourceName = source.getAliasName();
const innerCollector = new SelectValueCollector(this.tableColumnResolver, this.commonTables, true);
const innerSelected = innerCollector.collect(source.datasource.query);
innerSelected.forEach(item => {
this.addSelectValueFromSelectItem(item.name, new ColumnReference(sourceName ? [sourceName] : null, item.name));
});
return;
}
else if (source.datasource instanceof ParenSource) {
return this.visit(source.datasource.source);
}
}
visitFromClause(clause) {
if (clause) {
this.processFromClause(clause, true);
}
}
addSelectValueAsUnique(name, value) {
// Check if a select value with the same name already exists before adding
if (!this.selectValues.some(item => item.name === name)) {
this.selectValues.push({ name, value });
}
}
addSelectValue(name, value) {
this.selectValues.push({ name, value });
}
addSelectValueFromSelectItem(name, value) {
if (this.preserveDuplicateSelectItems) {
this.addSelectValue(name, value);
return;
}
this.addSelectValueAsUnique(name, value);
}
collectValuesFromCteQuery(query, commonTables) {
if (this.isSelectQuery(query)) {
const innerCollector = new SelectValueCollector(this.tableColumnResolver, commonTables, true);
return innerCollector.collect(query);
}
// Writable CTEs expose their output via RETURNING.
return this.collectValuesFromReturning(query);
}
collectValuesFromReturning(query) {
if (!(query instanceof InsertQuery || query instanceof UpdateQuery || query instanceof DeleteQuery || query instanceof MergeQuery)) {
return [];
}
if (!query.returningClause) {
return [];
}
return this.extractValuesFromReturningClause(query.returningClause);
}
extractValuesFromReturningClause(clause) {
var _a, _b;
const values = [];
for (const item of clause.items) {
const name = (_b = (_a = item.identifier) === null || _a === void 0 ? void 0 : _a.name) !== null && _b !== void 0 ? _b : this.extractSelectItemName(item);
if (name) {
values.push({ name, value: item.value });
}
}
return values;
}
extractSelectItemName(item) {
if (item.identifier) {
return item.identifier.name;
}
if (item.value instanceof ColumnReference) {
return item.value.column.name;
}
return null;
}
isSelectQuery(query) {
return '__selectQueryType' in query && query.__selectQueryType === 'SelectQuery';
}
}
//# sourceMappingURL=SelectValueCollector.js.map