@duckdbfan/drizzle-duckdb 0.0.7 → 1.3.2

package/src/sql/visitors/generate-series-alias.ts

@@ -0,0 +1,291 @@
+/**
+ * AST visitor to rewrite Postgres style generate_series aliases.
+ *
+ * Postgres lets you reference a generate_series alias as a column:
+ *   FROM generate_series(...) AS gs
+ *   SELECT gs::date
+ *
+ * DuckDB treats gs as a table alias, and the column is generate_series.
+ * This visitor rewrites unqualified column refs that match a
+ * generate_series alias to gs.generate_series.
+ */
+
+import type {
+  AST,
+  Binary,
+  ColumnRefItem,
+  ExpressionValue,
+  From,
+  Join,
+  Select,
+  OrderBy,
+  Column,
+} from 'node-sql-parser';
+
+function getColumnName(col: ColumnRefItem): string | null {
+  if (typeof col.column === 'string') {
+    return col.column;
+  }
+  if (col.column && 'expr' in col.column && col.column.expr?.value) {
+    return String(col.column.expr.value);
+  }
+  return null;
+}
+
+function isColumnRef(expr: ExpressionValue): expr is ColumnRefItem {
+  return (
+    typeof expr === 'object' &&
+    expr !== null &&
+    'type' in expr &&
+    expr.type === 'column_ref'
+  );
+}
+
+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 getGenerateSeriesAliases(from: Select['from']): Set<string> {
+  const aliases = new Set<string>();
+  if (!from || !Array.isArray(from)) return aliases;
+
+  for (const f of from) {
+    if ('expr' in f && f.expr && typeof f.expr === 'object') {
+      const exprObj = f.expr as Record<string, unknown>;
+      if (exprObj.type === 'function' && 'name' in exprObj) {
+        const nameObj = exprObj.name as Record<string, unknown> | undefined;
+        const nameParts = nameObj?.name as
+          | Array<Record<string, unknown>>
+          | undefined;
+        const fnName = nameParts?.[0]?.value;
+        if (
+          typeof fnName === 'string' &&
+          fnName.toLowerCase() === 'generate_series'
+        ) {
+          const alias = typeof f.as === 'string' ? f.as : null;
+          if (alias && !alias.includes('(')) {
+            aliases.add(alias);
+          }
+        }
+      }
+    }
+  }
+
+  return aliases;
+}
+
+function rewriteAliasColumnRef(col: ColumnRefItem, alias: string): void {
+  col.table = alias;
+  col.column = { expr: { type: 'default', value: 'generate_series' } };
+}
+
+function walkExpression(
+  expr: ExpressionValue | null | undefined,
+  aliases: Set<string>
+): boolean {
+  if (!expr || typeof expr !== 'object') return false;
+
+  let transformed = false;
+  const exprObj = expr as Record<string, unknown>;
+
+  if (isColumnRef(expr)) {
+    if (!('table' in expr) || !expr.table) {
+      const colName = getColumnName(expr);
+      if (colName && aliases.has(colName)) {
+        rewriteAliasColumnRef(expr, colName);
+        transformed = true;
+      }
+    }
+    return transformed;
+  }
+
+  if (isBinaryExpr(expr)) {
+    const binary = expr as Binary;
+    transformed =
+      walkExpression(binary.left as ExpressionValue, aliases) || transformed;
+    transformed =
+      walkExpression(binary.right as ExpressionValue, aliases) || transformed;
+    return transformed;
+  }
+
+  if (exprObj.type === 'unary_expr' && exprObj.expr) {
+    transformed =
+      walkExpression(exprObj.expr as ExpressionValue, aliases) || transformed;
+  }
+
+  if (exprObj.type === 'cast' && exprObj.expr) {
+    transformed =
+      walkExpression(exprObj.expr as ExpressionValue, aliases) || transformed;
+  }
+
+  if (exprObj.type === 'case') {
+    if (exprObj.expr) {
+      transformed =
+        walkExpression(exprObj.expr as ExpressionValue, aliases) || transformed;
+    }
+    if (Array.isArray(exprObj.args)) {
+      for (const whenClause of exprObj.args as Array<Record<string, unknown>>) {
+        if (whenClause.cond) {
+          transformed =
+            walkExpression(whenClause.cond as ExpressionValue, aliases) ||
+            transformed;
+        }
+        if (whenClause.result) {
+          transformed =
+            walkExpression(whenClause.result as ExpressionValue, aliases) ||
+            transformed;
+        }
+      }
+    }
+  }
+
+  if ('args' in exprObj && exprObj.args) {
+    const args = exprObj.args as Record<string, unknown>;
+    if (Array.isArray(args.value)) {
+      for (const arg of args.value as ExpressionValue[]) {
+        transformed = walkExpression(arg, aliases) || transformed;
+      }
+    } else if (args.expr) {
+      transformed =
+        walkExpression(args.expr as ExpressionValue, aliases) || transformed;
+    }
+  }
+
+  if ('over' in exprObj && exprObj.over && typeof exprObj.over === 'object') {
+    const over = exprObj.over as Record<string, unknown>;
+    if (Array.isArray(over.partition)) {
+      for (const part of over.partition as ExpressionValue[]) {
+        transformed = walkExpression(part, aliases) || transformed;
+      }
+    }
+    if (Array.isArray(over.orderby)) {
+      for (const order of over.orderby as ExpressionValue[]) {
+        transformed = walkExpression(order, aliases) || transformed;
+      }
+    }
+  }
+
+  if ('ast' in exprObj && exprObj.ast) {
+    const subAst = exprObj.ast as Select;
+    if (subAst.type === 'select') {
+      transformed = walkSelect(subAst) || transformed;
+    }
+  }
+
+  if (exprObj.type === 'expr_list' && Array.isArray(exprObj.value)) {
+    for (const item of exprObj.value as ExpressionValue[]) {
+      transformed = walkExpression(item, aliases) || transformed;
+    }
+  }
+
+  return transformed;
+}
+
+function walkFrom(from: Select['from'], aliases: Set<string>): boolean {
+  if (!from || !Array.isArray(from)) return false;
+
+  let transformed = false;
+
+  for (const f of from) {
+    if ('join' in f) {
+      const join = f as Join;
+      transformed =
+        walkExpression(join.on as ExpressionValue, aliases) || transformed;
+    }
+    if ('expr' in f && f.expr && 'ast' in f.expr) {
+      transformed = walkSelect(f.expr.ast as Select) || transformed;
+    }
+  }
+
+  return transformed;
+}
+
+function walkSelect(select: Select): boolean {
+  let transformed = false;
+  const aliases = getGenerateSeriesAliases(select.from);
+
+  if (select.with) {
+    for (const cte of select.with) {
+      const cteSelect = cte.stmt?.ast ?? cte.stmt;
+      if (cteSelect && cteSelect.type === 'select') {
+        transformed = walkSelect(cteSelect as Select) || transformed;
+      }
+    }
+  }
+
+  transformed = walkFrom(select.from, aliases) || transformed;
+
+  transformed = walkExpression(select.where, aliases) || transformed;
+
+  if (select.having) {
+    if (Array.isArray(select.having)) {
+      for (const h of select.having) {
+        transformed =
+          walkExpression(h as ExpressionValue, aliases) || transformed;
+      }
+    } else {
+      transformed =
+        walkExpression(select.having as ExpressionValue, aliases) ||
+        transformed;
+    }
+  }
+
+  if (Array.isArray(select.columns)) {
+    for (const col of select.columns as Column[]) {
+      if ('expr' in col) {
+        transformed =
+          walkExpression(col.expr as ExpressionValue, aliases) || transformed;
+      }
+    }
+  }
+
+  if (Array.isArray(select.groupby)) {
+    for (const g of select.groupby as ExpressionValue[]) {
+      transformed = walkExpression(g, aliases) || transformed;
+    }
+  }
+
+  if (Array.isArray(select.orderby)) {
+    for (const order of select.orderby as OrderBy[]) {
+      if (order.expr) {
+        transformed =
+          walkExpression(order.expr as ExpressionValue, aliases) || transformed;
+      }
+    }
+  }
+
+  if (select._orderby) {
+    for (const order of select._orderby as OrderBy[]) {
+      if (order.expr) {
+        transformed =
+          walkExpression(order.expr as ExpressionValue, aliases) || transformed;
+      }
+    }
+  }
+
+  if (select._next) {
+    transformed = walkSelect(select._next) || transformed;
+  }
+
+  return transformed;
+}
+
+export function rewriteGenerateSeriesAliases(ast: AST | AST[]): boolean {
+  const statements = Array.isArray(ast) ? ast : [ast];
+  let transformed = false;
+
+  for (const stmt of statements) {
+    if (stmt.type === 'select') {
+      transformed = walkSelect(stmt as Select) || transformed;
+    }
+  }
+
+  return transformed;
+}

package/src/sql/visitors/union-with-hoister.ts

@@ -0,0 +1,106 @@
+/**
+ * AST visitor to hoist WITH clauses out of UNION and other set operations.
+ *
+ * Drizzle can emit SQL like:
+ *   (with a as (...) select ...) union (with b as (...) select ...)
+ *
+ * DuckDB 1.4.x has an internal binder bug for this pattern.
+ * We merge per arm CTEs into a single top level WITH when names do not collide.
+ */
+
+import type { AST, Select, From } from 'node-sql-parser';
+
+function getCteName(cte: { name?: unknown }): string | null {
+  const nameObj = cte.name as Record<string, unknown> | undefined;
+  if (!nameObj) return null;
+  const value = nameObj.value;
+  if (typeof value === 'string') return value;
+  return null;
+}
+
+function hoistWithInSelect(select: Select): boolean {
+  if (!select.set_op || !select._next) return false;
+
+  const arms: Select[] = [];
+  let current: Select | null = select;
+  while (current && current.type === 'select') {
+    arms.push(current);
+    current = current._next as Select | null;
+  }
+
+  const mergedWith: NonNullable<Select['with']> = [];
+  const seen = new Set<string>();
+  let hasWithBeyondFirst = false;
+
+  for (const arm of arms) {
+    if (arm.with && arm.with.length > 0) {
+      if (arm !== arms[0]) {
+        hasWithBeyondFirst = true;
+      }
+      for (const cte of arm.with) {
+        const cteName = getCteName(cte);
+        if (!cteName) return false;
+        if (seen.has(cteName)) {
+          return false;
+        }
+        seen.add(cteName);
+        mergedWith.push(cte);
+      }
+    }
+  }
+
+  if (!hasWithBeyondFirst) return false;
+
+  arms[0].with = mergedWith;
+  if ('parentheses_symbol' in arms[0]) {
+    (arms[0] as Select & { parentheses_symbol?: boolean }).parentheses_symbol =
+      false;
+  }
+  for (let i = 1; i < arms.length; i++) {
+    arms[i].with = null;
+  }
+
+  return true;
+}
+
+function walkSelect(select: Select): boolean {
+  let transformed = false;
+
+  if (select.with) {
+    for (const cte of select.with) {
+      const cteSelect = cte.stmt?.ast ?? cte.stmt;
+      if (cteSelect && cteSelect.type === 'select') {
+        transformed = walkSelect(cteSelect as Select) || transformed;
+      }
+    }
+  }
+
+  if (Array.isArray(select.from)) {
+    for (const from of select.from as From[]) {
+      if ('expr' in from && from.expr && 'ast' in from.expr) {
+        transformed = walkSelect(from.expr.ast as Select) || transformed;
+      }
+    }
+  }
+
+  transformed = hoistWithInSelect(select) || transformed;
+
+  if (select._next) {
+    transformed = walkSelect(select._next) || transformed;
+  }
+
+  return transformed;
+}
+
+export function hoistUnionWith(ast: AST | AST[]): boolean {
+  const statements = Array.isArray(ast) ? ast : [ast];
+  let transformed = false;
+
+  for (const stmt of statements) {
+    if (stmt.type === 'select') {
+      transformed = walkSelect(stmt as Select) || transformed;
+    }
+  }
+
+  return transformed;
+}

package/src/utils.ts

@@ -1,222 +1,2 @@
-/* Marked as internal in the original source, so we need to copy it here */
-import {
-  type SelectedFieldsOrdered,
-  type AnyColumn,
-  type DriverValueDecoder,
-  is,
-  Column,
-  SQL,
-  getTableName,
-  sql,
-} from 'drizzle-orm';
-import { PgCustomColumn, type SelectedFields } from 'drizzle-orm/pg-core';
-
-// Need to get around "decoder" property being marked as internal
-type SQLInternal<T = unknown> = SQL<T> & {
-  decoder: DriverValueDecoder<T, any>;
-};
-
-export function mapResultRow<TResult>(
-  columns: SelectedFieldsOrdered<AnyColumn>,
-  row: unknown[],
-  joinsNotNullableMap: Record<string, boolean> | undefined
-): TResult {
-  // Key -> nested object key, value -> table name if all fields in the nested object are from the same table, false otherwise
-  const nullifyMap: Record<string, string | false> = {};
-
-  const result = columns.reduce<Record<string, any>>(
-    (result, { path, field }, columnIndex) => {
-      let decoder: DriverValueDecoder<unknown, unknown>;
-      if (is(field, Column)) {
-        decoder = field;
-      } else if (is(field, SQL)) {
-        decoder = (field as SQLInternal).decoder;
-      } else {
-        const col = field.sql.queryChunks.find((chunk) => is(chunk, Column));
-
-        if (is(col, PgCustomColumn)) {
-          decoder = col;
-        } else {
-          decoder = (field.sql as SQLInternal).decoder;
-        }
-      }
-      let node = result;
-      for (const [pathChunkIndex, pathChunk] of path.entries()) {
-        if (pathChunkIndex < path.length - 1) {
-          if (!(pathChunk in node)) {
-            node[pathChunk] = {};
-          }
-          node = node[pathChunk];
-          continue;
-        }
-
-        const rawValue = row[columnIndex]!;
-
-        const value = (node[pathChunk] =
-          rawValue === null ? null : decoder.mapFromDriverValue(rawValue));
-
-        if (joinsNotNullableMap && is(field, Column) && path.length === 2) {
-          const objectName = path[0]!;
-          if (!(objectName in nullifyMap)) {
-            nullifyMap[objectName] =
-              value === null ? getTableName(field.table) : false;
-          } else if (
-            typeof nullifyMap[objectName] === 'string' &&
-            nullifyMap[objectName] !== getTableName(field.table)
-          ) {
-            nullifyMap[objectName] = false;
-          }
-          continue;
-        }
-
-        // may need to add a condition for non-Aliased SQL
-        if (
-          joinsNotNullableMap &&
-          is(field, SQL.Aliased) &&
-          path.length === 2
-        ) {
-          const col = field.sql.queryChunks.find((chunk) => is(chunk, Column));
-          const tableName = col?.table && getTableName(col?.table);
-
-          if (!tableName) {
-            continue;
-          }
-
-          const objectName = path[0]!;
-
-          if (!(objectName in nullifyMap)) {
-            nullifyMap[objectName] = value === null ? tableName : false;
-            continue;
-          }
-
-          if (nullifyMap[objectName] && nullifyMap[objectName] !== tableName) {
-            nullifyMap[objectName] = false;
-          }
-          continue;
-        }
-      }
-      return result;
-    },
-    {}
-  );
-
-  // Nullify all nested objects from nullifyMap that are nullable
-  if (joinsNotNullableMap && Object.keys(nullifyMap).length > 0) {
-    for (const [objectName, tableName] of Object.entries(nullifyMap)) {
-      if (typeof tableName === 'string' && !joinsNotNullableMap[tableName]) {
-        result[objectName] = null;
-      }
-    }
-  }
-
-  return result as TResult;
-}
-
-export function aliasFields(
-  fields: SelectedFields,
-  fullJoin = false
-): SelectedFields {
-  return Object.fromEntries(
-    Object.entries(fields)
-      .filter(([key]) => key !== 'enableRLS')
-      .map(([key, value]) => {
-        if (fullJoin && is(value, Column)) {
-          return [
-            key,
-            sql`${value}`.as(`${getTableName(value.table)}.${value.name}`),
-          ];
-        }
-
-        if (fullJoin && is(value, SQL)) {
-          const col = value
-            .getSQL()
-            .queryChunks.find((chunk) => is(chunk, Column));
-
-          const tableName = col?.table && getTableName(col?.table);
-
-          return [key, value.as(tableName ? `${tableName}.${key}` : key)];
-        }
-
-        if (is(value, SQL) || is(value, Column)) {
-          return [key, (is(value, SQL) ? value : sql`${value}`).as(key)];
-        }
-
-        if (is(value, SQL.Aliased)) {
-          return [key, value];
-        }
-
-        // todo: should probably make this recursive?
-        if (typeof value === 'object') {
-          const parentKey = key;
-
-          return [
-            key,
-            Object.fromEntries(
-              Object.entries(value)
-                .filter(([childKey]) => childKey !== 'enableRLS')
-                .map(([childKey, childValue]) => [
-                  childKey,
-                  (is(childValue, SQL) ? childValue : sql`${childValue}`).as(
-                    `${parentKey}.${childKey}`
-                  ),
-                ])
-            ),
-          ];
-        }
-
-        return [key, value];
-      })
-  );
-}
-
-// DuckDB names returned variables differently than Postgres
-// so we need to remap them to match the Postgres names
-
-const selectionRegex = /select\s+(.+)\s+from/i;
-// const tableIdPropSelectionRegex = /("(.+)"\."(.+)")(\s+as\s+'?(.+?)'?\.'?(.+?)'?)?/i;
-const tableIdPropSelectionRegex = new RegExp(
-  [
-    `("(.+)"\\."(.+)")`, // table identifier + property
-    `(\\s+as\\s+'?(.+?)'?\\.'?(.+?)'?)?`, // optional AS clause
-  ].join(''),
-  'i'
-);
-const noTableIdPropSelectionRegex = /"(.+)"(\s+as\s+'?\1'?)?/i;
-
-const tablePropRegex = /"(.+)"\."(.+)"/i;
-const asClauseRegex = /as\s+(.+)$/i;
-const aliasRegex = /as\s+'?(.+)'?\.'?(.+)'?$/i;
-
-/* Takes an SQL query as a string, and adds or updates "AS" clauses
- * to the form: `AS 'table_name.column_name'`
- * instead of : `AS "table_name"."column_name"`
- */
-export function queryAdapter(query: string): string {
-  // Things to consider:
-  // - need to handle nested selects
-  // - what about full joins?
-  const selection = selectionRegex.exec(query);
-
-  if (selection?.length !== 2) {
-    return query;
-  }
-
-  const fields = selection[1].split(',').map((field) => {
-    const trimmedField = field.trim();
-
-    // - different scenarios:
-    //    - no table identifier + no AS clause -> ignore
-    //    - no table identifier + AS clause -> ensure AS clause format
-    //    - table identifier + no AS clause -> add AS clause
-    //    - table identifier + AS clause -> ensure AS clause format
-    const propSelection = tableIdPropSelectionRegex
-      .exec(trimmedField)
-      ?.filter(Boolean);
-
-    if (!propSelection) {
-      return trimmedField;
-    }
-  });
-
-  return query.replace(selection[1], fields.join(', '));
-}
+export { aliasFields } from './sql/selection.ts';
+export { mapResultRow } from './sql/result-mapper.ts';