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

package/dist/duckdb-introspect.mjs

@@ -1,12 +1,13 @@
 #!/usr/bin/env node
 
 // src/bin/duckdb-introspect.ts
-import { DuckDBInstance } from "@duckdb/node-api";
+import { DuckDBInstance as DuckDBInstance3 } from "@duckdb/node-api";
 import { mkdir, writeFile } from "node:fs/promises";
 import path from "node:path";
 import process from "node:process";
 
 // src/driver.ts
+import { DuckDBInstance as DuckDBInstance2 } from "@duckdb/node-api";
 import { entityKind as entityKind3 } from "drizzle-orm/entity";
 import { DefaultLogger } from "drizzle-orm/logger";
 import { PgDatabase } from "drizzle-orm/pg-core/db";
@@ -23,10 +24,6 @@ import { PgPreparedQuery, PgSession } from "drizzle-orm/pg-core/session";
 import { fillPlaceholders, sql } from "drizzle-orm/sql/sql";
 
 // src/sql/query-rewriters.ts
-var tableIdPropSelectionRegex = new RegExp([
-  `("(.+)"\\."(.+)")`,
-  `(\\s+as\\s+'?(.+?)'?\\.'?(.+?)'?)?`
-].join(""), "i");
 function adaptArrayOperators(query) {
   const operators = [
     { token: "@>", fn: "array_has_all" },
@@ -43,6 +40,8 @@ function adaptArrayOperators(query) {
     let inString = false;
     for (;idx >= 0; idx--) {
       const ch = source[idx];
+      if (ch === undefined)
+        break;
       if (ch === "'" && source[idx - 1] !== "\\") {
         inString = !inString;
       }
@@ -70,6 +69,8 @@ function adaptArrayOperators(query) {
     let inString = false;
     for (;idx < source.length; idx++) {
       const ch = source[idx];
+      if (ch === undefined)
+        break;
       if (ch === "'" && source[idx - 1] !== "\\") {
         inString = !inString;
       }
@@ -311,8 +312,85 @@ import { TransactionRollbackError } from "drizzle-orm/errors";
 
 // src/client.ts
 import {
-  listValue
+  listValue as listValue2,
+  timestampValue as timestampValue2
+} from "@duckdb/node-api";
+
+// src/value-wrappers.ts
+import {
+  listValue,
+  arrayValue,
+  structValue,
+  mapValue,
+  blobValue,
+  timestampValue,
+  timestampTZValue
 } from "@duckdb/node-api";
+
+// src/value-wrappers-core.ts
+var DUCKDB_VALUE_MARKER = Symbol.for("drizzle-duckdb:value");
+
+// src/value-wrappers.ts
+function dateToMicros(value) {
+  if (value instanceof Date) {
+    return BigInt(value.getTime()) * 1000n;
+  }
+  let normalized = value;
+  if (!value.includes("T") && value.includes(" ")) {
+    normalized = value.replace(" ", "T");
+  }
+  if (!normalized.endsWith("Z") && !/[+-]\d{2}:?\d{2}$/.test(normalized)) {
+    normalized += "Z";
+  }
+  const date = new Date(normalized);
+  if (isNaN(date.getTime())) {
+    throw new Error(`Invalid timestamp string: ${value}`);
+  }
+  return BigInt(date.getTime()) * 1000n;
+}
+function toUint8Array(data) {
+  return data instanceof Uint8Array && !(data instanceof Buffer) ? data : new Uint8Array(data);
+}
+function convertStructEntries(data, toValue) {
+  const entries = {};
+  for (const [key, val] of Object.entries(data)) {
+    entries[key] = toValue(val);
+  }
+  return entries;
+}
+function convertMapEntries(data, toValue) {
+  return Object.entries(data).map(([key, val]) => ({
+    key,
+    value: toValue(val)
+  }));
+}
+function wrapperToNodeApiValue(wrapper, toValue) {
+  switch (wrapper.kind) {
+    case "list":
+      return listValue(wrapper.data.map(toValue));
+    case "array":
+      return arrayValue(wrapper.data.map(toValue));
+    case "struct":
+      return structValue(convertStructEntries(wrapper.data, toValue));
+    case "map":
+      return mapValue(convertMapEntries(wrapper.data, toValue));
+    case "timestamp":
+      return wrapper.withTimezone ? timestampTZValue(dateToMicros(wrapper.data)) : timestampValue(dateToMicros(wrapper.data));
+    case "blob":
+      return blobValue(toUint8Array(wrapper.data));
+    case "json":
+      return JSON.stringify(wrapper.data);
+    default: {
+      const _exhaustive = wrapper;
+      throw new Error(`Unknown wrapper kind: ${_exhaustive.kind}`);
+    }
+  }
+}
+
+// src/client.ts
+function isPool(client) {
+  return typeof client.acquire === "function";
+}
 function isPgArrayLiteral(value) {
   return value.startsWith("{") && value.endsWith("}");
 }
@@ -324,15 +402,13 @@ function parsePgArrayLiteral(value) {
     return value;
   }
 }
-var warnedArrayLiteral = false;
 function prepareParams(params, options = {}) {
   return params.map((param) => {
     if (typeof param === "string" && isPgArrayLiteral(param)) {
       if (options.rejectStringArrayLiterals) {
         throw new Error("Stringified array literals are not supported. Use duckDbList()/duckDbArray() or pass native arrays.");
       }
-      if (!warnedArrayLiteral && options.warnOnStringArrayLiteral) {
-        warnedArrayLiteral = true;
+      if (options.warnOnStringArrayLiteral) {
         options.warnOnStringArrayLiteral();
       }
       return parsePgArrayLiteral(param);
@@ -341,30 +417,116 @@ function prepareParams(params, options = {}) {
   });
 }
 function toNodeApiValue(value) {
+  if (value == null)
+    return null;
+  const t = typeof value;
+  if (t === "string" || t === "number" || t === "bigint" || t === "boolean") {
+    return value;
+  }
+  if (t === "object" && DUCKDB_VALUE_MARKER in value) {
+    return wrapperToNodeApiValue(value, toNodeApiValue);
+  }
   if (Array.isArray(value)) {
-    return listValue(value.map((inner) => toNodeApiValue(inner)));
+    return listValue2(value.map((inner) => toNodeApiValue(inner)));
+  }
+  if (value instanceof Date) {
+    return timestampValue2(BigInt(value.getTime()) * 1000n);
   }
   return value;
 }
-async function executeOnClient(client, query, params) {
-  const values = params.length > 0 ? params.map((param) => toNodeApiValue(param)) : undefined;
-  const result = await client.run(query, values);
-  const rows = await result.getRowsJS();
-  const columns = result.columnNames();
+function deduplicateColumns(columns) {
   const seen = {};
-  const uniqueColumns = columns.map((col) => {
+  return columns.map((col) => {
     const count = seen[col] ?? 0;
     seen[col] = count + 1;
     return count === 0 ? col : `${col}_${count}`;
   });
-  return (rows ?? []).map((vals) => {
+}
+function mapRowsToObjects(columns, rows) {
+  return rows.map((vals) => {
     const obj = {};
-    uniqueColumns.forEach((col, idx) => {
+    columns.forEach((col, idx) => {
       obj[col] = vals[idx];
     });
     return obj;
   });
 }
+async function closeClientConnection(connection) {
+  if ("close" in connection && typeof connection.close === "function") {
+    await connection.close();
+    return;
+  }
+  if ("closeSync" in connection && typeof connection.closeSync === "function") {
+    connection.closeSync();
+    return;
+  }
+  if ("disconnectSync" in connection && typeof connection.disconnectSync === "function") {
+    connection.disconnectSync();
+  }
+}
+async function executeOnClient(client, query, params) {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      return await executeOnClient(connection, query, params);
+    } finally {
+      await client.release(connection);
+    }
+  }
+  const values = params.length > 0 ? params.map((param) => toNodeApiValue(param)) : undefined;
+  const result = await client.run(query, values);
+  const rows = await result.getRowsJS();
+  const columns = result.deduplicatedColumnNames?.() ?? result.columnNames();
+  const uniqueColumns = deduplicateColumns(columns);
+  return rows ? mapRowsToObjects(uniqueColumns, rows) : [];
+}
+async function* executeInBatches(client, query, params, options = {}) {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      yield* executeInBatches(connection, query, params, options);
+      return;
+    } finally {
+      await client.release(connection);
+    }
+  }
+  const rowsPerChunk = options.rowsPerChunk && options.rowsPerChunk > 0 ? options.rowsPerChunk : 1e5;
+  const values = params.length > 0 ? params.map((param) => toNodeApiValue(param)) : undefined;
+  const result = await client.stream(query, values);
+  const columns = result.deduplicatedColumnNames?.() ?? result.columnNames();
+  const uniqueColumns = deduplicateColumns(columns);
+  let buffer = [];
+  for await (const chunk of result.yieldRowsJs()) {
+    const objects = mapRowsToObjects(uniqueColumns, chunk);
+    for (const row of objects) {
+      buffer.push(row);
+      if (buffer.length >= rowsPerChunk) {
+        yield buffer;
+        buffer = [];
+      }
+    }
+  }
+  if (buffer.length > 0) {
+    yield buffer;
+  }
+}
+async function executeArrowOnClient(client, query, params) {
+  if (isPool(client)) {
+    const connection = await client.acquire();
+    try {
+      return await executeArrowOnClient(connection, query, params);
+    } finally {
+      await client.release(connection);
+    }
+  }
+  const values = params.length > 0 ? params.map((param) => toNodeApiValue(param)) : undefined;
+  const result = await client.run(query, values);
+  const maybeArrow = result.toArrow ?? result.getArrowTable;
+  if (typeof maybeArrow === "function") {
+    return await maybeArrow.call(result);
+  }
+  return result.getColumnsObjectJS();
+}
 
 // src/session.ts
 class DuckDBPreparedQuery extends PgPreparedQuery {
@@ -405,11 +567,7 @@ class DuckDBPreparedQuery extends PgPreparedQuery {
       this.logger.logQuery(`[duckdb] original query before array rewrite: ${this.queryString}`, params);
     }
     this.logger.logQuery(rewrittenQuery, params);
-    const {
-      fields,
-      joinsNotNullableMap,
-      customResultMapper
-    } = this;
+    const { fields, joinsNotNullableMap, customResultMapper } = this;
     const rows = await executeOnClient(this.client, rewrittenQuery, params);
     if (rows.length === 0 || !fields) {
       return rows;
@@ -449,16 +607,30 @@ class DuckDBSession extends PgSession {
     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);
   }
   async transaction(transaction) {
-    const session = new DuckDBSession(this.client, this.dialect, this.schema, this.options);
+    let pinnedConnection;
+    let pool;
+    let clientForTx = this.client;
+    if (isPool(this.client)) {
+      pool = this.client;
+      pinnedConnection = await pool.acquire();
+      clientForTx = pinnedConnection;
+    }
+    const session = new DuckDBSession(clientForTx, this.dialect, this.schema, this.options);
     const tx = new DuckDBTransaction(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;
+      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;
+      }
+    } finally {
+      if (pinnedConnection && pool) {
+        await pool.release(pinnedConnection);
+      }
     }
   }
   warnOnStringArrayLiteral = (query) => {
@@ -469,6 +641,32 @@ class DuckDBSession extends PgSession {
     this.logger.logQuery(`[duckdb] ${arrayLiteralWarning}
 query: ${query}`, []);
   };
+  executeBatches(query, options = {}) {
+    const builtQuery = this.dialect.sqlToQuery(query);
+    const params = prepareParams(builtQuery.params, {
+      rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+      warnOnStringArrayLiteral: this.rejectStringArrayLiterals ? undefined : () => this.warnOnStringArrayLiteral(builtQuery.sql)
+    });
+    const rewrittenQuery = this.rewriteArrays ? adaptArrayOperators(builtQuery.sql) : builtQuery.sql;
+    if (this.rewriteArrays && rewrittenQuery !== builtQuery.sql) {
+      this.logger.logQuery(`[duckdb] original query before array rewrite: ${builtQuery.sql}`, params);
+    }
+    this.logger.logQuery(rewrittenQuery, params);
+    return executeInBatches(this.client, rewrittenQuery, params, options);
+  }
+  async executeArrow(query) {
+    const builtQuery = this.dialect.sqlToQuery(query);
+    const params = prepareParams(builtQuery.params, {
+      rejectStringArrayLiterals: this.rejectStringArrayLiterals,
+      warnOnStringArrayLiteral: this.rejectStringArrayLiterals ? undefined : () => this.warnOnStringArrayLiteral(builtQuery.sql)
+    });
+    const rewrittenQuery = this.rewriteArrays ? adaptArrayOperators(builtQuery.sql) : builtQuery.sql;
+    if (this.rewriteArrays && rewrittenQuery !== builtQuery.sql) {
+      this.logger.logQuery(`[duckdb] original query before array rewrite: ${builtQuery.sql}`, params);
+    }
+    this.logger.logQuery(rewrittenQuery, params);
+    return executeArrowOnClient(this.client, rewrittenQuery, params);
+  }
 }
 
 class DuckDBTransaction extends PgTransaction {
@@ -492,6 +690,12 @@ class DuckDBTransaction extends PgTransaction {
   setTransaction(config) {
     return this.session.execute(sql`set transaction ${this.getTransactionConfigSQL(config)}`);
   }
+  executeBatches(query, options = {}) {
+    return this.session.executeBatches(query, options);
+  }
+  executeArrow(query) {
+    return this.session.executeArrow(query);
+  }
   async transaction(transaction) {
     const nestedTx = new DuckDBTransaction(this.dialect, this.session, this.schema, this.nestedIndex + 1);
     return transaction(nestedTx);
@@ -590,13 +794,7 @@ import { PgViewBase } from "drizzle-orm/pg-core/view-base";
 import { SQL as SQL4 } from "drizzle-orm/sql/sql";
 
 // src/sql/selection.ts
-import {
-  Column as Column2,
-  SQL as SQL3,
-  getTableName as getTableName2,
-  is as is3,
-  sql as sql3
-} from "drizzle-orm";
+import { Column as Column2, SQL as SQL3, getTableName as getTableName2, is as is3, sql as sql3 } from "drizzle-orm";
 function mapEntries(obj, prefix, fullJoin = false) {
   return Object.fromEntries(Object.entries(obj).filter(([key]) => key !== "enableRLS").map(([key, value]) => {
     const qualified = prefix ? `${prefix}.${key}` : key;
@@ -680,6 +878,72 @@ class DuckDBSelectBuilder extends PgSelectBuilder {
   }
 }
 
+// src/pool.ts
+import { DuckDBConnection } from "@duckdb/node-api";
+var POOL_PRESETS = {
+  pulse: 4,
+  standard: 6,
+  jumbo: 8,
+  mega: 12,
+  giga: 16,
+  local: 8,
+  memory: 4
+};
+function resolvePoolSize(pool) {
+  if (pool === false)
+    return false;
+  if (pool === undefined)
+    return 4;
+  if (typeof pool === "string")
+    return POOL_PRESETS[pool];
+  return pool.size ?? 4;
+}
+function createDuckDBConnectionPool(instance, options = {}) {
+  const size = options.size && options.size > 0 ? options.size : 4;
+  const idle = [];
+  const waiting = [];
+  let total = 0;
+  let closed = false;
+  const acquire = async () => {
+    if (closed) {
+      throw new Error("DuckDB connection pool is closed");
+    }
+    if (idle.length > 0) {
+      return idle.pop();
+    }
+    if (total < size) {
+      total += 1;
+      return await DuckDBConnection.create(instance);
+    }
+    return await new Promise((resolve) => {
+      waiting.push(resolve);
+    });
+  };
+  const release = async (connection) => {
+    if (closed) {
+      await closeClientConnection(connection);
+      return;
+    }
+    const waiter = waiting.shift();
+    if (waiter) {
+      waiter(connection);
+      return;
+    }
+    idle.push(connection);
+  };
+  const close = async () => {
+    closed = true;
+    const toClose = idle.splice(0, idle.length);
+    await Promise.all(toClose.map((conn) => closeClientConnection(conn)));
+  };
+  return {
+    acquire,
+    release,
+    close,
+    size
+  };
+}
+
 // src/driver.ts
 class DuckDBDriver {
   client;
@@ -699,7 +963,14 @@ class DuckDBDriver {
     });
   }
 }
-function drizzle(client, config = {}) {
+function isConfigObject(data) {
+  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;
+}
+function createFromClient(client, config = {}, instance) {
   const dialect = new DuckDBDialect;
   const logger = config.logger === true ? new DefaultLogger : config.logger || undefined;
   let schema;
@@ -717,17 +988,60 @@ function drizzle(client, config = {}) {
     rejectStringArrayLiterals: config.rejectStringArrayLiterals
   });
   const session = driver.createSession(schema);
-  return new DuckDBDatabase(dialect, session, schema);
+  const db = new DuckDBDatabase(dialect, session, schema, client, instance);
+  return db;
+}
+async function createFromConnectionString(path, instanceOptions, config = {}) {
+  const instance = await DuckDBInstance2.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);
+}
+function drizzle(clientOrConfigOrPath, config) {
+  if (typeof clientOrConfigOrPath === "string") {
+    return createFromConnectionString(clientOrConfigOrPath, undefined, config);
+  }
+  if (isConfigObject(clientOrConfigOrPath)) {
+    const configObj = clientOrConfigOrPath;
+    if ("connection" in configObj) {
+      const connConfig = configObj;
+      const { connection, ...restConfig } = connConfig;
+      if (typeof connection === "string") {
+        return createFromConnectionString(connection, undefined, restConfig);
+      }
+      return createFromConnectionString(connection.path, connection.options, restConfig);
+    }
+    if ("client" in configObj) {
+      const clientConfig = configObj;
+      const { client: clientValue, ...restConfig } = clientConfig;
+      return createFromClient(clientValue, restConfig);
+    }
+    throw new Error("Invalid drizzle config: either connection or client must be provided");
+  }
+  return createFromClient(clientOrConfigOrPath, config);
 }
 
 class DuckDBDatabase extends PgDatabase {
   dialect;
   session;
   static [entityKind3] = "DuckDBDatabase";
-  constructor(dialect, session, schema) {
+  $client;
+  $instance;
+  constructor(dialect, session, schema, client, instance) {
     super(dialect, session, schema);
     this.dialect = dialect;
     this.session = session;
+    this.$client = client;
+    this.$instance = instance;
+  }
+  async close() {
+    if (isPool(this.$client) && this.$client.close) {
+      await this.$client.close();
+    }
   }
   select(fields) {
     const selectedFields = fields ? aliasFields(fields) : undefined;
@@ -737,6 +1051,12 @@ class DuckDBDatabase extends PgDatabase {
       dialect: this.dialect
     });
   }
+  executeBatches(query, options = {}) {
+    return this.session.executeBatches(query, options);
+  }
+  executeArrow(query) {
+    return this.session.executeArrow(query);
+  }
   async transaction(transaction) {
     return await this.session.transaction(transaction);
   }
@@ -745,7 +1065,7 @@ class DuckDBDatabase extends PgDatabase {
 // src/introspect.ts
 import { sql as sql4 } from "drizzle-orm";
 var SYSTEM_SCHEMAS = new Set(["information_schema", "pg_catalog"]);
-var DEFAULT_IMPORT_BASE = "@leonardovida-md/drizzle-neo-duckdb";
+var DEFAULT_IMPORT_BASE = "@leonardovida-md/drizzle-neo-duckdb/helpers";
 async function introspect(db, opts = {}) {
   const database = await resolveDatabase(db, opts.database, opts.allDatabases);
   const schemas = await resolveSchemas(db, database, opts.schemas);
@@ -1046,7 +1366,9 @@ function mapDuckDbType(column, imports, options) {
   }
   if (upper === "BIGINT" || upper === "INT8" || upper === "UBIGINT") {
     imports.pgCore.add("bigint");
-    return { builder: `bigint(${columnName(column.name)})` };
+    return {
+      builder: `bigint(${columnName(column.name)}, { mode: 'number' })`
+    };
   }
   const decimalMatch = /^DECIMAL\((\d+),(\d+)\)/i.exec(upper);
   const numericMatch = /^NUMERIC\((\d+),(\d+)\)/i.exec(upper);
@@ -1079,6 +1401,22 @@ function mapDuckDbType(column, imports, options) {
     imports.pgCore.add("doublePrecision");
     return { builder: `doublePrecision(${columnName(column.name)})` };
   }
+  const arrayMatch = /^(.*)\[(\d+)\]$/.exec(upper);
+  if (arrayMatch) {
+    imports.local.add("duckDbArray");
+    const [, base, length] = arrayMatch;
+    return {
+      builder: `duckDbArray(${columnName(column.name)}, ${JSON.stringify(base)}, ${Number(length)})`
+    };
+  }
+  const listMatch = /^(.*)\[\]$/.exec(upper);
+  if (listMatch) {
+    imports.local.add("duckDbList");
+    const [, base] = listMatch;
+    return {
+      builder: `duckDbList(${columnName(column.name)}, ${JSON.stringify(base)})`
+    };
+  }
   if (upper.startsWith("CHAR(") || upper === "CHAR") {
     imports.pgCore.add("char");
     const length = column.characterLength;
@@ -1119,22 +1457,6 @@ function mapDuckDbType(column, imports, options) {
     imports.local.add("duckDbBlob");
     return { builder: `duckDbBlob(${columnName(column.name)})` };
   }
-  const arrayMatch = /^(.*)\[(\d+)\]$/.exec(upper);
-  if (arrayMatch) {
-    imports.local.add("duckDbArray");
-    const [, base, length] = arrayMatch;
-    return {
-      builder: `duckDbArray(${columnName(column.name)}, ${JSON.stringify(base)}, ${Number(length)})`
-    };
-  }
-  const listMatch = /^(.*)\[\]$/.exec(upper);
-  if (listMatch) {
-    imports.local.add("duckDbList");
-    const [, base] = listMatch;
-    return {
-      builder: `duckDbList(${columnName(column.name)}, ${JSON.stringify(base)})`
-    };
-  }
   if (upper.startsWith("STRUCT")) {
     imports.local.add("duckDbStruct");
     const inner = upper.replace(/^STRUCT\s*\(/i, "").replace(/\)$/, "");
@@ -1383,7 +1705,7 @@ async function main() {
     throw new Error("Missing required --url");
   }
   const instanceOptions = options.url.startsWith("md:") && process.env.MOTHERDUCK_TOKEN ? { motherduck_token: process.env.MOTHERDUCK_TOKEN } : undefined;
-  const instance = await DuckDBInstance.create(options.url, instanceOptions);
+  const instance = await DuckDBInstance3.create(options.url, instanceOptions);
   const connection = await instance.connect();
   const db = drizzle(connection);
   try {

package/dist/helpers.d.ts

@@ -0,0 +1 @@
+export { duckDbList, duckDbArray, duckDbMap, duckDbStruct, duckDbJson, duckDbBlob, duckDbInet, duckDbInterval, duckDbTimestamp, duckDbDate, duckDbTime, duckDbArrayContains, duckDbArrayContained, duckDbArrayOverlaps, } from './columns.ts';