@duckdbfan/drizzle-duckdb 0.0.7 → 1.3.1

package/src/columns.ts

@@ -1,5 +1,20 @@
 import { sql, type SQL } from 'drizzle-orm';
+import type { SQLWrapper } from 'drizzle-orm/sql/sql';
 import { customType } from 'drizzle-orm/pg-core';
+import {
+  wrapList,
+  wrapArray,
+  wrapMap,
+  wrapBlob,
+  wrapJson,
+  wrapTimestamp,
+  type ListValueWrapper,
+  type ArrayValueWrapper,
+  type MapValueWrapper,
+  type BlobValueWrapper,
+  type JsonValueWrapper,
+  type TimestampValueWrapper,
+} from './value-wrappers-core.ts';
 
 type IntColType =
   | 'SMALLINT'
@@ -46,70 +61,200 @@ type AnyColType =
 
 type ListColType = `${AnyColType}[]`;
 type ArrayColType = `${AnyColType}[${number}]`;
-/**
- * @example
- * const structColType: StructColType = 'STRUCT (name: STRING, age: INT)';
- */
 type StructColType = `STRUCT (${string})`;
 
+type Primitive = AnyColType | ListColType | ArrayColType | StructColType;
+
+export function coerceArrayString(value: string): unknown[] | undefined {
+  const trimmed = value.trim();
+  if (!trimmed) {
+    return [];
+  }
+  if (trimmed.startsWith('[')) {
+    try {
+      return JSON.parse(trimmed) as unknown[];
+    } catch {
+      return undefined;
+    }
+  }
+  if (trimmed.startsWith('{') && trimmed.endsWith('}')) {
+    try {
+      const json = trimmed.replace(/{/g, '[').replace(/}/g, ']');
+      return JSON.parse(json) as unknown[];
+    } catch {
+      return undefined;
+    }
+  }
+  return undefined;
+}
+
+export function formatLiteral(value: unknown, typeHint?: string): string {
+  if (value === null || value === undefined) {
+    return 'NULL';
+  }
+
+  const upperType = typeHint?.toUpperCase() ?? '';
+  if (value instanceof Date) {
+    return `'${value.toISOString()}'`;
+  }
+
+  if (typeof value === 'number' || typeof value === 'bigint') {
+    return value.toString();
+  }
+
+  if (typeof value === 'boolean') {
+    return value ? 'TRUE' : 'FALSE';
+  }
+
+  const str =
+    typeof value === 'string'
+      ? value
+      : (JSON.stringify(value) ?? String(value));
+
+  const escaped = str.replace(/'/g, "''");
+  // Simple quoting based on hint.
+  if (
+    upperType.includes('CHAR') ||
+    upperType.includes('TEXT') ||
+    upperType.includes('STRING') ||
+    upperType.includes('VARCHAR')
+  ) {
+    return `'${escaped}'`;
+  }
+
+  return `'${escaped}'`;
+}
+
+export function buildListLiteral(values: unknown[], elementType?: string): SQL {
+  if (values.length === 0) {
+    return sql`[]`;
+  }
+  const chunks = values.map((v) =>
+    typeof v === 'object' && !Array.isArray(v)
+      ? sql`${v as SQLWrapper}`
+      : sql.raw(formatLiteral(v, elementType))
+  );
+  return sql`list_value(${sql.join(chunks, sql.raw(', '))})`;
+}
+
+export function buildStructLiteral(
+  value: Record<string, unknown>,
+  schema?: Record<string, Primitive>
+): SQL {
+  const parts = Object.entries(value).map(([key, val]) => {
+    const typeHint = schema?.[key];
+    if (Array.isArray(val)) {
+      const inner =
+        typeof typeHint === 'string' && typeHint.endsWith('[]')
+          ? typeHint.slice(0, -2)
+          : undefined;
+
+      return sql`${sql.identifier(key)} := ${buildListLiteral(val, inner)}`;
+    }
+    return sql`${sql.identifier(key)} := ${val}`;
+  });
+  return sql`struct_pack(${sql.join(parts, sql.raw(', '))})`;
+}
+
+export function buildMapLiteral(
+  value: Record<string, unknown>,
+  valueType?: string
+): SQL {
+  const keys = Object.keys(value);
+  const vals = Object.values(value);
+  const keyList = buildListLiteral(keys, 'TEXT');
+  const valList = buildListLiteral(
+    vals,
+    valueType?.endsWith('[]') ? valueType.slice(0, -2) : valueType
+  );
+  return sql`map(${keyList}, ${valList})`;
+}
+
+export const duckDbList = <TData = unknown>(
+  name: string,
+  elementType: AnyColType
+) =>
+  customType<{
+    data: TData[];
+    driverData: ListValueWrapper | unknown[] | string;
+  }>({
+    dataType() {
+      return `${elementType}[]`;
+    },
+    toDriver(value: TData[]): ListValueWrapper {
+      return wrapList(value, elementType);
+    },
+    fromDriver(value: unknown[] | string | ListValueWrapper): TData[] {
+      if (Array.isArray(value)) {
+        return value as TData[];
+      }
+      if (typeof value === 'string') {
+        const parsed = coerceArrayString(value);
+        if (parsed !== undefined) {
+          return parsed as TData[];
+        }
+      }
+      return value as unknown as TData[];
+    },
+  })(name);
+
+export const duckDbArray = <TData = unknown>(
+  name: string,
+  elementType: AnyColType,
+  fixedLength?: number
+) =>
+  customType<{
+    data: TData[];
+    driverData: ArrayValueWrapper | unknown[] | string;
+  }>({
+    dataType() {
+      return fixedLength
+        ? `${elementType}[${fixedLength}]`
+        : `${elementType}[]`;
+    },
+    toDriver(value: TData[]): ArrayValueWrapper {
+      return wrapArray(value, elementType, fixedLength);
+    },
+    fromDriver(value: unknown[] | string | ArrayValueWrapper): TData[] {
+      if (Array.isArray(value)) {
+        return value as TData[];
+      }
+      if (typeof value === 'string') {
+        const parsed = coerceArrayString(value);
+        if (parsed !== undefined) {
+          return parsed as TData[];
+        }
+      }
+      return value as unknown as TData[];
+    },
+  })(name);
+
 export const duckDbMap = <TData extends Record<string, any>>(
   name: string,
   valueType: AnyColType | ListColType | ArrayColType
 ) =>
-  customType<{ data: TData; driverData: string }>({
+  customType<{ data: TData; driverData: MapValueWrapper | TData }>({
     dataType() {
-      console.log('dataType');
       return `MAP (STRING, ${valueType})`;
     },
     toDriver(value: TData) {
-      console.log('toDriver');
-      // todo: more sophisticated encoding based on data type
-      const valueFormatter = (value: any) => {
-        if (['STRING', 'TEXT', 'VARCHAR'].includes(valueType)) {
-          return `'${value}'`;
-        }
-
-        return JSON.stringify(value);
-      };
-
-      const values = Object.entries(value).map(([key, value]) => {
-        return sql.raw(`'${key}': ${valueFormatter(value)}`);
-      });
-
-      const sqlChunks: SQL[] = [];
-
-      for (const value of values) {
-        sqlChunks.push(value);
+      // Use SQL literals for empty maps due to DuckDB type inference issues
+      // with mapValue() when there are no entries to infer types from
+      if (Object.keys(value).length === 0) {
+        return buildMapLiteral(value, valueType);
       }
-
-      return sql`MAP {${sql.join(sqlChunks, sql.raw(', '))}}`;
+      return wrapMap(value, valueType);
     },
-    // ! this won't actually ever work because of how map values are returned
-    fromDriver(value: string): TData {
-      console.log('fromDriver');
-      // todo: more sophisticated decoding based on data type
-
-      const replacedValue = value.replaceAll(
-        /(?:^{)?([^=]+?)=(.+)(?:}$)?/g,
-        '"$1":"$2"'
-      );
-      const formattedValue = `{${replacedValue}}`;
-
-      const valueObj = JSON.parse(formattedValue);
-
-      return Object.fromEntries(
-        Object.entries(valueObj).map(([key, value]) => {
-          return [key, JSON.parse(value as string)];
-        })
-      ) as TData;
+    fromDriver(value: TData | MapValueWrapper): TData {
+      return value as TData;
     },
   })(name);
 
 export const duckDbStruct = <TData extends Record<string, any>>(
   name: string,
-  schema: Record<string, AnyColType | ListColType | ArrayColType>
+  schema: Record<string, Primitive>
 ) =>
-  customType<{ data: TData; driverData: string }>({
+  customType<{ data: TData; driverData: TData }>({
     dataType() {
       const fields = Object.entries(schema).map(
         ([key, type]) => `${key} ${type}`
@@ -118,38 +263,248 @@ export const duckDbStruct = <TData extends Record<string, any>>(
       return `STRUCT (${fields.join(', ')})`;
     },
     toDriver(value: TData) {
-      // todo: more sophisticated encoding based on data type
-      const valueFormatter = (value: any) =>
-        JSON.stringify(value).replaceAll(/(?<!\\)"/g, "'");
-
-      const values = Object.entries(value).map(([key, value]) => {
-        return sql.raw(`'${key}': ${valueFormatter(value)}`);
-      });
-
-      const sqlChunks: SQL[] = [];
-
-      for (const value of values) {
-        sqlChunks.push(value);
+      // Use SQL literals for structs due to DuckDB type inference issues
+      // with nested empty lists
+      return buildStructLiteral(value, schema);
+    },
+    fromDriver(value: TData | string): TData {
+      if (typeof value === 'string') {
+        try {
+          return JSON.parse(value) as TData;
+        } catch {
+          return value as unknown as TData;
+        }
       }
+      return value;
+    },
+  })(name);
 
-      return sql`(SELECT {${sql.join(sqlChunks, sql.raw(', '))}})`;
+/**
+ * JSON column type that wraps values and delays JSON.stringify() to binding time.
+ * This ensures consistent handling with other wrapped types.
+ *
+ * Note: DuckDB stores JSON as VARCHAR internally, so the final binding
+ * is always a stringified JSON value.
+ */
+export const duckDbJson = <TData = unknown>(name: string) =>
+  customType<{ data: TData; driverData: JsonValueWrapper | SQL | string }>({
+    dataType() {
+      return 'JSON';
     },
-    fromDriver(value: string): TData {
-      return value as unknown as TData;
+    toDriver(value: TData): JsonValueWrapper | SQL | string {
+      // Pass through strings directly
+      if (typeof value === 'string') {
+        return value;
+      }
+      // Pass through SQL objects (for raw SQL expressions)
+      if (
+        value !== null &&
+        typeof value === 'object' &&
+        'queryChunks' in (value as Record<string, unknown>)
+      ) {
+        return value as unknown as SQL;
+      }
+      // Wrap non-string values for delayed stringify at binding time
+      return wrapJson(value);
+    },
+    fromDriver(value: SQL | string | JsonValueWrapper) {
+      if (typeof value !== 'string') {
+        return value as unknown as TData;
+      }
+      const trimmed = value.trim();
+      if (!trimmed) {
+        return value as unknown as TData;
+      }
+      try {
+        return JSON.parse(trimmed) as TData;
+      } catch {
+        return value as unknown as TData;
+      }
     },
   })(name);
 
 export const duckDbBlob = customType<{
   data: Buffer;
+  driverData: BlobValueWrapper;
   default: false;
 }>({
   dataType() {
     return 'BLOB';
   },
-  toDriver(value: Buffer) {
-    const hexString = value.toString('hex');
-    return sql`from_hex(${hexString})`;
+  toDriver(value: Buffer): BlobValueWrapper {
+    return wrapBlob(value);
   },
 });
 
-// todo: date/time types
+export const duckDbInet = (name: string) =>
+  customType<{ data: string; driverData: string }>({
+    dataType() {
+      return 'INET';
+    },
+    toDriver(value: string) {
+      return value;
+    },
+  })(name);
+
+export const duckDbInterval = (name: string) =>
+  customType<{ data: string; driverData: string }>({
+    dataType() {
+      return 'INTERVAL';
+    },
+    toDriver(value: string) {
+      return value;
+    },
+  })(name);
+
+type TimestampMode = 'date' | 'string';
+
+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 | TimestampValueWrapper;
+  }>({
+    dataType() {
+      if (options.withTimezone) {
+        return 'TIMESTAMPTZ';
+      }
+      const precision = options.precision ? `(${options.precision})` : '';
+      return `TIMESTAMP${precision}`;
+    },
+    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 | 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');
+        }
+        return typeof value === 'string' ? value : value.toString();
+      }
+      if (value instanceof Date) {
+        return value;
+      }
+      const stringValue = typeof value === 'string' ? value : value.toString();
+      const hasOffset =
+        stringValue.endsWith('Z') || /[+-]\d{2}:?\d{2}$/.test(stringValue);
+      const normalized = hasOffset
+        ? stringValue.replace(' ', 'T')
+        : `${stringValue.replace(' ', 'T')}Z`;
+      return new Date(normalized);
+    },
+  })(name);
+
+export const duckDbDate = (name: string) =>
+  customType<{ data: string | Date; driverData: string | Date }>({
+    dataType() {
+      return 'DATE';
+    },
+    toDriver(value: string | Date) {
+      return value;
+    },
+    fromDriver(value: string | Date) {
+      const str =
+        value instanceof Date ? value.toISOString().slice(0, 10) : value;
+      return str;
+    },
+  })(name);
+
+export const duckDbTime = (name: string) =>
+  customType<{ data: string; driverData: string | bigint }>({
+    dataType() {
+      return 'TIME';
+    },
+    toDriver(value: string) {
+      return value;
+    },
+    fromDriver(value: string | bigint) {
+      if (typeof value === 'bigint') {
+        const totalMillis = Number(value) / 1000;
+        const date = new Date(totalMillis);
+        return date.toISOString().split('T')[1]!.replace('Z', '');
+      }
+      return value;
+    },
+  })(name);
+
+function toListValue(values: (unknown | SQLWrapper)[]): SQL {
+  return buildListLiteral(values);
+}
+
+export function duckDbArrayContains(
+  column: SQLWrapper,
+  values: unknown[] | SQLWrapper
+): SQL {
+  const rhs = Array.isArray(values) ? toListValue(values) : values;
+  return sql`array_has_all(${column}, ${rhs})`;
+}
+
+export function duckDbArrayContained(
+  column: SQLWrapper,
+  values: unknown[] | SQLWrapper
+): SQL {
+  const rhs = Array.isArray(values) ? toListValue(values) : values;
+  return sql`array_has_all(${rhs}, ${column})`;
+}
+
+export function duckDbArrayOverlaps(
+  column: SQLWrapper,
+  values: unknown[] | SQLWrapper
+): SQL {
+  const rhs = Array.isArray(values) ? toListValue(values) : values;
+  return sql`array_has_any(${column}, ${rhs})`;
+}

package/src/dialect.ts

@@ -13,41 +13,118 @@ import {
   PgTimestampString,
   PgUUID,
 } from 'drizzle-orm/pg-core';
-import { DuckDBSession } from './session';
 import {
   sql,
+  SQL,
   type DriverValueEncoder,
   type QueryTypingsValue,
 } from 'drizzle-orm';
+import type { QueryWithTypings } from 'drizzle-orm/sql/sql';
+
+import { transformSQL } from './sql/ast-transformer.ts';
+
+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."
+      );
+    }
+  }
+
+  /**
+   * 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,
-    config: MigrationConfig
+    config: MigrationConfig | string
   ): Promise<void> {
-    const migrationsSchema = config.migrationsSchema ?? 'drizzle';
+    const migrationConfig: MigrationConfig =
+      typeof config === 'string' ? { migrationsFolder: config } : config;
 
-    const migrationsTable = config.migrationsTable ?? '__drizzle_migrations';
+    const migrationsSchema = migrationConfig.migrationsSchema ?? 'drizzle';
+    const migrationsTable =
+      migrationConfig.migrationsTable ?? '__drizzle_migrations';
+    const migrationsSequence = `${migrationsTable}_id_seq`;
+    const legacySequence = 'migrations_pk_seq';
+
+    const escapeIdentifier = (value: string) => value.replace(/"/g, '""');
+    const sequenceLiteral = `"${escapeIdentifier(
+      migrationsSchema
+    )}"."${escapeIdentifier(migrationsSequence)}"`;
 
     const migrationTableCreate = sql`
-			CREATE TABLE IF NOT EXISTS ${sql.identifier(migrationsSchema)}.${sql.identifier(
-      migrationsTable
-    )} (
-				id integer PRIMARY KEY default nextval('migrations_pk_seq'),
-				hash text NOT NULL,
-				created_at bigint
-			)
-		`;
+      CREATE TABLE IF NOT EXISTS ${sql.identifier(migrationsSchema)}.${sql.identifier(
+        migrationsTable
+      )} (
+        id integer PRIMARY KEY default nextval('${sql.raw(sequenceLiteral)}'),
+        hash text NOT NULL,
+        created_at bigint
+      )
+    `;
 
     await session.execute(
-      sql.raw('CREATE SEQUENCE IF NOT EXISTS migrations_pk_seq')
+      sql`CREATE SCHEMA IF NOT EXISTS ${sql.identifier(migrationsSchema)}`
     );
     await session.execute(
-      sql`CREATE SCHEMA IF NOT EXISTS ${sql.identifier(migrationsSchema)}`
+      sql`CREATE SEQUENCE IF NOT EXISTS ${sql.identifier(
+        migrationsSchema
+      )}.${sql.identifier(migrationsSequence)}`
     );
+    if (legacySequence !== migrationsSequence) {
+      await session.execute(
+        sql`CREATE SEQUENCE IF NOT EXISTS ${sql.identifier(
+          migrationsSchema
+        )}.${sql.identifier(legacySequence)}`
+      );
+    }
     await session.execute(migrationTableCreate);
 
     const dbMigrations = await session.all<{
@@ -90,7 +167,10 @@ export class DuckDBDialect extends PgDialect {
     encoder: DriverValueEncoder<unknown, unknown>
   ): QueryTypingsValue {
     if (is(encoder, PgJsonb) || is(encoder, PgJson)) {
-      throw new Error('JSON and JSONB types are not supported in DuckDB');
+      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)) {
@@ -105,4 +185,20 @@ export class DuckDBDialect extends PgDialect {
       return 'none';
     }
   }
+
+  override sqlToQuery(
+    sqlObj: SQL,
+    invokeSource?: 'indexes' | undefined
+  ): QueryWithTypings {
+    // First, let the parent generate the SQL string
+    const result = super.sqlToQuery(sqlObj, invokeSource);
+
+    // Apply AST-based transformations for DuckDB compatibility
+    const transformed = transformSQL(result.sql);
+
+    return {
+      ...result,
+      sql: transformed.sql,
+    };
+  }
 }