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

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 { closeClientConnection, 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,46 @@ 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();
+    }
+    if (!isPool(this.$client)) {
+      await closeClientConnection(this.$client);
+    }
+    if (this.$instance) {
+      const maybeClosable = this.$instance as unknown as {
+        close?: () => Promise<void> | void;
+        closeSync?: () => void;
+      };
+      if (typeof maybeClosable.close === 'function') {
+        await maybeClosable.close();
+      } else if (typeof maybeClosable.closeSync === 'function') {
+        maybeClosable.closeSync();
+      }
+    }
   }
 
   select(): DuckDBSelectBuilder<undefined>;
@@ -119,6 +315,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,
@@ -534,7 +535,7 @@ function emitColumn(
   return builder;
 }
 
-function buildDefault(defaultValue: string | null): string {
+export function buildDefault(defaultValue: string | null): string {
   if (!defaultValue) {
     return '';
   }
@@ -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;
@@ -684,6 +707,22 @@ function mapDuckDbType(
     return { builder: `text(${columnName(column.name)}) /* JSON */` };
   }
 
+  if (upper.startsWith('ENUM')) {
+    imports.pgCore.add('text');
+    const enumLiteral = raw.replace(/^ENUM\s*/i, '').trim();
+    return {
+      builder: `text(${columnName(column.name)}) /* ENUM ${enumLiteral} */`,
+    };
+  }
+
+  if (upper.startsWith('UNION')) {
+    imports.pgCore.add('text');
+    const unionLiteral = raw.replace(/^UNION\s*/i, '').trim();
+    return {
+      builder: `text(${columnName(column.name)}) /* UNION ${unionLiteral} */`,
+    };
+  }
+
   if (upper === 'INET') {
     imports.local.add('duckDbInet');
     return { builder: `duckDbInet(${columnName(column.name)})` };
@@ -699,28 +738,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(/\)$/, '');
@@ -787,15 +804,16 @@ function mapDuckDbType(
   }
 
   // Fallback: keep as text to avoid runtime failures.
+  // Unknown types are mapped to text with a comment indicating the original type.
   imports.pgCore.add('text');
   return {
     builder: `text(${columnName(
       column.name
-    )}) /* TODO: verify type ${upper} */`,
+    )}) /* unsupported DuckDB type: ${upper} */`,
   };
 }
 
-function parseStructFields(
+export function parseStructFields(
   inner: string
 ): Array<{ name: string; type: string }> {
   const result: Array<{ name: string; type: string }> = [];
@@ -812,7 +830,7 @@ function parseStructFields(
   return result;
 }
 
-function parseMapValue(raw: string): string {
+export function parseMapValue(raw: string): string {
   const inner = raw.replace(/^MAP\(/i, '').replace(/\)$/, '');
   const parts = splitTopLevel(inner, ',');
   if (parts.length < 2) {
@@ -821,7 +839,7 @@ function parseMapValue(raw: string): string {
   return parts[1]?.trim() ?? 'TEXT';
 }
 
-function splitTopLevel(input: string, delimiter: string): string[] {
+export function splitTopLevel(input: string, delimiter: string): string[] {
   const parts: string[] = [];
   let depth = 0;
   let current = '';
@@ -846,7 +864,7 @@ function tableKey(schema: string, table: string): string {
   return `${schema}.${table}`;
 }
 
-function toIdentifier(name: string): string {
+export function toIdentifier(name: string): string {
   const cleaned = name.replace(/[^A-Za-z0-9_]/g, '_');
   const parts = cleaned.split('_').filter(Boolean);
   const base = parts

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!)