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

package/src/client.ts

@@ -2,6 +2,7 @@ import {
   listValue,
   timestampValue,
   type DuckDBConnection,
+  type DuckDBPreparedStatement,
   type DuckDBValue,
 } from '@duckdb/node-api';
 import {
@@ -9,6 +10,7 @@ import {
   wrapperToNodeApiValue,
   type AnyDuckDBValueWrapper,
 } from './value-wrappers.ts';
+import type { PreparedStatementCacheConfig } from './options.ts';
 
 export type DuckDBClientLike = DuckDBConnection | DuckDBConnectionPool;
 export type RowData = Record<string, unknown>;
@@ -25,6 +27,25 @@ export function isPool(
   return typeof (client as DuckDBConnectionPool).acquire === 'function';
 }
 
+export interface ExecuteClientOptions {
+  prepareCache?: PreparedStatementCacheConfig;
+}
+
+export type ExecuteArraysResult = { columns: string[]; rows: unknown[][] };
+
+type MaterializedRows = ExecuteArraysResult;
+
+type PreparedCacheEntry = {
+  statement: DuckDBPreparedStatement;
+};
+
+type PreparedStatementCache = {
+  size: number;
+  entries: Map<string, PreparedCacheEntry>;
+};
+
+const PREPARED_CACHE = Symbol.for('drizzle-duckdb:prepared-cache');
+
 export interface PrepareParamsOptions {
   rejectStringArrayLiterals?: boolean;
   warnOnStringArrayLiteral?: () => void;
@@ -49,19 +70,30 @@ export function prepareParams(
   options: PrepareParamsOptions = {}
 ): unknown[] {
   return params.map((param) => {
-    if (typeof param === 'string') {
-      const trimmed = param.trim();
-      if (trimmed && isPgArrayLiteral(trimmed)) {
-        if (options.rejectStringArrayLiterals) {
-          throw new Error(
-            'Stringified array literals are not supported. Use duckDbList()/duckDbArray() or pass native arrays.'
-          );
-        }
-
-        if (options.warnOnStringArrayLiteral) {
-          options.warnOnStringArrayLiteral();
+    if (typeof param === 'string' && param.length > 0) {
+      const firstChar = param[0];
+      const maybeArrayLiteral =
+        firstChar === '{' ||
+        firstChar === '[' ||
+        firstChar === ' ' ||
+        firstChar === '\t';
+
+      if (maybeArrayLiteral) {
+        const trimmed =
+          firstChar === '{' || firstChar === '[' ? param : param.trim();
+
+        if (trimmed && isPgArrayLiteral(trimmed)) {
+          if (options.rejectStringArrayLiterals) {
+            throw new Error(
+              'Stringified array literals are not supported. Use duckDbList()/duckDbArray() or pass native arrays.'
+            );
+          }
+
+          if (options.warnOnStringArrayLiteral) {
+            options.warnOnStringArrayLiteral();
+          }
+          return parsePgArrayLiteral(trimmed);
         }
-        return parsePgArrayLiteral(trimmed);
       }
     }
     return param;
@@ -106,14 +138,177 @@ function toNodeApiValue(value: unknown): DuckDBValue {
 }
 
 function deduplicateColumns(columns: string[]): string[] {
-  const seen: Record<string, number> = {};
-  return columns.map((col) => {
-    const count = seen[col] ?? 0;
-    seen[col] = count + 1;
-    return count === 0 ? col : `${col}_${count}`;
+  const counts = new Map<string, number>();
+  let hasDuplicates = false;
+
+  for (const column of columns) {
+    const next = (counts.get(column) ?? 0) + 1;
+    counts.set(column, next);
+    if (next > 1) {
+      hasDuplicates = true;
+      break;
+    }
+  }
+
+  if (!hasDuplicates) {
+    return columns;
+  }
+
+  counts.clear();
+  return columns.map((column) => {
+    const count = counts.get(column) ?? 0;
+    counts.set(column, count + 1);
+    return count === 0 ? column : `${column}_${count}`;
   });
 }
 
+function destroyPreparedStatement(entry: PreparedCacheEntry | undefined): void {
+  if (!entry) return;
+  try {
+    entry.statement.destroySync();
+  } catch {
+    // Ignore cleanup errors
+  }
+}
+
+function getPreparedCache(
+  connection: DuckDBConnection,
+  size: number
+): PreparedStatementCache {
+  const store = connection as unknown as Record<
+    symbol,
+    PreparedStatementCache | undefined
+  >;
+  const existing = store[PREPARED_CACHE];
+  if (existing) {
+    existing.size = size;
+    return existing;
+  }
+
+  const cache: PreparedStatementCache = { size, entries: new Map() };
+  store[PREPARED_CACHE] = cache;
+  return cache;
+}
+
+function evictOldest(cache: PreparedStatementCache): void {
+  const oldest = cache.entries.keys().next();
+  if (!oldest.done) {
+    const key = oldest.value as string;
+    const entry = cache.entries.get(key);
+    cache.entries.delete(key);
+    destroyPreparedStatement(entry);
+  }
+}
+
+function evictCacheEntry(cache: PreparedStatementCache, key: string): void {
+  const entry = cache.entries.get(key);
+  cache.entries.delete(key);
+  destroyPreparedStatement(entry);
+}
+
+async function getOrPrepareStatement(
+  connection: DuckDBConnection,
+  query: string,
+  cacheConfig: PreparedStatementCacheConfig
+): Promise<DuckDBPreparedStatement> {
+  const cache = getPreparedCache(connection, cacheConfig.size);
+  const cached = cache.entries.get(query);
+  if (cached) {
+    cache.entries.delete(query);
+    cache.entries.set(query, cached);
+    return cached.statement;
+  }
+
+  const statement = await connection.prepare(query);
+  cache.entries.set(query, { statement });
+
+  while (cache.entries.size > cache.size) {
+    evictOldest(cache);
+  }
+
+  return statement;
+}
+
+async function materializeResultRows(result: {
+  getRowsJS: () => Promise<unknown[][] | undefined>;
+  columnNames: () => string[];
+  deduplicatedColumnNames?: () => string[];
+}): Promise<MaterializedRows> {
+  const rows = (await result.getRowsJS()) ?? [];
+  const baseColumns =
+    typeof result.deduplicatedColumnNames === 'function'
+      ? result.deduplicatedColumnNames()
+      : result.columnNames();
+  const columns =
+    typeof result.deduplicatedColumnNames === 'function'
+      ? baseColumns
+      : deduplicateColumns(baseColumns);
+
+  return { columns, rows };
+}
+
+async function materializeRows(
+  client: DuckDBClientLike,
+  query: string,
+  params: unknown[],
+  options: ExecuteClientOptions = {}
+): Promise<MaterializedRows> {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      return await materializeRows(connection, query, params, options);
+    } finally {
+      await client.release(connection);
+    }
+  }
+
+  const values =
+    params.length > 0
+      ? (params.map((param) => toNodeApiValue(param)) as DuckDBValue[])
+      : undefined;
+
+  const connection = client as DuckDBConnection;
+
+  if (options.prepareCache && typeof connection.prepare === 'function') {
+    const cache = getPreparedCache(connection, options.prepareCache.size);
+    try {
+      const statement = await getOrPrepareStatement(
+        connection,
+        query,
+        options.prepareCache
+      );
+      if (values) {
+        statement.bind(values as DuckDBValue[]);
+      } else {
+        statement.clearBindings?.();
+      }
+      const result = await statement.run();
+      cache.entries.delete(query);
+      cache.entries.set(query, { statement });
+      return await materializeResultRows(result);
+    } catch (error) {
+      evictCacheEntry(cache, query);
+      throw error;
+    }
+  }
+
+  const result = await connection.run(query, values);
+  return await materializeResultRows(result);
+}
+
+function clearPreparedCache(connection: DuckDBConnection): void {
+  const store = connection as unknown as Record<
+    symbol,
+    PreparedStatementCache | undefined
+  >;
+  const cache = store[PREPARED_CACHE];
+  if (!cache) return;
+  for (const entry of cache.entries.values()) {
+    destroyPreparedStatement(entry);
+  }
+  cache.entries.clear();
+}
+
 function mapRowsToObjects(columns: string[], rows: unknown[][]): RowData[] {
   return rows.map((vals) => {
     const obj: Record<string, unknown> = {};
@@ -127,6 +322,8 @@ function mapRowsToObjects(columns: string[], rows: unknown[][]): RowData[] {
 export async function closeClientConnection(
   connection: DuckDBConnection
 ): Promise<void> {
+  clearPreparedCache(connection);
+
   if ('close' in connection && typeof connection.close === 'function') {
     await connection.close();
     return;
@@ -148,35 +345,41 @@ export async function closeClientConnection(
 export async function executeOnClient(
   client: DuckDBClientLike,
   query: string,
-  params: unknown[]
+  params: unknown[],
+  options: ExecuteClientOptions = {}
 ): Promise<RowData[]> {
-  if (isPool(client)) {
-    const connection = await client.acquire();
-    try {
-      return await executeOnClient(connection, query, params);
-    } finally {
-      await client.release(connection);
-    }
+  const { columns, rows } = await materializeRows(
+    client,
+    query,
+    params,
+    options
+  );
+
+  if (!rows || rows.length === 0) {
+    return [];
   }
 
-  const values =
-    params.length > 0
-      ? (params.map((param) => toNodeApiValue(param)) as DuckDBValue[])
-      : undefined;
-  const result = await client.run(query, values);
-  const rows = await result.getRowsJS();
-  const columns =
-    // prefer deduplicated names when available (Node API >=1.4.2)
-    result.deduplicatedColumnNames?.() ?? result.columnNames();
-  const uniqueColumns = deduplicateColumns(columns);
+  return mapRowsToObjects(columns, rows);
+}
 
-  return rows ? mapRowsToObjects(uniqueColumns, rows) : [];
+export async function executeArraysOnClient(
+  client: DuckDBClientLike,
+  query: string,
+  params: unknown[],
+  options: ExecuteClientOptions = {}
+): Promise<ExecuteArraysResult> {
+  return await materializeRows(client, query, params, options);
 }
 
 export interface ExecuteInBatchesOptions {
   rowsPerChunk?: number;
 }
 
+export interface ExecuteBatchesRawChunk {
+  columns: string[];
+  rows: unknown[][];
+}
+
 /**
  * Stream results from DuckDB in batches to avoid fully materializing rows in JS.
  */
@@ -206,15 +409,19 @@ export async function* executeInBatches(
       : undefined;
 
   const result = await client.stream(query, values);
+  const rawColumns =
+    typeof result.deduplicatedColumnNames === 'function'
+      ? result.deduplicatedColumnNames()
+      : result.columnNames();
   const columns =
-    // prefer deduplicated names when available (Node API >=1.4.2)
-    result.deduplicatedColumnNames?.() ?? result.columnNames();
-  const uniqueColumns = deduplicateColumns(columns);
+    typeof result.deduplicatedColumnNames === 'function'
+      ? rawColumns
+      : deduplicateColumns(rawColumns);
 
   let buffer: RowData[] = [];
 
   for await (const chunk of result.yieldRowsJs()) {
-    const objects = mapRowsToObjects(uniqueColumns, chunk);
+    const objects = mapRowsToObjects(columns, chunk);
     for (const row of objects) {
       buffer.push(row);
       if (buffer.length >= rowsPerChunk) {
@@ -229,6 +436,59 @@ export async function* executeInBatches(
   }
 }
 
+export async function* executeInBatchesRaw(
+  client: DuckDBClientLike,
+  query: string,
+  params: unknown[],
+  options: ExecuteInBatchesOptions = {}
+): AsyncGenerator<ExecuteBatchesRawChunk, void, void> {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      yield* executeInBatchesRaw(connection, query, params, options);
+      return;
+    } finally {
+      await client.release(connection);
+    }
+  }
+
+  const rowsPerChunk =
+    options.rowsPerChunk && options.rowsPerChunk > 0
+      ? options.rowsPerChunk
+      : 100_000;
+
+  const values =
+    params.length > 0
+      ? (params.map((param) => toNodeApiValue(param)) as DuckDBValue[])
+      : undefined;
+
+  const result = await client.stream(query, values);
+  const rawColumns =
+    typeof result.deduplicatedColumnNames === 'function'
+      ? result.deduplicatedColumnNames()
+      : result.columnNames();
+  const columns =
+    typeof result.deduplicatedColumnNames === 'function'
+      ? rawColumns
+      : deduplicateColumns(rawColumns);
+
+  let buffer: unknown[][] = [];
+
+  for await (const chunk of result.yieldRowsJs()) {
+    for (const row of chunk) {
+      buffer.push(row as unknown[]);
+      if (buffer.length >= rowsPerChunk) {
+        yield { columns, rows: buffer };
+        buffer = [];
+      }
+    }
+  }
+
+  if (buffer.length > 0) {
+    yield { columns, rows: buffer };
+  }
+}
+
 /**
  * Return columnar results when the underlying node-api exposes an Arrow/columnar API.
  * Falls back to column-major JS arrays when Arrow is unavailable.

package/src/columns.ts

@@ -7,11 +7,13 @@ import {
   wrapMap,
   wrapBlob,
   wrapJson,
+  wrapTimestamp,
   type ListValueWrapper,
   type ArrayValueWrapper,
   type MapValueWrapper,
   type BlobValueWrapper,
   type JsonValueWrapper,
+  type TimestampValueWrapper,
 } from './value-wrappers-core.ts';
 
 type IntColType =
@@ -355,12 +357,35 @@ interface TimestampOptions {
   withTimezone?: boolean;
   mode?: TimestampMode;
   precision?: number;
+  bindMode?: 'auto' | 'bind' | 'literal';
+}
+
+function shouldBindTimestamp(options: TimestampOptions): boolean {
+  const bindMode = options.bindMode ?? 'auto';
+  if (bindMode === 'bind') return true;
+  if (bindMode === 'literal') return false;
+
+  const isBun =
+    typeof process !== 'undefined' &&
+    typeof process.versions?.bun !== 'undefined';
+  if (isBun) return false;
+
+  const forceLiteral =
+    typeof process !== 'undefined'
+      ? process.env.DRIZZLE_DUCKDB_FORCE_LITERAL_TIMESTAMPS
+      : undefined;
+
+  if (forceLiteral && forceLiteral !== '0') {
+    return false;
+  }
+
+  return true;
 }
 
 export const duckDbTimestamp = (name: string, options: TimestampOptions = {}) =>
   customType<{
     data: Date | string;
-    driverData: SQL | string | Date;
+    driverData: SQL | string | Date | TimestampValueWrapper;
   }>({
     dataType() {
       if (options.withTimezone) {
@@ -369,14 +394,36 @@ export const duckDbTimestamp = (name: string, options: TimestampOptions = {}) =>
       const precision = options.precision ? `(${options.precision})` : '';
       return `TIMESTAMP${precision}`;
     },
-    toDriver(value: Date | string) {
-      // Use SQL literals for timestamps due to Bun/DuckDB bigint binding issues
+    toDriver(
+      value: Date | string
+    ): SQL | string | Date | TimestampValueWrapper {
+      if (shouldBindTimestamp(options)) {
+        return wrapTimestamp(
+          value,
+          options.withTimezone ?? false,
+          options.precision
+        );
+      }
+
       const iso = value instanceof Date ? value.toISOString() : value;
       const normalized = iso.replace('T', ' ').replace('Z', '+00');
       const typeKeyword = options.withTimezone ? 'TIMESTAMPTZ' : 'TIMESTAMP';
       return sql.raw(`${typeKeyword} '${normalized}'`);
     },
-    fromDriver(value: Date | string | SQL) {
+    fromDriver(value: Date | string | SQL | TimestampValueWrapper) {
+      if (
+        value &&
+        typeof value === 'object' &&
+        'kind' in value &&
+        (value as TimestampValueWrapper).kind === 'timestamp'
+      ) {
+        const wrapped = value as TimestampValueWrapper;
+        return wrapped.data instanceof Date
+          ? wrapped.data
+          : typeof wrapped.data === 'number' || typeof wrapped.data === 'bigint'
+            ? new Date(Number(wrapped.data) / 1000)
+            : wrapped.data;
+      }
       if (options.mode === 'string') {
         if (value instanceof Date) {
           return value.toISOString().replace('T', ' ').replace('Z', '+00');

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)) };
+}