@duckdbfan/drizzle-duckdb 0.0.6 → 1.3.1

package/src/session.ts

@@ -1,7 +1,6 @@
-import type { Connection, Database, RowData } from 'duckdb-async';
 import { entityKind } from 'drizzle-orm/entity';
-import { type Logger, NoopLogger } from 'drizzle-orm/logger';
-import type { PgDialect } from 'drizzle-orm/pg-core/dialect';
+import type { Logger } from 'drizzle-orm/logger';
+import { NoopLogger } from 'drizzle-orm/logger';
 import { PgTransaction } from 'drizzle-orm/pg-core';
 import type { SelectedFieldsOrdered } from 'drizzle-orm/pg-core/query-builders/select.types';
 import type {
@@ -16,28 +15,59 @@ import type {
 } from 'drizzle-orm/relations';
 import { fillPlaceholders, type Query, SQL, sql } from 'drizzle-orm/sql/sql';
 import type { Assume } from 'drizzle-orm/utils';
-import { mapResultRow } from './utils';
-import type { DuckDBDialect } from './dialect';
+import { mapResultRow } from './sql/result-mapper.ts';
 import { TransactionRollbackError } from 'drizzle-orm/errors';
-
-export type DuckDBClient = Database;
+import type { DuckDBDialect } from './dialect.ts';
+import type {
+  DuckDBClientLike,
+  DuckDBConnectionPool,
+  RowData,
+} from './client.ts';
+import {
+  executeArrowOnClient,
+  executeArraysOnClient,
+  executeInBatches,
+  executeInBatchesRaw,
+  executeOnClient,
+  prepareParams,
+  type ExecuteBatchesRawChunk,
+  type ExecuteInBatchesOptions,
+} from './client.ts';
+import { isPool } from './client.ts';
+import type { DuckDBConnection } from '@duckdb/node-api';
+import type { PreparedStatementCacheConfig } from './options.ts';
+
+export type { DuckDBClientLike, RowData } from './client.ts';
+
+function isSavepointSyntaxError(error: unknown): boolean {
+  if (!(error instanceof Error) || !error.message) {
+    return false;
+  }
+  return (
+    error.message.toLowerCase().includes('savepoint') &&
+    error.message.toLowerCase().includes('syntax error')
+  );
+}
 
 export class DuckDBPreparedQuery<
-  T extends PreparedQueryConfig
+  T extends PreparedQueryConfig,
 > extends PgPreparedQuery<T> {
   static readonly [entityKind]: string = 'DuckDBPreparedQuery';
 
-  // private rawQueryConfig: QueryOptions;
-  // private queryConfig: QueryOptions;
-
   constructor(
-    private client: DuckDBClient | Connection,
+    private client: DuckDBClientLike,
+    private dialect: DuckDBDialect,
     private queryString: string,
     private params: unknown[],
     private logger: Logger,
     private fields: SelectedFieldsOrdered | undefined,
     private _isResponseInArrayMode: boolean,
-    private customResultMapper?: (rows: unknown[][]) => T['execute']
+    private customResultMapper:
+      | ((rows: unknown[][]) => T['execute'])
+      | undefined,
+    private rejectStringArrayLiterals: boolean,
+    private prepareCache: PreparedStatementCacheConfig | undefined,
+    private warnOnStringArrayLiteral?: (sql: string) => void
   ) {
     super({ sql: queryString, params });
   }
@@ -45,31 +75,45 @@ export class DuckDBPreparedQuery<
   async execute(
     placeholderValues: Record<string, unknown> | undefined = {}
   ): Promise<T['execute']> {
-    const params = fillPlaceholders(this.params, placeholderValues);
-
+    this.dialect.assertNoPgJsonColumns();
+    const params = prepareParams(
+      fillPlaceholders(this.params, placeholderValues),
+      {
+        rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+        warnOnStringArrayLiteral: this.warnOnStringArrayLiteral
+          ? () => this.warnOnStringArrayLiteral?.(this.queryString)
+          : undefined,
+      }
+    );
     this.logger.logQuery(this.queryString, params);
 
-    const {
-      fields,
-      client,
-      joinsNotNullableMap,
-      customResultMapper,
-      queryString,
-    } = this as typeof this & { joinsNotNullableMap?: Record<string, boolean> };
-
-    const rows = (await client.all(queryString, ...params)) ?? [];
-
-    if (rows.length === 0 || !fields) {
-      return rows;
+    const { fields, joinsNotNullableMap, customResultMapper } =
+      this as typeof this & { joinsNotNullableMap?: Record<string, boolean> };
+
+    if (fields) {
+      const { rows } = await executeArraysOnClient(
+        this.client,
+        this.queryString,
+        params,
+        { prepareCache: this.prepareCache }
+      );
+
+      if (rows.length === 0) {
+        return [] as T['execute'];
+      }
+
+      return customResultMapper
+        ? customResultMapper(rows)
+        : rows.map((row) =>
+            mapResultRow<T['execute']>(fields, row, joinsNotNullableMap)
+          );
     }
 
-    const rowValues = rows.map((row) => Object.values(row));
+    const rows = await executeOnClient(this.client, this.queryString, params, {
+      prepareCache: this.prepareCache,
+    });
 
-    return customResultMapper
-      ? customResultMapper(rowValues)
-      : rowValues.map((row) =>
-          mapResultRow<T['execute']>(fields!, row, joinsNotNullableMap)
-        );
+    return rows as T['execute'];
   }
 
   all(
@@ -79,30 +123,44 @@ export class DuckDBPreparedQuery<
   }
 
   isResponseInArrayMode(): boolean {
-    return false;
+    return this._isResponseInArrayMode;
   }
 }
 
 export interface DuckDBSessionOptions {
   logger?: Logger;
+  rejectStringArrayLiterals?: boolean;
+  prepareCache?: PreparedStatementCacheConfig;
 }
 
 export class DuckDBSession<
   TFullSchema extends Record<string, unknown> = Record<string, never>,
-  TSchema extends TablesRelationalConfig = Record<string, never>
+  TSchema extends TablesRelationalConfig = Record<string, never>,
 > extends PgSession<DuckDBQueryResultHKT, TFullSchema, TSchema> {
   static readonly [entityKind]: string = 'DuckDBSession';
 
+  protected override dialect: DuckDBDialect;
   private logger: Logger;
+  private rejectStringArrayLiterals: boolean;
+  private prepareCache: PreparedStatementCacheConfig | undefined;
+  private hasWarnedArrayLiteral = false;
+  private rollbackOnly = false;
 
   constructor(
-    private client: DuckDBClient | Connection,
+    private client: DuckDBClientLike,
     dialect: DuckDBDialect,
     private schema: RelationalSchemaConfig<TSchema> | undefined,
     private options: DuckDBSessionOptions = {}
   ) {
     super(dialect);
+    this.dialect = dialect;
     this.logger = options.logger ?? new NoopLogger();
+    this.rejectStringArrayLiterals = options.rejectStringArrayLiterals ?? false;
+    this.prepareCache = options.prepareCache;
+    this.options = {
+      ...options,
+      prepareCache: this.prepareCache,
+    };
   }
 
   prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
@@ -112,25 +170,48 @@ export class DuckDBSession<
     isResponseInArrayMode: boolean,
     customResultMapper?: (rows: unknown[][]) => T['execute']
   ): PgPreparedQuery<T> {
+    void name; // DuckDB doesn't support prepared statement names but the signature must match.
     return new DuckDBPreparedQuery(
       this.client,
+      this.dialect,
       query.sql,
       query.params,
       this.logger,
       fields,
       isResponseInArrayMode,
-      customResultMapper
+      customResultMapper,
+      this.rejectStringArrayLiterals,
+      this.prepareCache,
+      this.rejectStringArrayLiterals ? undefined : this.warnOnStringArrayLiteral
     );
   }
 
+  override execute<T>(query: SQL): Promise<T> {
+    this.dialect.resetPgJsonFlag();
+    return super.execute(query);
+  }
+
+  override all<T = unknown>(query: SQL): Promise<T[]> {
+    this.dialect.resetPgJsonFlag();
+    return super.all(query);
+  }
+
   override async transaction<T>(
-    transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
+    transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>,
+    config?: PgTransactionConfig
   ): Promise<T> {
-    const connection =
-      'connect' in this.client ? await this.client.connect() : this.client;
+    let pinnedConnection: DuckDBConnection | undefined;
+    let pool: DuckDBConnectionPool | undefined;
+
+    let clientForTx: DuckDBClientLike = this.client;
+    if (isPool(this.client)) {
+      pool = this.client;
+      pinnedConnection = await pool.acquire();
+      clientForTx = pinnedConnection;
+    }
 
     const session = new DuckDBSession(
-      connection,
+      clientForTx,
       this.dialect,
       this.schema,
       this.options
@@ -142,24 +223,114 @@ export class DuckDBSession<
       this.schema
     );
 
-    await tx.execute(sql`BEGIN TRANSACTION;`);
-
     try {
-      const result = await transaction(tx);
-      await tx.execute(sql`commit`);
-      return result;
-    } catch (error) {
-      await tx.execute(sql`rollback`);
-      throw error;
+      await tx.execute(sql`BEGIN TRANSACTION;`);
+
+      if (config) {
+        await tx.setTransaction(config);
+      }
+
+      try {
+        const result = await transaction(tx);
+        if (session.isRollbackOnly()) {
+          await tx.execute(sql`rollback`);
+          throw new TransactionRollbackError();
+        }
+        await tx.execute(sql`commit`);
+        return result;
+      } catch (error) {
+        await tx.execute(sql`rollback`);
+        throw error;
+      }
     } finally {
-      await connection.close();
+      if (pinnedConnection && pool) {
+        await pool.release(pinnedConnection);
+      }
     }
   }
+
+  private warnOnStringArrayLiteral = (query: string) => {
+    if (this.hasWarnedArrayLiteral) {
+      return;
+    }
+    this.hasWarnedArrayLiteral = true;
+    this.logger.logQuery(
+      `[duckdb] ${arrayLiteralWarning}\nquery: ${query}`,
+      []
+    );
+  };
+
+  executeBatches<T extends RowData = RowData>(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<GenericRowData<T>[], void, void> {
+    this.dialect.resetPgJsonFlag();
+    const builtQuery = this.dialect.sqlToQuery(query);
+    this.dialect.assertNoPgJsonColumns();
+    const params = prepareParams(builtQuery.params, {
+      rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+      warnOnStringArrayLiteral: this.rejectStringArrayLiterals
+        ? undefined
+        : () => this.warnOnStringArrayLiteral(builtQuery.sql),
+    });
+
+    this.logger.logQuery(builtQuery.sql, params);
+
+    return executeInBatches(
+      this.client,
+      builtQuery.sql,
+      params,
+      options
+    ) as AsyncGenerator<GenericRowData<T>[], void, void>;
+  }
+
+  executeBatchesRaw(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<ExecuteBatchesRawChunk, void, void> {
+    this.dialect.resetPgJsonFlag();
+    const builtQuery = this.dialect.sqlToQuery(query);
+    this.dialect.assertNoPgJsonColumns();
+    const params = prepareParams(builtQuery.params, {
+      rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+      warnOnStringArrayLiteral: this.rejectStringArrayLiterals
+        ? undefined
+        : () => this.warnOnStringArrayLiteral(builtQuery.sql),
+    });
+
+    this.logger.logQuery(builtQuery.sql, params);
+
+    return executeInBatchesRaw(this.client, builtQuery.sql, params, options);
+  }
+
+  async executeArrow(query: SQL): Promise<unknown> {
+    this.dialect.resetPgJsonFlag();
+    const builtQuery = this.dialect.sqlToQuery(query);
+    this.dialect.assertNoPgJsonColumns();
+    const params = prepareParams(builtQuery.params, {
+      rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+      warnOnStringArrayLiteral: this.rejectStringArrayLiterals
+        ? undefined
+        : () => this.warnOnStringArrayLiteral(builtQuery.sql),
+    });
+
+    this.logger.logQuery(builtQuery.sql, params);
+
+    return executeArrowOnClient(this.client, builtQuery.sql, params);
+  }
+
+  markRollbackOnly(): void {
+    this.rollbackOnly = true;
+  }
+
+  isRollbackOnly(): boolean {
+    return this.rollbackOnly;
+  }
 }
 
 type PgTransactionInternals<
   TFullSchema extends Record<string, unknown> = Record<string, never>,
-  TSchema extends TablesRelationalConfig = Record<string, never>
+  TSchema extends TablesRelationalConfig = Record<string, never>,
 > = {
   dialect: DuckDBDialect;
   session: DuckDBSession<TFullSchema, TSchema>;
@@ -167,13 +338,13 @@ type PgTransactionInternals<
 
 type DuckDBTransactionWithInternals<
   TFullSchema extends Record<string, unknown> = Record<string, never>,
-  TSchema extends TablesRelationalConfig = Record<string, never>
+  TSchema extends TablesRelationalConfig = Record<string, never>,
 > = PgTransactionInternals<TFullSchema, TSchema> &
   DuckDBTransaction<TFullSchema, TSchema>;
 
 export class DuckDBTransaction<
   TFullSchema extends Record<string, unknown>,
-  TSchema extends TablesRelationalConfig
+  TSchema extends TablesRelationalConfig,
 > extends PgTransaction<DuckDBQueryResultHKT, TFullSchema, TSchema> {
   static readonly [entityKind]: string = 'DuckDBTransaction';
 
@@ -196,42 +367,111 @@ export class DuckDBTransaction<
   }
 
   setTransaction(config: PgTransactionConfig): Promise<void> {
-    // Need to work around omitted internal types from drizzle...
+    // Cast needed: PgTransaction doesn't expose dialect/session properties in public API
     type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
     return (this as unknown as Tx).session.execute(
       sql`set transaction ${this.getTransactionConfigSQL(config)}`
     );
   }
 
+  executeBatches<T extends RowData = RowData>(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<GenericRowData<T>[], void, void> {
+    // Cast needed: PgTransaction doesn't expose session property in public API
+    type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    return (this as unknown as Tx).session.executeBatches<T>(query, options);
+  }
+
+  executeBatchesRaw(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<ExecuteBatchesRawChunk, void, void> {
+    // Cast needed: PgTransaction doesn't expose session property in public API
+    type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    return (this as unknown as Tx).session.executeBatchesRaw(query, options);
+  }
+
+  executeArrow(query: SQL): Promise<unknown> {
+    // Cast needed: PgTransaction doesn't expose session property in public API
+    type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    return (this as unknown as Tx).session.executeArrow(query);
+  }
+
   override async transaction<T>(
     transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
   ): Promise<T> {
-    // Need to work around omitted internal types from drizzle...
+    // Cast needed: PgTransaction doesn't expose dialect/session properties in public API
     type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
-
-    const savepointName = `sp${this.nestedIndex + 1}`;
-    const tx = new DuckDBTransaction<TFullSchema, TSchema>(
-      (this as unknown as Tx).dialect,
-      (this as unknown as Tx).session,
+    const internals = this as unknown as Tx;
+    const savepoint = `drizzle_savepoint_${this.nestedIndex + 1}`;
+    const savepointSql = sql.raw(`savepoint ${savepoint}`);
+    const releaseSql = sql.raw(`release savepoint ${savepoint}`);
+    const rollbackSql = sql.raw(`rollback to savepoint ${savepoint}`);
+
+    const nestedTx = new DuckDBTransaction<TFullSchema, TSchema>(
+      internals.dialect,
+      internals.session,
       this.schema,
       this.nestedIndex + 1
     );
-    await tx.execute(sql.raw(`savepoint ${savepointName}`));
+
+    // Check dialect-level savepoint support (per-instance, not global)
+    if (internals.dialect.areSavepointsUnsupported()) {
+      return this.runNestedWithoutSavepoint(transaction, nestedTx, internals);
+    }
+
+    let createdSavepoint = false;
+    try {
+      await internals.session.execute(savepointSql);
+      internals.dialect.markSavepointsSupported();
+      createdSavepoint = true;
+    } catch (error) {
+      if (!isSavepointSyntaxError(error)) {
+        throw error;
+      }
+      internals.dialect.markSavepointsUnsupported();
+      return this.runNestedWithoutSavepoint(transaction, nestedTx, internals);
+    }
+
     try {
-      const result = await transaction(tx);
-      await tx.execute(sql.raw(`release savepoint ${savepointName}`));
+      const result = await transaction(nestedTx);
+      if (createdSavepoint) {
+        await internals.session.execute(releaseSql);
+      }
       return result;
-    } catch (err) {
-      await tx.execute(sql.raw(`rollback to savepoint ${savepointName}`));
-      throw err;
+    } catch (error) {
+      if (createdSavepoint) {
+        await internals.session.execute(rollbackSql);
+      }
+      (
+        internals.session as DuckDBSession<TFullSchema, TSchema>
+      ).markRollbackOnly();
+      throw error;
     }
   }
+
+  private runNestedWithoutSavepoint<T>(
+    transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>,
+    nestedTx: DuckDBTransaction<TFullSchema, TSchema>,
+    internals: DuckDBTransactionWithInternals<TFullSchema, TSchema>
+  ): Promise<T> {
+    return transaction(nestedTx).catch((error) => {
+      (
+        internals.session as DuckDBSession<TFullSchema, TSchema>
+      ).markRollbackOnly();
+      throw error;
+    });
+  }
 }
 
 export type GenericRowData<T extends RowData = RowData> = T;
 
 export type GenericTableData<T = RowData> = T[];
 
+const arrayLiteralWarning =
+  'Received a stringified Postgres-style array literal. Use duckDbList()/duckDbArray() or pass native arrays instead. You can also set rejectStringArrayLiterals=true to throw.';
+
 export interface DuckDBQueryResultHKT extends PgQueryResultHKT {
   type: GenericTableData<Assume<this['row'], RowData>>;
 }

package/src/sql/ast-transformer.ts

@@ -0,0 +1,170 @@
+/**
+ * AST-based SQL transformer for DuckDB compatibility.
+ *
+ * Transforms:
+ * - Array operators: @>, <@, && -> array_has_all(), array_has_any()
+ * - JOIN column qualification: "col" = "col" -> "left"."col" = "right"."col"
+ *
+ * Performance optimizations:
+ * - LRU cache for transformed queries (avoids re-parsing identical queries)
+ * - Smart heuristics to skip JOIN qualification when not needed
+ * - Early exit when no transformation is required
+ */
+
+import nodeSqlParser from 'node-sql-parser';
+const { Parser } = nodeSqlParser;
+import type { AST } from 'node-sql-parser';
+
+import { transformArrayOperators } from './visitors/array-operators.ts';
+import { qualifyJoinColumns } from './visitors/column-qualifier.ts';
+import { rewriteGenerateSeriesAliases } from './visitors/generate-series-alias.ts';
+import { hoistUnionWith } from './visitors/union-with-hoister.ts';
+
+const parser = new Parser();
+
+export type TransformResult = {
+  sql: string;
+  transformed: boolean;
+};
+
+// LRU cache for transformed SQL queries
+// Key: original SQL, Value: transformed result
+const CACHE_SIZE = 500;
+const transformCache = new Map<string, TransformResult>();
+
+function getCachedOrTransform(
+  query: string,
+  transform: () => TransformResult
+): TransformResult {
+  const cached = transformCache.get(query);
+  if (cached) {
+    // Move to end for LRU behavior
+    transformCache.delete(query);
+    transformCache.set(query, cached);
+    return cached;
+  }
+
+  const result = transform();
+
+  // Add to cache with LRU eviction
+  if (transformCache.size >= CACHE_SIZE) {
+    // Delete oldest entry (first key in Map iteration order)
+    const oldestKey = transformCache.keys().next().value;
+    if (oldestKey) {
+      transformCache.delete(oldestKey);
+    }
+  }
+  transformCache.set(query, result);
+
+  return result;
+}
+
+const DEBUG_ENV = 'DRIZZLE_DUCKDB_DEBUG_AST';
+
+function hasJoin(query: string): boolean {
+  return /\bjoin\b/i.test(query);
+}
+
+function debugLog(message: string, payload?: unknown): void {
+  if (process?.env?.[DEBUG_ENV]) {
+    // eslint-disable-next-line no-console
+    console.debug('[duckdb-ast]', message, payload ?? '');
+  }
+}
+
+export function transformSQL(query: string): TransformResult {
+  const needsArrayTransform =
+    query.includes('@>') || query.includes('<@') || query.includes('&&');
+  const needsJoinTransform =
+    hasJoin(query) || /\bupdate\b/i.test(query) || /\bdelete\b/i.test(query);
+  const needsUnionTransform =
+    /\bunion\b/i.test(query) ||
+    /\bintersect\b/i.test(query) ||
+    /\bexcept\b/i.test(query);
+  const needsGenerateSeriesTransform = /\bgenerate_series\b/i.test(query);
+
+  if (
+    !needsArrayTransform &&
+    !needsJoinTransform &&
+    !needsUnionTransform &&
+    !needsGenerateSeriesTransform
+  ) {
+    return { sql: query, transformed: false };
+  }
+
+  // Use cache for repeated queries
+  return getCachedOrTransform(query, () => {
+    try {
+      const ast = parser.astify(query, { database: 'PostgreSQL' });
+
+      let transformed = false;
+
+      if (needsArrayTransform) {
+        transformed = transformArrayOperators(ast) || transformed;
+      }
+
+      if (needsJoinTransform) {
+        transformed = qualifyJoinColumns(ast) || transformed;
+      }
+
+      if (needsGenerateSeriesTransform) {
+        transformed = rewriteGenerateSeriesAliases(ast) || transformed;
+      }
+
+      if (needsUnionTransform) {
+        transformed = hoistUnionWith(ast) || transformed;
+      }
+
+      if (!transformed) {
+        debugLog('AST parsed but no transformation applied', {
+          join: needsJoinTransform,
+        });
+        return { sql: query, transformed: false };
+      }
+
+      const transformedSql = parser.sqlify(ast, { database: 'PostgreSQL' });
+
+      return { sql: transformedSql, transformed: true };
+    } catch (err) {
+      debugLog('AST transform failed; returning original SQL', {
+        error: (err as Error).message,
+      });
+      return { sql: query, transformed: false };
+    }
+  });
+}
+
+/**
+ * Clear the transformation cache. Useful for testing or memory management.
+ */
+export function clearTransformCache(): void {
+  transformCache.clear();
+}
+
+/**
+ * Get current cache statistics for monitoring.
+ */
+export function getTransformCacheStats(): { size: number; maxSize: number } {
+  return { size: transformCache.size, maxSize: CACHE_SIZE };
+}
+
+export function needsTransformation(query: string): boolean {
+  const lower = query.toLowerCase();
+  return (
+    query.includes('@>') ||
+    query.includes('<@') ||
+    query.includes('&&') ||
+    lower.includes('join') ||
+    lower.includes('union') ||
+    lower.includes('intersect') ||
+    lower.includes('except') ||
+    lower.includes('generate_series') ||
+    lower.includes('update') ||
+    lower.includes('delete')
+  );
+}
+
+export { transformArrayOperators } from './visitors/array-operators.ts';
+export { qualifyJoinColumns } from './visitors/column-qualifier.ts';
+export { rewriteGenerateSeriesAliases } from './visitors/generate-series-alias.ts';
+export { hoistUnionWith } from './visitors/union-with-hoister.ts';