@leonardovida-md/drizzle-neo-duckdb 1.2.0 → 1.2.2

package/src/sql/visitors/column-qualifier.ts

@@ -1,5 +1,10 @@
 /**
  * AST visitor to qualify unqualified column references in JOIN ON clauses.
+ *
+ * Performance optimizations:
+ * - Early exit when no unqualified columns found in ON clause
+ * - Skip processing if all columns are already qualified
+ * - Minimal tree traversal when possible
  */
 
 import type {
@@ -17,6 +22,12 @@ import type {
 type TableSource = {
   name: string;
   alias: string | null;
+  schema: string | null;
+};
+
+type Qualifier = {
+  table: string;
+  schema: string | null;
 };
 
 function getTableSource(from: From): TableSource | null {
@@ -24,19 +35,24 @@ function getTableSource(from: From): TableSource | null {
     return {
       name: from.table,
       alias: from.as ?? null,
+      schema: 'db' in from ? (from.db ?? null) : null,
     };
   }
   if ('expr' in from && from.as) {
     return {
       name: from.as,
       alias: from.as,
+      schema: null,
     };
   }
   return null;
 }
 
-function getQualifier(source: TableSource): string {
-  return source.alias ?? source.name;
+function getQualifier(source: TableSource): Qualifier {
+  return {
+    table: source.alias ?? source.name,
+    schema: source.schema,
+  };
 }
 
 function isUnqualifiedColumnRef(expr: ExpressionValue): expr is ColumnRefItem {
@@ -45,7 +61,18 @@ function isUnqualifiedColumnRef(expr: ExpressionValue): expr is ColumnRefItem {
     expr !== null &&
     'type' in expr &&
     expr.type === 'column_ref' &&
-    !('table' in expr && expr.table)
+    (!('table' in expr) || !expr.table)
+  );
+}
+
+function isQualifiedColumnRef(expr: ExpressionValue): expr is ColumnRefItem {
+  return (
+    typeof expr === 'object' &&
+    expr !== null &&
+    'type' in expr &&
+    expr.type === 'column_ref' &&
+    'table' in expr &&
+    !!expr.table
   );
 }
 
@@ -59,51 +86,140 @@ function getColumnName(col: ColumnRefItem): string | null {
   return null;
 }
 
+function applyQualifier(col: ColumnRefItem, qualifier: Qualifier): void {
+  col.table = qualifier.table;
+  if (!('schema' in col) || !col.schema) {
+    (col as ColumnRefItem & { schema?: string | null }).schema =
+      qualifier.schema;
+  }
+}
+
+function unwrapColumnRef(
+  expr: ExpressionValue | undefined
+): ColumnRefItem | null {
+  if (!expr || typeof expr !== 'object') return null;
+  if ('type' in expr && expr.type === 'column_ref') {
+    return expr as ColumnRefItem;
+  }
+  if ('expr' in expr && expr.expr) {
+    return unwrapColumnRef(expr.expr as ExpressionValue);
+  }
+  if ('ast' in expr && expr.ast && typeof expr.ast === 'object') {
+    return null;
+  }
+  if ('args' in expr && expr.args) {
+    const args = expr.args as {
+      value?: ExpressionValue[];
+      expr?: ExpressionValue;
+    };
+    if (args.expr) {
+      return unwrapColumnRef(args.expr as ExpressionValue);
+    }
+    if (args.value && args.value.length === 1) {
+      return unwrapColumnRef(args.value[0] as ExpressionValue);
+    }
+  }
+  return null;
+}
+
+function isBinaryExpr(
+  expr: ExpressionValue | Binary | null | undefined
+): expr is Binary {
+  return (
+    !!expr &&
+    typeof expr === 'object' &&
+    'type' in expr &&
+    (expr as { type?: string }).type === 'binary_expr'
+  );
+}
+
 function walkOnClause(
-  expr: Binary | null | undefined,
-  leftSource: string,
-  rightSource: string,
+  expr: Binary | ExpressionValue | null | undefined,
+  leftQualifier: Qualifier,
+  rightQualifier: Qualifier,
   ambiguousColumns: Set<string>
 ): boolean {
   if (!expr || typeof expr !== 'object') return false;
 
   let transformed = false;
 
-  if (expr.type === 'binary_expr') {
-    if (expr.operator === '=') {
-      const left = expr.left as ExpressionValue;
-      const right = expr.right as ExpressionValue;
-
-      if (isUnqualifiedColumnRef(left) && isUnqualifiedColumnRef(right)) {
-        const leftColName = getColumnName(left);
-        const rightColName = getColumnName(right);
-
-        if (leftColName && rightColName && leftColName === rightColName) {
-          left.table = leftSource;
-          right.table = rightSource;
-
-          ambiguousColumns.add(leftColName);
-          transformed = true;
-        }
+  if (isBinaryExpr(expr)) {
+    const left = expr.left as ExpressionValue;
+    const right = expr.right as ExpressionValue;
+
+    const leftCol = unwrapColumnRef(left);
+    const rightCol = unwrapColumnRef(right);
+
+    const leftUnqualified = leftCol ? isUnqualifiedColumnRef(leftCol) : false;
+    const rightUnqualified = rightCol
+      ? isUnqualifiedColumnRef(rightCol)
+      : false;
+    const leftQualified = leftCol ? isQualifiedColumnRef(leftCol) : false;
+    const rightQualified = rightCol ? isQualifiedColumnRef(rightCol) : false;
+    const leftColName = leftCol ? getColumnName(leftCol) : null;
+    const rightColName = rightCol ? getColumnName(rightCol) : null;
+
+    if (
+      expr.operator === '=' &&
+      leftColName &&
+      rightColName &&
+      leftColName === rightColName
+    ) {
+      if (leftUnqualified && rightUnqualified) {
+        applyQualifier(leftCol!, leftQualifier);
+        applyQualifier(rightCol!, rightQualifier);
+        ambiguousColumns.add(leftColName);
+        transformed = true;
+      } else if (leftQualified && rightUnqualified) {
+        applyQualifier(rightCol!, rightQualifier);
+        ambiguousColumns.add(rightColName);
+        transformed = true;
+      } else if (leftUnqualified && rightQualified) {
+        applyQualifier(leftCol!, leftQualifier);
+        ambiguousColumns.add(leftColName);
+        transformed = true;
       }
     }
 
-    if (expr.operator === 'AND' || expr.operator === 'OR') {
-      transformed =
-        walkOnClause(
-          expr.left as Binary,
-          leftSource,
-          rightSource,
-          ambiguousColumns
-        ) || transformed;
-      transformed =
-        walkOnClause(
-          expr.right as Binary,
-          leftSource,
-          rightSource,
-          ambiguousColumns
-        ) || transformed;
+    if (
+      expr.operator === '=' &&
+      leftCol &&
+      rightCol &&
+      leftColName &&
+      rightColName &&
+      leftColName !== rightColName
+    ) {
+      if (leftQualified && rightUnqualified && !rightColName.includes('.')) {
+        applyQualifier(rightCol, rightQualifier);
+        transformed = true;
+      } else if (
+        leftUnqualified &&
+        rightQualified &&
+        !leftColName.includes('.')
+      ) {
+        applyQualifier(leftCol, leftQualifier);
+        transformed = true;
+      }
     }
+
+    transformed =
+      walkOnClause(
+        isBinaryExpr(expr.left as Binary)
+          ? (expr.left as Binary)
+          : (expr.left as ExpressionValue),
+        leftQualifier,
+        rightQualifier,
+        ambiguousColumns
+      ) || transformed;
+    transformed =
+      walkOnClause(
+        isBinaryExpr(expr.right as Binary)
+          ? (expr.right as Binary)
+          : (expr.right as ExpressionValue),
+        leftQualifier,
+        rightQualifier,
+        ambiguousColumns
+      ) || transformed;
   }
 
   return transformed;
@@ -111,7 +227,7 @@ function walkOnClause(
 
 function qualifyAmbiguousInExpression(
   expr: ExpressionValue | null | undefined,
-  defaultQualifier: string,
+  defaultQualifier: Qualifier,
   ambiguousColumns: Set<string>
 ): boolean {
   if (!expr || typeof expr !== 'object') return false;
@@ -121,13 +237,13 @@ function qualifyAmbiguousInExpression(
   if (isUnqualifiedColumnRef(expr)) {
     const colName = getColumnName(expr);
     if (colName && ambiguousColumns.has(colName)) {
-      expr.table = defaultQualifier;
+      applyQualifier(expr, defaultQualifier);
       transformed = true;
     }
     return transformed;
   }
 
-  if ('type' in expr && expr.type === 'binary_expr') {
+  if (isBinaryExpr(expr)) {
     const binary = expr as Binary;
     transformed =
       qualifyAmbiguousInExpression(
@@ -169,79 +285,186 @@ function qualifyAmbiguousInExpression(
     }
   }
 
+  if ('over' in expr && expr.over && typeof expr.over === 'object') {
+    const over = expr.over as {
+      partition?: ExpressionValue[];
+      orderby?: ExpressionValue[];
+    };
+    if (Array.isArray(over.partition)) {
+      for (const part of over.partition) {
+        transformed =
+          qualifyAmbiguousInExpression(
+            part,
+            defaultQualifier,
+            ambiguousColumns
+          ) || transformed;
+      }
+    }
+    if (Array.isArray(over.orderby)) {
+      for (const order of over.orderby) {
+        transformed =
+          qualifyAmbiguousInExpression(
+            order,
+            defaultQualifier,
+            ambiguousColumns
+          ) || transformed;
+      }
+    }
+  }
+
   return transformed;
 }
 
+/**
+ * Quick check if an ON clause has any unqualified column references.
+ * Used for early exit optimization.
+ */
+function hasUnqualifiedColumns(expr: Binary | null | undefined): boolean {
+  if (!expr || typeof expr !== 'object') return false;
+
+  if ('type' in expr && expr.type === 'binary_expr') {
+    const left = expr.left as ExpressionValue;
+    const right = expr.right as ExpressionValue;
+    const leftCol = unwrapColumnRef(left);
+    const rightCol = unwrapColumnRef(right);
+    if (
+      isUnqualifiedColumnRef(left) ||
+      isUnqualifiedColumnRef(right) ||
+      (leftCol && isUnqualifiedColumnRef(leftCol)) ||
+      (rightCol && isUnqualifiedColumnRef(rightCol))
+    ) {
+      return true;
+    }
+    if (
+      isBinaryExpr(expr.left as Binary) &&
+      hasUnqualifiedColumns(expr.left as Binary)
+    )
+      return true;
+    if (
+      isBinaryExpr(expr.right as Binary) &&
+      hasUnqualifiedColumns(expr.right as Binary)
+    )
+      return true;
+  }
+
+  if ('args' in expr && expr.args) {
+    const args = expr.args as {
+      value?: ExpressionValue[];
+      expr?: ExpressionValue;
+    };
+    if (args.expr && isUnqualifiedColumnRef(args.expr as ExpressionValue))
+      return true;
+    if (args.value) {
+      for (const arg of args.value) {
+        if (isUnqualifiedColumnRef(arg)) return true;
+      }
+    }
+  }
+
+  return false;
+}
+
 function walkSelect(select: Select): boolean {
   let transformed = false;
   const ambiguousColumns = new Set<string>();
 
   if (Array.isArray(select.from) && select.from.length >= 2) {
     const firstSource = getTableSource(select.from[0]);
-    const defaultQualifier = firstSource ? getQualifier(firstSource) : '';
+    const defaultQualifier = firstSource ? getQualifier(firstSource) : null;
     let prevSource = firstSource;
 
+    let hasAnyUnqualified = false;
     for (const from of select.from) {
       if ('join' in from) {
         const join = from as Join;
-        const currentSource = getTableSource(join);
-
-        if (join.on && prevSource && currentSource) {
-          const leftQualifier = getQualifier(prevSource);
-          const rightQualifier = getQualifier(currentSource);
-
-          transformed =
-            walkOnClause(
-              join.on,
-              leftQualifier,
-              rightQualifier,
-              ambiguousColumns
-            ) || transformed;
+        if (join.on && hasUnqualifiedColumns(join.on)) {
+          hasAnyUnqualified = true;
+          break;
         }
+      }
+    }
 
-        prevSource = currentSource;
-      } else {
-        const source = getTableSource(from);
-        if (source) {
-          prevSource = source;
+    if (!hasAnyUnqualified) {
+      for (const from of select.from) {
+        if ('expr' in from && from.expr && 'ast' in from.expr) {
+          transformed = walkSelect(from.expr.ast) || transformed;
         }
       }
+    } else {
+      for (const from of select.from) {
+        if ('join' in from) {
+          const join = from as Join;
+          const currentSource = getTableSource(join);
 
-      if ('expr' in from && from.expr && 'ast' in from.expr) {
-        transformed = walkSelect(from.expr.ast) || transformed;
-      }
-    }
+          if (join.on && prevSource && currentSource) {
+            const leftQualifier = getQualifier(prevSource);
+            const rightQualifier = getQualifier(currentSource);
 
-    if (ambiguousColumns.size > 0 && defaultQualifier) {
-      if (Array.isArray(select.columns)) {
-        for (const col of select.columns as Column[]) {
-          if ('expr' in col) {
             transformed =
-              qualifyAmbiguousInExpression(
-                col.expr,
-                defaultQualifier,
+              walkOnClause(
+                join.on,
+                leftQualifier,
+                rightQualifier,
                 ambiguousColumns
               ) || transformed;
           }
+
+          if (join.using && prevSource && currentSource) {
+            for (const usingCol of join.using) {
+              if (typeof usingCol === 'string') {
+                ambiguousColumns.add(usingCol);
+              } else if ('value' in usingCol) {
+                ambiguousColumns.add(
+                  String((usingCol as { value: unknown }).value)
+                );
+              }
+            }
+          }
+
+          prevSource = currentSource;
+        } else {
+          const source = getTableSource(from);
+          if (source) {
+            prevSource = source;
+          }
+        }
+
+        if ('expr' in from && from.expr && 'ast' in from.expr) {
+          transformed = walkSelect(from.expr.ast) || transformed;
         }
       }
 
-      transformed =
-        qualifyAmbiguousInExpression(
-          select.where,
-          defaultQualifier,
-          ambiguousColumns
-        ) || transformed;
+      if (ambiguousColumns.size > 0 && defaultQualifier) {
+        if (Array.isArray(select.columns)) {
+          for (const col of select.columns as Column[]) {
+            if ('expr' in col) {
+              transformed =
+                qualifyAmbiguousInExpression(
+                  col.expr,
+                  defaultQualifier,
+                  ambiguousColumns
+                ) || transformed;
+            }
+          }
+        }
 
-      if (Array.isArray(select.orderby)) {
-        for (const order of select.orderby as OrderBy[]) {
-          if (order.expr) {
-            transformed =
-              qualifyAmbiguousInExpression(
-                order.expr,
-                defaultQualifier,
-                ambiguousColumns
-              ) || transformed;
+        transformed =
+          qualifyAmbiguousInExpression(
+            select.where,
+            defaultQualifier,
+            ambiguousColumns
+          ) || transformed;
+
+        if (Array.isArray(select.orderby)) {
+          for (const order of select.orderby as OrderBy[]) {
+            if (order.expr) {
+              transformed =
+                qualifyAmbiguousInExpression(
+                  order.expr,
+                  defaultQualifier,
+                  ambiguousColumns
+                ) || transformed;
+            }
           }
         }
       }
@@ -271,6 +494,91 @@ export function qualifyJoinColumns(ast: AST | AST[]): boolean {
   for (const stmt of statements) {
     if (stmt.type === 'select') {
       transformed = walkSelect(stmt as Select) || transformed;
+    } else if (stmt.type === 'insert') {
+      const insert = stmt as unknown as { values?: unknown };
+      if (
+        insert.values &&
+        typeof insert.values === 'object' &&
+        'type' in insert.values &&
+        (insert.values as { type: string }).type === 'select'
+      ) {
+        transformed =
+          walkSelect(insert.values as unknown as Select) || transformed;
+      }
+    } else if (stmt.type === 'update') {
+      const update = stmt as unknown as {
+        table?: From[];
+        from?: From[];
+        where?: ExpressionValue;
+        returning?: ExpressionValue | ExpressionValue[];
+      };
+      const mainSource = update.table?.[0]
+        ? getTableSource(update.table[0] as From)
+        : null;
+      const defaultQualifier = mainSource ? getQualifier(mainSource) : null;
+      const fromSources = update.from ?? [];
+      const firstFrom = fromSources[0] ? getTableSource(fromSources[0]) : null;
+      if (update.where && defaultQualifier && firstFrom) {
+        const ambiguous = new Set<string>();
+        transformed =
+          walkOnClause(
+            update.where as Binary,
+            defaultQualifier,
+            getQualifier(firstFrom),
+            ambiguous
+          ) || transformed;
+        transformed =
+          qualifyAmbiguousInExpression(
+            update.where,
+            defaultQualifier,
+            ambiguous
+          ) || transformed;
+      }
+      if (Array.isArray(update.returning) && defaultQualifier) {
+        for (const ret of update.returning) {
+          transformed =
+            qualifyAmbiguousInExpression(
+              ret,
+              defaultQualifier,
+              new Set<string>()
+            ) || transformed;
+        }
+      }
+    } else if (stmt.type === 'delete') {
+      const del = stmt as unknown as {
+        table?: From[];
+        from?: From[];
+        where?: ExpressionValue;
+      };
+      const mainSource = del.table?.[0]
+        ? getTableSource(del.table[0] as From)
+        : null;
+      const defaultQualifier = mainSource ? getQualifier(mainSource) : null;
+      const fromSources = del.from ?? [];
+      const firstFrom = fromSources[0] ? getTableSource(fromSources[0]) : null;
+      if (del.where && defaultQualifier && firstFrom) {
+        const ambiguous = new Set<string>();
+        transformed =
+          walkOnClause(
+            del.where as Binary,
+            defaultQualifier,
+            getQualifier(firstFrom),
+            ambiguous
+          ) || transformed;
+        transformed =
+          qualifyAmbiguousInExpression(
+            del.where,
+            defaultQualifier,
+            ambiguous
+          ) || transformed;
+      } else if (del.where && defaultQualifier) {
+        transformed =
+          qualifyAmbiguousInExpression(
+            del.where,
+            defaultQualifier,
+            new Set<string>()
+          ) || transformed;
+      }
     }
   }