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

package/src/columns.ts

@@ -7,11 +7,13 @@ import {
   wrapMap,
   wrapBlob,
   wrapJson,
+  wrapTimestamp,
   type ListValueWrapper,
   type ArrayValueWrapper,
   type MapValueWrapper,
   type BlobValueWrapper,
   type JsonValueWrapper,
+  type TimestampValueWrapper,
 } from './value-wrappers-core.ts';
 
 type IntColType =
@@ -63,7 +65,7 @@ type StructColType = `STRUCT (${string})`;
 
 type Primitive = AnyColType | ListColType | ArrayColType | StructColType;
 
-function coerceArrayString(value: string): unknown[] | undefined {
+export function coerceArrayString(value: string): unknown[] | undefined {
   const trimmed = value.trim();
   if (!trimmed) {
     return [];
@@ -86,7 +88,7 @@ function coerceArrayString(value: string): unknown[] | undefined {
   return undefined;
 }
 
-function formatLiteral(value: unknown, typeHint?: string): string {
+export function formatLiteral(value: unknown, typeHint?: string): string {
   if (value === null || value === undefined) {
     return 'NULL';
   }
@@ -123,7 +125,7 @@ function formatLiteral(value: unknown, typeHint?: string): string {
   return `'${escaped}'`;
 }
 
-function buildListLiteral(values: unknown[], elementType?: string): SQL {
+export function buildListLiteral(values: unknown[], elementType?: string): SQL {
   if (values.length === 0) {
     return sql`[]`;
   }
@@ -135,7 +137,7 @@ function buildListLiteral(values: unknown[], elementType?: string): SQL {
   return sql`list_value(${sql.join(chunks, sql.raw(', '))})`;
 }
 
-function buildStructLiteral(
+export function buildStructLiteral(
   value: Record<string, unknown>,
   schema?: Record<string, Primitive>
 ): SQL {
@@ -154,7 +156,7 @@ function buildStructLiteral(
   return sql`struct_pack(${sql.join(parts, sql.raw(', '))})`;
 }
 
-function buildMapLiteral(
+export function buildMapLiteral(
   value: Record<string, unknown>,
   valueType?: string
 ): SQL {
@@ -188,11 +190,11 @@ export const duckDbList = <TData = unknown>(
       }
       if (typeof value === 'string') {
         const parsed = coerceArrayString(value);
-        if (parsed) {
+        if (parsed !== undefined) {
           return parsed as TData[];
         }
       }
-      return [] as TData[];
+      return value as unknown as TData[];
     },
   })(name);
 
@@ -219,11 +221,11 @@ export const duckDbArray = <TData = unknown>(
       }
       if (typeof value === 'string') {
         const parsed = coerceArrayString(value);
-        if (parsed) {
+        if (parsed !== undefined) {
           return parsed as TData[];
         }
       }
-      return [] as TData[];
+      return value as unknown as TData[];
     },
   })(name);
 
@@ -355,12 +357,35 @@ interface TimestampOptions {
   withTimezone?: boolean;
   mode?: TimestampMode;
   precision?: number;
+  bindMode?: 'auto' | 'bind' | 'literal';
+}
+
+function shouldBindTimestamp(options: TimestampOptions): boolean {
+  const bindMode = options.bindMode ?? 'auto';
+  if (bindMode === 'bind') return true;
+  if (bindMode === 'literal') return false;
+
+  const isBun =
+    typeof process !== 'undefined' &&
+    typeof process.versions?.bun !== 'undefined';
+  if (isBun) return false;
+
+  const forceLiteral =
+    typeof process !== 'undefined'
+      ? process.env.DRIZZLE_DUCKDB_FORCE_LITERAL_TIMESTAMPS
+      : undefined;
+
+  if (forceLiteral && forceLiteral !== '0') {
+    return false;
+  }
+
+  return true;
 }
 
 export const duckDbTimestamp = (name: string, options: TimestampOptions = {}) =>
   customType<{
     data: Date | string;
-    driverData: SQL | string | Date;
+    driverData: SQL | string | Date | TimestampValueWrapper;
   }>({
     dataType() {
       if (options.withTimezone) {
@@ -369,14 +394,36 @@ export const duckDbTimestamp = (name: string, options: TimestampOptions = {}) =>
       const precision = options.precision ? `(${options.precision})` : '';
       return `TIMESTAMP${precision}`;
     },
-    toDriver(value: Date | string) {
-      // Use SQL literals for timestamps due to Bun/DuckDB bigint binding issues
+    toDriver(
+      value: Date | string
+    ): SQL | string | Date | TimestampValueWrapper {
+      if (shouldBindTimestamp(options)) {
+        return wrapTimestamp(
+          value,
+          options.withTimezone ?? false,
+          options.precision
+        );
+      }
+
       const iso = value instanceof Date ? value.toISOString() : value;
       const normalized = iso.replace('T', ' ').replace('Z', '+00');
       const typeKeyword = options.withTimezone ? 'TIMESTAMPTZ' : 'TIMESTAMP';
       return sql.raw(`${typeKeyword} '${normalized}'`);
     },
-    fromDriver(value: Date | string | SQL) {
+    fromDriver(value: Date | string | SQL | TimestampValueWrapper) {
+      if (
+        value &&
+        typeof value === 'object' &&
+        'kind' in value &&
+        (value as TimestampValueWrapper).kind === 'timestamp'
+      ) {
+        const wrapped = value as TimestampValueWrapper;
+        return wrapped.data instanceof Date
+          ? wrapped.data
+          : typeof wrapped.data === 'number' || typeof wrapped.data === 'bigint'
+            ? new Date(Number(wrapped.data) / 1000)
+            : wrapped.data;
+      }
       if (options.mode === 'string') {
         if (value instanceof Date) {
           return value.toISOString().replace('T', ' ').replace('Z', '+00');

package/src/dialect.ts

@@ -19,18 +19,64 @@ import {
   type QueryTypingsValue,
 } from 'drizzle-orm';
 
+const enum SavepointSupport {
+  Unknown = 0,
+  Yes = 1,
+  No = 2,
+}
+
 export class DuckDBDialect extends PgDialect {
   static readonly [entityKind]: string = 'DuckDBPgDialect';
+  // Track if PG JSON columns were detected during the current query preparation.
+  // Reset before each query via DuckDBSession to keep detection per-query.
   private hasPgJsonColumn = false;
+  // Track savepoint support per-dialect instance to avoid cross-contamination
+  // when multiple database connections with different capabilities exist.
+  private savepointsSupported: SavepointSupport = SavepointSupport.Unknown;
+
+  /**
+   * Reset the PG JSON detection flag. Should be called before preparing a new query.
+   */
+  resetPgJsonFlag(): void {
+    this.hasPgJsonColumn = false;
+  }
+
+  /**
+   * Mark that a PG JSON/JSONB column was detected during query preparation.
+   */
+  markPgJsonDetected(): void {
+    this.hasPgJsonColumn = true;
+  }
 
   assertNoPgJsonColumns(): void {
     if (this.hasPgJsonColumn) {
       throw new Error(
-        'Pg JSON/JSONB columns are not supported in DuckDB. Replace them with duckDbJson() to use DuckDB’s native JSON type.'
+        "Pg JSON/JSONB columns are not supported in DuckDB. Replace them with duckDbJson() to use DuckDB's native JSON type."
       );
     }
   }
 
+  /**
+   * Check if savepoints are known to be unsupported for this dialect instance.
+   */
+  areSavepointsUnsupported(): boolean {
+    return this.savepointsSupported === SavepointSupport.No;
+  }
+
+  /**
+   * Mark that savepoints are supported for this dialect instance.
+   */
+  markSavepointsSupported(): void {
+    this.savepointsSupported = SavepointSupport.Yes;
+  }
+
+  /**
+   * Mark that savepoints are not supported for this dialect instance.
+   */
+  markSavepointsUnsupported(): void {
+    this.savepointsSupported = SavepointSupport.No;
+  }
+
   override async migrate(
     migrations: MigrationMeta[],
     session: PgSession,
@@ -117,8 +163,10 @@ export class DuckDBDialect extends PgDialect {
     encoder: DriverValueEncoder<unknown, unknown>
   ): QueryTypingsValue {
     if (is(encoder, PgJsonb) || is(encoder, PgJson)) {
-      this.hasPgJsonColumn = true;
-      return 'none';
+      this.markPgJsonDetected();
+      throw new Error(
+        "Pg JSON/JSONB columns are not supported in DuckDB. Replace them with duckDbJson() to use DuckDB's native JSON type."
+      );
     } else if (is(encoder, PgNumeric)) {
       return 'decimal';
     } else if (is(encoder, PgTime)) {

package/src/driver.ts

@@ -23,19 +23,32 @@ import { DuckDBSession } from './session.ts';
 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 type {
+  ExecuteBatchesRawChunk,
+  ExecuteInBatchesOptions,
+  RowData,
+} from './client.ts';
+import { closeClientConnection, isPool } from './client.ts';
 import {
   createDuckDBConnectionPool,
   resolvePoolSize,
   type DuckDBPoolConfig,
   type PoolPreset,
 } from './pool.ts';
+import {
+  resolvePrepareCacheOption,
+  resolveRewriteArraysOption,
+  type PreparedStatementCacheConfig,
+  type PrepareCacheOption,
+  type RewriteArraysMode,
+  type RewriteArraysOption,
+} from './options.ts';
 
 export interface PgDriverOptions {
   logger?: Logger;
-  rewriteArrays?: boolean;
+  rewriteArrays?: RewriteArraysMode;
   rejectStringArrayLiterals?: boolean;
+  prepareCache?: PreparedStatementCacheConfig;
 }
 
 export class DuckDBDriver {
@@ -52,8 +65,9 @@ export class DuckDBDriver {
   ): DuckDBSession<Record<string, unknown>, TablesRelationalConfig> {
     return new DuckDBSession(this.client, this.dialect, schema, {
       logger: this.options.logger,
-      rewriteArrays: this.options.rewriteArrays,
+      rewriteArrays: this.options.rewriteArrays ?? 'auto',
       rejectStringArrayLiterals: this.options.rejectStringArrayLiterals,
+      prepareCache: this.options.prepareCache,
     });
   }
 }
@@ -69,8 +83,9 @@ export interface DuckDBConnectionConfig {
 export interface DuckDBDrizzleConfig<
   TSchema extends Record<string, unknown> = Record<string, never>,
 > extends DrizzleConfig<TSchema> {
-  rewriteArrays?: boolean;
+  rewriteArrays?: RewriteArraysOption;
   rejectStringArrayLiterals?: boolean;
+  prepareCache?: PrepareCacheOption;
   /** Pool configuration. Use preset name, size config, or false to disable. */
   pool?: DuckDBPoolConfig | PoolPreset | false;
 }
@@ -111,6 +126,8 @@ function createFromClient<
   instance?: DuckDBInstance
 ): DuckDBDatabase<TSchema, ExtractTablesWithRelations<TSchema>> {
   const dialect = new DuckDBDialect();
+  const rewriteArraysMode = resolveRewriteArraysOption(config.rewriteArrays);
+  const prepareCache = resolvePrepareCacheOption(config.prepareCache);
 
   const logger =
     config.logger === true ? new DefaultLogger() : config.logger || undefined;
@@ -131,8 +148,9 @@ function createFromClient<
 
   const driver = new DuckDBDriver(client, dialect, {
     logger,
-    rewriteArrays: config.rewriteArrays,
+    rewriteArrays: rewriteArraysMode,
     rejectStringArrayLiterals: config.rejectStringArrayLiterals,
+    prepareCache,
   });
   const session = driver.createSession(schema);
 
@@ -290,7 +308,20 @@ export class DuckDBDatabase<
     if (isPool(this.$client) && this.$client.close) {
       await this.$client.close();
     }
-    // Note: DuckDBInstance doesn't have a close method in the current API
+    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>;
@@ -317,6 +348,13 @@ export class DuckDBDatabase<
     return this.session.executeBatches<T>(query, options);
   }
 
+  executeBatchesRaw(
+    query: SQL,
+    options: ExecuteInBatchesOptions = {}
+  ): AsyncGenerator<ExecuteBatchesRawChunk, void, void> {
+    return this.session.executeBatchesRaw(query, options);
+  }
+
   executeArrow(query: SQL): Promise<unknown> {
     return this.session.executeArrow(query);
   }

package/src/index.ts

@@ -7,3 +7,4 @@ export * from './client.ts';
 export * from './pool.ts';
 export * from './olap.ts';
 export * from './value-wrappers.ts';
+export * from './options.ts';

package/src/introspect.ts

@@ -535,7 +535,7 @@ function emitColumn(
   return builder;
 }
 
-function buildDefault(defaultValue: string | null): string {
+export function buildDefault(defaultValue: string | null): string {
   if (!defaultValue) {
     return '';
   }
@@ -707,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)})` };
@@ -788,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 }> = [];
@@ -813,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) {
@@ -822,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 = '';
@@ -847,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/options.ts

@@ -0,0 +1,40 @@
+export type RewriteArraysMode = 'auto' | 'always' | 'never';
+
+export type RewriteArraysOption = boolean | RewriteArraysMode;
+
+const DEFAULT_REWRITE_ARRAYS_MODE: RewriteArraysMode = 'auto';
+
+export function resolveRewriteArraysOption(
+  value?: RewriteArraysOption
+): RewriteArraysMode {
+  if (value === undefined) return DEFAULT_REWRITE_ARRAYS_MODE;
+  if (value === true) return 'auto';
+  if (value === false) return 'never';
+  return value;
+}
+
+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)) };
+}