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

package/src/client.ts

@@ -10,9 +10,21 @@ import {
   type AnyDuckDBValueWrapper,
 } from './value-wrappers.ts';
 
-export type DuckDBClientLike = DuckDBConnection;
+export type DuckDBClientLike = DuckDBConnection | DuckDBConnectionPool;
 export type RowData = Record<string, unknown>;
 
+export interface DuckDBConnectionPool {
+  acquire(): Promise<DuckDBConnection>;
+  release(connection: DuckDBConnection): void | Promise<void>;
+  close?(): Promise<void> | void;
+}
+
+export function isPool(
+  client: DuckDBClientLike
+): client is DuckDBConnectionPool {
+  return typeof (client as DuckDBConnectionPool).acquire === 'function';
+}
+
 export interface PrepareParamsOptions {
   rejectStringArrayLiterals?: boolean;
   warnOnStringArrayLiteral?: () => void;
@@ -32,8 +44,6 @@ function parsePgArrayLiteral(value: string): unknown {
   }
 }
 
-let warnedArrayLiteral = false;
-
 export function prepareParams(
   params: unknown[],
   options: PrepareParamsOptions = {}
@@ -46,8 +56,7 @@ export function prepareParams(
         );
       }
 
-      if (!warnedArrayLiteral && options.warnOnStringArrayLiteral) {
-        warnedArrayLiteral = true;
+      if (options.warnOnStringArrayLiteral) {
         options.warnOnStringArrayLiteral();
       }
       return parsePgArrayLiteral(param);
@@ -138,6 +147,15 @@ export async function executeOnClient(
   query: string,
   params: unknown[]
 ): Promise<RowData[]> {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      return await executeOnClient(connection, query, params);
+    } finally {
+      await client.release(connection);
+    }
+  }
+
   const values =
     params.length > 0
       ? (params.map((param) => toNodeApiValue(param)) as DuckDBValue[])
@@ -165,6 +183,16 @@ export async function* executeInBatches(
   params: unknown[],
   options: ExecuteInBatchesOptions = {}
 ): AsyncGenerator<RowData[], void, void> {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      yield* executeInBatches(connection, query, params, options);
+      return;
+    } finally {
+      await client.release(connection);
+    }
+  }
+
   const rowsPerChunk =
     options.rowsPerChunk && options.rowsPerChunk > 0
       ? options.rowsPerChunk
@@ -207,12 +235,22 @@ export async function executeArrowOnClient(
   query: string,
   params: unknown[]
 ): Promise<unknown> {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      return await executeArrowOnClient(connection, query, params);
+    } finally {
+      await client.release(connection);
+    }
+  }
+
   const values =
     params.length > 0
       ? (params.map((param) => toNodeApiValue(param)) as DuckDBValue[])
       : undefined;
   const result = await client.run(query, values);
 
+  // Runtime detection for Arrow API support (optional method, not in base type)
   const maybeArrow =
     (result as unknown as { toArrow?: () => Promise<unknown> }).toArrow ??
     (result as unknown as { getArrowTable?: () => Promise<unknown> })

package/src/columns.ts

@@ -12,7 +12,7 @@ import {
   type MapValueWrapper,
   type BlobValueWrapper,
   type JsonValueWrapper,
-} from './value-wrappers.ts';
+} from './value-wrappers-core.ts';
 
 type IntColType =
   | 'SMALLINT'

package/src/driver.ts

@@ -1,3 +1,4 @@
+import { DuckDBInstance } from '@duckdb/node-api';
 import { entityKind } from 'drizzle-orm/entity';
 import type { Logger } from 'drizzle-orm/logger';
 import { DefaultLogger } from 'drizzle-orm/logger';
@@ -23,6 +24,13 @@ 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 { isPool } from './client.ts';
+import {
+  createDuckDBConnectionPool,
+  resolvePoolSize,
+  type DuckDBPoolConfig,
+  type PoolPreset,
+} from './pool.ts';
 
 export interface PgDriverOptions {
   logger?: Logger;
@@ -50,18 +58,57 @@ export class DuckDBDriver {
   }
 }
 
+/** Connection configuration when using path-based connection */
+export interface DuckDBConnectionConfig {
+  /** Database path: ':memory:', './file.duckdb', 'md:', 'md:database' */
+  path: string;
+  /** DuckDB instance options (e.g., motherduck_token) */
+  options?: Record<string, string>;
+}
+
 export interface DuckDBDrizzleConfig<
   TSchema extends Record<string, unknown> = Record<string, never>,
 > extends DrizzleConfig<TSchema> {
   rewriteArrays?: boolean;
   rejectStringArrayLiterals?: boolean;
+  /** Pool configuration. Use preset name, size config, or false to disable. */
+  pool?: DuckDBPoolConfig | PoolPreset | false;
 }
 
-export function drizzle<
+export interface DuckDBDrizzleConfigWithConnection<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+> extends DuckDBDrizzleConfig<TSchema> {
+  /** Connection string or config object */
+  connection: string | DuckDBConnectionConfig;
+}
+
+export interface DuckDBDrizzleConfigWithClient<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+> extends DuckDBDrizzleConfig<TSchema> {
+  /** Explicit client (connection or pool) */
+  client: DuckDBClientLike;
+}
+
+/** Check if a value looks like a config object (not a client) */
+function isConfigObject(data: unknown): data is Record<string, unknown> {
+  if (typeof data !== 'object' || data === null) return false;
+  if (data.constructor?.name !== 'Object') return false;
+  return (
+    'connection' in data ||
+    'client' in data ||
+    'pool' in data ||
+    'schema' in data ||
+    'logger' in data
+  );
+}
+
+/** Internal: create database from a client (connection or pool) */
+function createFromClient<
   TSchema extends Record<string, unknown> = Record<string, never>,
 >(
   client: DuckDBClientLike,
-  config: DuckDBDrizzleConfig<TSchema> = {}
+  config: DuckDBDrizzleConfig<TSchema> = {},
+  instance?: DuckDBInstance
 ): DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>> {
   const dialect = new DuckDBDialect();
 
@@ -89,10 +136,125 @@ export function drizzle<
   });
   const session = driver.createSession(schema);
 
-  return new DuckDBDatabase(dialect, session, schema) as DuckDBDatabase<
-    TSchema,
-    ExtractTablesWithRelations<TSchema>
-  >;
+  const db = new DuckDBDatabase(dialect, session, schema, client, instance);
+  return db as DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>;
+}
+
+/** Internal: create database from a connection string */
+async function createFromConnectionString<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  path: string,
+  instanceOptions: Record<string, string> | undefined,
+  config: DuckDBDrizzleConfig<TSchema> = {}
+): Promise<DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>> {
+  const instance = await DuckDBInstance.create(path, instanceOptions);
+  const poolSize = resolvePoolSize(config.pool);
+
+  if (poolSize === false) {
+    const connection = await instance.connect();
+    return createFromClient(connection, config, instance);
+  }
+
+  const pool = createDuckDBConnectionPool(instance, { size: poolSize });
+  return createFromClient(pool, config, instance);
+}
+
+// Overload 1: Connection string (async, auto-pools)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  connectionString: string
+): Promise<DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>>;
+
+// Overload 2: Connection string + config (async, auto-pools)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  connectionString: string,
+  config: DuckDBDrizzleConfig<TSchema>
+): Promise<DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>>;
+
+// Overload 3: Config with connection (async, auto-pools)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  config: DuckDBDrizzleConfigWithConnection<TSchema>
+): Promise<DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>>;
+
+// Overload 4: Config with explicit client (sync)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  config: DuckDBDrizzleConfigWithClient<TSchema>
+): DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>;
+
+// Overload 5: Explicit client (sync, backward compatible)
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  client: DuckDBClientLike,
+  config?: DuckDBDrizzleConfig<TSchema>
+): DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>;
+
+// Implementation
+export function drizzle<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+>(
+  clientOrConfigOrPath:
+    | string
+    | DuckDBClientLike
+    | DuckDBDrizzleConfigWithConnection<TSchema>
+    | DuckDBDrizzleConfigWithClient<TSchema>,
+  config?: DuckDBDrizzleConfig<TSchema>
+):
+  | DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>
+  | Promise<DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>>> {
+  // String path -> async with auto-pool
+  if (typeof clientOrConfigOrPath === 'string') {
+    return createFromConnectionString(clientOrConfigOrPath, undefined, config);
+  }
+
+  // Config object with connection or client
+  if (isConfigObject(clientOrConfigOrPath)) {
+    const configObj = clientOrConfigOrPath as
+      | DuckDBDrizzleConfigWithConnection<TSchema>
+      | DuckDBDrizzleConfigWithClient<TSchema>;
+
+    if ('connection' in configObj) {
+      const connConfig =
+        configObj as DuckDBDrizzleConfigWithConnection<TSchema>;
+      const { connection, ...restConfig } = connConfig;
+      if (typeof connection === 'string') {
+        return createFromConnectionString(
+          connection,
+          undefined,
+          restConfig as DuckDBDrizzleConfig<TSchema>
+        );
+      }
+      return createFromConnectionString(
+        connection.path,
+        connection.options,
+        restConfig as DuckDBDrizzleConfig<TSchema>
+      );
+    }
+
+    if ('client' in configObj) {
+      const clientConfig = configObj as DuckDBDrizzleConfigWithClient<TSchema>;
+      const { client: clientValue, ...restConfig } = clientConfig;
+      return createFromClient(
+        clientValue,
+        restConfig as DuckDBDrizzleConfig<TSchema>
+      );
+    }
+
+    throw new Error(
+      'Invalid drizzle config: either connection or client must be provided'
+    );
+  }
+
+  // Direct client (backward compatible)
+  return createFromClient(clientOrConfigOrPath as DuckDBClientLike, config);
 }
 
 export class DuckDBDatabase<
@@ -102,12 +264,33 @@ export class DuckDBDatabase<
 > extends PgDatabase<DuckDBQueryResultHKT, TFullSchema, TSchema> {
   static readonly [entityKind]: string = 'DuckDBDatabase';
 
+  /** The underlying connection or pool */
+  readonly $client: DuckDBClientLike;
+
+  /** The DuckDB instance (when created from connection string) */
+  readonly $instance?: DuckDBInstance;
+
   constructor(
     readonly dialect: DuckDBDialect,
     readonly session: DuckDBSession<TFullSchema, TSchema>,
-    schema: RelationalSchemaConfig<TSchema> | undefined
+    schema: RelationalSchemaConfig<TSchema> | undefined,
+    client: DuckDBClientLike,
+    instance?: DuckDBInstance
   ) {
     super(dialect, session, schema);
+    this.$client = client;
+    this.$instance = instance;
+  }
+
+  /**
+   * Close the database connection pool and instance.
+   * Should be called when shutting down the application.
+   */
+  async close(): Promise<void> {
+    if (isPool(this.$client) && this.$client.close) {
+      await this.$client.close();
+    }
+    // Note: DuckDBInstance doesn't have a close method in the current API
   }
 
   select(): DuckDBSelectBuilder<undefined>;
@@ -119,6 +302,7 @@ export class DuckDBDatabase<
   ): DuckDBSelectBuilder<SelectedFields | undefined> {
     const selectedFields = fields ? aliasFields(fields) : undefined;
 
+    // Cast needed: DuckDBSession is compatible but types don't align exactly with PgSession
     return new DuckDBSelectBuilder({
       fields: selectedFields ?? undefined,
       session: this.session as unknown as PgSession<DuckDBQueryResultHKT>,

package/src/helpers.ts

@@ -0,0 +1,18 @@
+// Client-safe entrypoint exposing schema builder utilities without pulling
+// the DuckDB Node API bindings. Intended for generated schemas and browser use.
+export {
+  duckDbList,
+  duckDbArray,
+  duckDbMap,
+  duckDbStruct,
+  duckDbJson,
+  duckDbBlob,
+  duckDbInet,
+  duckDbInterval,
+  duckDbTimestamp,
+  duckDbDate,
+  duckDbTime,
+  duckDbArrayContains,
+  duckDbArrayContained,
+  duckDbArrayOverlaps,
+} from './columns.ts';

package/src/index.ts

@@ -4,5 +4,6 @@ export * from './columns.ts';
 export * from './migrator.ts';
 export * from './introspect.ts';
 export * from './client.ts';
+export * from './pool.ts';
 export * from './olap.ts';
 export * from './value-wrappers.ts';

package/src/introspect.ts

@@ -108,7 +108,8 @@ type ImportBuckets = {
   local: Set<string>;
 };
 
-const DEFAULT_IMPORT_BASE = '@leonardovida-md/drizzle-neo-duckdb';
+export const DEFAULT_IMPORT_BASE =
+  '@leonardovida-md/drizzle-neo-duckdb/helpers';
 
 export async function introspect(
   db: DuckDBDatabase,
@@ -649,6 +650,28 @@ function mapDuckDbType(
     return { builder: `doublePrecision(${columnName(column.name)})` };
   }
 
+  const arrayMatch = /^(.*)\[(\d+)\]$/.exec(upper);
+  if (arrayMatch) {
+    imports.local.add('duckDbArray');
+    const [, base, length] = arrayMatch;
+    return {
+      builder: `duckDbArray(${columnName(
+        column.name
+      )}, ${JSON.stringify(base)}, ${Number(length)})`,
+    };
+  }
+
+  const listMatch = /^(.*)\[\]$/.exec(upper);
+  if (listMatch) {
+    imports.local.add('duckDbList');
+    const [, base] = listMatch;
+    return {
+      builder: `duckDbList(${columnName(
+        column.name
+      )}, ${JSON.stringify(base)})`,
+    };
+  }
+
   if (upper.startsWith('CHAR(') || upper === 'CHAR') {
     imports.pgCore.add('char');
     const length = column.characterLength;
@@ -699,28 +722,6 @@ function mapDuckDbType(
     return { builder: `duckDbBlob(${columnName(column.name)})` };
   }
 
-  const arrayMatch = /^(.*)\[(\d+)\]$/.exec(upper);
-  if (arrayMatch) {
-    imports.local.add('duckDbArray');
-    const [, base, length] = arrayMatch;
-    return {
-      builder: `duckDbArray(${columnName(
-        column.name
-      )}, ${JSON.stringify(base)}, ${Number(length)})`,
-    };
-  }
-
-  const listMatch = /^(.*)\[\]$/.exec(upper);
-  if (listMatch) {
-    imports.local.add('duckDbList');
-    const [, base] = listMatch;
-    return {
-      builder: `duckDbList(${columnName(
-        column.name
-      )}, ${JSON.stringify(base)})`,
-    };
-  }
-
   if (upper.startsWith('STRUCT')) {
     imports.local.add('duckDbStruct');
     const inner = upper.replace(/^STRUCT\s*\(/i, '').replace(/\)$/, '');

package/src/migrator.ts

@@ -14,9 +14,9 @@ export async function migrate<TSchema extends Record<string, unknown>>(
 
   const migrations = readMigrationFiles(migrationConfig);
 
+  // Cast needed: Drizzle's internal PgSession type differs from exported type
   await db.dialect.migrate(
     migrations,
-    // Need to work around omitted internal types from drizzle...
     db.session as unknown as PgSession,
     migrationConfig
   );

package/src/olap.ts

@@ -169,6 +169,7 @@ export class OlapBuilder {
 
     Object.assign(selection, this.measureMap);
 
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any -- Drizzle's query builder types don't allow reassignment after groupBy
     let query: any = this.db
       .select(selection as SelectedFields)
       .from(this.source!)

package/src/pool.ts

@@ -0,0 +1,104 @@
+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;
+}
+
+export function createDuckDBConnectionPool(
+  instance: DuckDBInstance,
+  options: DuckDBConnectionPoolOptions = {}
+): DuckDBConnectionPool & { size: number } {
+  const size = options.size && options.size > 0 ? options.size : 4;
+  const idle: DuckDBConnection[] = [];
+  const waiting: Array<(conn: DuckDBConnection) => void> = [];
+  let total = 0;
+  let closed = 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;
+    }
+
+    if (total < size) {
+      total += 1;
+      return await DuckDBConnection.create(instance);
+    }
+
+    return await new Promise((resolve) => {
+      waiting.push(resolve);
+    });
+  };
+
+  const release = async (connection: DuckDBConnection): Promise<void> => {
+    if (closed) {
+      await closeClientConnection(connection);
+      return;
+    }
+
+    const waiter = waiting.shift();
+    if (waiter) {
+      waiter(connection);
+      return;
+    }
+
+    idle.push(connection);
+  };
+
+  const close = async (): Promise<void> => {
+    closed = true;
+    const toClose = idle.splice(0, idle.length);
+    await Promise.all(toClose.map((conn) => closeClientConnection(conn)));
+  };
+
+  return {
+    acquire,
+    release,
+    close,
+    size,
+  };
+}

package/src/session.ts

@@ -19,7 +19,11 @@ 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 {
+  DuckDBClientLike,
+  DuckDBConnectionPool,
+  RowData,
+} from './client.ts';
 import {
   executeArrowOnClient,
   executeInBatches,
@@ -27,6 +31,8 @@ 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';
 
@@ -165,8 +171,18 @@ export class DuckDBSession<
   override async transaction<T>(
     transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
   ): 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 +194,21 @@ 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;`);
+
+      try {
+        const result = await transaction(tx);
+        await tx.execute(sql`commit`);
+        return result;
+      } catch (error) {
+        await tx.execute(sql`rollback`);
+        throw error;
+      }
+    } finally {
+      if (pinnedConnection && pool) {
+        await pool.release(pinnedConnection);
+      }
     }
   }
 
@@ -297,6 +319,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 +330,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,6 +344,7 @@ 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 nestedTx = new DuckDBTransaction<TFullSchema, TSchema>(
       (this as unknown as Tx).dialect,