rawsql-ts
Version: 
[beta]High-performance SQL parser and AST analyzer written in TypeScript. Provides fast parsing and advanced transformation capabilities.
435 lines • 17.7 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.TableSourceCollector = void 0;
const Clause_1 = require("../models/Clause");
const SelectQuery_1 = require("../models/SelectQuery");
const ValueComponent_1 = require("../models/ValueComponent");
const CTECollector_1 = require("./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.
 */
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(SelectQuery_1.SimpleSelectQuery.kind, (expr) => this.visitSimpleSelectQuery(expr));
        this.handlers.set(SelectQuery_1.BinarySelectQuery.kind, (expr) => this.visitBinarySelectQuery(expr));
        this.handlers.set(SelectQuery_1.ValuesQuery.kind, (expr) => this.visitValuesQuery(expr));
        // WITH clause and common tables
        this.handlers.set(Clause_1.WithClause.kind, (expr) => this.visitWithClause(expr));
        this.handlers.set(Clause_1.CommonTable.kind, (expr) => this.visitCommonTable(expr));
        // Handlers for FROM and JOIN components
        this.handlers.set(Clause_1.FromClause.kind, (expr) => this.visitFromClause(expr));
        this.handlers.set(Clause_1.JoinClause.kind, (expr) => this.visitJoinClause(expr));
        this.handlers.set(Clause_1.JoinOnClause.kind, (expr) => this.visitJoinOnClause(expr));
        this.handlers.set(Clause_1.JoinUsingClause.kind, (expr) => this.visitJoinUsingClause(expr));
        // Source components
        this.handlers.set(Clause_1.SourceExpression.kind, (expr) => this.visitSourceExpression(expr));
        this.handlers.set(Clause_1.TableSource.kind, (expr) => this.visitTableSource(expr));
        this.handlers.set(Clause_1.FunctionSource.kind, (expr) => this.visitFunctionSource(expr));
        this.handlers.set(Clause_1.ParenSource.kind, (expr) => this.visitParenSource(expr));
        this.handlers.set(Clause_1.SubQuerySource.kind, (expr) => this.visitSubQuerySource(expr));
        this.handlers.set(ValueComponent_1.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(Clause_1.WhereClause.kind, (expr) => this.visitWhereClause(expr));
            this.handlers.set(Clause_1.GroupByClause.kind, (expr) => this.visitGroupByClause(expr));
            this.handlers.set(Clause_1.HavingClause.kind, (expr) => this.visitHavingClause(expr));
            this.handlers.set(Clause_1.OrderByClause.kind, (expr) => this.visitOrderByClause(expr));
            this.handlers.set(Clause_1.WindowFrameClause.kind, (expr) => this.visitWindowFrameClause(expr));
            this.handlers.set(Clause_1.LimitClause.kind, (expr) => this.visitLimitClause(expr));
            this.handlers.set(Clause_1.OffsetClause.kind, (expr) => this.visitOffsetClause(expr));
            this.handlers.set(Clause_1.FetchClause.kind, (expr) => this.visitFetchClause(expr));
            this.handlers.set(Clause_1.ForClause.kind, (expr) => this.visitForClause(expr));
            this.handlers.set(Clause_1.OrderByItem.kind, (expr) => this.visitOrderByItem(expr));
            this.handlers.set(Clause_1.SelectClause.kind, (expr) => this.visitSelectClause(expr));
            this.handlers.set(Clause_1.SelectItem.kind, (expr) => this.visitSelectItem(expr));
            // Value components that might contain table references
            this.handlers.set(ValueComponent_1.ParenExpression.kind, (expr) => this.visitParenExpression(expr));
            this.handlers.set(ValueComponent_1.BinaryExpression.kind, (expr) => this.visitBinaryExpression(expr));
            this.handlers.set(ValueComponent_1.UnaryExpression.kind, (expr) => this.visitUnaryExpression(expr));
            this.handlers.set(ValueComponent_1.CaseExpression.kind, (expr) => this.visitCaseExpression(expr));
            this.handlers.set(ValueComponent_1.CaseKeyValuePair.kind, (expr) => this.visitCaseKeyValuePair(expr));
            this.handlers.set(ValueComponent_1.SwitchCaseArgument.kind, (expr) => this.visitSwitchCaseArgument(expr));
            this.handlers.set(ValueComponent_1.BetweenExpression.kind, (expr) => this.visitBetweenExpression(expr));
            this.handlers.set(ValueComponent_1.FunctionCall.kind, (expr) => this.visitFunctionCall(expr));
            this.handlers.set(ValueComponent_1.ArrayExpression.kind, (expr) => this.visitArrayExpression(expr));
            this.handlers.set(ValueComponent_1.ArrayQueryExpression.kind, (expr) => this.visitArrayQueryExpression(expr));
            this.handlers.set(ValueComponent_1.TupleExpression.kind, (expr) => this.visitTupleExpression(expr));
            this.handlers.set(ValueComponent_1.CastExpression.kind, (expr) => this.visitCastExpression(expr));
            this.handlers.set(ValueComponent_1.ValueList.kind, (expr) => this.visitValueList(expr));
            this.handlers.set(ValueComponent_1.StringSpecifierExpression.kind, (expr) => this.visitStringSpecifierExpression(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 ValueComponent_1.RawString ? source.qualifiedName.name.value : source.qualifiedName.name.name);
        }
        else {
            return source.qualifiedName.name instanceof ValueComponent_1.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_1.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);
        }
    }
    visitFunctionSource(source) {
        // Function sources are not regular table sources, but may contain subqueries in their arguments
        if (source.argument) {
            // Special handling for function arguments to ensure we traverse nested structures
            this.visitValueComponent(source.argument);
        }
        // Function sources themselves are not collected as table sources
    }
    /**
     * Helper method to visit value components, handling special cases like TupleExpression, ParenExpression, InlineQuery, and ArrayQueryExpression
     * even in selectableOnly mode when they appear in function arguments
     */
    visitValueComponent(value) {
        // Always use the normal accept pattern - let handlers deal with the logic
        value.accept(this);
    }
    /**
     * 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);
    }
    visitArrayQueryExpression(expr) {
        expr.query.accept(this);
    }
    visitTupleExpression(expr) {
        for (const value of expr.values) {
            value.accept(this);
        }
    }
    visitCastExpression(expr) {
        expr.input.accept(this);
        expr.castType.accept(this);
    }
    visitValueList(valueList) {
        // Process all values in the list, this may include InlineQuery and other table-referencing components
        for (const value of valueList.values) {
            value.accept(this);
        }
    }
    // Handle StringSpecifierExpression (PostgreSQL E-strings)
    visitStringSpecifierExpression(_expr) {
        // StringSpecifierExpression is just a literal string with an escape specifier
        // It doesn't contain table references, so we don't need to visit any children
        // This is a no-op method to prevent "No handler" errors
    }
}
exports.TableSourceCollector = TableSourceCollector;
//# sourceMappingURL=TableSourceCollector.js.map