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

package/dist/index.cjs

@@ -46,6 +46,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 {
@@ -375,9 +381,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 +395,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 +420,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 +438,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 +448,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 +488,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 +498,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 +576,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 +585,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 +616,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 +631,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 +684,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 +722,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 +738,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

@@ -40,6 +40,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 {
@@ -369,9 +375,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 +389,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 +414,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 +432,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 +442,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 +482,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 +492,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 +570,8 @@ function drizzleStorage({
   db,
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
-  schema,
-  idColumn = "id",
+  schema: _schema,
+  idColumn,
   migrate: migrateOptions
 }) {
   return defineIndexerPlugin((indexer) => {
@@ -569,15 +579,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 +610,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 +625,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 +678,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 +716,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 +732,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.21",
   "type": "module",
   "files": [
     "dist",
@@ -45,8 +45,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.21",
+    "@apibara/protocol": "2.1.0-beta.21",
     "postgres-range": "^1.1.4"
   }
 }

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["*"];
+};