rawsql-ts 0.26.1 → 0.27.0

package/dist/esm/transformers/StaticPredicatePlacementOptimizer.js

@@ -1,9 +1,10 @@
-import { SubQuerySource, TableSource, WhereClause } from "../models/Clause";
+import { DistinctOn, SubQuerySource, TableSource, WhereClause } from "../models/Clause";
 import { BinarySelectQuery, SimpleSelectQuery } from "../models/SelectQuery";
 import { ArrayExpression, ArrayQueryExpression, BetweenExpression, BinaryExpression, CaseExpression, CastExpression, ColumnReference, FunctionCall, InlineQuery, JsonPredicateExpression, ParameterExpression, ParenExpression, TupleExpression, TypeValue, UnaryExpression, ValueList } from "../models/ValueComponent";
 import { SelectQueryParser } from "../parsers/SelectQueryParser";
 import { ValueParser } from "../parsers/ValueParser";
-import { formatSqlComponent } from "./SqlComponentFormatter";
+import { formatSqlComponent, hasSqlComponentFormatOverride } from "./SqlComponentFormatter";
+import { SelectOutputCollector } from "./SelectOutputCollector";
 const VOLATILE_OR_UNSUPPORTED_FUNCTION_REASON = "Predicate contains a function call; volatile and expression predicates are not moved in the safe-only implementation.";
 const normalizeIdentifier = (value) => value.trim().toLowerCase();
 const unwrapParens = (expression) => {
@@ -44,8 +45,8 @@ const rebuildWhereWithoutTerms = (query, termsToRemove) => {
     }
     query.whereClause = new WhereClause(rebuilt);
 };
-const cloneValueComponent = (expression) => {
-    return ValueParser.parse(formatSqlComponent(expression));
+const cloneValueComponent = (expression, options) => {
+    return ValueParser.parse(formatSqlComponent(expression, options));
 };
 const cloneColumnReference = (reference) => {
     var _a, _b;
@@ -68,11 +69,12 @@ const appendUnique = (items, value) => {
 export class StaticPredicatePlacementOptimizer {
     plan(input, options = {}) {
         var _a, _b, _c;
-        const parsed = this.parseInput(input);
+        const parsed = this.parseInput(input, options);
         const warnings = [...parsed.warnings];
         const errors = [...parsed.errors];
         if (!parsed.query) {
             return this.buildResult({
+                query: null,
                 sql: parsed.sql,
                 applied: [],
                 skipped: [],
@@ -88,6 +90,7 @@ export class StaticPredicatePlacementOptimizer {
                 message: "Static predicate placement currently supports only SimpleSelectQuery roots."
             });
             return this.buildResult({
+                query: parsed.query,
                 sql: parsed.sql,
                 applied: [],
                 skipped: [],
@@ -102,39 +105,74 @@ export class StaticPredicatePlacementOptimizer {
         const skipped = [];
         const movedTerms = [];
         for (const term of query.whereClause ? collectTopLevelAndTerms(query.whereClause.condition) : []) {
-            const candidate = this.analyzeCandidate(query, term);
+            const candidate = this.analyzeCandidate(query, term, options);
             if (!candidate) {
                 continue;
             }
             if ("code" in candidate) {
-                skipped.push(this.makeSkipped(term, candidate));
+                skipped.push(this.makeSkipped(term, candidate, options));
                 continue;
             }
             const target = this.resolveTarget(query, candidate);
             if ("code" in target) {
-                skipped.push(this.makeSkipped(term, target));
+                skipped.push(this.makeSkipped(term, target, options));
                 continue;
             }
-            const targetColumns = this.resolveTargetColumns(target.query, candidate.references);
-            if ("code" in targetColumns) {
-                skipped.push(this.makeSkipped(term, targetColumns));
-                continue;
+            let appliedReason = "";
+            if (target.kind === "simple") {
+                const targetColumns = this.resolveTargetColumns(query, target.query, candidate.references);
+                if ("code" in targetColumns) {
+                    skipped.push(this.makeSkipped(term, targetColumns, options));
+                    continue;
+                }
+                const placement = this.resolveTargetPlacement(target.query, targetColumns);
+                if ("code" in placement) {
+                    skipped.push(this.makeSkipped(term, placement, options));
+                    continue;
+                }
+                const movedPredicate = this.rebasePredicate(candidate.expression, targetColumns, options);
+                target.query.appendWhere(movedPredicate);
+                appliedReason = placement.reason;
+            }
+            else {
+                const placements = [];
+                let skip = null;
+                for (const branch of target.branches) {
+                    const placement = this.resolveTargetPlacement(branch.query, branch.targetColumns);
+                    if ("code" in placement) {
+                        skip = placement;
+                        break;
+                    }
+                    placements.push(placement);
+                }
+                if (skip) {
+                    skipped.push(this.makeSkipped(term, skip, options));
+                    continue;
+                }
+                for (const branch of target.branches) {
+                    const movedPredicate = this.rebasePredicate(candidate.expression, branch.targetColumns, options);
+                    branch.query.appendWhere(movedPredicate);
+                }
+                appliedReason = placements.some(item => /group by/i.test(item.reason))
+                    ? "Predicate is distributed to every UNION branch by output column position; grouped branches only receive GROUP BY-key predicates."
+                    : "Predicate is distributed to every UNION branch by output column position before unsafe query boundaries.";
             }
-            const movedPredicate = this.rebasePredicate(candidate.expression, targetColumns);
-            target.query.appendWhere(movedPredicate);
             movedTerms.push(term);
             applied.push({
                 kind: "move_static_predicate",
                 predicateSql: candidate.predicateSql,
                 fromScopeId: "scope:root",
                 toScopeId: target.scopeId,
-                reason: "All outer references resolve to direct upstream outputs before unsafe query boundaries.",
+                reason: appliedReason,
                 columnReferences: candidate.references.map(columnReferenceText)
             });
         }
         rebuildWhereWithoutTerms(query, new Set(movedTerms));
-        const sql = applied.length > 0 ? formatSqlComponent(query) : parsed.sql;
+        const sql = applied.length > 0 || hasSqlComponentFormatOverride(options)
+            ? formatSqlComponent(query, options)
+            : parsed.sql;
         return this.buildResult({
+            query,
             sql,
             applied,
             skipped,
@@ -148,13 +186,22 @@ export class StaticPredicatePlacementOptimizer {
         var _a;
         return this.plan(input, Object.assign(Object.assign({}, options), { dryRun: (_a = options.dryRun) !== null && _a !== void 0 ? _a : false }));
     }
-    parseInput(input) {
+    parseInput(input, options) {
         const warnings = [];
         const errors = [];
         try {
             const sourceSql = typeof input === "string"
                 ? input
-                : formatSqlComponent(input);
+                : formatSqlComponent(input, options);
+            if (typeof input !== "string" && options.cloneInput === false) {
+                return {
+                    query: input,
+                    sql: sourceSql,
+                    formatterGeneratedSource: false,
+                    warnings,
+                    errors
+                };
+            }
             if (typeof input !== "string") {
                 warnings.push({
                     code: "AST_INPUT_FORMATTED",
@@ -185,7 +232,7 @@ export class StaticPredicatePlacementOptimizer {
             };
         }
     }
-    analyzeCandidate(root, expression) {
+    analyzeCandidate(root, expression, options) {
         if (this.collectParameterNames(expression).length > 0) {
             return null;
         }
@@ -202,7 +249,7 @@ export class StaticPredicatePlacementOptimizer {
         }
         return {
             expression,
-            predicateSql: formatSqlComponent(expression),
+            predicateSql: formatSqlComponent(expression, options),
             references
         };
     }
@@ -263,13 +310,6 @@ export class StaticPredicatePlacementOptimizer {
             }
             const candidate = unwrapParens(value);
             if (candidate instanceof BinaryExpression) {
-                if (candidate.operator.value.trim().toLowerCase() === "or") {
-                    found = {
-                        code: "OR_PREDICATE_UNSUPPORTED",
-                        reason: "Predicate contains OR predicates, which are not moved in the first safe-only implementation."
-                    };
-                    return;
-                }
                 visit(candidate.left);
                 visit(candidate.right);
                 return;
@@ -306,13 +346,6 @@ export class StaticPredicatePlacementOptimizer {
                 };
                 return;
             }
-            if (candidate instanceof BetweenExpression) {
-                found = {
-                    code: "BETWEEN_PREDICATE_UNSUPPORTED",
-                    reason: "BETWEEN predicates are not moved in the first safe-only implementation."
-                };
-                return;
-            }
             if (candidate instanceof UnaryExpression) {
                 visit(candidate.expression);
                 return;
@@ -345,7 +378,7 @@ export class StaticPredicatePlacementOptimizer {
         return found;
     }
     resolveTarget(root, candidate) {
-        const boundary = this.findCurrentQueryBoundary(root);
+        const boundary = this.findRootQueryBoundary(root);
         if (boundary) {
             return boundary;
         }
@@ -357,7 +390,7 @@ export class StaticPredicatePlacementOptimizer {
         }
         const bindings = [];
         for (const reference of candidate.references) {
-            const binding = this.resolveSourceBinding(root, reference);
+            const binding = this.resolveSourceBinding(root, root, reference);
             if ("code" in binding) {
                 return binding;
             }
@@ -376,27 +409,17 @@ export class StaticPredicatePlacementOptimizer {
         if (nullableSide) {
             return nullableSide;
         }
-        const upstream = this.resolveUpstreamQuery(root, binding);
+        const upstream = this.resolveUpstreamQuery(root, binding, candidate.references);
         if ("code" in upstream) {
             return upstream;
         }
-        const targetBoundary = this.findTargetQueryBoundary(upstream.query);
-        if (targetBoundary) {
-            return targetBoundary;
-        }
         return upstream;
     }
-    findCurrentQueryBoundary(query) {
-        if (query.selectClause.distinct) {
+    findRootQueryBoundary(query) {
+        if (this.hasDistinctOnBoundary(query)) {
             return {
                 code: "DISTINCT_BOUNDARY",
-                reason: "Predicate crosses DISTINCT boundary; moving it may change semantics."
-            };
-        }
-        if (query.groupByClause || query.havingClause) {
-            return {
-                code: "GROUP_BY_BOUNDARY",
-                reason: "Predicate crosses GROUP BY boundary; moving it may change semantics."
+                reason: "Predicate crosses DISTINCT ON boundary; moving it may change semantics."
             };
         }
         if (this.hasWindowUsage(query)) {
@@ -407,10 +430,19 @@ export class StaticPredicatePlacementOptimizer {
         }
         return null;
     }
-    findTargetQueryBoundary(query) {
-        const commonBoundary = this.findCurrentQueryBoundary(query);
-        if (commonBoundary) {
-            return commonBoundary;
+    resolveTargetPlacement(query, targetColumns) {
+        const hasOrdinaryDistinct = this.hasOrdinaryDistinct(query);
+        if (this.hasDistinctOnBoundary(query)) {
+            return {
+                code: "DISTINCT_BOUNDARY",
+                reason: "Predicate crosses DISTINCT ON boundary; moving it may change semantics."
+            };
+        }
+        if (this.hasWindowUsage(query)) {
+            return {
+                code: "WINDOW_BOUNDARY",
+                reason: "Predicate crosses WINDOW boundary; moving it may change semantics."
+            };
         }
         if (query.limitClause || query.offsetClause || query.fetchClause) {
             return {
@@ -424,10 +456,35 @@ export class StaticPredicatePlacementOptimizer {
                 reason: "Target query contains an OUTER JOIN boundary that is not moved across in the safe-only implementation."
             };
         }
-        return null;
+        if (query.groupByClause) {
+            const allReferencesAreGroupKeys = targetColumns.every(item => this.isGroupKeyColumn(query, item.targetColumn));
+            if (!allReferencesAreGroupKeys) {
+                return {
+                    code: "GROUP_BY_BOUNDARY",
+                    reason: "Predicate references a target column that is not proven to be a GROUP BY key."
+                };
+            }
+            return {
+                reason: "Predicate references only GROUP BY keys; it is moved into pre-aggregation WHERE."
+            };
+        }
+        if (query.havingClause) {
+            return {
+                code: "GROUP_BY_BOUNDARY",
+                reason: "Predicate crosses HAVING aggregation boundary; moving it may change semantics."
+            };
+        }
+        if (hasOrdinaryDistinct) {
+            return {
+                reason: "Predicate references direct ordinary DISTINCT output columns; it is moved into the DISTINCT input WHERE."
+            };
+        }
+        return {
+            reason: "All outer references resolve to direct upstream outputs before unsafe query boundaries."
+        };
     }
-    resolveSourceBinding(root, column) {
-        const fromClause = root.fromClause;
+    resolveSourceBinding(contextRoot, query, column) {
+        const fromClause = query.fromClause;
         if (!fromClause) {
             return {
                 code: "NO_FROM_CLAUSE",
@@ -445,10 +502,7 @@ export class StaticPredicatePlacementOptimizer {
                     reason: `Column source alias '${column.getNamespace()}' is not uniquely resolvable.`
                 };
             }
-            if (this.resolveSourceQueryForColumns(root, matches[0].source) instanceof BinarySelectQuery) {
-                return matches[0];
-            }
-            const matchCount = this.getOutputColumnMatchCount(root, matches[0], columnName);
+            const matchCount = this.getOutputColumnMatchCount(contextRoot, matches[0], columnName);
             if (matchCount === 0) {
                 return {
                     code: "COLUMN_NOT_AVAILABLE_UPSTREAM",
@@ -464,13 +518,8 @@ export class StaticPredicatePlacementOptimizer {
             return matches[0];
         }
         const matches = [];
-        let unionSource = null;
         for (const binding of bindings) {
-            if (this.resolveSourceQueryForColumns(root, binding.source) instanceof BinarySelectQuery) {
-                unionSource !== null && unionSource !== void 0 ? unionSource : (unionSource = binding);
-                continue;
-            }
-            const matchCount = this.getOutputColumnMatchCount(root, binding, columnName);
+            const matchCount = this.getOutputColumnMatchCount(contextRoot, binding, columnName);
             if (matchCount > 1) {
                 return {
                     code: "AMBIGUOUS_COLUMN_REFERENCE",
@@ -481,12 +530,6 @@ export class StaticPredicatePlacementOptimizer {
                 matches.push(binding);
             }
         }
-        if (matches.length === 0 && unionSource) {
-            return {
-                code: "UNION_BOUNDARY",
-                reason: "Predicate would need distribution into a UNION branch, which is unsupported."
-            };
-        }
         if (matches.length !== 1) {
             return {
                 code: "AMBIGUOUS_COLUMN_REFERENCE",
@@ -497,16 +540,20 @@ export class StaticPredicatePlacementOptimizer {
         }
         return matches[0];
     }
-    resolveUpstreamQuery(root, binding) {
+    resolveUpstreamQuery(root, binding, references) {
         const source = binding.source.datasource;
         if (source instanceof SubQuerySource) {
             if (source.query instanceof SimpleSelectQuery) {
                 return {
+                    kind: "simple",
                     query: source.query,
                     scopeId: `subquery:${binding.alias}`,
                     sourceBinding: binding
                 };
             }
+            if (source.query instanceof BinarySelectQuery) {
+                return this.resolveUnionTarget(root, source.query, `subquery:${binding.alias}`, references);
+            }
             return {
                 code: "UNION_BOUNDARY",
                 reason: "Predicate would need distribution into a UNION or non-simple subquery, which is unsupported."
@@ -534,10 +581,7 @@ export class StaticPredicatePlacementOptimizer {
             };
         }
         if (commonTable.query instanceof BinarySelectQuery) {
-            return {
-                code: "UNION_BOUNDARY",
-                reason: "Predicate would need distribution into a UNION branch, which is unsupported."
-            };
+            return this.resolveUnionTarget(root, commonTable.query, `cte:${commonTable.getSourceAliasName()}`, references);
         }
         if (!(commonTable.query instanceof SimpleSelectQuery)) {
             return {
@@ -546,16 +590,48 @@ export class StaticPredicatePlacementOptimizer {
             };
         }
         return {
+            kind: "simple",
             query: commonTable.query,
             scopeId: `cte:${commonTable.getSourceAliasName()}`,
             sourceBinding: binding,
             cteName: commonTable.getSourceAliasName()
         };
     }
-    resolveTargetColumns(query, references) {
+    resolveUnionTarget(root, query, scopeId, references) {
+        const branches = this.collectUnionBranches(query);
+        if ("code" in branches) {
+            return branches;
+        }
+        const outputIndexes = new Map();
+        for (const reference of references) {
+            const outputIndex = this.resolveUnionOutputIndex(root, branches[0], reference.column.name);
+            if ("code" in outputIndex) {
+                return outputIndex;
+            }
+            outputIndexes.set(reference, outputIndex.index);
+        }
+        const branchTargets = [];
+        for (const branch of branches) {
+            const targetColumns = [];
+            for (const [reference, outputIndex] of outputIndexes.entries()) {
+                const targetColumn = this.resolveTargetColumnByOutputIndex(root, branch, outputIndex, reference);
+                if ("code" in targetColumn) {
+                    return targetColumn;
+                }
+                targetColumns.push(targetColumn);
+            }
+            branchTargets.push(this.resolveDeepestBranchTarget(root, branch, targetColumns));
+        }
+        return {
+            kind: "union",
+            scopeId,
+            branches: branchTargets
+        };
+    }
+    resolveTargetColumns(root, query, references) {
         const resolved = [];
         for (const reference of references) {
-            const matches = query.selectClause.items.filter(item => { var _a; return normalizeIdentifier((_a = this.getSelectItemOutputName(item)) !== null && _a !== void 0 ? _a : "") === normalizeIdentifier(reference.column.name); });
+            const matches = this.collectSelectOutputs(root, query).filter(item => normalizeIdentifier(item.name) === normalizeIdentifier(reference.column.name));
             if (matches.length !== 1) {
                 return {
                     code: "AMBIGUOUS_TARGET_COLUMN",
@@ -582,6 +658,50 @@ export class StaticPredicatePlacementOptimizer {
         }
         return resolved;
     }
+    resolveDeepestBranchTarget(contextRoot, query, targetColumns, visited = new Set()) {
+        if (visited.has(query) || targetColumns.length === 0) {
+            return { query, targetColumns: [...targetColumns] };
+        }
+        const nextVisited = new Set(visited);
+        nextVisited.add(query);
+        const bindings = [];
+        for (const item of targetColumns) {
+            const binding = this.resolveSourceBinding(contextRoot, query, item.targetColumn);
+            if ("code" in binding) {
+                return { query, targetColumns: [...targetColumns] };
+            }
+            if (!bindings.some(existing => existing.source === binding.source)) {
+                bindings.push(binding);
+            }
+        }
+        if (bindings.length !== 1) {
+            return { query, targetColumns: [...targetColumns] };
+        }
+        const binding = bindings[0];
+        const nullableSide = query.fromClause
+            ? this.findNullableSideBoundary(query.fromClause, binding)
+            : null;
+        if (nullableSide) {
+            return { query, targetColumns: [...targetColumns] };
+        }
+        const upstream = this.resolveUpstreamQuery(contextRoot, binding, targetColumns.map(item => item.targetColumn));
+        if ("code" in upstream || upstream.kind === "union") {
+            return { query, targetColumns: [...targetColumns] };
+        }
+        const upstreamColumns = this.resolveTargetColumns(contextRoot, upstream.query, targetColumns.map(item => item.targetColumn));
+        if ("code" in upstreamColumns) {
+            return { query, targetColumns: [...targetColumns] };
+        }
+        const placement = this.resolveTargetPlacement(upstream.query, upstreamColumns);
+        if ("code" in placement) {
+            return { query, targetColumns: [...targetColumns] };
+        }
+        const rebasedColumns = upstreamColumns.map((item, index) => ({
+            sourceColumn: targetColumns[index].sourceColumn,
+            targetColumn: item.targetColumn
+        }));
+        return this.resolveDeepestBranchTarget(contextRoot, upstream.query, rebasedColumns, nextVisited);
+    }
     verifyColumnResolvableInQuery(query, column) {
         if (!query.fromClause) {
             return {
@@ -608,8 +728,42 @@ export class StaticPredicatePlacementOptimizer {
             reason: `Unqualified target column '${column.column.name}' is ambiguous in a multi-source destination query.`
         };
     }
-    rebasePredicate(expression, targetColumns) {
-        const cloned = cloneValueComponent(expression);
+    isGroupKeyColumn(query, column) {
+        const groupBy = query.groupByClause;
+        if (!groupBy || groupBy.mode || groupBy.grouping.length === 0) {
+            return false;
+        }
+        return groupBy.grouping.some(grouping => {
+            const candidate = unwrapParens(grouping);
+            return candidate instanceof ColumnReference
+                && this.sameResolvableColumnInQuery(query, candidate, column);
+        });
+    }
+    sameResolvableColumnInQuery(query, left, right) {
+        if (sameColumnReference(left, right)) {
+            return true;
+        }
+        if (normalizeIdentifier(left.column.name) !== normalizeIdentifier(right.column.name)) {
+            return false;
+        }
+        const leftSource = this.resolveColumnSourceAlias(query, left);
+        const rightSource = this.resolveColumnSourceAlias(query, right);
+        return leftSource !== null && leftSource === rightSource;
+    }
+    resolveColumnSourceAlias(query, column) {
+        if (!query.fromClause) {
+            return null;
+        }
+        const bindings = this.getSourceBindings(query.fromClause);
+        const namespace = normalizeIdentifier(column.getNamespace());
+        if (namespace) {
+            const matches = bindings.filter(binding => normalizeIdentifier(binding.alias) === namespace);
+            return matches.length === 1 ? normalizeIdentifier(matches[0].alias) : null;
+        }
+        return bindings.length === 1 ? normalizeIdentifier(bindings[0].alias) : null;
+    }
+    rebasePredicate(expression, targetColumns, options) {
+        const cloned = cloneValueComponent(expression, options);
         const visitSelect = (query, inheritedLocalAliases) => {
             var _a, _b, _c, _d;
             if (!(query instanceof SimpleSelectQuery)) {
@@ -794,12 +948,103 @@ export class StaticPredicatePlacementOptimizer {
             return joinType.includes("left") || joinType.includes("right") || joinType.includes("full") || joinType.includes("outer");
         });
     }
+    hasDistinctOnBoundary(query) {
+        return query.selectClause.distinct instanceof DistinctOn;
+    }
+    hasOrdinaryDistinct(query) {
+        return query.selectClause.distinct !== null && !this.hasDistinctOnBoundary(query);
+    }
     getOutputColumnMatchCount(root, binding, columnName) {
         const target = this.resolveSourceQueryForColumns(root, binding.source);
-        if (!target || !(target instanceof SimpleSelectQuery)) {
+        if (!target) {
+            return 0;
+        }
+        if (target instanceof BinarySelectQuery) {
+            const branches = this.collectUnionBranches(target);
+            if ("code" in branches) {
+                return 0;
+            }
+            return this.collectSelectOutputs(root, branches[0]).filter(item => normalizeIdentifier(item.name) === normalizeIdentifier(columnName)).length;
+        }
+        if (!(target instanceof SimpleSelectQuery)) {
             return 0;
         }
-        return target.selectClause.items.filter(item => { var _a; return normalizeIdentifier((_a = this.getSelectItemOutputName(item)) !== null && _a !== void 0 ? _a : "") === normalizeIdentifier(columnName); }).length;
+        return this.collectSelectOutputs(root, target).filter(item => normalizeIdentifier(item.name) === normalizeIdentifier(columnName)).length;
+    }
+    collectUnionBranches(query) {
+        const branches = [];
+        const visit = (select) => {
+            var _a;
+            if (select instanceof SimpleSelectQuery) {
+                branches.push(select);
+                return null;
+            }
+            if (!(select instanceof BinarySelectQuery)) {
+                return {
+                    code: "UNION_BOUNDARY",
+                    reason: "Predicate would need distribution into a non-simple UNION branch, which is unsupported."
+                };
+            }
+            const operator = select.operator.value.trim().toLowerCase();
+            if (operator !== "union" && operator !== "union all") {
+                return {
+                    code: "UNION_BOUNDARY",
+                    reason: `Predicate would need distribution through '${select.operator.value}', which is unsupported.`
+                };
+            }
+            return (_a = visit(select.left)) !== null && _a !== void 0 ? _a : visit(select.right);
+        };
+        const skip = visit(query);
+        return skip !== null && skip !== void 0 ? skip : branches;
+    }
+    resolveUnionOutputIndex(root, firstBranch, outputColumnName) {
+        const matches = [];
+        this.collectSelectOutputs(root, firstBranch).forEach((item, index) => {
+            if (normalizeIdentifier(item.name) === normalizeIdentifier(outputColumnName)) {
+                matches.push(index);
+            }
+        });
+        if (matches.length !== 1) {
+            return {
+                code: "AMBIGUOUS_TARGET_COLUMN",
+                reason: matches.length === 0
+                    ? `UNION output does not expose '${outputColumnName}' from its first branch.`
+                    : `UNION first branch exposes multiple '${outputColumnName}' columns.`
+            };
+        }
+        return { index: matches[0] };
+    }
+    resolveTargetColumnByOutputIndex(root, query, outputIndex, sourceColumn) {
+        const output = this.collectSelectOutputs(root, query)[outputIndex];
+        if (!output) {
+            return {
+                code: "UNION_COLUMN_MISMATCH",
+                reason: `UNION branch does not expose column '${sourceColumn.column.name}' at output position ${outputIndex + 1}.`
+            };
+        }
+        if (!(output.value instanceof ColumnReference)) {
+            return {
+                code: "EXPRESSION_OUTPUT_UNSUPPORTED",
+                reason: `UNION branch output at position ${outputIndex + 1} for '${sourceColumn.column.name}' is an expression, not a direct column reference.`
+            };
+        }
+        const sourceResolution = this.verifyColumnResolvableInQuery(query, output.value);
+        if (sourceResolution) {
+            return sourceResolution;
+        }
+        return {
+            sourceColumn,
+            targetColumn: output.value
+        };
+    }
+    collectSelectOutputs(root, query) {
+        var _a, _b, _c, _d;
+        const commonTables = [
+            ...((_b = (_a = query.withClause) === null || _a === void 0 ? void 0 : _a.tables) !== null && _b !== void 0 ? _b : []),
+            ...((_d = (_c = root.withClause) === null || _c === void 0 ? void 0 : _c.tables) !== null && _d !== void 0 ? _d : [])
+        ];
+        const collector = new SelectOutputCollector(null, commonTables.length > 0 ? commonTables : null);
+        return collector.collect(query);
     }
     resolveSourceQueryForColumns(root, source) {
         var _a;
@@ -814,15 +1059,6 @@ export class StaticPredicatePlacementOptimizer {
         }
         return null;
     }
-    getSelectItemOutputName(item) {
-        if (item.identifier) {
-            return item.identifier.name;
-        }
-        if (item.value instanceof ColumnReference) {
-            return item.value.column.name;
-        }
-        return null;
-    }
     findCte(root, name) {
         var _a, _b;
         const normalized = normalizeIdentifier(name);
@@ -1168,13 +1404,14 @@ export class StaticPredicatePlacementOptimizer {
         visit(expression);
         return names;
     }
-    makeSkipped(expression, draft) {
-        return Object.assign({ predicateSql: formatSqlComponent(expression), scopeId: "scope:root" }, draft);
+    makeSkipped(expression, draft, options) {
+        return Object.assign({ predicateSql: formatSqlComponent(expression, options), scopeId: "scope:root" }, draft);
     }
     buildResult(params) {
         return {
             ok: params.errors.length === 0,
             sql: params.sql,
+            query: params.query,
             applied: params.applied,
             skipped: params.skipped,
             warnings: params.warnings,