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

package/src/pool.ts

@@ -43,6 +43,14 @@ export function resolvePoolSize(
 export interface DuckDBConnectionPoolOptions {
   /** Maximum concurrent connections. Defaults to 4. */
   size?: number;
+  /** Timeout in milliseconds to wait for a connection. Defaults to 30000 (30s). */
+  acquireTimeout?: number;
+  /** Maximum number of requests waiting for a connection. Defaults to 100. */
+  maxWaitingRequests?: number;
+  /** Max time (ms) a connection may live before being recycled. */
+  maxLifetimeMs?: number;
+  /** Max idle time (ms) before an idle connection is discarded. */
+  idleTimeoutMs?: number;
 }
 
 export function createDuckDBConnectionPool(
@@ -50,49 +58,211 @@ export function createDuckDBConnectionPool(
   options: DuckDBConnectionPoolOptions = {}
 ): DuckDBConnectionPool & { size: number } {
   const size = options.size && options.size > 0 ? options.size : 4;
-  const idle: DuckDBConnection[] = [];
-  const waiting: Array<(conn: DuckDBConnection) => void> = [];
+  const acquireTimeout = options.acquireTimeout ?? 30_000;
+  const maxWaitingRequests = options.maxWaitingRequests ?? 100;
+  const maxLifetimeMs = options.maxLifetimeMs;
+  const idleTimeoutMs = options.idleTimeoutMs;
+  const metadata = new WeakMap<
+    DuckDBConnection,
+    { createdAt: number; lastUsedAt: number }
+  >();
+
+  type PooledConnection = {
+    connection: DuckDBConnection;
+    createdAt: number;
+    lastUsedAt: number;
+  };
+
+  const idle: PooledConnection[] = [];
+  const waiting: Array<{
+    resolve: (conn: DuckDBConnection) => void;
+    reject: (error: Error) => void;
+    timeoutId: ReturnType<typeof setTimeout>;
+  }> = [];
   let total = 0;
   let closed = false;
+  // Track pending acquires to handle race conditions during close
+  let pendingAcquires = 0;
+
+  const shouldRecycle = (conn: PooledConnection, now: number): boolean => {
+    if (maxLifetimeMs !== undefined && now - conn.createdAt >= maxLifetimeMs) {
+      return true;
+    }
+    if (idleTimeoutMs !== undefined && now - conn.lastUsedAt >= idleTimeoutMs) {
+      return true;
+    }
+    return false;
+  };
 
   const acquire = async (): Promise<DuckDBConnection> => {
     if (closed) {
       throw new Error('DuckDB connection pool is closed');
     }
 
-    if (idle.length > 0) {
-      return idle.pop() as DuckDBConnection;
+    while (idle.length > 0) {
+      const pooled = idle.pop() as PooledConnection;
+      const now = Date.now();
+      if (shouldRecycle(pooled, now)) {
+        await closeClientConnection(pooled.connection);
+        total = Math.max(0, total - 1);
+        metadata.delete(pooled.connection);
+        continue;
+      }
+      pooled.lastUsedAt = now;
+      metadata.set(pooled.connection, {
+        createdAt: pooled.createdAt,
+        lastUsedAt: pooled.lastUsedAt,
+      });
+      return pooled.connection;
     }
 
     if (total < size) {
+      pendingAcquires += 1;
       total += 1;
-      return await DuckDBConnection.create(instance);
+      try {
+        const connection = await DuckDBConnection.create(instance);
+        // Check if pool was closed during async connection creation
+        if (closed) {
+          await closeClientConnection(connection);
+          total -= 1;
+          throw new Error('DuckDB connection pool is closed');
+        }
+        const now = Date.now();
+        metadata.set(connection, { createdAt: now, lastUsedAt: now });
+        return connection;
+      } catch (error) {
+        total -= 1;
+        throw error;
+      } finally {
+        pendingAcquires -= 1;
+      }
     }
 
-    return await new Promise((resolve) => {
-      waiting.push(resolve);
+    // Check queue limit before waiting
+    if (waiting.length >= maxWaitingRequests) {
+      throw new Error(
+        `DuckDB connection pool queue is full (max ${maxWaitingRequests} waiting requests)`
+      );
+    }
+
+    return await new Promise((resolve, reject) => {
+      const timeoutId = setTimeout(() => {
+        // Remove this waiter from the queue
+        const idx = waiting.findIndex((w) => w.timeoutId === timeoutId);
+        if (idx !== -1) {
+          waiting.splice(idx, 1);
+        }
+        reject(
+          new Error(
+            `DuckDB connection pool acquire timeout after ${acquireTimeout}ms`
+          )
+        );
+      }, acquireTimeout);
+
+      waiting.push({ resolve, reject, timeoutId });
     });
   };
 
   const release = async (connection: DuckDBConnection): Promise<void> => {
+    const waiter = waiting.shift();
+    if (waiter) {
+      clearTimeout(waiter.timeoutId);
+      const now = Date.now();
+      const meta =
+        metadata.get(connection) ??
+        ({ createdAt: now, lastUsedAt: now } as {
+          createdAt: number;
+          lastUsedAt: number;
+        });
+
+      const expired =
+        maxLifetimeMs !== undefined && now - meta.createdAt >= maxLifetimeMs;
+
+      if (closed) {
+        await closeClientConnection(connection);
+        total = Math.max(0, total - 1);
+        metadata.delete(connection);
+        waiter.reject(new Error('DuckDB connection pool is closed'));
+        return;
+      }
+
+      if (expired) {
+        await closeClientConnection(connection);
+        total = Math.max(0, total - 1);
+        metadata.delete(connection);
+        try {
+          const replacement = await acquire();
+          waiter.resolve(replacement);
+        } catch (error) {
+          waiter.reject(error as Error);
+        }
+        return;
+      }
+
+      meta.lastUsedAt = now;
+      metadata.set(connection, meta);
+      waiter.resolve(connection);
+      return;
+    }
+
     if (closed) {
       await closeClientConnection(connection);
+      metadata.delete(connection);
+      total = Math.max(0, total - 1);
       return;
     }
 
-    const waiter = waiting.shift();
-    if (waiter) {
-      waiter(connection);
+    const now = Date.now();
+    const existingMeta =
+      metadata.get(connection) ??
+      ({ createdAt: now, lastUsedAt: now } as {
+        createdAt: number;
+        lastUsedAt: number;
+      });
+    existingMeta.lastUsedAt = now;
+    metadata.set(connection, existingMeta);
+
+    if (
+      maxLifetimeMs !== undefined &&
+      now - existingMeta.createdAt >= maxLifetimeMs
+    ) {
+      await closeClientConnection(connection);
+      total -= 1;
+      metadata.delete(connection);
       return;
     }
 
-    idle.push(connection);
+    idle.push({
+      connection,
+      createdAt: existingMeta.createdAt,
+      lastUsedAt: existingMeta.lastUsedAt,
+    });
   };
 
   const close = async (): Promise<void> => {
     closed = true;
+
+    // Clear all waiting requests with their timeouts
+    const waiters = waiting.splice(0, waiting.length);
+    for (const waiter of waiters) {
+      clearTimeout(waiter.timeoutId);
+      waiter.reject(new Error('DuckDB connection pool is closed'));
+    }
+
+    // Close all idle connections (use allSettled to ensure all are attempted)
     const toClose = idle.splice(0, idle.length);
-    await Promise.all(toClose.map((conn) => closeClientConnection(conn)));
+    await Promise.allSettled(
+      toClose.map((item) => closeClientConnection(item.connection))
+    );
+    total = Math.max(0, total - toClose.length);
+    toClose.forEach((item) => metadata.delete(item.connection));
+
+    // Wait for pending acquires to complete (with a reasonable timeout)
+    const maxWait = 5000;
+    const start = Date.now();
+    while (pendingAcquires > 0 && Date.now() - start < maxWait) {
+      await new Promise((r) => setTimeout(r, 10));
+    }
   };
 
   return {

package/src/session.ts

@@ -17,8 +17,8 @@ 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 { mapResultRow } from './sql/result-mapper.ts';
-import type { DuckDBDialect } from './dialect.ts';
 import { TransactionRollbackError } from 'drizzle-orm/errors';
+import type { DuckDBDialect } from './dialect.ts';
 import type {
   DuckDBClientLike,
   DuckDBConnectionPool,
@@ -26,16 +26,45 @@ 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';
 
+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')
+  );
+}
+
+function rewriteQuery(
+  mode: RewriteArraysMode,
+  query: string
+): { sql: string; rewritten: boolean } {
+  if (mode === 'never') {
+    return { sql: query, rewritten: false };
+  }
+
+  const rewritten = adaptArrayOperators(query);
+  return { sql: rewritten, rewritten: rewritten !== query };
+}
+
 export class DuckDBPreparedQuery<
   T extends PreparedQueryConfig,
 > extends PgPreparedQuery<T> {
@@ -52,8 +81,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 });
@@ -72,11 +102,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
@@ -88,19 +119,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 || !fields) {
-      return rows as T['execute'];
+      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, rewrittenQuery, params, {
+      prepareCache: this.prepareCache,
+    });
 
-    return customResultMapper
-      ? customResultMapper(rowValues)
-      : rowValues.map((row) =>
-          mapResultRow<T['execute']>(fields, row, joinsNotNullableMap)
-        );
+    return rows as T['execute'];
   }
 
   all(
@@ -116,8 +158,9 @@ export class DuckDBPreparedQuery<
 
 export interface DuckDBSessionOptions {
   logger?: Logger;
-  rewriteArrays?: boolean;
+  rewriteArrays?: RewriteArraysMode;
   rejectStringArrayLiterals?: boolean;
+  prepareCache?: PreparedStatementCacheConfig;
 }
 
 export class DuckDBSession<
@@ -128,9 +171,11 @@ 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;
 
   constructor(
     private client: DuckDBClientLike,
@@ -141,8 +186,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>(
@@ -162,14 +213,26 @@ export class DuckDBSession<
       fields,
       isResponseInArrayMode,
       customResultMapper,
-      this.rewriteArrays,
+      this.rewriteArraysMode,
       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> {
     let pinnedConnection: DuckDBConnection | undefined;
     let pool: DuckDBConnectionPool | undefined;
@@ -197,8 +260,16 @@ export class DuckDBSession<
     try {
       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) {
@@ -227,18 +298,21 @@ export class DuckDBSession<
     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),
     });
-    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
@@ -255,19 +329,52 @@ 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);
+    this.dialect.assertNoPgJsonColumns();
     const params = prepareParams(builtQuery.params, {
       rejectStringArrayLiterals: this.rejectStringArrayLiterals,
       warnOnStringArrayLiteral: this.rejectStringArrayLiterals
         ? 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
@@ -278,6 +385,14 @@ export class DuckDBSession<
 
     return executeArrowOnClient(this.client, rewrittenQuery, params);
   }
+
+  markRollbackOnly(): void {
+    this.rollbackOnly = true;
+  }
+
+  isRollbackOnly(): boolean {
+    return this.rollbackOnly;
+  }
 }
 
 type PgTransactionInternals<
@@ -335,6 +450,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>;
@@ -346,14 +470,65 @@ export class DuckDBTransaction<
   ): Promise<T> {
     // Cast needed: PgTransaction doesn't expose dialect/session properties in public API
     type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    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>(
-      (this as unknown as Tx).dialect,
-      (this as unknown as Tx).session,
+      internals.dialect,
+      internals.session,
       this.schema,
       this.nestedIndex + 1
     );
 
-    return transaction(nestedTx);
+    // 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(nestedTx);
+      if (createdSavepoint) {
+        await internals.session.execute(releaseSql);
+      }
+      return result;
+    } 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;
+    });
   }
 }