@apibara/plugin-drizzle 2.1.0-beta.4 → 2.1.0-beta.41

package/dist/index.cjs

@@ -4,8 +4,9 @@ const indexer = require('@apibara/indexer');
 const plugins = require('@apibara/indexer/plugins');
 const internal = require('@apibara/indexer/internal');
 const plugins$1 = require('@apibara/indexer/internal/plugins');
-const protocol = require('@apibara/protocol');
+const constants = require('./shared/plugin-drizzle.cae20704.cjs');
 const drizzleOrm = require('drizzle-orm');
+const protocol = require('@apibara/protocol');
 const pgCore = require('drizzle-orm/pg-core');
 
 class DrizzleStorageError extends Error {
@@ -35,16 +36,89 @@ 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 {
+    connectionString = process.env["POSTGRES_CONNECTION_STRING"] ?? "memory://",
+    schema,
+    type = "pglite",
+    config,
+    poolConfig
+  } = options ?? {};
+  if (isPgliteConnectionString(connectionString) && type === "pglite") {
+    const { drizzle: drizzlePGLite } = require("drizzle-orm/pglite");
+    return drizzlePGLite({
+      schema,
+      connection: {
+        dataDir: connectionString || "memory://pglite"
+      },
+      ...config || {}
+    });
+  }
+  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 = isDrizzleKind(db, "PgliteDatabase");
+  try {
+    if (isPglite) {
+      const { migrate: migratePGLite } = require("drizzle-orm/pglite/migrator");
+      await migratePGLite(db, options);
+    } else {
+      const {
+        migrate: migrateNode
+      } = require("drizzle-orm/node-postgres/migrator");
+      await migrateNode(db, options);
+    }
+  } catch (error) {
+    throw new DrizzleStorageError(
+      "Failed to apply migrations! Please check if you have generated migrations using drizzle:generate",
+      {
+        cause: error
+      }
+    );
+  }
+}
+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 = "__indexer_checkpoints";
-const FILTERS_TABLE_NAME = "__indexer_filters";
-const SCHEMA_VERSION_TABLE_NAME = "__indexer_schema_version";
-const checkpoints = pgCore.pgTable(CHECKPOINTS_TABLE_NAME, {
+const CHECKPOINTS_TABLE_NAME = "checkpoints";
+const FILTERS_TABLE_NAME = "filters";
+const SCHEMA_VERSION_TABLE_NAME = "schema_version";
+const schema$1 = pgCore.pgSchema(constants.SCHEMA_NAME);
+const checkpoints = schema$1.table(CHECKPOINTS_TABLE_NAME, {
   id: pgCore.text("id").notNull().primaryKey(),
   orderKey: pgCore.integer("order_key").notNull(),
   uniqueKey: pgCore.text("unique_key")
 });
-const filters = pgCore.pgTable(
+const filters = schema$1.table(
   FILTERS_TABLE_NAME,
   {
     id: pgCore.text("id").notNull(),
@@ -58,7 +132,16 @@ const filters = pgCore.pgTable(
     }
   ]
 );
-const schemaVersion = pgCore.pgTable(SCHEMA_VERSION_TABLE_NAME, {
+const chainReorganizations = schema$1.table("chain_reorganizations", {
+  id: pgCore.serial("id").primaryKey(),
+  indexerId: pgCore.text("indexer_id").notNull(),
+  oldHeadOrderKey: pgCore.integer("old_head_order_key"),
+  oldHeadUniqueKey: pgCore.text("old_head_unique_key").$type().default(null),
+  newHeadOrderKey: pgCore.integer("new_head_order_key").notNull(),
+  newHeadUniqueKey: pgCore.text("new_head_unique_key").$type().default(null),
+  recordedAt: pgCore.timestamp("recorded_at").defaultNow().notNull()
+});
+const schemaVersion = schema$1.table(SCHEMA_VERSION_TABLE_NAME, {
   k: pgCore.integer("k").notNull().primaryKey(),
   version: pgCore.integer("version").notNull()
 });
@@ -69,12 +152,38 @@ const MIGRATIONS = [
   // Add more migration arrays for future versions
 ];
 async function initializePersistentState(tx) {
-  await tx.execute(`
-    CREATE TABLE IF NOT EXISTS ${SCHEMA_VERSION_TABLE_NAME} (
+  await tx.execute(
+    drizzleOrm.sql.raw(`
+      CREATE SCHEMA IF NOT EXISTS ${constants.SCHEMA_NAME};
+  `)
+  );
+  await tx.execute(
+    drizzleOrm.sql.raw(`
+    CREATE TABLE IF NOT EXISTS ${constants.SCHEMA_NAME}.${SCHEMA_VERSION_TABLE_NAME} (
       k INTEGER PRIMARY KEY,
       version INTEGER NOT NULL
     );
-  `);
+  `)
+  );
+  await tx.execute(
+    drizzleOrm.sql.raw(`
+      CREATE TABLE IF NOT EXISTS ${constants.SCHEMA_NAME}.chain_reorganizations (
+        id SERIAL PRIMARY KEY,
+        indexer_id TEXT NOT NULL,
+        old_head_order_key INTEGER,
+        old_head_unique_key TEXT DEFAULT NULL,
+        new_head_order_key INTEGER NOT NULL,
+        new_head_unique_key TEXT DEFAULT NULL,
+        recorded_at TIMESTAMP NOT NULL DEFAULT NOW()
+      );
+    `)
+  );
+  await tx.execute(
+    drizzleOrm.sql.raw(`
+      CREATE INDEX IF NOT EXISTS idx_chain_reorgs_indexer_id 
+      ON ${constants.SCHEMA_NAME}.chain_reorganizations(indexer_id);
+    `)
+  );
   const versionRows = await tx.select().from(schemaVersion).where(drizzleOrm.eq(schemaVersion.k, 0));
   const storedVersion = versionRows[0]?.version ?? -1;
   if (storedVersion > CURRENT_SCHEMA_VERSION) {
@@ -84,22 +193,26 @@ async function initializePersistentState(tx) {
   }
   try {
     if (storedVersion === -1) {
-      await tx.execute(`
-        CREATE TABLE IF NOT EXISTS ${CHECKPOINTS_TABLE_NAME} (
+      await tx.execute(
+        drizzleOrm.sql.raw(`
+        CREATE TABLE IF NOT EXISTS ${constants.SCHEMA_NAME}.${CHECKPOINTS_TABLE_NAME} (
           id TEXT PRIMARY KEY,
           order_key INTEGER NOT NULL,
           unique_key TEXT
         );
-      `);
-      await tx.execute(`
-        CREATE TABLE IF NOT EXISTS ${FILTERS_TABLE_NAME} (
+      `)
+      );
+      await tx.execute(
+        drizzleOrm.sql.raw(`
+        CREATE TABLE IF NOT EXISTS ${constants.SCHEMA_NAME}.${FILTERS_TABLE_NAME} (
           id TEXT NOT NULL,
           filter TEXT NOT NULL,
           from_block INTEGER NOT NULL,
           to_block INTEGER DEFAULT NULL,
           PRIMARY KEY (id, from_block)
         );
-      `);
+      `)
+      );
       await tx.insert(schemaVersion).values({
         k: 0,
         version: CURRENT_SCHEMA_VERSION
@@ -122,6 +235,22 @@ async function initializePersistentState(tx) {
     );
   }
 }
+async function recordChainReorganization(props) {
+  const { tx, indexerId, oldHead, newHead } = props;
+  try {
+    await tx.insert(chainReorganizations).values({
+      indexerId,
+      oldHeadOrderKey: oldHead ? Number(oldHead.orderKey) : null,
+      oldHeadUniqueKey: oldHead?.uniqueKey ? oldHead.uniqueKey : null,
+      newHeadOrderKey: Number(newHead.orderKey),
+      newHeadUniqueKey: newHead.uniqueKey ? newHead.uniqueKey : null
+    });
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to record chain reorganization", {
+      cause: error
+    });
+  }
+}
 async function persistState(props) {
   const { tx, endCursor, filter, indexerId } = props;
   try {
@@ -134,7 +263,9 @@ async function persistState(props) {
         target: checkpoints.id,
         set: {
           orderKey: Number(endCursor.orderKey),
-          uniqueKey: endCursor.uniqueKey
+          // Explicitly set the unique key to `null` to indicate that it has been deleted
+          // Otherwise drizzle will not update its value.
+          uniqueKey: endCursor.uniqueKey ? endCursor.uniqueKey : null
         }
       });
       if (filter) {
@@ -180,6 +311,10 @@ async function getState(props) {
 async function invalidateState(props) {
   const { tx, cursor, indexerId } = props;
   try {
+    await tx.update(checkpoints).set({
+      orderKey: Number(cursor.orderKey),
+      uniqueKey: cursor.uniqueKey ? cursor.uniqueKey : null
+    }).where(drizzleOrm.eq(checkpoints.id, indexerId));
     await tx.delete(filters).where(
       drizzleOrm.and(
         drizzleOrm.eq(filters.id, indexerId),
@@ -213,11 +348,24 @@ async function finalizeState(props) {
     });
   }
 }
+async function resetPersistence(props) {
+  const { tx, indexerId } = props;
+  try {
+    await tx.delete(checkpoints).where(drizzleOrm.eq(checkpoints.id, indexerId));
+    await tx.delete(filters).where(drizzleOrm.eq(filters.id, indexerId));
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to reset persistence state", {
+      cause: error
+    });
+  }
+}
 
+const ROLLBACK_TABLE_NAME = "reorg_rollback";
+const schema = pgCore.pgSchema(constants.SCHEMA_NAME);
 function getReorgTriggerName(table, indexerId) {
   return `${table}_reorg_${indexerId}`;
 }
-pgCore.pgTable("__reorg_rollback", {
+schema.table(ROLLBACK_TABLE_NAME, {
   n: pgCore.serial("n").primaryKey(),
   op: pgCore.char("op", { length: 1 }).$type().notNull(),
   table_name: pgCore.text("table_name").notNull(),
@@ -228,9 +376,12 @@ pgCore.pgTable("__reorg_rollback", {
 });
 async function initializeReorgRollbackTable(tx, indexerId) {
   try {
+    await tx.execute(`
+    CREATE SCHEMA IF NOT EXISTS ${constants.SCHEMA_NAME};
+    `);
     await tx.execute(
       drizzleOrm.sql.raw(`
-        CREATE TABLE IF NOT EXISTS __reorg_rollback(
+        CREATE TABLE IF NOT EXISTS ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(
           n SERIAL PRIMARY KEY,
           op CHAR(1) NOT NULL,
           table_name TEXT NOT NULL,
@@ -243,7 +394,7 @@ async function initializeReorgRollbackTable(tx, indexerId) {
     );
     await tx.execute(
       drizzleOrm.sql.raw(`
-        CREATE INDEX IF NOT EXISTS idx_reorg_rollback_indexer_id_cursor ON __reorg_rollback(indexer_id, cursor);
+        CREATE INDEX IF NOT EXISTS idx_reorg_rollback_indexer_id_cursor ON ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(indexer_id, cursor);
       `)
     );
   } catch (error) {
@@ -254,24 +405,25 @@ async function initializeReorgRollbackTable(tx, indexerId) {
   try {
     await tx.execute(
       drizzleOrm.sql.raw(`
-      CREATE OR REPLACE FUNCTION reorg_checkpoint()
+      CREATE OR REPLACE FUNCTION ${constants.SCHEMA_NAME}.reorg_checkpoint()
       RETURNS TRIGGER AS $$
       DECLARE
-        id_col TEXT := TG_ARGV[0]::TEXT;
-        order_key INTEGER := TG_ARGV[1]::INTEGER;
-        indexer_id TEXT := TG_ARGV[2]::TEXT;
+        table_name TEXT := TG_ARGV[0]::TEXT;
+        id_col TEXT := TG_ARGV[1]::TEXT;
+        order_key INTEGER := TG_ARGV[2]::INTEGER;
+        indexer_id TEXT := TG_ARGV[3]::TEXT;
         new_id_value TEXT := row_to_json(NEW.*)->>id_col;
         old_id_value TEXT := row_to_json(OLD.*)->>id_col;
       BEGIN
         IF (TG_OP = 'DELETE') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
-            SELECT 'D', TG_TABLE_NAME, order_key, old_id_value, row_to_json(OLD.*), indexer_id;
+          INSERT INTO ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
+            SELECT 'D', table_name, order_key, old_id_value, row_to_json(OLD.*), indexer_id;
         ELSIF (TG_OP = 'UPDATE') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
-            SELECT 'U', TG_TABLE_NAME, order_key, new_id_value, row_to_json(OLD.*), indexer_id;
+          INSERT INTO ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
+            SELECT 'U', table_name, order_key, new_id_value, row_to_json(OLD.*), indexer_id;
         ELSIF (TG_OP = 'INSERT') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
-            SELECT 'I', TG_TABLE_NAME, order_key, new_id_value, null, indexer_id;
+          INSERT INTO ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
+            SELECT 'I', table_name, order_key, new_id_value, null, indexer_id;
         END IF;
         RETURN NULL;
       END;
@@ -287,9 +439,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};`
@@ -300,7 +453,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 reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
+          FOR EACH ROW EXECUTE FUNCTION ${constants.SCHEMA_NAME}.reorg_checkpoint('${table}', '${tableIdColumn}', ${Number(endCursor.orderKey)}, '${indexerId}');
         `)
       );
     }
@@ -325,11 +478,11 @@ 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 (
-        DELETE FROM __reorg_rollback
+        DELETE FROM ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}
         WHERE cursor > ${Number(cursor.orderKey)}
         AND indexer_id = '${indexerId}'
         RETURNING *
@@ -343,6 +496,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 {
@@ -352,7 +506,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) {
@@ -392,7 +546,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}
@@ -400,7 +556,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) {
@@ -422,7 +578,7 @@ async function finalize(tx, cursor, indexerId) {
   try {
     await tx.execute(
       drizzleOrm.sql.raw(`
-      DELETE FROM __reorg_rollback
+      DELETE FROM ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}
       WHERE cursor <= ${Number(cursor.orderKey)}
       AND indexer_id = '${indexerId}'
     `)
@@ -433,52 +589,134 @@ async function finalize(tx, cursor, indexerId) {
     });
   }
 }
+async function cleanupStorage(tx, tables, indexerId) {
+  try {
+    for (const table of tables) {
+      await tx.execute(
+        drizzleOrm.sql.raw(
+          `DROP TRIGGER IF EXISTS ${getReorgTriggerName(table, indexerId)} ON ${table};`
+        )
+      );
+    }
+    await tx.execute(
+      drizzleOrm.sql.raw(`
+        DELETE FROM ${constants.SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}
+        WHERE indexer_id = '${indexerId}'
+      `)
+    );
+    for (const table of tables) {
+      try {
+        await tx.execute(drizzleOrm.sql.raw(`TRUNCATE TABLE ${table} CASCADE;`));
+      } catch (error) {
+        throw new DrizzleStorageError(`Failed to truncate table ${table}`, {
+          cause: error
+        });
+      }
+    }
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to clean up storage", {
+      cause: error
+    });
+  }
+}
 
-const DRIZZLE_PROPERTY = "_drizzle";
 const MAX_RETRIES = 5;
 function useDrizzleStorage(_db) {
   const context = indexer.useIndexerContext();
-  if (!context[DRIZZLE_PROPERTY]) {
+  if (!context[constants.DRIZZLE_PROPERTY]) {
     throw new DrizzleStorageError(
       "drizzle storage is not available. Did you register the plugin?"
     );
   }
-  return context[DRIZZLE_PROPERTY];
+  return context[constants.DRIZZLE_PROPERTY];
+}
+function useTestDrizzleStorage() {
+  const context = indexer.useIndexerContext();
+  if (!context[constants.DRIZZLE_STORAGE_DB_PROPERTY]) {
+    throw new DrizzleStorageError(
+      "drizzle storage db is not available. Did you register the plugin?"
+    );
+  }
+  return context[constants.DRIZZLE_STORAGE_DB_PROPERTY];
 }
 function drizzleStorage({
   db,
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
-  schema,
-  idColumn = "id"
+  schema: _schema,
+  idColumn,
+  migrate: migrateOptions,
+  recordChainReorganizations = false
 }) {
-  return plugins.defineIndexerPlugin((indexer) => {
+  return plugins.defineIndexerPlugin((indexer$1) => {
     let tableNames = [];
     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
       });
     }
-    indexer.hooks.hook("run:before", async () => {
-      const { indexerName: indexerFileName, availableIndexers } = plugins$1.useInternalContext();
+    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("plugins:init", async () => {
+      const internalContext = plugins$1.useInternalContext();
+      const context = indexer.useIndexerContext();
+      const logger = plugins.useLogger();
+      context[constants.DRIZZLE_STORAGE_DB_PROPERTY] = db;
+      const { indexerName: indexerFileName, availableIndexers } = internalContext;
       indexerId = internal.generateIndexerId(indexerFileName, identifier);
       let retries = 0;
+      let migrationsApplied = false;
+      let cleanupApplied = false;
       while (retries <= MAX_RETRIES) {
         try {
+          if (migrateOptions && !migrationsApplied) {
+            await migrate(db, migrateOptions);
+            migrationsApplied = true;
+            logger.success("Migrations applied");
+          }
           await withTransaction(db, async (tx) => {
             await initializeReorgRollbackTable(tx, indexerId);
             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) {
           if (retries === MAX_RETRIES) {
+            if (error instanceof DrizzleStorageError) {
+              throw error;
+            }
             throw new DrizzleStorageError(
               "Initialization failed after 5 retries",
               {
@@ -491,7 +729,7 @@ function drizzleStorage({
         }
       }
     });
-    indexer.hooks.hook("connect:before", async ({ request }) => {
+    indexer$1.hooks.hook("connect:before", async ({ request }) => {
       if (!enablePersistence) {
         return;
       }
@@ -508,19 +746,19 @@ function drizzleStorage({
         }
       });
     });
-    indexer.hooks.hook("connect:after", async ({ request }) => {
+    indexer$1.hooks.hook("connect:after", async ({ request }) => {
       const cursor = request.startingCursor;
       if (!cursor) {
         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 });
         }
       });
     });
-    indexer.hooks.hook("connect:factory", async ({ request, endCursor }) => {
+    indexer$1.hooks.hook("connect:factory", async ({ request, endCursor }) => {
       if (!enablePersistence) {
         return;
       }
@@ -534,8 +772,8 @@ function drizzleStorage({
         });
       }
     });
-    indexer.hooks.hook("message:finalize", async ({ message }) => {
-      const { cursor } = message.finalize;
+    indexer$1.hooks.hook("message:finalize", async ({ message }) => {
+      const { cursor } = message;
       if (!cursor) {
         throw new DrizzleStorageError("Finalized Cursor is undefined");
       }
@@ -546,59 +784,79 @@ function drizzleStorage({
         }
       });
     });
-    indexer.hooks.hook("message:invalidate", async ({ message }) => {
-      const { cursor } = message.invalidate;
+    indexer$1.hooks.hook("message:invalidate", async ({ message }) => {
+      const { cursor } = message;
       if (!cursor) {
         throw new DrizzleStorageError("Invalidate Cursor is undefined");
       }
       await withTransaction(db, async (tx) => {
-        await invalidate(tx, cursor, idColumn, indexerId);
+        let oldHead;
+        if (recordChainReorganizations) {
+          const { cursor: currentCursor } = await getState({
+            tx,
+            indexerId
+          });
+          oldHead = currentCursor;
+          await recordChainReorganization({
+            tx,
+            indexerId,
+            oldHead,
+            newHead: cursor
+          });
+        }
+        await invalidate(tx, cursor, idColumnMap, indexerId);
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
         }
       });
     });
-    indexer.hooks.hook("handler:middleware", async ({ use }) => {
+    indexer$1.hooks.hook("handler:middleware", async ({ use }) => {
       use(async (context, next) => {
         try {
-          const { endCursor, finality } = context;
+          const { endCursor, finality, cursor } = context;
           if (!endCursor) {
             throw new DrizzleStorageError("End Cursor is undefined");
           }
           await withTransaction(db, async (tx) => {
-            context[DRIZZLE_PROPERTY] = { db: tx };
+            context[constants.DRIZZLE_PROPERTY] = { db: tx };
+            if (prevFinality === "pending") {
+              await invalidate(tx, cursor, idColumnMap, indexerId);
+            }
             if (finality !== "finalized") {
               await registerTriggers(
                 tx,
                 tableNames,
                 endCursor,
-                idColumn,
+                idColumnMap,
                 indexerId
               );
             }
             await next();
-            delete context[DRIZZLE_PROPERTY];
-            if (enablePersistence) {
+            delete context[constants.DRIZZLE_PROPERTY];
+            if (enablePersistence && finality !== "pending") {
               await persistState({
                 tx,
                 endCursor,
                 indexerId
               });
             }
+            prevFinality = finality;
           });
           if (finality !== "finalized") {
             await removeTriggers(db, tableNames, indexerId);
           }
         } catch (error) {
           await removeTriggers(db, tableNames, indexerId);
-          throw new DrizzleStorageError("Failed to run handler:middleware", {
-            cause: error
-          });
+          throw error;
         }
       });
     });
   });
 }
 
+exports.drizzle = drizzle;
 exports.drizzleStorage = drizzleStorage;
+exports.migrate = migrate;
 exports.useDrizzleStorage = useDrizzleStorage;
+exports.useTestDrizzleStorage = useTestDrizzleStorage;
+//# sourceMappingURL=index.cjs.map