@leonardovida-md/drizzle-neo-duckdb 1.1.1 → 1.1.3

package/src/driver.ts

@@ -23,7 +23,11 @@ import { DuckDBSession } from './session.ts';
 import { DuckDBDialect } from './dialect.ts';
 import { DuckDBSelectBuilder } from './select-builder.ts';
 import { aliasFields } from './sql/selection.ts';
-import type { ExecuteInBatchesOptions, RowData } from './client.ts';
+import type {
+  ExecuteBatchesRawChunk,
+  ExecuteInBatchesOptions,
+  RowData,
+} from './client.ts';
 import { closeClientConnection, isPool } from './client.ts';
 import {
   createDuckDBConnectionPool,
@@ -31,11 +35,20 @@ import {
   type DuckDBPoolConfig,
   type PoolPreset,
 } from './pool.ts';
+import {
+  resolvePrepareCacheOption,
+  resolveRewriteArraysOption,
+  type PreparedStatementCacheConfig,
+  type PrepareCacheOption,
+  type RewriteArraysMode,
+  type RewriteArraysOption,
+} from './options.ts';
 
 export interface PgDriverOptions {
   logger?: Logger;
-  rewriteArrays?: boolean;
+  rewriteArrays?: RewriteArraysMode;
   rejectStringArrayLiterals?: boolean;
+  prepareCache?: PreparedStatementCacheConfig;
 }
 
 export class DuckDBDriver {
@@ -52,8 +65,9 @@ export class DuckDBDriver {
   ): DuckDBSession<Record<string, unknown>, TablesRelationalConfig> {
     return new DuckDBSession(this.client, this.dialect, schema, {
       logger: this.options.logger,
-      rewriteArrays: this.options.rewriteArrays,
+      rewriteArrays: this.options.rewriteArrays ?? 'auto',
       rejectStringArrayLiterals: this.options.rejectStringArrayLiterals,
+      prepareCache: this.options.prepareCache,
     });
   }
 }
@@ -69,8 +83,9 @@ export interface DuckDBConnectionConfig {
 export interface DuckDBDrizzleConfig<
   TSchema extends Record<string, unknown> = Record<string, never>,
 > extends DrizzleConfig<TSchema> {
-  rewriteArrays?: boolean;
+  rewriteArrays?: RewriteArraysOption;
   rejectStringArrayLiterals?: boolean;
+  prepareCache?: PrepareCacheOption;
   /** Pool configuration. Use preset name, size config, or false to disable. */
   pool?: DuckDBPoolConfig | PoolPreset | false;
 }
@@ -111,6 +126,8 @@ function createFromClient<
   instance?: DuckDBInstance
 ): DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>> {
   const dialect = new DuckDBDialect();
+  const rewriteArraysMode = resolveRewriteArraysOption(config.rewriteArrays);
+  const prepareCache = resolvePrepareCacheOption(config.prepareCache);
 
   const logger =
     config.logger === true ? new DefaultLogger() : config.logger || undefined;
@@ -131,8 +148,9 @@ function createFromClient<
 
   const driver = new DuckDBDriver(client, dialect, {
     logger,
-    rewriteArrays: config.rewriteArrays,
+    rewriteArrays: rewriteArraysMode,
     rejectStringArrayLiterals: config.rejectStringArrayLiterals,
+    prepareCache,
   });
   const session = driver.createSession(schema);
 
@@ -330,6 +348,13 @@ export class DuckDBDatabase<
     return this.session.executeBatches<T>(query, options);
   }
 
+  executeBatchesRaw(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<ExecuteBatchesRawChunk, void, void> {
+    return this.session.executeBatchesRaw(query, options);
+  }
+
   executeArrow(query: SQL): Promise<unknown> {
     return this.session.executeArrow(query);
   }

package/src/index.ts

@@ -7,3 +7,4 @@ export * from './client.ts';
 export * from './pool.ts';
 export * from './olap.ts';
 export * from './value-wrappers.ts';
+export * from './options.ts';

package/src/options.ts

@@ -0,0 +1,40 @@
+export type RewriteArraysMode = 'auto' | 'always' | 'never';
+
+export type RewriteArraysOption = boolean | RewriteArraysMode;
+
+const DEFAULT_REWRITE_ARRAYS_MODE: RewriteArraysMode = 'auto';
+
+export function resolveRewriteArraysOption(
+  value?: RewriteArraysOption
+): RewriteArraysMode {
+  if (value === undefined) return DEFAULT_REWRITE_ARRAYS_MODE;
+  if (value === true) return 'auto';
+  if (value === false) return 'never';
+  return value;
+}
+
+export type PrepareCacheOption = boolean | number | { size?: number };
+
+export interface PreparedStatementCacheConfig {
+  size: number;
+}
+
+const DEFAULT_PREPARED_CACHE_SIZE = 32;
+
+export function resolvePrepareCacheOption(
+  option?: PrepareCacheOption
+): PreparedStatementCacheConfig | undefined {
+  if (!option) return undefined;
+
+  if (option === true) {
+    return { size: DEFAULT_PREPARED_CACHE_SIZE };
+  }
+
+  if (typeof option === 'number') {
+    const size = Math.max(1, Math.floor(option));
+    return { size };
+  }
+
+  const size = option.size ?? DEFAULT_PREPARED_CACHE_SIZE;
+  return { size: Math.max(1, Math.floor(size)) };
+}

package/src/session.ts

@@ -15,7 +15,10 @@ 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 { adaptArrayOperators } from './sql/query-rewriters.ts';
+import {
+  adaptArrayOperators,
+  qualifyJoinColumns,
+} from './sql/query-rewriters.ts';
 import { mapResultRow } from './sql/result-mapper.ts';
 import { TransactionRollbackError } from 'drizzle-orm/errors';
 import type { DuckDBDialect } from './dialect.ts';
@@ -26,13 +29,20 @@ import type {
 } 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,
+  RewriteArraysMode,
+} from './options.ts';
 
 export type { DuckDBClientLike, RowData } from './client.ts';
 
@@ -46,6 +56,34 @@ function isSavepointSyntaxError(error: unknown): boolean {
   );
 }
 
+function rewriteQuery(
+  mode: RewriteArraysMode,
+  query: string
+): { sql: string; rewritten: boolean } {
+  if (mode === 'never') {
+    return { sql: query, rewritten: false };
+  }
+
+  let result = query;
+  let wasRewritten = false;
+
+  // Rewrite Postgres array operators to DuckDB functions
+  const arrayRewritten = adaptArrayOperators(result);
+  if (arrayRewritten !== result) {
+    result = arrayRewritten;
+    wasRewritten = true;
+  }
+
+  // Qualify unqualified column references in JOIN ON clauses
+  const joinQualified = qualifyJoinColumns(result);
+  if (joinQualified !== result) {
+    result = joinQualified;
+    wasRewritten = true;
+  }
+
+  return { sql: result, rewritten: wasRewritten };
+}
+
 export class DuckDBPreparedQuery<
   T extends PreparedQueryConfig,
 > extends PgPreparedQuery<T> {
@@ -62,8 +100,9 @@ export class DuckDBPreparedQuery<
     private customResultMapper:
       | ((rows: unknown[][]) => T['execute'])
       | undefined,
-    private rewriteArrays: boolean,
+    private rewriteArraysMode: RewriteArraysMode,
     private rejectStringArrayLiterals: boolean,
+    private prepareCache: PreparedStatementCacheConfig | undefined,
     private warnOnStringArrayLiteral?: (sql: string) => void
   ) {
     super({ sql: queryString, params });
@@ -82,11 +121,12 @@ export class DuckDBPreparedQuery<
           : undefined,
       }
     );
-    const rewrittenQuery = this.rewriteArrays
-      ? adaptArrayOperators(this.queryString)
-      : this.queryString;
+    const { sql: rewrittenQuery, rewritten: didRewrite } = rewriteQuery(
+      this.rewriteArraysMode,
+      this.queryString
+    );
 
-    if (this.rewriteArrays && rewrittenQuery !== this.queryString) {
+    if (didRewrite) {
       this.logger.logQuery(
         `[duckdb] original query before array rewrite: ${this.queryString}`,
         params
@@ -98,19 +138,30 @@ export class DuckDBPreparedQuery<
     const { fields, joinsNotNullableMap, customResultMapper } =
       this as typeof this & { joinsNotNullableMap?: Record<string, boolean> };
 
-    const rows = await executeOnClient(this.client, rewrittenQuery, params);
+    if (fields) {
+      const { rows } = await executeArraysOnClient(
+        this.client,
+        rewrittenQuery,
+        params,
+        { prepareCache: this.prepareCache }
+      );
+
+      if (rows.length === 0) {
+        return [] as T['execute'];
+      }
 
-    if (rows.length === 0 || !fields) {
-      return rows 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, rewrittenQuery, params, {
+      prepareCache: this.prepareCache,
+    });
 
-    return customResultMapper
-      ? customResultMapper(rowValues)
-      : rowValues.map((row) =>
-          mapResultRow<T['execute']>(fields, row, joinsNotNullableMap)
-        );
+    return rows as T['execute'];
   }
 
   all(
@@ -126,8 +177,9 @@ export class DuckDBPreparedQuery<
 
 export interface DuckDBSessionOptions {
   logger?: Logger;
-  rewriteArrays?: boolean;
+  rewriteArrays?: RewriteArraysMode;
   rejectStringArrayLiterals?: boolean;
+  prepareCache?: PreparedStatementCacheConfig;
 }
 
 export class DuckDBSession<
@@ -138,8 +190,9 @@ export class DuckDBSession<
 
   protected override dialect: DuckDBDialect;
   private logger: Logger;
-  private rewriteArrays: boolean;
+  private rewriteArraysMode: RewriteArraysMode;
   private rejectStringArrayLiterals: boolean;
+  private prepareCache: PreparedStatementCacheConfig | undefined;
   private hasWarnedArrayLiteral = false;
   private rollbackOnly = false;
 
@@ -152,8 +205,14 @@ export class DuckDBSession<
     super(dialect);
     this.dialect = dialect;
     this.logger = options.logger ?? new NoopLogger();
-    this.rewriteArrays = options.rewriteArrays ?? true;
+    this.rewriteArraysMode = options.rewriteArrays ?? 'auto';
     this.rejectStringArrayLiterals = options.rejectStringArrayLiterals ?? false;
+    this.prepareCache = options.prepareCache;
+    this.options = {
+      ...options,
+      rewriteArrays: this.rewriteArraysMode,
+      prepareCache: this.prepareCache,
+    };
   }
 
   prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
@@ -173,8 +232,9 @@ export class DuckDBSession<
       fields,
       isResponseInArrayMode,
       customResultMapper,
-      this.rewriteArrays,
+      this.rewriteArraysMode,
       this.rejectStringArrayLiterals,
+      this.prepareCache,
       this.rejectStringArrayLiterals ? undefined : this.warnOnStringArrayLiteral
     );
   }
@@ -266,11 +326,12 @@ export class DuckDBSession<
         ? undefined
         : () => this.warnOnStringArrayLiteral(builtQuery.sql),
     });
-    const rewrittenQuery = this.rewriteArrays
-      ? adaptArrayOperators(builtQuery.sql)
-      : builtQuery.sql;
+    const { sql: rewrittenQuery, rewritten: didRewrite } = rewriteQuery(
+      this.rewriteArraysMode,
+      builtQuery.sql
+    );
 
-    if (this.rewriteArrays && rewrittenQuery !== builtQuery.sql) {
+    if (didRewrite) {
       this.logger.logQuery(
         `[duckdb] original query before array rewrite: ${builtQuery.sql}`,
         params
@@ -287,6 +348,36 @@ export class DuckDBSession<
     ) 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),
+    });
+    const { sql: rewrittenQuery, rewritten: didRewrite } = rewriteQuery(
+      this.rewriteArraysMode,
+      builtQuery.sql
+    );
+
+    if (didRewrite) {
+      this.logger.logQuery(
+        `[duckdb] original query before array rewrite: ${builtQuery.sql}`,
+        params
+      );
+    }
+
+    this.logger.logQuery(rewrittenQuery, params);
+
+    return executeInBatchesRaw(this.client, rewrittenQuery, params, options);
+  }
+
   async executeArrow(query: SQL): Promise<unknown> {
     this.dialect.resetPgJsonFlag();
     const builtQuery = this.dialect.sqlToQuery(query);
@@ -297,11 +388,12 @@ export class DuckDBSession<
         ? undefined
         : () => this.warnOnStringArrayLiteral(builtQuery.sql),
     });
-    const rewrittenQuery = this.rewriteArrays
-      ? adaptArrayOperators(builtQuery.sql)
-      : builtQuery.sql;
+    const { sql: rewrittenQuery, rewritten: didRewrite } = rewriteQuery(
+      this.rewriteArraysMode,
+      builtQuery.sql
+    );
 
-    if (this.rewriteArrays && rewrittenQuery !== builtQuery.sql) {
+    if (didRewrite) {
       this.logger.logQuery(
         `[duckdb] original query before array rewrite: ${builtQuery.sql}`,
         params
@@ -377,6 +469,15 @@ export class DuckDBTransaction<
     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>;