@rawsql-ts/ztd-cli 0.19.0 → 0.20.2

package/dist/query/execute.js

@@ -0,0 +1,784 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.executeQueryPipeline = executeQueryPipeline;
+const node_fs_1 = require("node:fs");
+const node_path_1 = __importDefault(require("node:path"));
+const rawsql_ts_1 = require("rawsql-ts");
+const analysis_1 = require("./analysis");
+const planner_1 = require("./planner");
+const SUPPORTED_COMPARISON_OPERATORS = new Set(['=', '!=', '<>', '>', '>=', '<', '<=']);
+/**
+ * Execute a decomposed query pipeline inside one DB session and reuse prior stage outputs.
+ */
+async function executeQueryPipeline(sessionFactory, options) {
+    const source = preparePipelineSource(options.sqlFile);
+    const plan = (0, planner_1.buildQueryPipelinePlan)(options.sqlFile, options.metadata);
+    const session = await sessionFactory.openSession();
+    const materializedCtes = [];
+    const createdTempTables = [];
+    const steps = [];
+    let executionError;
+    try {
+        for (const step of plan.steps) {
+            if (step.kind === 'materialize') {
+                const stage = await buildPipelineStageQuery(source, session, {
+                    cte: step.target,
+                    runtimeParams: options.params,
+                    materializedCtes,
+                    scalarFilterColumns: plan.metadata.scalarFilterColumns
+                });
+                steps.push(...stage.scalarSteps);
+                const sql = `create temp table ${quoteIdentifier(step.target)} as ${stage.sql.trim()}`;
+                const stageParams = normalizeParamsForSql(sql, stage.params, options.params);
+                const result = normalizePipelineQueryResult(await session.query(sql, stageParams));
+                materializedCtes.push(step.target);
+                createdTempTables.push(step.target);
+                steps.push({
+                    kind: step.kind,
+                    target: step.target,
+                    sql,
+                    params: stageParams,
+                    rowCount: result.rowCount
+                });
+                continue;
+            }
+            const finalStage = await buildPipelineStageQuery(source, session, {
+                final: true,
+                runtimeParams: options.params,
+                materializedCtes,
+                scalarFilterColumns: plan.metadata.scalarFilterColumns
+            });
+            steps.push(...finalStage.scalarSteps);
+            const finalParams = normalizeParamsForSql(finalStage.sql, finalStage.params, options.params);
+            const result = normalizePipelineQueryResult(await session.query(finalStage.sql, finalParams));
+            steps.push({
+                kind: step.kind,
+                target: step.target,
+                sql: finalStage.sql,
+                params: finalParams,
+                rowCount: result.rowCount
+            });
+            return {
+                plan,
+                final: {
+                    rows: result.rows,
+                    rowCount: result.rowCount,
+                    sql: finalStage.sql,
+                    params: finalParams
+                },
+                steps
+            };
+        }
+        throw new Error('Query pipeline plan did not include a final-query step.');
+    }
+    catch (error) {
+        executionError = error;
+        throw error;
+    }
+    finally {
+        await cleanupPipelineSession(session, createdTempTables, executionError);
+    }
+}
+async function buildPipelineStageQuery(source, session, options) {
+    validateStageOptions(options);
+    return options.cte
+        ? buildTargetStageQuery(source, session, options)
+        : buildFinalStageQuery(source, session, options);
+}
+async function buildTargetStageQuery(source, session, options) {
+    var _a, _b, _c, _d;
+    const parsed = (0, analysis_1.assertSupportedStatement)(rawsql_ts_1.SqlParser.parse(source.sql), 'executeQueryPipeline');
+    const analysis = (0, analysis_1.analyzeStatement)(parsed);
+    const targetName = options.cte;
+    const materializedStopSet = new Set(options.materializedCtes.filter((name) => name !== targetName));
+    const includedNames = (0, analysis_1.collectDependencyClosure)(targetName, analysis.dependencyMap, materializedStopSet);
+    if (!includedNames.includes(targetName)) {
+        throw new Error(`CTE not found in query: ${targetName}`);
+    }
+    const includedCtes = buildStageCtes(analysis.ctes, includedNames, targetName, options.materializedCtes);
+    const targetCte = includedCtes.find((cte) => cte.aliasExpression.table.name === targetName);
+    if (!targetCte) {
+        throw new Error(`CTE not found in query: ${targetName}`);
+    }
+    const dependencyCtes = includedCtes.filter((cte) => cte.aliasExpression.table.name !== targetName);
+    const bindingContext = createScalarBindingContext(options, includedNames, source.sql, (_b = (_a = getWithClause(parsed)) === null || _a === void 0 ? void 0 : _a.recursive) !== null && _b !== void 0 ? _b : false);
+    const scalarSteps = await bindScalarFilterPredicatesInCtes(includedCtes, bindingContext, session);
+    const formatter = createPipelineFormatter(options.runtimeParams);
+    const targetQuery = assertSelectQuery(targetCte.query);
+    const withComponent = dependencyCtes.length > 0 ? new rawsql_ts_1.WithClause((_d = (_c = getWithClause(parsed)) === null || _c === void 0 ? void 0 : _c.recursive) !== null && _d !== void 0 ? _d : false, dependencyCtes) : null;
+    const withResult = withComponent
+        ? formatter.format(withComponent)
+        : { formattedSql: '', params: emptyFormatterParams(options.runtimeParams) };
+    let mainResult;
+    if (options.limit !== undefined) {
+        if (targetQuery instanceof rawsql_ts_1.SimpleSelectQuery) {
+            targetQuery.limitClause = new rawsql_ts_1.LimitClause(new rawsql_ts_1.LiteralValue(options.limit));
+            mainResult = formatter.format(targetQuery);
+        }
+        else {
+            mainResult = formatter.format(buildWrappedLimitQuery(targetQuery, options.limit));
+        }
+    }
+    else {
+        mainResult = formatter.format(targetQuery);
+    }
+    const mergedParams = mergeFormatterParams([withResult.params, mainResult.params], options.runtimeParams);
+    const sql = withResult.formattedSql ? `${withResult.formattedSql} ${mainResult.formattedSql}` : mainResult.formattedSql;
+    return {
+        sql: `${sql}\n`,
+        params: normalizeParamsForSql(sql, mergedParams, options.runtimeParams),
+        scalarSteps
+    };
+}
+async function buildFinalStageQuery(source, session, options) {
+    var _a, _b, _c, _d;
+    const parsed = (0, analysis_1.assertSupportedStatement)(rawsql_ts_1.SqlParser.parse(source.sql), 'executeQueryPipeline');
+    const analysis = (0, analysis_1.analyzeStatement)(parsed);
+    const stopSet = new Set(options.materializedCtes);
+    const includedNames = [...(0, analysis_1.collectReachableCtes)(analysis.rootDependencies, analysis.dependencyMap, stopSet)];
+    const includedCtes = buildStageCtes(analysis.ctes, includedNames, null, options.materializedCtes);
+    const bindingContext = createScalarBindingContext(options, includedNames, source.sql, (_b = (_a = getWithClause(parsed)) === null || _a === void 0 ? void 0 : _a.recursive) !== null && _b !== void 0 ? _b : false);
+    applyMinimalWithClause(parsed, includedCtes, (_d = (_c = getWithClause(parsed)) === null || _c === void 0 ? void 0 : _c.recursive) !== null && _d !== void 0 ? _d : false);
+    const scalarSteps = [
+        ...(await bindScalarFilterPredicatesInCtes(includedCtes, bindingContext, session)),
+        ...(await bindScalarFilterPredicates(parsed, bindingContext, session))
+    ];
+    const formatter = createPipelineFormatter(options.runtimeParams);
+    if (options.limit !== undefined) {
+        if (parsed instanceof rawsql_ts_1.SimpleSelectQuery) {
+            parsed.limitClause = new rawsql_ts_1.LimitClause(new rawsql_ts_1.LiteralValue(options.limit));
+        }
+        else if (parsed instanceof rawsql_ts_1.ValuesQuery || parsed instanceof rawsql_ts_1.BinarySelectQuery) {
+            const wrapped = buildWrappedLimitQuery(parsed, options.limit);
+            const { formattedSql, params } = formatter.format(wrapped);
+            return {
+                sql: `${formattedSql}\n`,
+                params: normalizeParamsForSql(formattedSql, mergeFormatterParams(params, options.runtimeParams), options.runtimeParams),
+                scalarSteps
+            };
+        }
+        else {
+            throw new Error('--limit is only supported for SELECT final slices or --cte slices.');
+        }
+    }
+    const { formattedSql, params } = formatter.format(parsed);
+    return {
+        sql: `${formattedSql}\n`,
+        params: normalizeParamsForSql(formattedSql, mergeFormatterParams(params, options.runtimeParams), options.runtimeParams),
+        scalarSteps
+    };
+}
+function createScalarBindingContext(options, includedNames, sourceSql, recursive) {
+    return {
+        runtimeParams: options.runtimeParams,
+        scalarFilterColumns: new Set(options.scalarFilterColumns),
+        materializedCtes: new Set(options.materializedCtes),
+        sourceSql,
+        recursive,
+        nextOrdinal: 1
+    };
+}
+function buildStageCtes(ctes, includedNames, currentTarget, materializedCtes) {
+    const includedSet = new Set(includedNames);
+    const materializedSet = new Set(materializedCtes);
+    return ctes.filter((cte) => {
+        const name = cte.aliasExpression.table.name;
+        if (!includedSet.has(name)) {
+            return false;
+        }
+        return name === currentTarget || !materializedSet.has(name);
+    });
+}
+async function bindScalarFilterPredicatesInCtes(ctes, context, session) {
+    const steps = [];
+    for (const cte of ctes) {
+        steps.push(...await bindScalarFilterPredicates(assertSelectQuery(cte.query), context, session));
+    }
+    return steps;
+}
+async function bindScalarFilterPredicates(statement, context, session) {
+    if (context.scalarFilterColumns.size === 0) {
+        return [];
+    }
+    if (statement instanceof rawsql_ts_1.SimpleSelectQuery) {
+        return bindScalarFilterPredicatesInSelect(statement, context, session);
+    }
+    if (statement instanceof rawsql_ts_1.BinarySelectQuery) {
+        const steps = [];
+        steps.push(...await bindScalarFilterPredicatesInSelectBranch(assertSelectQuery(statement.left), context, session));
+        steps.push(...await bindScalarFilterPredicatesInSelectBranch(assertSelectQuery(statement.right), context, session));
+        return steps;
+    }
+    if (statement instanceof rawsql_ts_1.ValuesQuery) {
+        return [];
+    }
+    if (statement instanceof rawsql_ts_1.InsertQuery) {
+        return statement.selectQuery
+            ? bindScalarFilterPredicatesInSelectBranch(assertSelectQuery(statement.selectQuery), context, session)
+            : [];
+    }
+    if (statement instanceof rawsql_ts_1.UpdateQuery || statement instanceof rawsql_ts_1.DeleteQuery) {
+        return [];
+    }
+    return [];
+}
+async function bindScalarFilterPredicatesInSelectBranch(statement, context, session) {
+    if (statement instanceof rawsql_ts_1.SimpleSelectQuery) {
+        return bindScalarFilterPredicatesInSelect(statement, context, session);
+    }
+    if (statement instanceof rawsql_ts_1.BinarySelectQuery) {
+        const steps = [];
+        steps.push(...await bindScalarFilterPredicatesInSelectBranch(assertSelectQuery(statement.left), context, session));
+        steps.push(...await bindScalarFilterPredicatesInSelectBranch(assertSelectQuery(statement.right), context, session));
+        return steps;
+    }
+    return [];
+}
+async function bindScalarFilterPredicatesInSelect(selectQuery, context, session) {
+    if (!selectQuery.whereClause) {
+        return [];
+    }
+    return rewritePredicateExpression(selectQuery.whereClause.condition, context, session);
+}
+async function rewritePredicateExpression(expression, context, session) {
+    if (!(expression instanceof rawsql_ts_1.BinaryExpression)) {
+        return [];
+    }
+    const binaryExpression = expression;
+    const candidate = findScalarBindingCandidate(binaryExpression, context.scalarFilterColumns);
+    if (candidate) {
+        const execution = await executeScalarFilterBinding(candidate, context, session);
+        if (execution) {
+            candidate.replace(new rawsql_ts_1.ParameterExpression(execution.paramName, execution.value));
+            return [execution.step];
+        }
+    }
+    const steps = [];
+    steps.push(...await rewritePredicateExpression(binaryExpression.left, context, session));
+    steps.push(...await rewritePredicateExpression(binaryExpression.right, context, session));
+    return steps;
+}
+function findScalarBindingCandidate(expression, scalarFilterColumns) {
+    const operator = extractOperator(expression.operator);
+    if (!SUPPORTED_COMPARISON_OPERATORS.has(operator)) {
+        return null;
+    }
+    const leftMatches = containsTargetColumn(expression.left, scalarFilterColumns);
+    const rightMatches = containsTargetColumn(expression.right, scalarFilterColumns);
+    const leftInline = unwrapInlineQuery(expression.left);
+    const rightInline = unwrapInlineQuery(expression.right);
+    if (leftMatches && rightInline) {
+        return {
+            columnName: leftMatches,
+            inlineQuery: rightInline,
+            replace: (parameter) => {
+                expression.right = parameter;
+            }
+        };
+    }
+    if (rightMatches && leftInline) {
+        return {
+            columnName: rightMatches,
+            inlineQuery: leftInline,
+            replace: (parameter) => {
+                expression.left = parameter;
+            }
+        };
+    }
+    return null;
+}
+async function executeScalarFilterBinding(candidate, context, session) {
+    const selectQuery = candidate.inlineQuery.selectQuery;
+    if (!(selectQuery instanceof rawsql_ts_1.SimpleSelectQuery)) {
+        throw new Error(`Scalar filter binding for column "${candidate.columnName}" requires a simple SELECT subquery.`);
+    }
+    if (isCorrelatedScalarSubquery(selectQuery)) {
+        return null;
+    }
+    assertSingleColumnScalarSelect(selectQuery, candidate.columnName);
+    const paramName = `__scalar_filter_${candidate.columnName}_${context.nextOrdinal}`;
+    const { sql, params: queryParams } = buildScalarBindingQuery(selectQuery, context);
+    const scalarParams = normalizeParamsForSql(sql, queryParams, context.runtimeParams);
+    const result = normalizePipelineQueryResult(await session.query(sql, scalarParams));
+    const value = extractScalarFilterValue(result.rows, candidate.columnName);
+    const step = {
+        kind: 'scalar-filter-bind',
+        target: candidate.columnName,
+        sql,
+        params: scalarParams,
+        rowCount: result.rowCount
+    };
+    context.nextOrdinal += 1;
+    return { paramName, value, step };
+}
+function buildScalarBindingQuery(selectQuery, context) {
+    const sourceStatement = (0, analysis_1.assertSupportedStatement)(rawsql_ts_1.SqlParser.parse(context.sourceSql), 'executeQueryPipeline');
+    const sourceAnalysis = (0, analysis_1.analyzeStatement)(sourceStatement);
+    const scalarRoots = collectScalarQueryCteRoots(selectQuery, sourceAnalysis.ctes, context.materializedCtes);
+    const dependencyNames = scalarRoots.length > 0
+        ? [...(0, analysis_1.collectReachableCtes)(scalarRoots, sourceAnalysis.dependencyMap, context.materializedCtes)]
+        : [];
+    const scalarCtes = buildStageCtes(sourceAnalysis.ctes, dependencyNames, null, [...context.materializedCtes]);
+    const formatter = createPipelineFormatter(context.runtimeParams);
+    const withComponent = scalarCtes.length > 0 ? new rawsql_ts_1.WithClause(context.recursive, scalarCtes) : null;
+    const withResult = withComponent
+        ? formatter.format(withComponent)
+        : { formattedSql: '', params: emptyFormatterParams(context.runtimeParams) };
+    const mainResult = formatter.format(selectQuery);
+    const mergedParams = mergeFormatterParams([withResult.params, mainResult.params], context.runtimeParams);
+    const sql = withResult.formattedSql ? `${withResult.formattedSql} ${mainResult.formattedSql}` : mainResult.formattedSql;
+    return {
+        sql: `${sql}
+`,
+        params: normalizeParamsForSql(sql, mergedParams, context.runtimeParams)
+    };
+}
+function collectScalarQueryCteRoots(selectQuery, ctes, materializedCtes) {
+    const cteNames = new Set(ctes.map((cte) => cte.aliasExpression.table.name));
+    const roots = new Set();
+    walkAst(selectQuery, (current) => {
+        if (!(current instanceof rawsql_ts_1.TableSource)) {
+            return;
+        }
+        const sourceName = extractQualifiedNameLeaf(current.qualifiedName.name);
+        if (cteNames.has(sourceName) && !materializedCtes.has(sourceName)) {
+            roots.add(sourceName);
+        }
+    });
+    return [...roots];
+}
+function assertSingleColumnScalarSelect(selectQuery, columnName) {
+    if (selectQuery.selectClause.items.length !== 1) {
+        throw new Error(`Scalar filter binding for column "${columnName}" requires a subquery that statically exposes exactly one column.`);
+    }
+    const [item] = selectQuery.selectClause.items;
+    if ((item === null || item === void 0 ? void 0 : item.value) instanceof rawsql_ts_1.RawString && item.value.value.trim() === '*') {
+        throw new Error(`Scalar filter binding for column "${columnName}" requires a subquery that statically exposes exactly one column.`);
+    }
+}
+function extractScalarFilterValue(rows, columnName) {
+    var _a;
+    if (rows.length !== 1) {
+        throw new Error(`Scalar filter binding for column "${columnName}" must return exactly one row.`);
+    }
+    const row = (_a = rows[0]) !== null && _a !== void 0 ? _a : {};
+    const columns = Object.keys(row);
+    if (columns.length !== 1) {
+        throw new Error(`Scalar filter binding for column "${columnName}" must return exactly one column.`);
+    }
+    return row[columns[0]];
+}
+function isCorrelatedScalarSubquery(selectQuery) {
+    const localNames = collectLocalRelationNames(selectQuery);
+    const columnReferences = collectColumnReferences(selectQuery);
+    return columnReferences.some((reference) => {
+        var _a, _b;
+        const qualifierParts = (_b = (_a = reference.qualifiedName.namespaces) === null || _a === void 0 ? void 0 : _a.map((namespace) => namespace.name)) !== null && _b !== void 0 ? _b : [];
+        if (qualifierParts.length === 0) {
+            return false;
+        }
+        const qualifier = qualifierParts[qualifierParts.length - 1];
+        return !localNames.has(qualifier);
+    });
+}
+function collectLocalRelationNames(selectQuery) {
+    var _a, _b, _c, _d;
+    const localNames = new Set();
+    const sources = (_b = (_a = selectQuery.fromClause) === null || _a === void 0 ? void 0 : _a.getSources()) !== null && _b !== void 0 ? _b : [];
+    for (const source of sources) {
+        const aliasName = (_d = (_c = source.aliasExpression) === null || _c === void 0 ? void 0 : _c.table) === null || _d === void 0 ? void 0 : _d.name;
+        if (aliasName) {
+            localNames.add(aliasName);
+        }
+        if (source.datasource instanceof rawsql_ts_1.TableSource) {
+            localNames.add(extractQualifiedNameLeaf(source.datasource.qualifiedName.name));
+        }
+    }
+    return localNames;
+}
+function collectColumnReferences(node) {
+    const matches = [];
+    walkAst(node, (current) => {
+        if (current instanceof rawsql_ts_1.ColumnReference) {
+            matches.push(current);
+        }
+    });
+    return matches;
+}
+function containsTargetColumn(node, scalarFilterColumns) {
+    let matched = null;
+    walkAst(node, (current) => {
+        if (matched || !(current instanceof rawsql_ts_1.ColumnReference)) {
+            return;
+        }
+        const columnName = extractQualifiedNameLeaf(current.qualifiedName.name);
+        if (scalarFilterColumns.has(columnName)) {
+            matched = columnName;
+        }
+    });
+    return matched;
+}
+function walkAst(node, visit) {
+    if (!node || typeof node !== 'object') {
+        return;
+    }
+    visit(node);
+    for (const value of Object.values(node)) {
+        if (!value) {
+            continue;
+        }
+        if (Array.isArray(value)) {
+            for (const item of value) {
+                walkAst(item, visit);
+            }
+            continue;
+        }
+        if (typeof value === 'object') {
+            walkAst(value, visit);
+        }
+    }
+}
+function unwrapInlineQuery(expression) {
+    if (expression instanceof rawsql_ts_1.InlineQuery) {
+        return expression;
+    }
+    if (expression instanceof rawsql_ts_1.ParenExpression) {
+        return unwrapInlineQuery(expression.expression);
+    }
+    return null;
+}
+function extractOperator(operator) {
+    if (operator instanceof rawsql_ts_1.RawString) {
+        return operator.value.trim();
+    }
+    return '';
+}
+function extractQualifiedNameLeaf(name) {
+    return name instanceof rawsql_ts_1.RawString ? name.value : name.name;
+}
+function preparePipelineSource(sqlFile) {
+    const sql = (0, node_fs_1.readFileSync)(node_path_1.default.resolve(sqlFile), 'utf8');
+    return {
+        sql
+    };
+}
+function createPipelineFormatter(runtimeParams) {
+    if (runtimeParams && !Array.isArray(runtimeParams)) {
+        return new rawsql_ts_1.SqlFormatter({
+            identifierEscape: { start: '"', end: '"' },
+            parameterStyle: 'named',
+            parameterSymbol: ':'
+        });
+    }
+    return new rawsql_ts_1.SqlFormatter({ preset: 'postgres' });
+}
+function emptyFormatterParams(runtimeParams) {
+    return Array.isArray(runtimeParams) ? [] : {};
+}
+function mergeFormatterParams(formatterParams, runtimeParams) {
+    const parts = Array.isArray(formatterParams) && formatterParams.some((part) => Array.isArray(part) || isPlainObject(part))
+        ? formatterParams
+        : [formatterParams];
+    if ((runtimeParams && !Array.isArray(runtimeParams)) || parts.some((part) => isPlainObject(part))) {
+        const merged = isPlainObject(runtimeParams) ? { ...runtimeParams } : {};
+        // Preserve original runtime params while layering formatter-generated named values on top.
+        for (const part of parts) {
+            if (!isPlainObject(part)) {
+                continue;
+            }
+            for (const [key, value] of Object.entries(part)) {
+                if (value === undefined || (value === null && key in merged)) {
+                    continue;
+                }
+                merged[key] = value;
+            }
+        }
+        return Object.keys(merged).length > 0 ? merged : undefined;
+    }
+    const merged = Array.isArray(runtimeParams) ? [...runtimeParams] : [];
+    for (const part of parts) {
+        if (!Array.isArray(part)) {
+            continue;
+        }
+        // Formatter arrays reuse positional indexes, so keep runtime slots in place and only fill missing/generated entries.
+        part.forEach((value, index) => {
+            if ((value === undefined || value === null) && index < merged.length) {
+                return;
+            }
+            merged[index] = value;
+        });
+    }
+    return merged.length > 0 ? merged : undefined;
+}
+function isPlainObject(value) {
+    return typeof value === 'object' && value !== null && !Array.isArray(value);
+}
+function normalizeParamsForSql(sql, params, runtimeParams) {
+    if (!Array.isArray(params)) {
+        return params;
+    }
+    const maxSlot = findHighestPositionalPlaceholder(sql);
+    if (maxSlot === 0) {
+        return undefined;
+    }
+    const finalized = params.slice(0, maxSlot).map((value, index) => {
+        if ((value === undefined || value === null) && Array.isArray(runtimeParams) && index < runtimeParams.length) {
+            return runtimeParams[index];
+        }
+        return value;
+    });
+    return finalized;
+}
+function findHighestPositionalPlaceholder(sql) {
+    let highest = 0;
+    let index = 0;
+    let dollarQuoteTag = null;
+    while (index < sql.length) {
+        const current = sql[index];
+        const next = sql[index + 1];
+        if (dollarQuoteTag) {
+            if (sql.startsWith(dollarQuoteTag, index)) {
+                index += dollarQuoteTag.length;
+                dollarQuoteTag = null;
+                continue;
+            }
+            index += 1;
+            continue;
+        }
+        if (current === "'") {
+            index = skipQuotedString(sql, index, "'");
+            continue;
+        }
+        if (current === '"') {
+            index = skipQuotedString(sql, index, '"');
+            continue;
+        }
+        if (current === '-' && next === '-') {
+            index = skipLineComment(sql, index + 2);
+            continue;
+        }
+        if (current === '/' && next === '*') {
+            index = skipBlockComment(sql, index + 2);
+            continue;
+        }
+        const dollarTag = readDollarQuoteTag(sql, index);
+        if (dollarTag) {
+            dollarQuoteTag = dollarTag;
+            index += dollarTag.length;
+            continue;
+        }
+        if (current === '$' && isAsciiDigit(next)) {
+            let end = index + 1;
+            while (end < sql.length && isAsciiDigit(sql[end])) {
+                end += 1;
+            }
+            highest = Math.max(highest, Number(sql.slice(index + 1, end)));
+            index = end;
+            continue;
+        }
+        index += 1;
+    }
+    return highest;
+}
+function skipQuotedString(sql, start, quote) {
+    let index = start + 1;
+    while (index < sql.length) {
+        if (sql[index] === quote) {
+            if (sql[index + 1] === quote) {
+                index += 2;
+                continue;
+            }
+            return index + 1;
+        }
+        index += 1;
+    }
+    return index;
+}
+function skipLineComment(sql, start) {
+    let index = start;
+    while (index < sql.length && sql[index] !== '\n') {
+        index += 1;
+    }
+    return index;
+}
+function skipBlockComment(sql, start) {
+    let index = start;
+    while (index < sql.length - 1) {
+        if (sql[index] === '*' && sql[index + 1] === '/') {
+            return index + 2;
+        }
+        index += 1;
+    }
+    return sql.length;
+}
+function readDollarQuoteTag(sql, start) {
+    if (sql[start] !== '$') {
+        return null;
+    }
+    let index = start + 1;
+    while (index < sql.length && sql[index] !== '$') {
+        const current = sql[index];
+        if (!isDollarQuoteTagChar(current)) {
+            return null;
+        }
+        index += 1;
+    }
+    if (sql[index] !== '$') {
+        return null;
+    }
+    return sql.slice(start, index + 1);
+}
+function isDollarQuoteTagChar(char) {
+    return char !== undefined && /[A-Za-z0-9_]/.test(char);
+}
+function isAsciiDigit(value) {
+    return value !== undefined && value >= '0' && value <= '9';
+}
+function normalizePipelineQueryResult(result) {
+    if (Array.isArray(result)) {
+        return { rows: result };
+    }
+    return {
+        rows: result.rows,
+        rowCount: result.rowCount
+    };
+}
+async function cleanupPipelineSession(session, createdTempTables, executionError) {
+    let cleanupError;
+    // Keep dropping later tables even if one cleanup statement fails.
+    for (const tableName of [...createdTempTables].reverse()) {
+        try {
+            await session.query(`drop table if exists ${quoteIdentifier(tableName)}`);
+        }
+        catch (error) {
+            if (!cleanupError) {
+                cleanupError = error;
+            }
+        }
+    }
+    try {
+        await closePipelineSession(session);
+    }
+    catch (error) {
+        if (!executionError && !cleanupError) {
+            throw error;
+        }
+    }
+    if (!executionError && cleanupError) {
+        throw cleanupError;
+    }
+}
+async function closePipelineSession(session) {
+    if (typeof session.release === 'function') {
+        await session.release();
+        return;
+    }
+    if (typeof session.end === 'function') {
+        await session.end();
+    }
+}
+function quoteIdentifier(name) {
+    return `"${name.replace(/"/g, '""')}"`;
+}
+function validateStageOptions(options) {
+    const hasTarget = typeof options.cte === 'string' && options.cte.trim() !== '';
+    const hasFinal = options.final === true;
+    if (hasTarget === hasFinal) {
+        throw new Error('Specify exactly one stage target or final query mode.');
+    }
+}
+function buildSelectFromTargetQuery(targetName, limit) {
+    return new rawsql_ts_1.SimpleSelectQuery({
+        selectClause: new rawsql_ts_1.SelectClause([new rawsql_ts_1.SelectItem(new rawsql_ts_1.RawString('*'))]),
+        fromClause: new rawsql_ts_1.FromClause(new rawsql_ts_1.SourceExpression(new rawsql_ts_1.TableSource(null, targetName), null), null),
+        limitClause: limit === undefined ? null : new rawsql_ts_1.LimitClause(new rawsql_ts_1.LiteralValue(limit))
+    });
+}
+function buildWrappedLimitQuery(statement, limit) {
+    return new rawsql_ts_1.SimpleSelectQuery({
+        selectClause: new rawsql_ts_1.SelectClause([new rawsql_ts_1.SelectItem(new rawsql_ts_1.RawString('*'))]),
+        fromClause: new rawsql_ts_1.FromClause(new rawsql_ts_1.SourceExpression(new rawsql_ts_1.SubQuerySource(statement), new rawsql_ts_1.SourceAliasExpression('final_slice', null)), null),
+        limitClause: new rawsql_ts_1.LimitClause(new rawsql_ts_1.LiteralValue(limit))
+    });
+}
+function applyMinimalWithClause(statement, ctes, recursive) {
+    const nextWithClause = ctes.length > 0 ? new rawsql_ts_1.WithClause(recursive, ctes) : null;
+    if (statement instanceof rawsql_ts_1.SimpleSelectQuery ||
+        statement instanceof rawsql_ts_1.ValuesQuery ||
+        statement instanceof rawsql_ts_1.UpdateQuery ||
+        statement instanceof rawsql_ts_1.DeleteQuery) {
+        statement.withClause = nextWithClause;
+        return;
+    }
+    if (statement instanceof rawsql_ts_1.InsertQuery) {
+        if (!statement.selectQuery) {
+            return;
+        }
+        // INSERT ... SELECT stores the CTEs on the nested selectQuery, not on InsertQuery itself.
+        applyMinimalWithClauseToSelect(assertSelectQuery(statement.selectQuery), nextWithClause);
+        return;
+    }
+    applyMinimalWithClauseToSelect(statement, nextWithClause);
+}
+function applyMinimalWithClauseToSelect(statement, withClause) {
+    if (statement instanceof rawsql_ts_1.SimpleSelectQuery || statement instanceof rawsql_ts_1.ValuesQuery) {
+        statement.withClause = withClause;
+        return;
+    }
+    let current = statement;
+    while (current instanceof rawsql_ts_1.BinarySelectQuery) {
+        if (current.left instanceof rawsql_ts_1.BinarySelectQuery) {
+            current = current.left;
+            continue;
+        }
+        if (current.left instanceof rawsql_ts_1.SimpleSelectQuery || current.left instanceof rawsql_ts_1.ValuesQuery) {
+            current.left.withClause = withClause;
+            return;
+        }
+        break;
+    }
+    throw new Error('Unable to apply rewritten WITH clause to the final query.');
+}
+function getWithClause(statement) {
+    var _a;
+    if (statement instanceof rawsql_ts_1.SimpleSelectQuery ||
+        statement instanceof rawsql_ts_1.ValuesQuery ||
+        statement instanceof rawsql_ts_1.UpdateQuery ||
+        statement instanceof rawsql_ts_1.DeleteQuery) {
+        return (_a = statement.withClause) !== null && _a !== void 0 ? _a : null;
+    }
+    if (statement instanceof rawsql_ts_1.InsertQuery) {
+        return statement.selectQuery ? getSelectWithClause(assertSelectQuery(statement.selectQuery)) : null;
+    }
+    return getSelectWithClause(statement);
+}
+function assertSelectQuery(statement) {
+    if (statement instanceof rawsql_ts_1.SimpleSelectQuery ||
+        statement instanceof rawsql_ts_1.BinarySelectQuery ||
+        statement instanceof rawsql_ts_1.ValuesQuery) {
+        return statement;
+    }
+    throw new Error('Expected a SELECT-compatible statement for query execution.');
+}
+function getSelectWithClause(statement) {
+    var _a, _b;
+    if (statement instanceof rawsql_ts_1.SimpleSelectQuery || statement instanceof rawsql_ts_1.ValuesQuery) {
+        return (_a = statement.withClause) !== null && _a !== void 0 ? _a : null;
+    }
+    let current = statement;
+    while (current instanceof rawsql_ts_1.BinarySelectQuery) {
+        if (current.left instanceof rawsql_ts_1.BinarySelectQuery) {
+            current = current.left;
+            continue;
+        }
+        if (current.left instanceof rawsql_ts_1.SimpleSelectQuery || current.left instanceof rawsql_ts_1.ValuesQuery) {
+            return (_b = current.left.withClause) !== null && _b !== void 0 ? _b : null;
+        }
+        break;
+    }
+    return null;
+}
+//# sourceMappingURL=execute.js.map