@leonardovida-md/drizzle-neo-duckdb 1.0.0

package/src/migrator.ts

@@ -0,0 +1,25 @@
+import type { MigrationConfig } from 'drizzle-orm/migrator';
+import { readMigrationFiles } from 'drizzle-orm/migrator';
+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: DuckDbMigrationConfig
+) {
+  const migrationConfig: MigrationConfig =
+    typeof config === 'string'
+      ? { migrationsFolder: config }
+      : config;
+
+  const migrations = readMigrationFiles(migrationConfig);
+
+  await db.dialect.migrate(
+    migrations,
+    // Need to work around omitted internal types from drizzle...
+    db.session as unknown as PgSession,
+    migrationConfig
+  );
+}

package/src/select-builder.ts

@@ -0,0 +1,114 @@
+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;
+  }
+}

package/src/session.ts

@@ -0,0 +1,274 @@
+import { entityKind } from 'drizzle-orm/entity';
+import type { Logger } from 'drizzle-orm/logger';
+import { NoopLogger } from 'drizzle-orm/logger';
+import { PgTransaction } from 'drizzle-orm/pg-core';
+import type { SelectedFieldsOrdered } from 'drizzle-orm/pg-core/query-builders/select.types';
+import type {
+  PgTransactionConfig,
+  PreparedQueryConfig,
+  PgQueryResultHKT,
+} from 'drizzle-orm/pg-core/session';
+import { PgPreparedQuery, PgSession } from 'drizzle-orm/pg-core/session';
+import type {
+  RelationalSchemaConfig,
+  TablesRelationalConfig,
+} from 'drizzle-orm/relations';
+import { fillPlaceholders, type Query, SQL, sql } from 'drizzle-orm/sql/sql';
+import type { Assume } from 'drizzle-orm/utils';
+import { adaptArrayOperators } from './sql/query-rewriters.ts';
+import { mapResultRow } from './sql/result-mapper.ts';
+import type { DuckDBDialect } from './dialect.ts';
+import { TransactionRollbackError } from 'drizzle-orm/errors';
+import type { DuckDBClientLike, RowData } from './client.ts';
+import { executeOnClient, prepareParams } from './client.ts';
+
+export type { DuckDBClientLike, RowData } from './client.ts';
+
+export class DuckDBPreparedQuery<
+  T extends PreparedQueryConfig
+> extends PgPreparedQuery<T> {
+  static readonly [entityKind]: string = 'DuckDBPreparedQuery';
+
+  constructor(
+    private client: DuckDBClientLike,
+    private dialect: DuckDBDialect,
+    private queryString: string,
+    private params: unknown[],
+    private logger: Logger,
+    private fields: SelectedFieldsOrdered | undefined,
+    private _isResponseInArrayMode: boolean,
+    private customResultMapper: ((rows: unknown[][]) => T['execute']) | undefined,
+    private rewriteArrays: boolean,
+    private rejectStringArrayLiterals: boolean,
+    private warnOnStringArrayLiteral?: (sql: string) => void
+  ) {
+    super({ sql: queryString, params });
+  }
+
+  async execute(
+    placeholderValues: Record<string, unknown> | undefined = {}
+  ): Promise<T['execute']> {
+    this.dialect.assertNoPgJsonColumns();
+    const params = prepareParams(
+      fillPlaceholders(this.params, placeholderValues),
+      {
+        rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+        warnOnStringArrayLiteral: this.warnOnStringArrayLiteral
+          ? () => this.warnOnStringArrayLiteral?.(this.queryString)
+          : undefined,
+      }
+    );
+    const rewrittenQuery = this.rewriteArrays
+      ? adaptArrayOperators(this.queryString)
+      : this.queryString;
+
+    if (this.rewriteArrays && rewrittenQuery !== this.queryString) {
+      this.logger.logQuery(
+        `[duckdb] original query before array rewrite: ${this.queryString}`,
+        params
+      );
+    }
+
+    this.logger.logQuery(rewrittenQuery, params);
+
+    const {
+      fields,
+      joinsNotNullableMap,
+      customResultMapper,
+    } = this as typeof this & { joinsNotNullableMap?: Record<string, boolean> };
+
+    const rows = await executeOnClient(
+      this.client,
+      rewrittenQuery,
+      params
+    );
+
+    if (rows.length === 0 || !fields) {
+      return rows as T['execute'];
+    }
+
+    const rowValues = rows.map((row) => Object.values(row));
+
+    return customResultMapper
+      ? customResultMapper(rowValues)
+      : rowValues.map((row) =>
+          mapResultRow<T['execute']>(fields, row, joinsNotNullableMap)
+        );
+  }
+
+  all(
+    placeholderValues: Record<string, unknown> | undefined = {}
+  ): Promise<T['all']> {
+    return this.execute(placeholderValues);
+  }
+
+  isResponseInArrayMode(): boolean {
+    return this._isResponseInArrayMode;
+  }
+}
+
+export interface DuckDBSessionOptions {
+  logger?: Logger;
+  rewriteArrays?: boolean;
+  rejectStringArrayLiterals?: boolean;
+}
+
+export class DuckDBSession<
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends TablesRelationalConfig = Record<string, never>
+> extends PgSession<DuckDBQueryResultHKT, TFullSchema, TSchema> {
+  static readonly [entityKind]: string = 'DuckDBSession';
+
+  protected override dialect: DuckDBDialect;
+  private logger: Logger;
+  private rewriteArrays: boolean;
+  private rejectStringArrayLiterals: boolean;
+  private hasWarnedArrayLiteral = false;
+
+  constructor(
+    private client: DuckDBClientLike,
+    dialect: DuckDBDialect,
+    private schema: RelationalSchemaConfig<TSchema> | undefined,
+    private options: DuckDBSessionOptions = {}
+  ) {
+    super(dialect);
+    this.dialect = dialect;
+    this.logger = options.logger ?? new NoopLogger();
+    this.rewriteArrays = options.rewriteArrays ?? true;
+    this.rejectStringArrayLiterals = options.rejectStringArrayLiterals ?? false;
+  }
+
+  prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
+    query: Query,
+    fields: SelectedFieldsOrdered | undefined,
+    name: string | undefined,
+    isResponseInArrayMode: boolean,
+    customResultMapper?: (rows: unknown[][]) => T['execute']
+  ): PgPreparedQuery<T> {
+    void name; // DuckDB doesn't support prepared statement names but the signature must match.
+    return new DuckDBPreparedQuery(
+      this.client,
+      this.dialect,
+      query.sql,
+      query.params,
+      this.logger,
+      fields,
+      isResponseInArrayMode,
+      customResultMapper,
+      this.rewriteArrays,
+      this.rejectStringArrayLiterals,
+      this.rejectStringArrayLiterals ? undefined : this.warnOnStringArrayLiteral
+    );
+  }
+
+  override async transaction<T>(
+    transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
+  ): Promise<T> {
+    const session = new DuckDBSession(
+      this.client,
+      this.dialect,
+      this.schema,
+      this.options
+    );
+
+    const tx = new DuckDBTransaction<TFullSchema, TSchema>(
+      this.dialect,
+      session,
+      this.schema
+    );
+
+    await tx.execute(sql`BEGIN TRANSACTION;`);
+
+    try {
+      const result = await transaction(tx);
+      await tx.execute(sql`commit`);
+      return result;
+    } catch (error) {
+      await tx.execute(sql`rollback`);
+      throw error;
+    }
+  }
+
+  private warnOnStringArrayLiteral = (query: string) => {
+    if (this.hasWarnedArrayLiteral) {
+      return;
+    }
+    this.hasWarnedArrayLiteral = true;
+    this.logger.logQuery(
+      `[duckdb] ${arrayLiteralWarning}\nquery: ${query}`,
+      []
+    );
+  };
+}
+
+type PgTransactionInternals<
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends TablesRelationalConfig = Record<string, never>
+> = {
+  dialect: DuckDBDialect;
+  session: DuckDBSession<TFullSchema, TSchema>;
+};
+
+type DuckDBTransactionWithInternals<
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends TablesRelationalConfig = Record<string, never>
+> = PgTransactionInternals<TFullSchema, TSchema> &
+  DuckDBTransaction<TFullSchema, TSchema>;
+
+export class DuckDBTransaction<
+  TFullSchema extends Record<string, unknown>,
+  TSchema extends TablesRelationalConfig
+> extends PgTransaction<DuckDBQueryResultHKT, TFullSchema, TSchema> {
+  static readonly [entityKind]: string = 'DuckDBTransaction';
+
+  rollback(): never {
+    throw new TransactionRollbackError();
+  }
+
+  getTransactionConfigSQL(config: PgTransactionConfig): SQL {
+    const chunks: string[] = [];
+    if (config.isolationLevel) {
+      chunks.push(`isolation level ${config.isolationLevel}`);
+    }
+    if (config.accessMode) {
+      chunks.push(config.accessMode);
+    }
+    if (typeof config.deferrable === 'boolean') {
+      chunks.push(config.deferrable ? 'deferrable' : 'not deferrable');
+    }
+    return sql.raw(chunks.join(' '));
+  }
+
+  setTransaction(config: PgTransactionConfig): Promise<void> {
+    type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    return (this as unknown as Tx).session.execute(
+      sql`set transaction ${this.getTransactionConfigSQL(config)}`
+    );
+  }
+
+  override async transaction<T>(
+    transaction: (tx: DuckDBTransaction<TFullSchema, TSchema>) => Promise<T>
+  ): Promise<T> {
+    type Tx = DuckDBTransactionWithInternals<TFullSchema, TSchema>;
+    const nestedTx = new DuckDBTransaction<TFullSchema, TSchema>(
+      (this as unknown as Tx).dialect,
+      (this as unknown as Tx).session,
+      this.schema,
+      this.nestedIndex + 1
+    );
+
+    return transaction(nestedTx);
+  }
+}
+
+export type GenericRowData<T extends RowData = RowData> = T;
+
+export type GenericTableData<T = RowData> = T[];
+
+const arrayLiteralWarning =
+  'Received a stringified Postgres-style array literal. Use duckDbList()/duckDbArray() or pass native arrays instead. You can also set rejectStringArrayLiterals=true to throw.';
+
+
+export interface DuckDBQueryResultHKT extends PgQueryResultHKT {
+  type: GenericTableData<Assume<this['row'], RowData>>;
+}

package/src/sql/query-rewriters.ts

@@ -0,0 +1,147 @@
+const selectionRegex = /select\s+(.+)\s+from/i;
+const tableIdPropSelectionRegex = new RegExp(
+  [
+    `("(.+)"\\."(.+)")`, // table identifier + property
+    `(\\s+as\\s+'?(.+?)'?\\.'?(.+?)'?)?`, // optional AS clause
+  ].join(''),
+  'i'
+);
+const noTableIdPropSelectionRegex = /"(.+)"(\s+as\s+'?\1'?)?/i;
+
+export function adaptArrayOperators(query: string): string {
+  type ArrayOperator = {
+    token: '@>' | '<@' | '&&';
+    fn: 'array_has_all' | 'array_has_any';
+    swap?: boolean;
+  };
+
+  const operators: ArrayOperator[] = [
+    { token: '@>', fn: 'array_has_all' },
+    { token: '<@', fn: 'array_has_all', swap: true },
+    { token: '&&', fn: 'array_has_any' },
+  ];
+
+  const isWhitespace = (char: string | undefined) =>
+    char !== undefined && /\s/.test(char);
+
+  const walkLeft = (source: string, start: number): [number, string] => {
+    let idx = start;
+    while (idx >= 0 && isWhitespace(source[idx])) {
+      idx--;
+    }
+
+    let depth = 0;
+    let inString = false;
+    for (; idx >= 0; idx--) {
+      const ch = source[idx];
+      if (ch === "'" && source[idx - 1] !== '\\') {
+        inString = !inString;
+      }
+      if (inString) continue;
+      if (ch === ')' || ch === ']') {
+        depth++;
+      } else if (ch === '(' || ch === '[') {
+        depth--;
+        if (depth < 0) {
+          return [idx + 1, source.slice(idx + 1, start + 1)];
+        }
+      } else if (depth === 0 && isWhitespace(ch)) {
+        return [idx + 1, source.slice(idx + 1, start + 1)];
+      }
+    }
+    return [0, source.slice(0, start + 1)];
+  };
+
+  const walkRight = (source: string, start: number): [number, string] => {
+    let idx = start;
+    while (idx < source.length && isWhitespace(source[idx])) {
+      idx++;
+    }
+
+    let depth = 0;
+    let inString = false;
+    for (; idx < source.length; idx++) {
+      const ch = source[idx];
+      if (ch === "'" && source[idx - 1] !== '\\') {
+        inString = !inString;
+      }
+      if (inString) continue;
+      if (ch === '(' || ch === '[') {
+        depth++;
+      } else if (ch === ')' || ch === ']') {
+        depth--;
+        if (depth < 0) {
+          return [idx, source.slice(start, idx)];
+        }
+      } else if (depth === 0 && isWhitespace(ch)) {
+        return [idx, source.slice(start, idx)];
+      }
+    }
+    return [source.length, source.slice(start)];
+  };
+
+  let rewritten = query;
+  for (const { token, fn, swap } of operators) {
+    let idx = rewritten.indexOf(token);
+    while (idx !== -1) {
+      const [leftStart, leftExpr] = walkLeft(rewritten, idx - 1);
+      const [rightEnd, rightExpr] = walkRight(
+        rewritten,
+        idx + token.length
+      );
+
+      const left = leftExpr.trim();
+      const right = rightExpr.trim();
+
+      const replacement = `${fn}(${swap ? right : left}, ${
+        swap ? left : right
+      })`;
+
+      rewritten =
+        rewritten.slice(0, leftStart) +
+        replacement +
+        rewritten.slice(rightEnd);
+
+      idx = rewritten.indexOf(token, leftStart + replacement.length);
+    }
+  }
+
+  return rewritten;
+}
+
+export function queryAdapter(query: string): string {
+  const selection = selectionRegex.exec(query);
+
+  if (selection?.length !== 2) {
+    return query;
+  }
+
+  const fields = selection[1]
+    .split(',')
+    .map((field) => {
+      const trimmedField = field.trim();
+      const tableProp = tableIdPropSelectionRegex.exec(trimmedField);
+      if (tableProp) {
+        const [, identifier, table, column, , aliasTable, aliasColumn] =
+          tableProp;
+
+        const asAlias = `'${aliasTable ?? table}.${aliasColumn ?? column}'`;
+        if (tableProp[4]) {
+          return trimmedField.replace(tableProp[4], ` as ${asAlias}`);
+        }
+        return `${identifier} as ${asAlias}`;
+      }
+
+      const noTableProp = noTableIdPropSelectionRegex.exec(trimmedField);
+      if (noTableProp) {
+        const [, column, alias] = noTableProp;
+        const asAlias = ` as '${column}'`;
+        return alias ? trimmedField.replace(alias, asAlias) : `${trimmedField}${asAlias}`;
+      }
+
+      return trimmedField;
+    })
+    .filter(Boolean) as string[];
+
+  return query.replace(selection[1], fields.join(', '));
+}