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

package/src/pool.ts

@@ -0,0 +1,274 @@
+import { DuckDBConnection, DuckDBInstance } from '@duckdb/node-api';
+import { closeClientConnection, type DuckDBConnectionPool } from './client.ts';
+
+/** Pool size presets for different MotherDuck instance types */
+export type PoolPreset =
+  | 'pulse'
+  | 'standard'
+  | 'jumbo'
+  | 'mega'
+  | 'giga'
+  | 'local'
+  | 'memory';
+
+/** Pool sizes optimized for each MotherDuck instance type */
+export const POOL_PRESETS: Record<PoolPreset, number> = {
+  pulse: 4, // Auto-scaling, ad-hoc analytics
+  standard: 6, // Balanced ETL/ELT workloads
+  jumbo: 8, // Complex queries, high-volume
+  mega: 12, // Large-scale transformations
+  giga: 16, // Maximum parallelism
+  local: 8, // Local DuckDB file
+  memory: 4, // In-memory testing
+};
+
+export interface DuckDBPoolConfig {
+  /** Maximum concurrent connections. Defaults to 4. */
+  size?: number;
+}
+
+/**
+ * Resolve pool configuration to a concrete size.
+ * Returns false if pooling is disabled.
+ */
+export function resolvePoolSize(
+  pool: DuckDBPoolConfig | PoolPreset | false | undefined
+): number | false {
+  if (pool === false) return false;
+  if (pool === undefined) return 4;
+  if (typeof pool === 'string') return POOL_PRESETS[pool];
+  return pool.size ?? 4;
+}
+
+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(
+  instance: DuckDBInstance,
+  options: DuckDBConnectionPoolOptions = {}
+): DuckDBConnectionPool & { size: number } {
+  const size = options.size && options.size > 0 ? options.size : 4;
+  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');
+    }
+
+    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;
+      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;
+      }
+    }
+
+    // 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 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,
+      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.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 {
+    acquire,
+    release,
+    close,
+    size,
+  };
+}

package/src/session.ts

@@ -17,9 +17,13 @@ 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 { DuckDBClientLike, RowData } from './client.ts';
+import type { DuckDBDialect } from './dialect.ts';
+import type {
+  DuckDBClientLike,
+  DuckDBConnectionPool,
+  RowData,
+} from './client.ts';
 import {
   executeArrowOnClient,
   executeInBatches,
@@ -27,9 +31,21 @@ import {
   prepareParams,
   type ExecuteInBatchesOptions,
 } from './client.ts';
+import { isPool } from './client.ts';
+import type { DuckDBConnection } from '@duckdb/node-api';
 
 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,
 > extends PgPreparedQuery<T> {
@@ -125,6 +141,7 @@ export class DuckDBSession<
   private rewriteArrays: boolean;
   private rejectStringArrayLiterals: boolean;
   private hasWarnedArrayLiteral = false;
+  private rollbackOnly = false;
 
   constructor(
     private client: DuckDBClientLike,
@@ -162,11 +179,32 @@ export class DuckDBSession<
     );
   }
 
+  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;
+
+    let clientForTx: DuckDBClientLike = this.client;
+    if (isPool(this.client)) {
+      pool = this.client;
+      pinnedConnection = await pool.acquire();
+      clientForTx = pinnedConnection;
+    }
+
     const session = new DuckDBSession(
-      this.client,
+      clientForTx,
       this.dialect,
       this.schema,
       this.options
@@ -178,15 +216,29 @@ 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 {
+      if (pinnedConnection && pool) {
+        await pool.release(pinnedConnection);
+      }
     }
   }
 
@@ -205,7 +257,9 @@ 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
@@ -234,7 +288,9 @@ export class DuckDBSession<
   }
 
   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
@@ -256,6 +312,14 @@ export class DuckDBSession<
 
     return executeArrowOnClient(this.client, rewrittenQuery, params);
   }
+
+  markRollbackOnly(): void {
+    this.rollbackOnly = true;
+  }
+
+  isRollbackOnly(): boolean {
+    return this.rollbackOnly;
+  }
 }
 
 type PgTransactionInternals<
@@ -297,6 +361,7 @@ export class DuckDBTransaction<
   }
 
   setTransaction(config: PgTransactionConfig): Promise<void> {
+    // 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)}`
@@ -307,11 +372,13 @@ export class DuckDBTransaction<
     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);
   }
 
   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);
   }
@@ -319,15 +386,67 @@ export class DuckDBTransaction<
   override async transaction<T>(
     transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
   ): 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;
+    });
   }
 }