@duckdbfan/drizzle-duckdb 0.0.6 → 1.3.1

package/src/migrator.ts

@@ -1,18 +1,23 @@
 import type { MigrationConfig } from 'drizzle-orm/migrator';
 import { readMigrationFiles } from 'drizzle-orm/migrator';
-import type { DuckDBDatabase } from './driver';
+import type { DuckDBDatabase } from './driver.ts';
 import type { PgSession } from 'drizzle-orm/pg-core/session';
 
+export type DuckDbMigrationConfig = MigrationConfig | string;
+
 export async function migrate<TSchema extends Record<string, unknown>>(
   db: DuckDBDatabase<TSchema>,
-  config: MigrationConfig
+  config: DuckDbMigrationConfig
 ) {
-  const migrations = readMigrationFiles(config);
+  const migrationConfig: MigrationConfig =
+    typeof config === 'string' ? { migrationsFolder: config } : config;
+
+  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,
-    config
+    migrationConfig
   );
 }

package/src/olap.ts

@@ -0,0 +1,190 @@
+import { is } from 'drizzle-orm/entity';
+import { sql, Subquery, type SQLWrapper } from 'drizzle-orm';
+import type { AnyPgColumn, PgTable } from 'drizzle-orm/pg-core';
+import type { PgViewBase } from 'drizzle-orm/pg-core/view-base';
+import type { SelectedFields } from 'drizzle-orm/pg-core/query-builders';
+import { SQL } from 'drizzle-orm/sql/sql';
+import { Column, getTableName } from 'drizzle-orm';
+import type { DuckDBDatabase } from './driver.ts';
+
+export const countN = (expr: SQLWrapper = sql`*`) =>
+  sql<number>`count(${expr})`.mapWith(Number);
+
+export const sumN = (expr: SQLWrapper) =>
+  sql<number>`sum(${expr})`.mapWith(Number);
+
+export const avgN = (expr: SQLWrapper) =>
+  sql<number>`avg(${expr})`.mapWith(Number);
+
+export const sumDistinctN = (expr: SQLWrapper) =>
+  sql<number>`sum(distinct ${expr})`.mapWith(Number);
+
+export const percentileCont = (p: number, expr: SQLWrapper) =>
+  sql<number>`percentile_cont(${p}) within group (order by ${expr})`.mapWith(
+    Number
+  );
+
+export const median = (expr: SQLWrapper) => percentileCont(0.5, expr);
+
+export const anyValue = <T = unknown>(expr: SQLWrapper) =>
+  sql<T>`any_value(${expr})`;
+
+type PartitionOrder =
+  | {
+      partitionBy?: SQLWrapper | SQLWrapper[];
+      orderBy?: SQLWrapper | SQLWrapper[];
+    }
+  | undefined;
+
+function normalizeArray<T>(value?: T | T[]): T[] {
+  if (!value) return [];
+  return Array.isArray(value) ? value : [value];
+}
+
+function overClause(options?: PartitionOrder) {
+  const partitions = normalizeArray(options?.partitionBy);
+  const orders = normalizeArray(options?.orderBy);
+
+  const chunks: SQLWrapper[] = [];
+
+  if (partitions.length > 0) {
+    chunks.push(sql`partition by ${sql.join(partitions, sql`, `)}`);
+  }
+
+  if (orders.length > 0) {
+    chunks.push(sql`order by ${sql.join(orders, sql`, `)}`);
+  }
+
+  if (chunks.length === 0) {
+    return sql``;
+  }
+
+  return sql`over (${sql.join(chunks, sql` `)})`;
+}
+
+export const rowNumber = (options?: PartitionOrder) =>
+  sql<number>`row_number() ${overClause(options)}`.mapWith(Number);
+
+export const rank = (options?: PartitionOrder) =>
+  sql<number>`rank() ${overClause(options)}`.mapWith(Number);
+
+export const denseRank = (options?: PartitionOrder) =>
+  sql<number>`dense_rank() ${overClause(options)}`.mapWith(Number);
+
+export const lag = <T = unknown>(
+  expr: SQLWrapper,
+  offset = 1,
+  defaultValue?: SQLWrapper,
+  options?: PartitionOrder
+) =>
+  defaultValue
+    ? sql<T>`lag(${expr}, ${offset}, ${defaultValue}) ${overClause(options)}`
+    : sql<T>`lag(${expr}, ${offset}) ${overClause(options)}`;
+
+export const lead = <T = unknown>(
+  expr: SQLWrapper,
+  offset = 1,
+  defaultValue?: SQLWrapper,
+  options?: PartitionOrder
+) =>
+  defaultValue
+    ? sql<T>`lead(${expr}, ${offset}, ${defaultValue}) ${overClause(options)}`
+    : sql<T>`lead(${expr}, ${offset}) ${overClause(options)}`;
+
+type ValueExpr = SQL | SQL.Aliased | AnyPgColumn;
+type GroupKey = ValueExpr;
+type MeasureMap = Record<string, ValueExpr>;
+type NonAggMap = Record<string, ValueExpr>;
+
+function keyAlias(key: SQLWrapper, fallback: string): string {
+  if (is(key, SQL.Aliased)) {
+    return key.fieldAlias ?? fallback;
+  }
+  if (is(key, Column)) {
+    return `${getTableName(key.table)}.${key.name}`;
+  }
+  return fallback;
+}
+
+export class OlapBuilder {
+  private source?: PgTable | Subquery | PgViewBase | SQL;
+  private keys: GroupKey[] = [];
+  private measureMap: MeasureMap = {};
+  private nonAggregates: NonAggMap = {};
+  private wrapNonAggWithAnyValue = false;
+  private orderByClauses: ValueExpr[] = [];
+
+  constructor(private db: DuckDBDatabase) {}
+
+  from(source: PgTable | Subquery | PgViewBase | SQL): this {
+    this.source = source;
+    return this;
+  }
+
+  groupBy(keys: GroupKey[]): this {
+    this.keys = keys;
+    return this;
+  }
+
+  measures(measures: MeasureMap): this {
+    this.measureMap = measures;
+    return this;
+  }
+
+  selectNonAggregates(
+    fields: NonAggMap,
+    options: { anyValue?: boolean } = {}
+  ): this {
+    this.nonAggregates = fields;
+    this.wrapNonAggWithAnyValue = options.anyValue ?? false;
+    return this;
+  }
+
+  orderBy(...clauses: ValueExpr[]): this {
+    this.orderByClauses = clauses;
+    return this;
+  }
+
+  build() {
+    if (!this.source) {
+      throw new Error('olap: .from() is required');
+    }
+    if (this.keys.length === 0) {
+      throw new Error('olap: .groupBy() is required');
+    }
+    if (Object.keys(this.measureMap).length === 0) {
+      throw new Error('olap: .measures() is required');
+    }
+
+    const selection: Record<string, ValueExpr> = {};
+
+    this.keys.forEach((key, idx) => {
+      const alias = keyAlias(key, `key_${idx}`);
+      selection[alias] = key;
+    });
+
+    Object.entries(this.nonAggregates).forEach(([alias, expr]) => {
+      selection[alias] = this.wrapNonAggWithAnyValue ? anyValue(expr) : expr;
+    });
+
+    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!)
+      .groupBy(...this.keys);
+
+    if (this.orderByClauses.length > 0) {
+      query = query.orderBy(...this.orderByClauses);
+    }
+
+    return query;
+  }
+
+  run() {
+    return this.build();
+  }
+}
+
+export const olap = (db: DuckDBDatabase) => new OlapBuilder(db);

package/src/operators.ts

@@ -0,0 +1,27 @@
+/**
+ * DuckDB-native array operators. Generate DuckDB-compatible SQL directly
+ * without query rewriting.
+ */
+
+import { sql, type SQL, type SQLWrapper } from 'drizzle-orm';
+
+export function arrayHasAll<T>(
+  column: SQLWrapper,
+  values: T[] | SQLWrapper
+): SQL {
+  return sql`array_has_all(${column}, ${values})`;
+}
+
+export function arrayHasAny<T>(
+  column: SQLWrapper,
+  values: T[] | SQLWrapper
+): SQL {
+  return sql`array_has_any(${column}, ${values})`;
+}
+
+export function arrayContainedBy<T>(
+  column: SQLWrapper,
+  values: T[] | SQLWrapper
+): SQL {
+  return sql`array_has_all(${values}, ${column})`;
+}

package/src/options.ts

@@ -0,0 +1,25 @@
+export type PrepareCacheOption = boolean | number | { size?: number };
+
+export interface PreparedStatementCacheConfig {
+  size: number;
+}
+
+const DEFAULT_PREPARED_CACHE_SIZE = 32;
+
+export function resolvePrepareCacheOption(
+  option?: PrepareCacheOption
+): PreparedStatementCacheConfig | undefined {
+  if (!option) return undefined;
+
+  if (option === true) {
+    return { size: DEFAULT_PREPARED_CACHE_SIZE };
+  }
+
+  if (typeof option === 'number') {
+    const size = Math.max(1, Math.floor(option));
+    return { size };
+  }
+
+  const size = option.size ?? DEFAULT_PREPARED_CACHE_SIZE;
+  return { size: Math.max(1, Math.floor(size)) };
+}

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/select-builder.ts

@@ -0,0 +1,110 @@
+import { is } from 'drizzle-orm/entity';
+import {
+  PgSelectBase,
+  PgSelectBuilder,
+  type CreatePgSelectFromBuilderMode,
+  type SelectedFields,
+  type TableLikeHasEmptySelection,
+} from 'drizzle-orm/pg-core/query-builders';
+import { PgColumn, PgTable, type PgSession } from 'drizzle-orm/pg-core';
+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 { SQL, type ColumnsSelection } from 'drizzle-orm/sql/sql';
+import { aliasFields } from './sql/selection.ts';
+import type { DuckDBDialect } from './dialect.ts';
+import { getTableColumns, type DrizzleTypeError } from 'drizzle-orm/utils';
+
+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: DuckDBDialect;
+    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;
+  }
+
+  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)) {
+      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);
+    }
+
+    return new PgSelectBase({
+      table: src,
+      fields,
+      isPartialSelect,
+      session: this._session,
+      dialect: this._dialect,
+      withList: this._withList,
+      distinct: this._distinct,
+    }) as any;
+  }
+}