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

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,
@@ -36,6 +36,16 @@ 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> {
@@ -131,6 +141,7 @@ export class DuckDBSession<
   private rewriteArrays: boolean;
   private rejectStringArrayLiterals: boolean;
   private hasWarnedArrayLiteral = false;
+  private rollbackOnly = false;
 
   constructor(
     private client: DuckDBClientLike,
@@ -168,8 +179,19 @@ 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;
@@ -197,8 +219,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,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
@@ -256,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
@@ -278,6 +312,14 @@ export class DuckDBSession<
 
     return executeArrowOnClient(this.client, rewrittenQuery, params);
   }
+
+  markRollbackOnly(): void {
+    this.rollbackOnly = true;
+  }
+
+  isRollbackOnly(): boolean {
+    return this.rollbackOnly;
+  }
 }
 
 type PgTransactionInternals<
@@ -346,14 +388,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;
+    });
   }
 }
 

package/src/sql/query-rewriters.ts

@@ -1,96 +1,204 @@
-export function adaptArrayOperators(query: string): string {
-  type ArrayOperator = {
-    token: '@>' | '<@' | '&&';
-    fn: 'array_has_all' | 'array_has_any';
-    swap?: boolean;
-  };
-
-  const operators: ArrayOperator[] = [
-    { token: '@>', fn: 'array_has_all' },
-    { token: '<@', fn: 'array_has_all', swap: true },
-    { token: '&&', fn: 'array_has_any' },
-  ];
-
-  const isWhitespace = (char: string | undefined) =>
-    char !== undefined && /\s/.test(char);
-
-  const walkLeft = (source: string, start: number): [number, string] => {
-    let idx = start;
-    while (idx >= 0 && isWhitespace(source[idx])) {
-      idx--;
+type ArrayOperator = {
+  token: '@>' | '<@' | '&&';
+  fn: 'array_has_all' | 'array_has_any';
+  swap?: boolean;
+};
+
+const OPERATORS: ArrayOperator[] = [
+  { token: '@>', fn: 'array_has_all' },
+  { token: '<@', fn: 'array_has_all', swap: true },
+  { token: '&&', fn: 'array_has_any' },
+];
+
+const isWhitespace = (char: string | undefined) =>
+  char !== undefined && /\s/.test(char);
+
+export function scrubForRewrite(query: string): string {
+  let scrubbed = '';
+  type State = 'code' | 'single' | 'double' | 'lineComment' | 'blockComment';
+  let state: State = 'code';
+
+  for (let i = 0; i < query.length; i += 1) {
+    const char = query[i]!;
+    const next = query[i + 1];
+
+    if (state === 'code') {
+      if (char === "'") {
+        scrubbed += "'";
+        state = 'single';
+        continue;
+      }
+      if (char === '"') {
+        scrubbed += '"';
+        state = 'double';
+        continue;
+      }
+      if (char === '-' && next === '-') {
+        scrubbed += '  ';
+        i += 1;
+        state = 'lineComment';
+        continue;
+      }
+      if (char === '/' && next === '*') {
+        scrubbed += '  ';
+        i += 1;
+        state = 'blockComment';
+        continue;
+      }
+
+      scrubbed += char;
+      continue;
     }
 
-    let depth = 0;
-    let inString = false;
-    for (; idx >= 0; idx--) {
-      const ch = source[idx];
-      if (ch === undefined) break;
-      if (ch === "'" && source[idx - 1] !== '\\') {
-        inString = !inString;
+    if (state === 'single') {
+      if (char === "'" && next === "'") {
+        scrubbed += "''";
+        i += 1;
+        continue;
       }
-      if (inString) continue;
-      if (ch === ')' || ch === ']') {
-        depth++;
-      } else if (ch === '(' || ch === '[') {
-        depth--;
-        if (depth < 0) {
-          return [idx + 1, source.slice(idx + 1, start + 1)];
-        }
-      } else if (depth === 0 && isWhitespace(ch)) {
-        return [idx + 1, source.slice(idx + 1, start + 1)];
+      // Preserve quote for boundary detection but mask inner chars with a
+      // non-whitespace placeholder to avoid false positives on operators.
+      scrubbed += char === "'" ? "'" : '.';
+      if (char === "'") {
+        state = 'code';
+      }
+      continue;
+    }
+
+    if (state === 'double') {
+      if (char === '"' && next === '"') {
+        scrubbed += '""';
+        i += 1;
+        continue;
+      }
+      scrubbed += char === '"' ? '"' : '.';
+      if (char === '"') {
+        state = 'code';
+      }
+      continue;
+    }
+
+    if (state === 'lineComment') {
+      scrubbed += char === '\n' ? '\n' : ' ';
+      if (char === '\n') {
+        state = 'code';
+      }
+      continue;
+    }
+
+    if (state === 'blockComment') {
+      if (char === '*' && next === '/') {
+        scrubbed += '  ';
+        i += 1;
+        state = 'code';
+      } else {
+        scrubbed += ' ';
       }
     }
-    return [0, source.slice(0, start + 1)];
-  };
+  }
+
+  return scrubbed;
+}
 
-  const walkRight = (source: string, start: number): [number, string] => {
-    let idx = start;
-    while (idx < source.length && isWhitespace(source[idx])) {
-      idx++;
+function findNextOperator(
+  scrubbed: string,
+  start: number
+): { index: number; operator: ArrayOperator } | null {
+  for (let idx = start; idx < scrubbed.length; idx += 1) {
+    for (const operator of OPERATORS) {
+      if (scrubbed.startsWith(operator.token, idx)) {
+        return { index: idx, operator };
+      }
     }
+  }
+  return null;
+}
 
-    let depth = 0;
-    let inString = false;
-    for (; idx < source.length; idx++) {
-      const ch = source[idx];
-      if (ch === undefined) break;
-      if (ch === "'" && source[idx - 1] !== '\\') {
-        inString = !inString;
+function walkLeft(
+  source: string,
+  scrubbed: string,
+  start: number
+): [number, string] {
+  let idx = start;
+  while (idx >= 0 && isWhitespace(scrubbed[idx])) {
+    idx -= 1;
+  }
+
+  let depth = 0;
+  for (; idx >= 0; idx -= 1) {
+    const ch = scrubbed[idx];
+    if (ch === ')' || ch === ']') {
+      depth += 1;
+    } else if (ch === '(' || ch === '[') {
+      if (depth === 0) {
+        return [idx + 1, source.slice(idx + 1, start + 1)];
       }
-      if (inString) continue;
-      if (ch === '(' || ch === '[') {
-        depth++;
-      } else if (ch === ')' || ch === ']') {
-        depth--;
-        if (depth < 0) {
-          return [idx, source.slice(start, idx)];
-        }
-      } else if (depth === 0 && isWhitespace(ch)) {
+      depth = Math.max(0, depth - 1);
+    } else if (depth === 0 && isWhitespace(ch)) {
+      return [idx + 1, source.slice(idx + 1, start + 1)];
+    }
+  }
+
+  return [0, source.slice(0, start + 1)];
+}
+
+function walkRight(
+  source: string,
+  scrubbed: string,
+  start: number
+): [number, string] {
+  let idx = start;
+  while (idx < scrubbed.length && isWhitespace(scrubbed[idx])) {
+    idx += 1;
+  }
+
+  let depth = 0;
+  for (; idx < scrubbed.length; idx += 1) {
+    const ch = scrubbed[idx];
+    if (ch === '(' || ch === '[') {
+      depth += 1;
+    } else if (ch === ')' || ch === ']') {
+      if (depth === 0) {
         return [idx, source.slice(start, idx)];
       }
+      depth = Math.max(0, depth - 1);
+    } else if (depth === 0 && isWhitespace(ch)) {
+      return [idx, source.slice(start, idx)];
     }
-    return [source.length, source.slice(start)];
-  };
+  }
 
+  return [scrubbed.length, source.slice(start)];
+}
+
+export function adaptArrayOperators(query: string): string {
   let rewritten = query;
-  for (const { token, fn, swap } of operators) {
-    let idx = rewritten.indexOf(token);
-    while (idx !== -1) {
-      const [leftStart, leftExpr] = walkLeft(rewritten, idx - 1);
-      const [rightEnd, rightExpr] = walkRight(rewritten, idx + token.length);
+  let scrubbed = scrubForRewrite(query);
+  let searchStart = 0;
 
-      const left = leftExpr.trim();
-      const right = rightExpr.trim();
+  // Re-run after each replacement to keep indexes aligned with the current string
+  while (true) {
+    const next = findNextOperator(scrubbed, searchStart);
+    if (!next) break;
 
-      const replacement = `${fn}(${swap ? right : left}, ${
-        swap ? left : right
-      })`;
+    const { index, operator } = next;
+    const [leftStart, leftExpr] = walkLeft(rewritten, scrubbed, index - 1);
+    const [rightEnd, rightExpr] = walkRight(
+      rewritten,
+      scrubbed,
+      index + operator.token.length
+    );
 
-      rewritten =
-        rewritten.slice(0, leftStart) + replacement + rewritten.slice(rightEnd);
+    const left = leftExpr.trim();
+    const right = rightExpr.trim();
 
-      idx = rewritten.indexOf(token, leftStart + replacement.length);
-    }
+    const replacement = `${operator.fn}(${operator.swap ? right : left}, ${
+      operator.swap ? left : right
+    })`;
+
+    rewritten =
+      rewritten.slice(0, leftStart) + replacement + rewritten.slice(rightEnd);
+    scrubbed = scrubForRewrite(rewritten);
+    searchStart = leftStart + replacement.length;
   }
 
   return rewritten;