@apibara/plugin-drizzle 2.1.0-beta.5 → 2.1.0-beta.50

package/src/index.ts

@@ -1,5 +1,5 @@
 import { useIndexerContext } from "@apibara/indexer";
-import { defineIndexerPlugin } from "@apibara/indexer/plugins";
+import { defineIndexerPlugin, useLogger } from "@apibara/indexer/plugins";
 
 import type {
   ExtractTablesWithRelations,
@@ -14,23 +14,37 @@ import type {
   PgQueryResultHKT,
   PgTransaction,
 } from "drizzle-orm/pg-core";
+import { DRIZZLE_PROPERTY, DRIZZLE_STORAGE_DB_PROPERTY } from "./constants";
+import { type MigrateOptions, migrate } from "./helper";
 import {
   finalizeState,
   getState,
   initializePersistentState,
   invalidateState,
   persistState,
+  recordChainReorganization,
+  resetPersistence,
 } from "./persistence";
 import {
+  cleanupStorage,
   finalize,
   initializeReorgRollbackTable,
   invalidate,
   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 DRIZZLE_PROPERTY = "_drizzle";
 const MAX_RETRIES = 5;
 
 export type DrizzleStorage<
@@ -61,17 +75,82 @@ export function useDrizzleStorage<
   return context[DRIZZLE_PROPERTY];
 }
 
+export function useTestDrizzleStorage<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(): PgDatabase<TQueryResult, TFullSchema, TSchema> {
+  const context = useIndexerContext();
+
+  if (!context[DRIZZLE_STORAGE_DB_PROPERTY]) {
+    throw new DrizzleStorageError(
+      "drizzle storage db is not available. Did you register the plugin?",
+    );
+  }
+
+  return context[DRIZZLE_STORAGE_DB_PROPERTY];
+}
+
 export interface DrizzleStorageOptions<
   TQueryResult extends PgQueryResultHKT,
   TFullSchema extends Record<string, unknown> = Record<string, never>,
   TSchema extends
     TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
 > {
+  /**
+   * The Drizzle database instance.
+   */
   db: PgDatabase<TQueryResult, TFullSchema, TSchema>;
+  /**
+   * Whether to persist the indexer's state. Defaults to true.
+   */
   persistState?: boolean;
+  /**
+   * The name of the indexer. Default value is 'default'.
+   */
   indexerName?: string;
+  /**
+   * The schema of the database.
+   */
   schema?: Record<string, unknown>;
-  idColumn?: string;
+  /**
+   * 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 | Partial<IdColumnMap>;
+  /**
+   * The options for the database migration. When provided, the database will automatically run migrations before the indexer runs.
+   */
+  migrate?: MigrateOptions;
+
+  /**
+   * Whether to record chain reorganizations in the database.
+   * @default false
+   */
+  recordChainReorganizations?: boolean;
 }
 
 /**
@@ -83,6 +162,8 @@ export interface DrizzleStorageOptions<
  * @param options.indexerName - The name of the indexer. Defaults value is 'default'.
  * @param options.schema - The schema of the database.
  * @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.
+ * @param options.recordChainReorganizations - Whether to record chain reorganizations in the database. Defaults to false.
  */
 export function drizzleStorage<
   TFilter,
@@ -95,42 +176,102 @@ export function drizzleStorage<
   db,
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
-  schema,
-  idColumn = "id",
+  schema: _schema,
+  idColumn,
+  migrate: migrateOptions,
+  recordChainReorganizations = false,
 }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>) {
   return defineIndexerPlugin<TFilter, TBlock>((indexer) => {
     let tableNames: string[] = [];
     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,
       });
     }
 
-    indexer.hooks.hook("run:before", async () => {
+    // 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("plugins:init", async () => {
+      const internalContext = useInternalContext();
+      const context = useIndexerContext();
+      const logger = useLogger();
+
+      // For testing purposes using vcr.
+      context[DRIZZLE_STORAGE_DB_PROPERTY] = db;
+
       const { indexerName: indexerFileName, availableIndexers } =
-        useInternalContext();
+        internalContext;
 
       indexerId = generateIndexerId(indexerFileName, identifier);
 
       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 {
+          if (migrateOptions && !migrationsApplied) {
+            // @ts-ignore type mismatch for db
+            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",
               {
@@ -177,7 +318,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 });
@@ -204,7 +346,7 @@ export function drizzleStorage<
     });
 
     indexer.hooks.hook("message:finalize", async ({ message }) => {
-      const { cursor } = message.finalize;
+      const { cursor } = message;
 
       if (!cursor) {
         throw new DrizzleStorageError("Finalized Cursor is undefined");
@@ -220,14 +362,36 @@ export function drizzleStorage<
     });
 
     indexer.hooks.hook("message:invalidate", async ({ message }) => {
-      const { cursor } = message.invalidate;
+      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: Cursor | undefined;
+
+        if (recordChainReorganizations) {
+          const { cursor: currentCursor } = await getState<
+            TFilter,
+            TQueryResult,
+            TFullSchema,
+            TSchema
+          >({
+            tx,
+            indexerId,
+          });
+          oldHead = currentCursor;
+
+          await recordChainReorganization({
+            tx,
+            indexerId,
+            oldHead,
+            newHead: cursor,
+          });
+        }
+
+        await invalidate(tx, cursor, idColumnMap, indexerId);
 
         if (enablePersistence) {
           await invalidateState({ tx, cursor, indexerId });
@@ -238,7 +402,8 @@ export function drizzleStorage<
     indexer.hooks.hook("handler:middleware", async ({ use }) => {
       use(async (context, next) => {
         try {
-          const { endCursor, finality } = context as {
+          const { endCursor, finality, cursor } = context as {
+            cursor: Cursor;
             endCursor: Cursor;
             finality: DataFinality;
           };
@@ -254,12 +419,17 @@ export function drizzleStorage<
               TSchema
             >;
 
+            if (prevFinality === "pending") {
+              // invalidate if previous block's finality was "pending"
+              await invalidate(tx, cursor, idColumnMap, indexerId);
+            }
+
             if (finality !== "finalized") {
               await registerTriggers(
                 tx,
                 tableNames,
                 endCursor,
-                idColumn,
+                idColumnMap,
                 indexerId,
               );
             }
@@ -267,13 +437,15 @@ export function drizzleStorage<
             await next();
             delete context[DRIZZLE_PROPERTY];
 
-            if (enablePersistence) {
+            if (enablePersistence && finality !== "pending") {
               await persistState({
                 tx,
                 endCursor,
                 indexerId,
               });
             }
+
+            prevFinality = finality;
           });
 
           if (finality !== "finalized") {
@@ -283,9 +455,7 @@ export function drizzleStorage<
         } catch (error) {
           await removeTriggers(db, tableNames, indexerId);
 
-          throw new DrizzleStorageError("Failed to run handler:middleware", {
-            cause: error,
-          });
+          throw error;
         }
       });
     });

package/src/persistence.ts

@@ -1,24 +1,36 @@
 import { type Cursor, normalizeCursor } from "@apibara/protocol";
-import { and, eq, gt, isNull, lt } from "drizzle-orm";
+import { and, eq, gt, isNull, lt, sql } from "drizzle-orm";
 import type {
   ExtractTablesWithRelations,
   TablesRelationalConfig,
 } from "drizzle-orm";
 import type { PgQueryResultHKT, PgTransaction } from "drizzle-orm/pg-core";
-import { integer, pgTable, primaryKey, text } from "drizzle-orm/pg-core";
+import {
+  integer,
+  pgSchema,
+  primaryKey,
+  serial,
+  text,
+  timestamp,
+} from "drizzle-orm/pg-core";
+import { SCHEMA_NAME } from "./constants";
 import { DrizzleStorageError, deserialize, serialize } from "./utils";
 
-const CHECKPOINTS_TABLE_NAME = "__indexer_checkpoints";
-const FILTERS_TABLE_NAME = "__indexer_filters";
-const SCHEMA_VERSION_TABLE_NAME = "__indexer_schema_version";
+const CHECKPOINTS_TABLE_NAME = "checkpoints";
+const FILTERS_TABLE_NAME = "filters";
+const SCHEMA_VERSION_TABLE_NAME = "schema_version";
 
-export const checkpoints = pgTable(CHECKPOINTS_TABLE_NAME, {
+const schema = pgSchema(SCHEMA_NAME);
+
+/** This table is not used for migrations, its only used for ease of internal operations with drizzle. */
+export const checkpoints = schema.table(CHECKPOINTS_TABLE_NAME, {
   id: text("id").notNull().primaryKey(),
   orderKey: integer("order_key").notNull(),
   uniqueKey: text("unique_key"),
 });
 
-export const filters = pgTable(
+/** This table is not used for migrations, its only used for ease of internal operations with drizzle. */
+export const filters = schema.table(
   FILTERS_TABLE_NAME,
   {
     id: text("id").notNull(),
@@ -33,7 +45,23 @@ export const filters = pgTable(
   ],
 );
 
-export const schemaVersion = pgTable(SCHEMA_VERSION_TABLE_NAME, {
+/** Table for recording chain reorganizations */
+export const chainReorganizations = schema.table("chain_reorganizations", {
+  id: serial("id").primaryKey(),
+  indexerId: text("indexer_id").notNull(),
+  oldHeadOrderKey: integer("old_head_order_key"),
+  oldHeadUniqueKey: text("old_head_unique_key")
+    .$type<string | null>()
+    .default(null),
+  newHeadOrderKey: integer("new_head_order_key").notNull(),
+  newHeadUniqueKey: text("new_head_unique_key")
+    .$type<string | null>()
+    .default(null),
+  recordedAt: timestamp("recorded_at").defaultNow().notNull(),
+});
+
+/** This table is not used for migrations, its only used for ease of internal operations with drizzle. */
+export const schemaVersion = schema.table(SCHEMA_VERSION_TABLE_NAME, {
   k: integer("k").notNull().primaryKey(),
   version: integer("version").notNull(),
 });
@@ -53,13 +81,43 @@ export async function initializePersistentState<
   TSchema extends
     TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
 >(tx: PgTransaction<TQueryResult, TFullSchema, TSchema>) {
+  // Create schema if it doesn't exist
+  await tx.execute(
+    sql.raw(`
+      CREATE SCHEMA IF NOT EXISTS ${SCHEMA_NAME};
+  `),
+  );
+
   // Create schema version table
-  await tx.execute(`
-    CREATE TABLE IF NOT EXISTS ${SCHEMA_VERSION_TABLE_NAME} (
+  await tx.execute(
+    sql.raw(`
+    CREATE TABLE IF NOT EXISTS ${SCHEMA_NAME}.${SCHEMA_VERSION_TABLE_NAME} (
       k INTEGER PRIMARY KEY,
       version INTEGER NOT NULL
     );
-  `);
+  `),
+  );
+
+  await tx.execute(
+    sql.raw(`
+      CREATE TABLE IF NOT EXISTS ${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(
+    sql.raw(`
+      CREATE INDEX IF NOT EXISTS idx_chain_reorgs_indexer_id 
+      ON ${SCHEMA_NAME}.chain_reorganizations(indexer_id);
+    `),
+  );
 
   // Get current schema version
   const versionRows = await tx
@@ -80,23 +138,27 @@ export async function initializePersistentState<
   try {
     if (storedVersion === -1) {
       // First time initialization
-      await tx.execute(`
-        CREATE TABLE IF NOT EXISTS ${CHECKPOINTS_TABLE_NAME} (
+      await tx.execute(
+        sql.raw(`
+        CREATE TABLE IF NOT EXISTS ${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(
+        sql.raw(`
+        CREATE TABLE IF NOT EXISTS ${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)
         );
-      `);
+      `),
+      );
 
       // Set initial schema version
       await tx.insert(schemaVersion).values({
@@ -128,6 +190,34 @@ export async function initializePersistentState<
   }
 }
 
+export async function recordChainReorganization<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(props: {
+  tx: PgTransaction<TQueryResult, TFullSchema, TSchema>;
+  indexerId: string;
+  oldHead: Cursor | undefined;
+  newHead: Cursor;
+}) {
+  const { tx, indexerId, oldHead, newHead } = props;
+
+  try {
+    await tx.insert(chainReorganizations).values({
+      indexerId: 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,
+    });
+  }
+}
+
 export async function persistState<
   TFilter,
   TQueryResult extends PgQueryResultHKT,
@@ -155,7 +245,9 @@ export async function persistState<
           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,
           },
         });
 
@@ -245,6 +337,14 @@ export async function invalidateState<
   const { tx, cursor, indexerId } = props;
 
   try {
+    await tx
+      .update(checkpoints)
+      .set({
+        orderKey: Number(cursor.orderKey),
+        uniqueKey: cursor.uniqueKey ? cursor.uniqueKey : null,
+      })
+      .where(eq(checkpoints.id, indexerId));
+
     await tx
       .delete(filters)
       .where(
@@ -297,3 +397,24 @@ export async function finalizeState<
     });
   }
 }
+
+export async function resetPersistence<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(props: {
+  tx: PgTransaction<TQueryResult, TFullSchema, TSchema>;
+  indexerId: string;
+}) {
+  const { tx, indexerId } = props;
+
+  try {
+    await tx.delete(checkpoints).where(eq(checkpoints.id, indexerId));
+    await tx.delete(filters).where(eq(filters.id, indexerId));
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to reset persistence state", {
+      cause: error,
+    });
+  }
+}