@leonardovida-md/drizzle-neo-duckdb 1.0.1 → 1.0.3

package/dist/duckdb-introspect.mjs

@@ -311,8 +311,78 @@ 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";
+var DUCKDB_VALUE_MARKER = Symbol.for("drizzle-duckdb:value");
+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 isPgArrayLiteral(value) {
   return value.startsWith("{") && value.endsWith("}");
 }
@@ -341,30 +411,78 @@ 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 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.deduplicatedColumnNames?.() ?? result.columnNames();
+  const uniqueColumns = deduplicateColumns(columns);
+  return rows ? mapRowsToObjects(uniqueColumns, rows) : [];
+}
+async function* executeInBatches(client, query, params, options = {}) {
+  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) {
+  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 +523,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;
@@ -469,6 +583,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 +632,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 +736,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;
@@ -737,6 +877,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);
   }
@@ -747,12 +893,13 @@ 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";
 async function introspect(db, opts = {}) {
-  const schemas = await resolveSchemas(db, opts.schemas);
+  const database = await resolveDatabase(db, opts.database, opts.allDatabases);
+  const schemas = await resolveSchemas(db, database, opts.schemas);
   const includeViews = opts.includeViews ?? false;
-  const tables = await loadTables(db, schemas, includeViews);
-  const columns = await loadColumns(db, schemas);
-  const constraints = await loadConstraints(db, schemas);
-  const indexes = await loadIndexes(db, schemas);
+  const tables = await loadTables(db, database, schemas, includeViews);
+  const columns = await loadColumns(db, database, schemas);
+  const constraints = await loadConstraints(db, database, schemas);
+  const indexes = await loadIndexes(db, database, schemas);
   const grouped = buildTables(tables, columns, constraints, indexes);
   const schemaTs = emitSchema(grouped, {
     useCustomTimeTypes: opts.useCustomTimeTypes ?? true,
@@ -766,27 +913,41 @@ async function introspect(db, opts = {}) {
     }
   };
 }
-async function resolveSchemas(db, targetSchemas) {
+async function resolveDatabase(db, targetDatabase, allDatabases) {
+  if (allDatabases) {
+    return null;
+  }
+  if (targetDatabase) {
+    return targetDatabase;
+  }
+  const rows = await db.execute(sql4`SELECT current_database() as current_database`);
+  return rows[0]?.current_database ?? null;
+}
+async function resolveSchemas(db, database, targetSchemas) {
   if (targetSchemas?.length) {
     return targetSchemas;
   }
-  const rows = await db.execute(sql4`select schema_name from information_schema.schemata`);
+  const databaseFilter = database ? sql4`catalog_name = ${database}` : sql4`1 = 1`;
+  const rows = await db.execute(sql4`SELECT schema_name FROM information_schema.schemata WHERE ${databaseFilter}`);
   return rows.map((row) => row.schema_name).filter((name) => !SYSTEM_SCHEMAS.has(name));
 }
-async function loadTables(db, schemas, includeViews) {
+async function loadTables(db, database, schemas, includeViews) {
   const schemaFragments = schemas.map((schema) => sql4`${schema}`);
+  const databaseFilter = database ? sql4`table_catalog = ${database}` : sql4`1 = 1`;
   return await db.execute(sql4`
-      select table_schema as schema_name, table_name, table_type
-      from information_schema.tables
-      where table_schema in (${sql4.join(schemaFragments, sql4.raw(", "))})
-      and ${includeViews ? sql4`1 = 1` : sql4`table_type = 'BASE TABLE'`}
-      order by table_schema, table_name
+      SELECT table_schema as schema_name, table_name, table_type
+      FROM information_schema.tables
+      WHERE ${databaseFilter}
+      AND table_schema IN (${sql4.join(schemaFragments, sql4.raw(", "))})
+      AND ${includeViews ? sql4`1 = 1` : sql4`table_type = 'BASE TABLE'`}
+      ORDER BY table_schema, table_name
     `);
 }
-async function loadColumns(db, schemas) {
+async function loadColumns(db, database, schemas) {
   const schemaFragments = schemas.map((schema) => sql4`${schema}`);
+  const databaseFilter = database ? sql4`database_name = ${database}` : sql4`1 = 1`;
   return await db.execute(sql4`
-      select
+      SELECT
         schema_name,
         table_name,
         column_name,
@@ -798,15 +959,17 @@ async function loadColumns(db, schemas) {
         numeric_precision,
         numeric_scale,
         internal
-      from duckdb_columns()
-      where schema_name in (${sql4.join(schemaFragments, sql4.raw(", "))})
-      order by schema_name, table_name, column_index
+      FROM duckdb_columns()
+      WHERE ${databaseFilter}
+      AND schema_name IN (${sql4.join(schemaFragments, sql4.raw(", "))})
+      ORDER BY schema_name, table_name, column_index
     `);
 }
-async function loadConstraints(db, schemas) {
+async function loadConstraints(db, database, schemas) {
   const schemaFragments = schemas.map((schema) => sql4`${schema}`);
+  const databaseFilter = database ? sql4`database_name = ${database}` : sql4`1 = 1`;
   return await db.execute(sql4`
-      select
+      SELECT
         schema_name,
         table_name,
         constraint_name,
@@ -815,23 +978,26 @@ async function loadConstraints(db, schemas) {
         constraint_column_names,
         referenced_table,
         referenced_column_names
-      from duckdb_constraints()
-      where schema_name in (${sql4.join(schemaFragments, sql4.raw(", "))})
-      order by schema_name, table_name, constraint_index
+      FROM duckdb_constraints()
+      WHERE ${databaseFilter}
+      AND schema_name IN (${sql4.join(schemaFragments, sql4.raw(", "))})
+      ORDER BY schema_name, table_name, constraint_index
     `);
 }
-async function loadIndexes(db, schemas) {
+async function loadIndexes(db, database, schemas) {
   const schemaFragments = schemas.map((schema) => sql4`${schema}`);
+  const databaseFilter = database ? sql4`database_name = ${database}` : sql4`1 = 1`;
   return await db.execute(sql4`
-      select
+      SELECT
         schema_name,
         table_name,
         index_name,
         is_unique,
         expressions
-      from duckdb_indexes()
-      where schema_name in (${sql4.join(schemaFragments, sql4.raw(", "))})
-      order by schema_name, table_name, index_name
+      FROM duckdb_indexes()
+      WHERE ${databaseFilter}
+      AND schema_name IN (${sql4.join(schemaFragments, sql4.raw(", "))})
+      ORDER BY schema_name, table_name, index_name
     `);
 }
 function buildTables(tables, columns, constraints, indexes) {
@@ -1026,7 +1192,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);
@@ -1278,6 +1446,7 @@ function renderImports(imports, importBasePath) {
 function parseArgs(argv) {
   const options = {
     outFile: path.resolve(process.cwd(), "drizzle/schema.ts"),
+    allDatabases: false,
     includeViews: false,
     useCustomTimeTypes: true
   };
@@ -1287,6 +1456,13 @@ function parseArgs(argv) {
       case "--url":
         options.url = argv[++i];
         break;
+      case "--database":
+      case "--db":
+        options.database = argv[++i];
+        break;
+      case "--all-databases":
+        options.allDatabases = true;
+        break;
       case "--schema":
       case "--schemas":
         options.schemas = argv[++i]?.split(",").map((s) => s.trim()).filter(Boolean);
@@ -1325,11 +1501,27 @@ Usage:
 
 Options:
   --url            DuckDB database path (e.g. :memory:, ./local.duckdb, md:)
+  --database, --db Database/catalog to introspect (default: current database)
+  --all-databases  Introspect all attached databases (not just current)
   --schema         Comma separated schema list (defaults to all non-system schemas)
   --out            Output file (default: ./drizzle/schema.ts)
   --include-views  Include views in the generated schema
   --use-pg-time    Use pg-core timestamp/date/time instead of DuckDB custom helpers
   --import-base    Override import path for duckdb helpers (default: package name)
+
+Database Filtering:
+  By default, only tables from the current database are introspected. This prevents
+  returning tables from all attached databases in MotherDuck workspaces.
+
+  Use --database to specify a different database, or --all-databases to introspect
+  all attached databases.
+
+Examples:
+  # Local DuckDB file
+  bun x duckdb-introspect --url ./my-database.duckdb --out ./schema.ts
+
+  # MotherDuck (requires MOTHERDUCK_TOKEN env var)
+  MOTHERDUCK_TOKEN=xxx bun x duckdb-introspect --url md: --database my_cloud_db --out ./schema.ts
 `);
 }
 async function main() {
@@ -1344,6 +1536,8 @@ async function main() {
   const db = drizzle(connection);
   try {
     const result = await introspect(db, {
+      database: options.database,
+      allDatabases: options.allDatabases,
       schemas: options.schemas,
       includeViews: options.includeViews,
       useCustomTimeTypes: options.useCustomTimeTypes,

package/dist/index.d.ts

@@ -3,3 +3,6 @@ export * from './session.ts';
 export * from './columns.ts';
 export * from './migrator.ts';
 export * from './introspect.ts';
+export * from './client.ts';
+export * from './olap.ts';
+export * from './value-wrappers.ts';