@apibara/plugin-drizzle 2.1.0-beta.20 → 2.1.0-beta.22

package/dist/index.cjs

@@ -5,20 +5,10 @@ const plugins = require('@apibara/indexer/plugins');
 const internal = require('@apibara/indexer/internal');
 const plugins$1 = require('@apibara/indexer/internal/plugins');
 const constants = require('./shared/plugin-drizzle.cae20704.cjs');
-const pglite$1 = require('@electric-sql/pglite');
-const nodePostgres = require('drizzle-orm/node-postgres');
-const migrator$1 = require('drizzle-orm/node-postgres/migrator');
-const pglite = require('drizzle-orm/pglite');
-const migrator = require('drizzle-orm/pglite/migrator');
-const pg = require('pg');
-const protocol = require('@apibara/protocol');
 const drizzleOrm = require('drizzle-orm');
+const protocol = require('@apibara/protocol');
 const pgCore = require('drizzle-orm/pg-core');
 
-function _interopDefaultCompat (e) { return e && typeof e === 'object' && 'default' in e ? e.default : e; }
-
-const pg__default = /*#__PURE__*/_interopDefaultCompat(pg);
-
 class DrizzleStorageError extends Error {
   constructor(message, options) {
     super(message, options);
@@ -46,6 +36,12 @@ function serialize(obj) {
 function sleep(ms) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
+const getIdColumnForTable = (tableName, idColumn) => {
+  if (idColumn[tableName]) {
+    return idColumn[tableName];
+  }
+  return idColumn["*"];
+};
 
 function drizzle(options) {
   const {
@@ -55,15 +51,9 @@ function drizzle(options) {
     config,
     poolConfig
   } = options ?? {};
-  if (isPostgresConnectionString(connectionString) || type === "node-postgres") {
-    const pool = new pg__default.Pool({
-      connectionString,
-      ...poolConfig || {}
-    });
-    return nodePostgres.drizzle(pool, { schema, ...config || {} });
-  }
-  if (type === "pglite") {
-    return pglite.drizzle({
+  if (isPgliteConnectionString(connectionString) && type === "pglite") {
+    const { drizzle: drizzlePGLite } = require("drizzle-orm/pglite");
+    return drizzlePGLite({
       schema,
       connection: {
         dataDir: connectionString || "memory://pglite"
@@ -71,15 +61,25 @@ function drizzle(options) {
       ...config || {}
     });
   }
-  throw new Error("Invalid database type");
+  const { Pool } = require("pg");
+  const { drizzle: drizzleNode } = require("drizzle-orm/node-postgres");
+  const pool = new Pool({
+    connectionString,
+    ...poolConfig || {}
+  });
+  return drizzleNode(pool, { schema, ...config || {} });
 }
 async function migrate(db, options) {
-  const isPglite = !!("$client" in db && db.$client instanceof pglite$1.PGlite);
+  const isPglite = isDrizzleKind(db, "PgliteDatabase");
   try {
     if (isPglite) {
-      await migrator.migrate(db, options);
+      const { migrate: migratePGLite } = require("drizzle-orm/pglite/migrator");
+      await migratePGLite(db, options);
     } else {
-      await migrator$1.migrate(db, options);
+      const {
+        migrate: migrateNode
+      } = require("drizzle-orm/node-postgres/migrator");
+      await migrateNode(db, options);
     }
   } catch (error) {
     throw new DrizzleStorageError(
@@ -90,8 +90,23 @@ async function migrate(db, options) {
     );
   }
 }
-function isPostgresConnectionString(conn) {
-  return conn.startsWith("postgres://") || conn.startsWith("postgresql://");
+function isPgliteConnectionString(conn) {
+  return conn.startsWith("memory://") || conn.startsWith("file://") || conn.startsWith("idb://");
+}
+function isDrizzleKind(value, entityKindValue) {
+  if (!value || typeof value !== "object") {
+    return false;
+  }
+  let cls = Object.getPrototypeOf(value).constructor;
+  if (cls) {
+    while (cls) {
+      if (drizzleOrm.entityKind in cls && cls[drizzleOrm.entityKind] === entityKindValue) {
+        return true;
+      }
+      cls = Object.getPrototypeOf(cls);
+    }
+  }
+  return false;
 }
 
 const CHECKPOINTS_TABLE_NAME = "checkpoints";
@@ -375,9 +390,10 @@ async function initializeReorgRollbackTable(tx, indexerId) {
     );
   }
 }
-async function registerTriggers(tx, tables, endCursor, idColumn, indexerId) {
+async function registerTriggers(tx, tables, endCursor, idColumnMap, indexerId) {
   try {
     for (const table of tables) {
+      const tableIdColumn = getIdColumnForTable(table, idColumnMap);
       await tx.execute(
         drizzleOrm.sql.raw(
           `DROP TRIGGER IF EXISTS ${getReorgTriggerName(table, indexerId)} ON ${table};`
@@ -388,7 +404,7 @@ async function registerTriggers(tx, tables, endCursor, idColumn, indexerId) {
           CREATE CONSTRAINT TRIGGER ${getReorgTriggerName(table, indexerId)}
           AFTER INSERT OR UPDATE OR DELETE ON ${table}
           DEFERRABLE INITIALLY DEFERRED
-          FOR EACH ROW EXECUTE FUNCTION ${constants.SCHEMA_NAME}.reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
+          FOR EACH ROW EXECUTE FUNCTION ${constants.SCHEMA_NAME}.reorg_checkpoint('${tableIdColumn}', ${Number(endCursor.orderKey)}, '${indexerId}');
         `)
       );
     }
@@ -413,7 +429,7 @@ async function removeTriggers(db, tables, indexerId) {
     });
   }
 }
-async function invalidate(tx, cursor, idColumn, indexerId) {
+async function invalidate(tx, cursor, idColumnMap, indexerId) {
   const { rows: result } = await tx.execute(
     drizzleOrm.sql.raw(`
       WITH deleted AS (
@@ -431,6 +447,7 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
     );
   }
   for (const op of result) {
+    const tableIdColumn = getIdColumnForTable(op.table_name, idColumnMap);
     switch (op.op) {
       case "I":
         try {
@@ -440,7 +457,7 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
           await tx.execute(
             drizzleOrm.sql.raw(`
                 DELETE FROM ${op.table_name}
-                WHERE ${idColumn} = '${op.row_id}'
+                WHERE ${tableIdColumn} = '${op.row_id}'
               `)
           );
         } catch (error) {
@@ -480,7 +497,9 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
             );
           }
           const rowValue = typeof op.row_value === "string" ? JSON.parse(op.row_value) : op.row_value;
-          const nonIdKeys = Object.keys(rowValue).filter((k) => k !== idColumn);
+          const nonIdKeys = Object.keys(rowValue).filter(
+            (k) => k !== tableIdColumn
+          );
           const fields = nonIdKeys.map((c) => `${c} = prev.${c}`).join(", ");
           const query = drizzleOrm.sql.raw(`
               UPDATE ${op.table_name}
@@ -488,7 +507,7 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
               FROM (
                 SELECT * FROM json_populate_record(null::${op.table_name}, '${JSON.stringify(op.row_value)}'::json)
               ) as prev
-              WHERE ${op.table_name}.${idColumn} = '${op.row_id}'
+              WHERE ${op.table_name}.${tableIdColumn} = '${op.row_id}'
               `);
           await tx.execute(query);
         } catch (error) {
@@ -566,8 +585,8 @@ function drizzleStorage({
   db,
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
-  schema,
-  idColumn = "id",
+  schema: _schema,
+  idColumn,
   migrate: migrateOptions
 }) {
   return plugins.defineIndexerPlugin((indexer$1) => {
@@ -575,15 +594,30 @@ function drizzleStorage({
     let indexerId = "";
     const alwaysReindex = process.env["APIBARA_ALWAYS_REINDEX"] === "true";
     let prevFinality;
+    const schema = _schema ?? db._.schema ?? {};
+    const idColumnMap = {
+      "*": typeof idColumn === "string" ? idColumn : "id",
+      ...typeof idColumn === "object" ? idColumn : {}
+    };
     try {
-      tableNames = Object.values(schema ?? db._.schema ?? {}).map(
-        (table) => table.dbName
-      );
+      tableNames = Object.values(schema).map((table) => table.dbName);
     } catch (error) {
       throw new DrizzleStorageError("Failed to get table names from schema", {
         cause: error
       });
     }
+    for (const table of Object.values(schema)) {
+      const columns = table.columns;
+      const tableIdColumn = getIdColumnForTable(table.dbName, idColumnMap);
+      const columnExists = Object.values(columns).some(
+        (column) => column.name === tableIdColumn
+      );
+      if (!columnExists) {
+        throw new DrizzleStorageError(
+          `Column \`"${tableIdColumn}"\` does not exist in table \`"${table.dbName}"\`. Make sure the table has the specified column or provide a valid \`idColumn\` mapping to \`drizzleStorage\`.`
+        );
+      }
+    }
     indexer$1.hooks.hook("run:before", async () => {
       const internalContext = plugins$1.useInternalContext();
       const context = indexer.useIndexerContext();
@@ -591,20 +625,9 @@ function drizzleStorage({
       context[constants.DRIZZLE_STORAGE_DB_PROPERTY] = db;
       const { indexerName: indexerFileName, availableIndexers } = internalContext;
       indexerId = internal.generateIndexerId(indexerFileName, identifier);
-      if (alwaysReindex) {
-        logger.warn(
-          `Reindexing: Deleting all data from tables - ${tableNames.join(", ")}`
-        );
-        await withTransaction(db, async (tx) => {
-          await cleanupStorage(tx, tableNames, indexerId);
-          if (enablePersistence) {
-            await resetPersistence({ tx, indexerId });
-          }
-          logger.success("Tables have been cleaned up for reindexing");
-        });
-      }
       let retries = 0;
       let migrationsApplied = false;
+      let cleanupApplied = false;
       while (retries <= MAX_RETRIES) {
         try {
           if (migrateOptions && !migrationsApplied) {
@@ -617,6 +640,17 @@ function drizzleStorage({
             if (enablePersistence) {
               await initializePersistentState(tx);
             }
+            if (alwaysReindex && !cleanupApplied) {
+              logger.warn(
+                `Reindexing: Deleting all data from tables - ${tableNames.join(", ")}`
+              );
+              await cleanupStorage(tx, tableNames, indexerId);
+              if (enablePersistence) {
+                await resetPersistence({ tx, indexerId });
+              }
+              cleanupApplied = true;
+              logger.success("Tables have been cleaned up for reindexing");
+            }
           });
           break;
         } catch (error) {
@@ -659,7 +693,7 @@ function drizzleStorage({
         return;
       }
       await withTransaction(db, async (tx) => {
-        await invalidate(tx, cursor, idColumn, indexerId);
+        await invalidate(tx, cursor, idColumnMap, indexerId);
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
         }
@@ -697,7 +731,7 @@ function drizzleStorage({
         throw new DrizzleStorageError("Invalidate Cursor is undefined");
       }
       await withTransaction(db, async (tx) => {
-        await invalidate(tx, cursor, idColumn, indexerId);
+        await invalidate(tx, cursor, idColumnMap, indexerId);
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
         }
@@ -713,14 +747,14 @@ function drizzleStorage({
           await withTransaction(db, async (tx) => {
             context[constants.DRIZZLE_PROPERTY] = { db: tx };
             if (prevFinality === "pending") {
-              await invalidate(tx, cursor, idColumn, indexerId);
+              await invalidate(tx, cursor, idColumnMap, indexerId);
             }
             if (finality !== "finalized") {
               await registerTriggers(
                 tx,
                 tableNames,
                 endCursor,
-                idColumn,
+                idColumnMap,
                 indexerId
               );
             }

package/dist/index.d.cts

@@ -113,6 +113,13 @@ type MigrateOptions = MigrationConfig;
  */
 declare function migrate<TSchema extends Record<string, unknown>>(db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>, options: MigrateOptions): Promise<void>;
 
+interface IdColumnMap extends Record<string, string> {
+    /**
+     * Wildcard mapping for all tables.
+     */
+    "*": string;
+}
+
 type DrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>> = {
     db: PgTransaction<TQueryResult, TFullSchema, TSchema>;
 };
@@ -135,9 +142,32 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
      */
     schema?: Record<string, unknown>;
     /**
-     * The column to use as the id. Defaults to 'id'.
+     * The column to use as the primary identifier for each table.
+     *
+     * This identifier is used for tracking changes during reorgs and rollbacks.
+     *
+     * Can be specified in two ways:
+     *
+     * 1. As a single string that applies to all tables:
+     * ```ts
+     * idColumn: "_id" // Uses "_id" column for all tables
+     * ```
+     *
+     * 2. As an object mapping table names to their ID columns:
+     * ```ts
+     * idColumn: {
+     *   transfers: "transaction_hash",    // Use "transaction_hash" for transfers table
+     *   blocks: "block_number",           // Use "block_number" for blocks table
+     *   "*": "_id"                        // Use "_id" for all other tables | defaults to "id"
+     * }
+     * ```
+     *
+     * The special "*" key acts as a fallback for any tables not explicitly mapped.
+     *
+     * @default "id"
+     * @type {string | Partial<IdColumnMap>}
      */
-    idColumn?: string;
+    idColumn?: string | Partial<IdColumnMap>;
     /**
      * The options for the database migration. When provided, the database will automatically run migrations before the indexer runs.
      */
@@ -154,6 +184,6 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
  * @param options.idColumn - The column to use as the id. Defaults to 'id'.
  * @param options.migrate - The options for the database migration. when provided, the database will automatically run migrations before the indexer runs.
  */
-declare function drizzleStorage<TFilter, TBlock, TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>({ db, persistState: enablePersistence, indexerName: identifier, schema, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
+declare function drizzleStorage<TFilter, TBlock, TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>({ db, persistState: enablePersistence, indexerName: identifier, schema: _schema, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
 
-export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };
+export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type IdColumnMap, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/index.d.mts

@@ -113,6 +113,13 @@ type MigrateOptions = MigrationConfig;
  */
 declare function migrate<TSchema extends Record<string, unknown>>(db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>, options: MigrateOptions): Promise<void>;
 
+interface IdColumnMap extends Record<string, string> {
+    /**
+     * Wildcard mapping for all tables.
+     */
+    "*": string;
+}
+
 type DrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>> = {
     db: PgTransaction<TQueryResult, TFullSchema, TSchema>;
 };
@@ -135,9 +142,32 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
      */
     schema?: Record<string, unknown>;
     /**
-     * The column to use as the id. Defaults to 'id'.
+     * The column to use as the primary identifier for each table.
+     *
+     * This identifier is used for tracking changes during reorgs and rollbacks.
+     *
+     * Can be specified in two ways:
+     *
+     * 1. As a single string that applies to all tables:
+     * ```ts
+     * idColumn: "_id" // Uses "_id" column for all tables
+     * ```
+     *
+     * 2. As an object mapping table names to their ID columns:
+     * ```ts
+     * idColumn: {
+     *   transfers: "transaction_hash",    // Use "transaction_hash" for transfers table
+     *   blocks: "block_number",           // Use "block_number" for blocks table
+     *   "*": "_id"                        // Use "_id" for all other tables | defaults to "id"
+     * }
+     * ```
+     *
+     * The special "*" key acts as a fallback for any tables not explicitly mapped.
+     *
+     * @default "id"
+     * @type {string | Partial<IdColumnMap>}
      */
-    idColumn?: string;
+    idColumn?: string | Partial<IdColumnMap>;
     /**
      * The options for the database migration. When provided, the database will automatically run migrations before the indexer runs.
      */
@@ -154,6 +184,6 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
  * @param options.idColumn - The column to use as the id. Defaults to 'id'.
  * @param options.migrate - The options for the database migration. when provided, the database will automatically run migrations before the indexer runs.
  */
-declare function drizzleStorage<TFilter, TBlock, TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>({ db, persistState: enablePersistence, indexerName: identifier, schema, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
+declare function drizzleStorage<TFilter, TBlock, TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>({ db, persistState: enablePersistence, indexerName: identifier, schema: _schema, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
 
-export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };
+export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type IdColumnMap, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/index.d.ts

@@ -113,6 +113,13 @@ type MigrateOptions = MigrationConfig;
  */
 declare function migrate<TSchema extends Record<string, unknown>>(db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>, options: MigrateOptions): Promise<void>;
 
+interface IdColumnMap extends Record<string, string> {
+    /**
+     * Wildcard mapping for all tables.
+     */
+    "*": string;
+}
+
 type DrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>> = {
     db: PgTransaction<TQueryResult, TFullSchema, TSchema>;
 };
@@ -135,9 +142,32 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
      */
     schema?: Record<string, unknown>;
     /**
-     * The column to use as the id. Defaults to 'id'.
+     * The column to use as the primary identifier for each table.
+     *
+     * This identifier is used for tracking changes during reorgs and rollbacks.
+     *
+     * Can be specified in two ways:
+     *
+     * 1. As a single string that applies to all tables:
+     * ```ts
+     * idColumn: "_id" // Uses "_id" column for all tables
+     * ```
+     *
+     * 2. As an object mapping table names to their ID columns:
+     * ```ts
+     * idColumn: {
+     *   transfers: "transaction_hash",    // Use "transaction_hash" for transfers table
+     *   blocks: "block_number",           // Use "block_number" for blocks table
+     *   "*": "_id"                        // Use "_id" for all other tables | defaults to "id"
+     * }
+     * ```
+     *
+     * The special "*" key acts as a fallback for any tables not explicitly mapped.
+     *
+     * @default "id"
+     * @type {string | Partial<IdColumnMap>}
      */
-    idColumn?: string;
+    idColumn?: string | Partial<IdColumnMap>;
     /**
      * The options for the database migration. When provided, the database will automatically run migrations before the indexer runs.
      */
@@ -154,6 +184,6 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
  * @param options.idColumn - The column to use as the id. Defaults to 'id'.
  * @param options.migrate - The options for the database migration. when provided, the database will automatically run migrations before the indexer runs.
  */
-declare function drizzleStorage<TFilter, TBlock, TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>({ db, persistState: enablePersistence, indexerName: identifier, schema, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
+declare function drizzleStorage<TFilter, TBlock, TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>({ db, persistState: enablePersistence, indexerName: identifier, schema: _schema, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
 
-export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };
+export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type IdColumnMap, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/index.mjs

@@ -3,14 +3,8 @@ import { defineIndexerPlugin, useLogger } from '@apibara/indexer/plugins';
 import { generateIndexerId } from '@apibara/indexer/internal';
 import { useInternalContext } from '@apibara/indexer/internal/plugins';
 import { S as SCHEMA_NAME, D as DRIZZLE_PROPERTY, a as DRIZZLE_STORAGE_DB_PROPERTY } from './shared/plugin-drizzle.2d226351.mjs';
-import { PGlite } from '@electric-sql/pglite';
-import { drizzle as drizzle$1 } from 'drizzle-orm/node-postgres';
-import { migrate as migrate$2 } from 'drizzle-orm/node-postgres/migrator';
-import { drizzle as drizzle$2 } from 'drizzle-orm/pglite';
-import { migrate as migrate$1 } from 'drizzle-orm/pglite/migrator';
-import pg from 'pg';
+import { entityKind, sql, eq, and, isNull, gt, lt } from 'drizzle-orm';
 import { normalizeCursor } from '@apibara/protocol';
-import { sql, eq, and, isNull, gt, lt } from 'drizzle-orm';
 import { pgSchema, text, integer, primaryKey, serial, char, jsonb } from 'drizzle-orm/pg-core';
 
 class DrizzleStorageError extends Error {
@@ -40,6 +34,12 @@ function serialize(obj) {
 function sleep(ms) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
+const getIdColumnForTable = (tableName, idColumn) => {
+  if (idColumn[tableName]) {
+    return idColumn[tableName];
+  }
+  return idColumn["*"];
+};
 
 function drizzle(options) {
   const {
@@ -49,15 +49,9 @@ function drizzle(options) {
     config,
     poolConfig
   } = options ?? {};
-  if (isPostgresConnectionString(connectionString) || type === "node-postgres") {
-    const pool = new pg.Pool({
-      connectionString,
-      ...poolConfig || {}
-    });
-    return drizzle$1(pool, { schema, ...config || {} });
-  }
-  if (type === "pglite") {
-    return drizzle$2({
+  if (isPgliteConnectionString(connectionString) && type === "pglite") {
+    const { drizzle: drizzlePGLite } = require("drizzle-orm/pglite");
+    return drizzlePGLite({
       schema,
       connection: {
         dataDir: connectionString || "memory://pglite"
@@ -65,15 +59,25 @@ function drizzle(options) {
       ...config || {}
     });
   }
-  throw new Error("Invalid database type");
+  const { Pool } = require("pg");
+  const { drizzle: drizzleNode } = require("drizzle-orm/node-postgres");
+  const pool = new Pool({
+    connectionString,
+    ...poolConfig || {}
+  });
+  return drizzleNode(pool, { schema, ...config || {} });
 }
 async function migrate(db, options) {
-  const isPglite = !!("$client" in db && db.$client instanceof PGlite);
+  const isPglite = isDrizzleKind(db, "PgliteDatabase");
   try {
     if (isPglite) {
-      await migrate$1(db, options);
+      const { migrate: migratePGLite } = require("drizzle-orm/pglite/migrator");
+      await migratePGLite(db, options);
     } else {
-      await migrate$2(db, options);
+      const {
+        migrate: migrateNode
+      } = require("drizzle-orm/node-postgres/migrator");
+      await migrateNode(db, options);
     }
   } catch (error) {
     throw new DrizzleStorageError(
@@ -84,8 +88,23 @@ async function migrate(db, options) {
     );
   }
 }
-function isPostgresConnectionString(conn) {
-  return conn.startsWith("postgres://") || conn.startsWith("postgresql://");
+function isPgliteConnectionString(conn) {
+  return conn.startsWith("memory://") || conn.startsWith("file://") || conn.startsWith("idb://");
+}
+function isDrizzleKind(value, entityKindValue) {
+  if (!value || typeof value !== "object") {
+    return false;
+  }
+  let cls = Object.getPrototypeOf(value).constructor;
+  if (cls) {
+    while (cls) {
+      if (entityKind in cls && cls[entityKind] === entityKindValue) {
+        return true;
+      }
+      cls = Object.getPrototypeOf(cls);
+    }
+  }
+  return false;
 }
 
 const CHECKPOINTS_TABLE_NAME = "checkpoints";
@@ -369,9 +388,10 @@ async function initializeReorgRollbackTable(tx, indexerId) {
     );
   }
 }
-async function registerTriggers(tx, tables, endCursor, idColumn, indexerId) {
+async function registerTriggers(tx, tables, endCursor, idColumnMap, indexerId) {
   try {
     for (const table of tables) {
+      const tableIdColumn = getIdColumnForTable(table, idColumnMap);
       await tx.execute(
         sql.raw(
           `DROP TRIGGER IF EXISTS ${getReorgTriggerName(table, indexerId)} ON ${table};`
@@ -382,7 +402,7 @@ async function registerTriggers(tx, tables, endCursor, idColumn, indexerId) {
           CREATE CONSTRAINT TRIGGER ${getReorgTriggerName(table, indexerId)}
           AFTER INSERT OR UPDATE OR DELETE ON ${table}
           DEFERRABLE INITIALLY DEFERRED
-          FOR EACH ROW EXECUTE FUNCTION ${SCHEMA_NAME}.reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
+          FOR EACH ROW EXECUTE FUNCTION ${SCHEMA_NAME}.reorg_checkpoint('${tableIdColumn}', ${Number(endCursor.orderKey)}, '${indexerId}');
         `)
       );
     }
@@ -407,7 +427,7 @@ async function removeTriggers(db, tables, indexerId) {
     });
   }
 }
-async function invalidate(tx, cursor, idColumn, indexerId) {
+async function invalidate(tx, cursor, idColumnMap, indexerId) {
   const { rows: result } = await tx.execute(
     sql.raw(`
       WITH deleted AS (
@@ -425,6 +445,7 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
     );
   }
   for (const op of result) {
+    const tableIdColumn = getIdColumnForTable(op.table_name, idColumnMap);
     switch (op.op) {
       case "I":
         try {
@@ -434,7 +455,7 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
           await tx.execute(
             sql.raw(`
                 DELETE FROM ${op.table_name}
-                WHERE ${idColumn} = '${op.row_id}'
+                WHERE ${tableIdColumn} = '${op.row_id}'
               `)
           );
         } catch (error) {
@@ -474,7 +495,9 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
             );
           }
           const rowValue = typeof op.row_value === "string" ? JSON.parse(op.row_value) : op.row_value;
-          const nonIdKeys = Object.keys(rowValue).filter((k) => k !== idColumn);
+          const nonIdKeys = Object.keys(rowValue).filter(
+            (k) => k !== tableIdColumn
+          );
           const fields = nonIdKeys.map((c) => `${c} = prev.${c}`).join(", ");
           const query = sql.raw(`
               UPDATE ${op.table_name}
@@ -482,7 +505,7 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
               FROM (
                 SELECT * FROM json_populate_record(null::${op.table_name}, '${JSON.stringify(op.row_value)}'::json)
               ) as prev
-              WHERE ${op.table_name}.${idColumn} = '${op.row_id}'
+              WHERE ${op.table_name}.${tableIdColumn} = '${op.row_id}'
               `);
           await tx.execute(query);
         } catch (error) {
@@ -560,8 +583,8 @@ function drizzleStorage({
   db,
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
-  schema,
-  idColumn = "id",
+  schema: _schema,
+  idColumn,
   migrate: migrateOptions
 }) {
   return defineIndexerPlugin((indexer) => {
@@ -569,15 +592,30 @@ function drizzleStorage({
     let indexerId = "";
     const alwaysReindex = process.env["APIBARA_ALWAYS_REINDEX"] === "true";
     let prevFinality;
+    const schema = _schema ?? db._.schema ?? {};
+    const idColumnMap = {
+      "*": typeof idColumn === "string" ? idColumn : "id",
+      ...typeof idColumn === "object" ? idColumn : {}
+    };
     try {
-      tableNames = Object.values(schema ?? db._.schema ?? {}).map(
-        (table) => table.dbName
-      );
+      tableNames = Object.values(schema).map((table) => table.dbName);
     } catch (error) {
       throw new DrizzleStorageError("Failed to get table names from schema", {
         cause: error
       });
     }
+    for (const table of Object.values(schema)) {
+      const columns = table.columns;
+      const tableIdColumn = getIdColumnForTable(table.dbName, idColumnMap);
+      const columnExists = Object.values(columns).some(
+        (column) => column.name === tableIdColumn
+      );
+      if (!columnExists) {
+        throw new DrizzleStorageError(
+          `Column \`"${tableIdColumn}"\` does not exist in table \`"${table.dbName}"\`. Make sure the table has the specified column or provide a valid \`idColumn\` mapping to \`drizzleStorage\`.`
+        );
+      }
+    }
     indexer.hooks.hook("run:before", async () => {
       const internalContext = useInternalContext();
       const context = useIndexerContext();
@@ -585,20 +623,9 @@ function drizzleStorage({
       context[DRIZZLE_STORAGE_DB_PROPERTY] = db;
       const { indexerName: indexerFileName, availableIndexers } = internalContext;
       indexerId = generateIndexerId(indexerFileName, identifier);
-      if (alwaysReindex) {
-        logger.warn(
-          `Reindexing: Deleting all data from tables - ${tableNames.join(", ")}`
-        );
-        await withTransaction(db, async (tx) => {
-          await cleanupStorage(tx, tableNames, indexerId);
-          if (enablePersistence) {
-            await resetPersistence({ tx, indexerId });
-          }
-          logger.success("Tables have been cleaned up for reindexing");
-        });
-      }
       let retries = 0;
       let migrationsApplied = false;
+      let cleanupApplied = false;
       while (retries <= MAX_RETRIES) {
         try {
           if (migrateOptions && !migrationsApplied) {
@@ -611,6 +638,17 @@ function drizzleStorage({
             if (enablePersistence) {
               await initializePersistentState(tx);
             }
+            if (alwaysReindex && !cleanupApplied) {
+              logger.warn(
+                `Reindexing: Deleting all data from tables - ${tableNames.join(", ")}`
+              );
+              await cleanupStorage(tx, tableNames, indexerId);
+              if (enablePersistence) {
+                await resetPersistence({ tx, indexerId });
+              }
+              cleanupApplied = true;
+              logger.success("Tables have been cleaned up for reindexing");
+            }
           });
           break;
         } catch (error) {
@@ -653,7 +691,7 @@ function drizzleStorage({
         return;
       }
       await withTransaction(db, async (tx) => {
-        await invalidate(tx, cursor, idColumn, indexerId);
+        await invalidate(tx, cursor, idColumnMap, indexerId);
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
         }
@@ -691,7 +729,7 @@ function drizzleStorage({
         throw new DrizzleStorageError("Invalidate Cursor is undefined");
       }
       await withTransaction(db, async (tx) => {
-        await invalidate(tx, cursor, idColumn, indexerId);
+        await invalidate(tx, cursor, idColumnMap, indexerId);
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
         }
@@ -707,14 +745,14 @@ function drizzleStorage({
           await withTransaction(db, async (tx) => {
             context[DRIZZLE_PROPERTY] = { db: tx };
             if (prevFinality === "pending") {
-              await invalidate(tx, cursor, idColumn, indexerId);
+              await invalidate(tx, cursor, idColumnMap, indexerId);
             }
             if (finality !== "finalized") {
               await registerTriggers(
                 tx,
                 tableNames,
                 endCursor,
-                idColumn,
+                idColumnMap,
                 indexerId
               );
             }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@apibara/plugin-drizzle",
-  "version": "2.1.0-beta.20",
+  "version": "2.1.0-beta.22",
   "type": "module",
   "files": [
     "dist",
@@ -32,9 +32,18 @@
     "test:ci": "vitest run"
   },
   "peerDependencies": {
+    "@electric-sql/pglite": ">=0.2.0",
     "drizzle-orm": "<1",
     "pg": ">=8"
   },
+  "peerDependenciesMeta": {
+    "@electric-sql/pglite": {
+      "optional": true
+    },
+    "pg": {
+      "optional": true
+    }
+  },
   "devDependencies": {
     "@electric-sql/pglite": "^0.2.17",
     "@types/node": "^20.14.0",
@@ -45,8 +54,8 @@
     "vitest": "^1.6.0"
   },
   "dependencies": {
-    "@apibara/indexer": "2.1.0-beta.20",
-    "@apibara/protocol": "2.1.0-beta.20",
+    "@apibara/indexer": "2.1.0-beta.22",
+    "@apibara/protocol": "2.1.0-beta.22",
     "postgres-range": "^1.1.4"
   }
 }

package/src/helper.ts

@@ -1,18 +1,10 @@
-import { PGlite, type PGliteOptions } from "@electric-sql/pglite";
+import type { PGlite, PGliteOptions } from "@electric-sql/pglite";
 import type { DrizzleConfig } from "drizzle-orm";
+import { entityKind } from "drizzle-orm";
 import type { MigrationConfig } from "drizzle-orm/migrator";
-import {
-  type NodePgDatabase as OriginalNodePgDatabase,
-  drizzle as drizzleNode,
-} from "drizzle-orm/node-postgres";
-import { migrate as migrateNode } from "drizzle-orm/node-postgres/migrator";
-import {} from "drizzle-orm/pg-core";
-import {
-  type PgliteDatabase as OriginalPgliteDatabase,
-  drizzle as drizzlePGLite,
-} from "drizzle-orm/pglite";
-import { migrate as migratePGLite } from "drizzle-orm/pglite/migrator";
-import pg from "pg";
+import type { NodePgDatabase as OriginalNodePgDatabase } from "drizzle-orm/node-postgres";
+import type { PgliteDatabase as OriginalPgliteDatabase } from "drizzle-orm/pglite";
+import type pg from "pg";
 import { DrizzleStorageError } from "./utils";
 
 /**
@@ -129,21 +121,9 @@ export function drizzle<
     poolConfig,
   } = options ?? {};
 
-  if (
-    isPostgresConnectionString(connectionString) ||
-    type === "node-postgres"
-  ) {
-    const pool = new pg.Pool({
-      connectionString,
-      ...(poolConfig || {}),
-    });
-    return drizzleNode(pool, { schema, ...(config || {}) }) as Database<
-      TOptions,
-      TSchema
-    >;
-  }
+  if (isPgliteConnectionString(connectionString) && type === "pglite") {
+    const { drizzle: drizzlePGLite } = require("drizzle-orm/pglite");
 
-  if (type === "pglite") {
     return drizzlePGLite({
       schema: schema as TSchema,
       connection: {
@@ -153,7 +133,16 @@ export function drizzle<
     }) as Database<TOptions, TSchema>;
   }
 
-  throw new Error("Invalid database type");
+  const { Pool } = require("pg");
+  const { drizzle: drizzleNode } = require("drizzle-orm/node-postgres");
+  const pool = new Pool({
+    connectionString,
+    ...(poolConfig || {}),
+  });
+  return drizzleNode(pool, { schema, ...(config || {}) }) as Database<
+    TOptions,
+    TSchema
+  >;
 }
 
 /**
@@ -179,12 +168,16 @@ export async function migrate<TSchema extends Record<string, unknown>>(
   db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>,
   options: MigrateOptions,
 ) {
-  const isPglite = !!("$client" in db && db.$client instanceof PGlite);
+  const isPglite = isDrizzleKind(db, "PgliteDatabase");
 
   try {
     if (isPglite) {
+      const { migrate: migratePGLite } = require("drizzle-orm/pglite/migrator");
       await migratePGLite(db as PgliteDatabase<TSchema>, options);
     } else {
+      const {
+        migrate: migrateNode,
+      } = require("drizzle-orm/node-postgres/migrator");
       await migrateNode(db as NodePgDatabase<TSchema>, options);
     }
   } catch (error) {
@@ -197,6 +190,30 @@ export async function migrate<TSchema extends Record<string, unknown>>(
   }
 }
 
-function isPostgresConnectionString(conn: string) {
-  return conn.startsWith("postgres://") || conn.startsWith("postgresql://");
+function isPgliteConnectionString(conn: string) {
+  return (
+    conn.startsWith("memory://") ||
+    conn.startsWith("file://") ||
+    conn.startsWith("idb://")
+  );
+}
+
+function isDrizzleKind(value: unknown, entityKindValue: string) {
+  if (!value || typeof value !== "object") {
+    return false;
+  }
+  // https://github.com/drizzle-team/drizzle-orm/blob/f39f885779800982e90dd3c89aba6df3217a6fd2/drizzle-orm/src/entity.ts#L29-L41
+  let cls = Object.getPrototypeOf(value).constructor;
+  if (cls) {
+    // Traverse the prototype chain to find the entityKind
+    while (cls) {
+      // https://github.com/drizzle-team/drizzle-orm/blob/f39f885779800982e90dd3c89aba6df3217a6fd2/drizzle-orm/src/pglite/driver.ts#L41
+      if (entityKind in cls && cls[entityKind] === entityKindValue) {
+        return true;
+      }
+      cls = Object.getPrototypeOf(cls);
+    }
+  }
+
+  return false;
 }

package/src/index.ts

@@ -32,10 +32,18 @@ import {
   registerTriggers,
   removeTriggers,
 } from "./storage";
-import { DrizzleStorageError, sleep, withTransaction } from "./utils";
+import {
+  DrizzleStorageError,
+  type IdColumnMap,
+  getIdColumnForTable,
+  sleep,
+  withTransaction,
+} from "./utils";
 
 export * from "./helper";
 
+export type { IdColumnMap };
+
 const MAX_RETRIES = 5;
 
 export type DrizzleStorage<
@@ -89,9 +97,32 @@ export interface DrizzleStorageOptions<
    */
   schema?: Record<string, unknown>;
   /**
-   * The column to use as the id. Defaults to 'id'.
+   * The column to use as the primary identifier for each table.
+   *
+   * This identifier is used for tracking changes during reorgs and rollbacks.
+   *
+   * Can be specified in two ways:
+   *
+   * 1. As a single string that applies to all tables:
+   * ```ts
+   * idColumn: "_id" // Uses "_id" column for all tables
+   * ```
+   *
+   * 2. As an object mapping table names to their ID columns:
+   * ```ts
+   * idColumn: {
+   *   transfers: "transaction_hash",    // Use "transaction_hash" for transfers table
+   *   blocks: "block_number",           // Use "block_number" for blocks table
+   *   "*": "_id"                        // Use "_id" for all other tables | defaults to "id"
+   * }
+   * ```
+   *
+   * The special "*" key acts as a fallback for any tables not explicitly mapped.
+   *
+   * @default "id"
+   * @type {string | Partial<IdColumnMap>}
    */
-  idColumn?: string;
+  idColumn?: string | Partial<IdColumnMap>;
   /**
    * The options for the database migration. When provided, the database will automatically run migrations before the indexer runs.
    */
@@ -120,8 +151,8 @@ export function drizzleStorage<
   db,
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
-  schema,
-  idColumn = "id",
+  schema: _schema,
+  idColumn,
   migrate: migrateOptions,
 }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>) {
   return defineIndexerPlugin<TFilter, TBlock>((indexer) => {
@@ -129,17 +160,37 @@ export function drizzleStorage<
     let indexerId = "";
     const alwaysReindex = process.env["APIBARA_ALWAYS_REINDEX"] === "true";
     let prevFinality: DataFinality | undefined;
+    const schema: TSchema = (_schema as TSchema) ?? db._.schema ?? {};
+    const idColumnMap: IdColumnMap = {
+      "*": typeof idColumn === "string" ? idColumn : "id",
+      ...(typeof idColumn === "object" ? idColumn : {}),
+    };
 
     try {
-      tableNames = Object.values((schema as TSchema) ?? db._.schema ?? {}).map(
-        (table) => table.dbName,
-      );
+      tableNames = Object.values(schema).map((table) => table.dbName);
     } catch (error) {
       throw new DrizzleStorageError("Failed to get table names from schema", {
         cause: error,
       });
     }
 
+    // Check if specified idColumn exists in all the tables in schema
+    for (const table of Object.values(schema)) {
+      const columns = table.columns;
+      const tableIdColumn = getIdColumnForTable(table.dbName, idColumnMap);
+
+      const columnExists = Object.values(columns).some(
+        (column) => column.name === tableIdColumn,
+      );
+
+      if (!columnExists) {
+        throw new DrizzleStorageError(
+          `Column \`"${tableIdColumn}"\` does not exist in table \`"${table.dbName}"\`. ` +
+            "Make sure the table has the specified column or provide a valid `idColumn` mapping to `drizzleStorage`.",
+        );
+      }
+    }
+
     indexer.hooks.hook("run:before", async () => {
       const internalContext = useInternalContext();
       const context = useIndexerContext();
@@ -153,25 +204,11 @@ export function drizzleStorage<
 
       indexerId = generateIndexerId(indexerFileName, identifier);
 
-      if (alwaysReindex) {
-        logger.warn(
-          `Reindexing: Deleting all data from tables - ${tableNames.join(", ")}`,
-        );
-        await withTransaction(db, async (tx) => {
-          await cleanupStorage(tx, tableNames, indexerId);
-
-          if (enablePersistence) {
-            await resetPersistence({ tx, indexerId });
-          }
-
-          logger.success("Tables have been cleaned up for reindexing");
-        });
-      }
-
       let retries = 0;
 
       // incase the migrations are already applied, we don't want to run them again
       let migrationsApplied = false;
+      let cleanupApplied = false;
 
       while (retries <= MAX_RETRIES) {
         try {
@@ -186,6 +223,22 @@ export function drizzleStorage<
             if (enablePersistence) {
               await initializePersistentState(tx);
             }
+
+            if (alwaysReindex && !cleanupApplied) {
+              logger.warn(
+                `Reindexing: Deleting all data from tables - ${tableNames.join(", ")}`,
+              );
+
+              await cleanupStorage(tx, tableNames, indexerId);
+
+              if (enablePersistence) {
+                await resetPersistence({ tx, indexerId });
+              }
+
+              cleanupApplied = true;
+
+              logger.success("Tables have been cleaned up for reindexing");
+            }
           });
           break;
         } catch (error) {
@@ -239,7 +292,8 @@ export function drizzleStorage<
       }
 
       await withTransaction(db, async (tx) => {
-        await invalidate(tx, cursor, idColumn, indexerId);
+        // Use the appropriate idColumn for each table when calling invalidate
+        await invalidate(tx, cursor, idColumnMap, indexerId);
 
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
@@ -289,7 +343,8 @@ export function drizzleStorage<
       }
 
       await withTransaction(db, async (tx) => {
-        await invalidate(tx, cursor, idColumn, indexerId);
+        // Use the appropriate idColumn for each table when calling invalidate
+        await invalidate(tx, cursor, idColumnMap, indexerId);
 
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
@@ -319,7 +374,7 @@ export function drizzleStorage<
 
             if (prevFinality === "pending") {
               // invalidate if previous block's finality was "pending"
-              await invalidate(tx, cursor, idColumn, indexerId);
+              await invalidate(tx, cursor, idColumnMap, indexerId);
             }
 
             if (finality !== "finalized") {
@@ -327,7 +382,7 @@ export function drizzleStorage<
                 tx,
                 tableNames,
                 endCursor,
-                idColumn,
+                idColumnMap,
                 indexerId,
               );
             }

package/src/storage.ts

@@ -16,7 +16,11 @@ import {
   text,
 } from "drizzle-orm/pg-core";
 import { SCHEMA_NAME } from "./constants";
-import { DrizzleStorageError } from "./utils";
+import {
+  DrizzleStorageError,
+  type IdColumnMap,
+  getIdColumnForTable,
+} from "./utils";
 
 const ROLLBACK_TABLE_NAME = "reorg_rollback";
 
@@ -125,11 +129,14 @@ export async function registerTriggers<
   tx: PgTransaction<TQueryResult, TFullSchema, TSchema>,
   tables: string[],
   endCursor: Cursor,
-  idColumn: string,
+  idColumnMap: IdColumnMap,
   indexerId: string,
 ) {
   try {
     for (const table of tables) {
+      // Determine the column ID for this specific table
+      const tableIdColumn = getIdColumnForTable(table, idColumnMap);
+
       await tx.execute(
         sql.raw(
           `DROP TRIGGER IF EXISTS ${getReorgTriggerName(table, indexerId)} ON ${table};`,
@@ -140,7 +147,7 @@ export async function registerTriggers<
           CREATE CONSTRAINT TRIGGER ${getReorgTriggerName(table, indexerId)}
           AFTER INSERT OR UPDATE OR DELETE ON ${table}
           DEFERRABLE INITIALLY DEFERRED
-          FOR EACH ROW EXECUTE FUNCTION ${SCHEMA_NAME}.reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
+          FOR EACH ROW EXECUTE FUNCTION ${SCHEMA_NAME}.reorg_checkpoint('${tableIdColumn}', ${Number(endCursor.orderKey)}, '${indexerId}');
         `),
       );
     }
@@ -184,7 +191,7 @@ export async function invalidate<
 >(
   tx: PgTransaction<TQueryResult, TFullSchema, TSchema>,
   cursor: Cursor,
-  idColumn: string,
+  idColumnMap: IdColumnMap,
   indexerId: string,
 ) {
   // Get and delete operations after cursor in one query, ordered by newest first
@@ -208,6 +215,9 @@ export async function invalidate<
 
   // Process each operation in reverse order
   for (const op of result) {
+    // Determine the column ID for this specific table
+    const tableIdColumn = getIdColumnForTable(op.table_name, idColumnMap);
+
     switch (op.op) {
       case "I":
         try {
@@ -218,7 +228,7 @@ export async function invalidate<
           await tx.execute(
             sql.raw(`
                 DELETE FROM ${op.table_name}
-                WHERE ${idColumn} = '${op.row_id}'
+                WHERE ${tableIdColumn} = '${op.row_id}'
               `),
           );
         } catch (error) {
@@ -271,7 +281,9 @@ export async function invalidate<
               ? JSON.parse(op.row_value)
               : op.row_value;
 
-          const nonIdKeys = Object.keys(rowValue).filter((k) => k !== idColumn);
+          const nonIdKeys = Object.keys(rowValue).filter(
+            (k) => k !== tableIdColumn,
+          );
 
           const fields = nonIdKeys.map((c) => `${c} = prev.${c}`).join(", ");
 
@@ -281,7 +293,7 @@ export async function invalidate<
               FROM (
                 SELECT * FROM json_populate_record(null::${op.table_name}, '${JSON.stringify(op.row_value)}'::json)
               ) as prev
-              WHERE ${op.table_name}.${idColumn} = '${op.row_id}'
+              WHERE ${op.table_name}.${tableIdColumn} = '${op.row_id}'
               `);
 
           await tx.execute(query);

package/src/utils.ts

@@ -48,3 +48,22 @@ export function serialize<T>(obj: T): string {
 export function sleep(ms: number) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
+
+export interface IdColumnMap extends Record<string, string> {
+  /**
+   * Wildcard mapping for all tables.
+   */
+  "*": string;
+}
+
+export const getIdColumnForTable = (
+  tableName: string,
+  idColumn: IdColumnMap,
+): string => {
+  // If there's a specific mapping for this table, use it
+  if (idColumn[tableName]) {
+    return idColumn[tableName];
+  }
+  // Default fallback
+  return idColumn["*"];
+};