@apibara/plugin-drizzle 2.1.0-beta.1 → 2.1.0-beta.11

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,14 +14,18 @@ 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,
+  resetPersistence,
 } from "./persistence";
 import {
+  cleanupStorage,
   finalize,
   initializeReorgRollbackTable,
   invalidate,
@@ -30,7 +34,8 @@ import {
 } from "./storage";
 import { DrizzleStorageError, sleep, withTransaction } from "./utils";
 
-const DRIZZLE_PROPERTY = "_drizzle";
+export * from "./helper";
+
 const MAX_RETRIES = 5;
 
 export type DrizzleStorage<
@@ -67,11 +72,30 @@ export interface DrizzleStorageOptions<
   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>;
+  /**
+   * The column to use as the id. Defaults to 'id'.
+   */
   idColumn?: string;
+  /**
+   * The options for the database migration. When provided, the database will automatically run migrations before the indexer runs.
+   */
+  migrate?: MigrateOptions;
 }
 
 /**
@@ -83,6 +107,7 @@ 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.
  */
 export function drizzleStorage<
   TFilter,
@@ -97,10 +122,12 @@ export function drizzleStorage<
   indexerName: identifier = "default",
   schema,
   idColumn = "id",
+  migrate: migrateOptions,
 }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>) {
   return defineIndexerPlugin<TFilter, TBlock>((indexer) => {
     let tableNames: string[] = [];
     let indexerId = "";
+    const alwaysReindex = process.env["APIBARA_ALWAYS_REINDEX"] === "true";
 
     try {
       tableNames = Object.values((schema as TSchema) ?? db._.schema ?? {}).map(
@@ -113,15 +140,46 @@ export function drizzleStorage<
     }
 
     indexer.hooks.hook("run:before", 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);
 
+      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;
+
       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) {
@@ -131,6 +189,9 @@ export function drizzleStorage<
           break;
         } catch (error) {
           if (retries === MAX_RETRIES) {
+            if (error instanceof DrizzleStorageError) {
+              throw error;
+            }
             throw new DrizzleStorageError(
               "Initialization failed after 5 retries",
               {

package/src/persistence.ts

@@ -297,3 +297,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,
+    });
+  }
+}

package/src/storage.ts

@@ -316,3 +316,45 @@ export async function finalize<
     });
   }
 }
+
+export async function cleanupStorage<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(
+  tx: PgTransaction<TQueryResult, TFullSchema, TSchema>,
+  tables: string[],
+  indexerId: string,
+) {
+  try {
+    for (const table of tables) {
+      await tx.execute(
+        sql.raw(
+          `DROP TRIGGER IF EXISTS ${getReorgTriggerName(table, indexerId)} ON ${table};`,
+        ),
+      );
+    }
+
+    await tx.execute(
+      sql.raw(`
+        DELETE FROM __reorg_rollback
+        WHERE indexer_id = '${indexerId}'
+      `),
+    );
+
+    for (const table of tables) {
+      try {
+        await tx.execute(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,
+    });
+  }
+}

package/src/testing.ts

@@ -0,0 +1,13 @@
+import type { VcrResult } from "@apibara/indexer/testing";
+import type { PgDatabase, PgQueryResultHKT } from "drizzle-orm/pg-core";
+import { DRIZZLE_STORAGE_DB_PROPERTY } from "./constants";
+
+export function getTestDatabase(context: VcrResult) {
+  const db = context[DRIZZLE_STORAGE_DB_PROPERTY];
+
+  if (!db) {
+    throw new Error("Drizzle database not found in context");
+  }
+
+  return db as PgDatabase<PgQueryResultHKT>;
+}