@duckdbfan/drizzle-duckdb 0.0.7 → 1.3.2

package/src/driver.ts

@@ -1,8 +1,10 @@
-import { entityKind, is } from 'drizzle-orm/entity';
+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';
 import { PgDatabase } from 'drizzle-orm/pg-core/db';
-import { PgDialect } from 'drizzle-orm/pg-core/dialect';
+import type { SelectedFields } from 'drizzle-orm/pg-core/query-builders';
+import type { PgSession } from 'drizzle-orm/pg-core';
 import {
   createTableRelationsHelpers,
   extractTablesRelationalConfig,
@@ -10,104 +12,115 @@ import {
   type RelationalSchemaConfig,
   type TablesRelationalConfig,
 } from 'drizzle-orm/relations';
-import {
-  getTableColumns,
-  type DrizzleConfig,
-  type DrizzleTypeError,
-} from 'drizzle-orm/utils';
+import { type DrizzleConfig } from 'drizzle-orm/utils';
+import type { SQL } from 'drizzle-orm/sql/sql';
 import type {
-  DuckDBClient,
+  DuckDBClientLike,
   DuckDBQueryResultHKT,
   DuckDBTransaction,
-} from './session';
+} from './session.ts';
 import { DuckDBSession } from './session.ts';
 import { DuckDBDialect } from './dialect.ts';
-
+import { DuckDBSelectBuilder } from './select-builder.ts';
+import { aliasFields } from './sql/selection.ts';
+import type {
+  ExecuteBatchesRawChunk,
+  ExecuteInBatchesOptions,
+  RowData,
+} from './client.ts';
+import { closeClientConnection, isPool } from './client.ts';
 import {
-  PgSelectBase,
-  PgSelectBuilder,
-  type CreatePgSelectFromBuilderMode,
-  type SelectedFields,
-  type TableLikeHasEmptySelection,
-} from 'drizzle-orm/pg-core/query-builders';
+  createDuckDBConnectionPool,
+  resolvePoolSize,
+  type DuckDBPoolConfig,
+  type PoolPreset,
+} from './pool.ts';
 import {
-  PgColumn,
-  PgTable,
-  type PgSession,
-  type PgTransactionConfig,
-} from 'drizzle-orm/pg-core';
-import { sql, SQL, type ColumnsSelection } from 'drizzle-orm/sql/sql';
-import { Column } from 'drizzle-orm/column';
-import { Subquery, ViewBaseConfig, type SQLWrapper } from 'drizzle-orm';
-import { PgViewBase } from 'drizzle-orm/pg-core/view-base';
-import type {
-  GetSelectTableName,
-  GetSelectTableSelection,
-} from 'drizzle-orm/query-builders/select.types';
-import { aliasFields } from './utils.ts';
+  resolvePrepareCacheOption,
+  type PreparedStatementCacheConfig,
+  type PrepareCacheOption,
+} from './options.ts';
 
 export interface PgDriverOptions {
   logger?: Logger;
+  rejectStringArrayLiterals?: boolean;
+  prepareCache?: PreparedStatementCacheConfig;
 }
 
-const selectionRegex = /select\s+(.+)\s+from/i;
-const aliasRegex = /as\s+(.+)$/i;
-
 export class DuckDBDriver {
   static readonly [entityKind]: string = 'DuckDBDriver';
 
   constructor(
-    private client: DuckDBClient,
+    private client: DuckDBClientLike,
     private dialect: DuckDBDialect,
     private options: PgDriverOptions = {}
-  ) {
-    // const clientProxy = new Proxy(client, {
-    //   get(target, prop, receiver) {
-    //     if (prop === 'all') {
-    //       return (query: string, params: Parameters<typeof target.all>) => {
-    //         // DuckDB names returned variables differently than Postgres
-    //         // so we need to remap them to match the Postgres names
-    //         const selection = selectionRegex.exec(query);
-    //         if (selection?.length !== 2) {
-    //           return target.all(query, params);
-    //         }
-    //         const fields = selection[1].split(',').map((field) => {});
-    //         return target.all(query, params);
-    //       };
-    //     }
-    //     return Reflect.get(target, prop, receiver);
-    //   },
-    // });
-  }
+  ) {}
 
   createSession(
     schema: RelationalSchemaConfig<TablesRelationalConfig> | undefined
   ): DuckDBSession<Record<string, unknown>, TablesRelationalConfig> {
     return new DuckDBSession(this.client, this.dialect, schema, {
       logger: this.options.logger,
+      rejectStringArrayLiterals: this.options.rejectStringArrayLiterals,
+      prepareCache: this.options.prepareCache,
     });
   }
 }
 
-// Need to work around omitted internal types from drizzle...
-// interface DuckDBDatabaseInternal<
-//   TSchema extends Record<string, unknown> = Record<string, never>
-// > extends PgDatabase<DuckDBQueryResultHKT, TSchema> {
-//   dialect: PgDialect;
-//   session: DuckDBSession<Record<string, unknown>, TablesRelationalConfig>;
-// }
+/** 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 type DuckDBDatabase<
-//   TSchema extends Record<string, unknown> = Record<string, never>
-// > = DuckDBDatabaseInternal<TSchema>;
+export interface DuckDBDrizzleConfig<
+  TSchema extends Record<string, unknown> = Record<string, never>,
+> extends DrizzleConfig<TSchema> {
+  rejectStringArrayLiterals?: boolean;
+  prepareCache?: PrepareCacheOption;
+  /** Pool configuration. Use preset name, size config, or false to disable. */
+  pool?: DuckDBPoolConfig | PoolPreset | false;
+}
 
-export function drizzle<
-  TSchema extends Record<string, unknown> = Record<string, never>
+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: DuckDBClient,
-  config: DrizzleConfig<TSchema> = {}
+  client: DuckDBClientLike,
+  config: DuckDBDrizzleConfig<TSchema> = {},
+  instance?: DuckDBInstance
 ): DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>> {
   const dialect = new DuckDBDialect();
+  const prepareCache = resolvePrepareCacheOption(config.prepareCache);
 
   const logger =
     config.logger === true ? new DefaultLogger() : config.logger || undefined;
@@ -126,148 +139,221 @@ export function drizzle<
     };
   }
 
-  const driver = new DuckDBDriver(client, dialect, { logger });
+  const driver = new DuckDBDriver(client, dialect, {
+    logger,
+    rejectStringArrayLiterals: config.rejectStringArrayLiterals,
+    prepareCache,
+  });
   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<
   TFullSchema extends Record<string, unknown> = Record<string, never>,
-  TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>
+  TSchema extends TablesRelationalConfig =
+    ExtractTablesWithRelations<TFullSchema>,
 > 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;
   }
 
-  select(): PgSelectBuilder<undefined>;
-  select<TSelection extends SelectedFields>(
-    fields: TSelection
-  ): PgSelectBuilder<TSelection>;
-  select(fields?: SelectedFields): PgSelectBuilder<SelectedFields | undefined> {
-    if (!fields) {
-      // iterate over all fields and do the same as below (may have to extend/override `.from()`)
-
-      return new DuckDBSelectBuilder({
-        fields: fields ?? undefined,
-        session: this.session as unknown as PgSession<DuckDBQueryResultHKT>,
-        dialect: this.dialect,
-      });
+  /**
+   * 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();
+      }
     }
+  }
 
-    const aliasedFields: SelectedFields = aliasFields(fields);
+  select(): DuckDBSelectBuilder<undefined>;
+  select<TSelection extends SelectedFields>(
+    fields: TSelection
+  ): DuckDBSelectBuilder<TSelection>;
+  select(
+    fields?: SelectedFields
+  ): 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: aliasedFields,
+      fields: selectedFields ?? undefined,
       session: this.session as unknown as PgSession<DuckDBQueryResultHKT>,
       dialect: this.dialect,
     });
   }
 
-  override async transaction<T>(
-    transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
-  ): Promise<T> {
-    return await this.session.transaction<T>(transaction);
+  executeBatches<T extends RowData = RowData>(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<T[], void, void> {
+    return this.session.executeBatches<T>(query, options);
   }
-}
-
-interface PgViewBaseInternal<
-  TName extends string = string,
-  TExisting extends boolean = boolean,
-  TSelectedFields extends ColumnsSelection = ColumnsSelection
-> extends PgViewBase<TName, TExisting, TSelectedFields> {
-  [ViewBaseConfig]?: {
-    selectedFields: SelectedFields;
-  };
-}
 
-export class DuckDBSelectBuilder<
-  TSelection extends SelectedFields | undefined,
-  TBuilderMode extends 'db' | 'qb' = 'db'
-> extends PgSelectBuilder<TSelection, TBuilderMode> {
-  private _fields: TSelection;
-  private _session: PgSession | undefined;
-  private _dialect: DuckDBDialect;
-  private _withList: Subquery[] = [];
-  private _distinct:
-    | boolean
-    | {
-        on: (PgColumn | SQLWrapper)[];
-      }
-    | undefined;
-
-  constructor(config: {
-    fields: TSelection;
-    session: PgSession | undefined;
-    dialect: PgDialect;
-    withList?: Subquery[];
-    distinct?:
-      | boolean
-      | {
-          on: (PgColumn | SQLWrapper)[];
-        };
-  }) {
-    super(config);
-    this._fields = config.fields;
-    this._session = config.session;
-    this._dialect = config.dialect;
-    if (config.withList) {
-      this._withList = config.withList;
-    }
-    this._distinct = config.distinct;
+  executeBatchesRaw(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<ExecuteBatchesRawChunk, void, void> {
+    return this.session.executeBatchesRaw(query, options);
   }
 
-  from<TFrom extends PgTable | Subquery | PgViewBaseInternal | SQL>(
-    source: TableLikeHasEmptySelection<TFrom> extends true
-      ? DrizzleTypeError<"Cannot reference a data-modifying statement subquery if it doesn't contain a `returning` clause">
-      : TFrom
-  ): CreatePgSelectFromBuilderMode<
-    TBuilderMode,
-    GetSelectTableName<TFrom>,
-    TSelection extends undefined ? GetSelectTableSelection<TFrom> : TSelection,
-    TSelection extends undefined ? 'single' : 'partial'
-  > {
-    const isPartialSelect = !!this._fields;
-    const src = source as TFrom;
-
-    let fields: SelectedFields;
-    if (this._fields) {
-      fields = this._fields;
-    } else if (is(src, Subquery)) {
-      // This is required to use the proxy handler to get the correct field values from the subquery
-      fields = Object.fromEntries(
-        Object.keys(src._.selectedFields).map((key) => [
-          key,
-          src[
-            key as unknown as keyof typeof src
-          ] as unknown as SelectedFields[string],
-        ])
-      );
-    } else if (is(src, PgViewBase)) {
-      fields = src[ViewBaseConfig]?.selectedFields as SelectedFields;
-    } else if (is(src, SQL)) {
-      fields = {};
-    } else {
-      fields = aliasFields(getTableColumns<PgTable>(src), !isPartialSelect);
-    }
+  executeArrow(query: SQL): Promise<unknown> {
+    return this.session.executeArrow(query);
+  }
 
-    return new PgSelectBase({
-      table: src,
-      fields,
-      isPartialSelect,
-      session: this._session,
-      dialect: this._dialect,
-      withList: this._withList,
-      distinct: this._distinct,
-    }) as any;
+  override async transaction<T>(
+    transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
+  ): Promise<T> {
+    return await this.session.transaction<T>(transaction);
   }
 }

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

@@ -1,4 +1,11 @@
 export * from './driver.ts';
 export * from './session.ts';
 export * from './columns.ts';
-export * from './migrator.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';
+export * from './options.ts';
+export * from './operators.ts';