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

package/src/helper.ts

@@ -0,0 +1,202 @@
+import { PGlite, type PGliteOptions } from "@electric-sql/pglite";
+import type { DrizzleConfig } from "drizzle-orm";
+import type { MigrationConfig } from "drizzle-orm/migrator";
+import {
+  type NodePgDatabase as OriginalNodePgDatabase,
+  drizzle as drizzleNode,
+} from "drizzle-orm/node-postgres";
+import { migrate as migrateNode } from "drizzle-orm/node-postgres/migrator";
+import {} from "drizzle-orm/pg-core";
+import {
+  type PgliteDatabase as OriginalPgliteDatabase,
+  drizzle as drizzlePGLite,
+} from "drizzle-orm/pglite";
+import { migrate as migratePGLite } from "drizzle-orm/pglite/migrator";
+import pg from "pg";
+import { DrizzleStorageError } from "./utils";
+
+/**
+ * Union type of all possible drizzle database options
+ */
+export type DrizzleOptions = PgliteDrizzleOptions | NodePgDrizzleOptions;
+
+/**
+ * Configuration options for Node-Postgres database connection
+ */
+export type NodePgDrizzleOptions = {
+  /**
+   * Type of database to use -
+   * - "pglite" - PGLite database
+   * - "node-postgres" - Node-Postgres database
+   * @default "pglite"
+   */
+  type: "node-postgres";
+  /**
+   * Connection string to use for the database
+   * @default ""
+   */
+  connectionString?: string;
+  /**
+   * Pool configuration options for Node-Postgres
+   */
+  poolConfig?: pg.PoolConfig;
+  /**
+   * Additional drizzle configuration options
+   */
+  config?: Omit<DrizzleConfig, "schema">;
+};
+
+/**
+ * Configuration options for PGLite database connection
+ */
+export type PgliteDrizzleOptions = {
+  /**
+   * Type of database to use -
+   * - "pglite" - PGLite database
+   * - "node-postgres" - Node-Postgres database
+   */
+  type?: "pglite";
+  /**
+   * Connection string to use for the database
+   * @default process.env["POSTGRES_CONNECTION_STRING"] ?? "memory://pglite"
+   */
+  connectionString?: string;
+  /**
+   * Pool configuration is not supported for PGLite
+   */
+  poolConfig?: never;
+  /**
+   * Additional drizzle configuration options with PGLite specific connection options
+   */
+  config?: Omit<DrizzleConfig, "schema"> & {
+    connection?:
+      | (PGliteOptions & {
+          dataDir?: string;
+        })
+      | string;
+  };
+};
+
+/**
+ * Extended PGLite database type with client information
+ */
+export type PgliteDatabase<TSchema extends Record<string, unknown>> =
+  OriginalPgliteDatabase<TSchema> & {
+    $client: PGlite;
+  };
+
+/**
+ * Extended Node-Postgres database type with client information
+ */
+export type NodePgDatabase<TSchema extends Record<string, unknown>> =
+  OriginalNodePgDatabase<TSchema> & {
+    $client: pg.Pool;
+  };
+
+export type Database<
+  TOptions extends DrizzleOptions,
+  TSchema extends Record<string, unknown>,
+> = TOptions extends PgliteDrizzleOptions
+  ? PgliteDatabase<TSchema>
+  : NodePgDatabase<TSchema>;
+
+/**
+ * Creates a new Drizzle database instance based on the provided options
+ *
+ * @important connectionString defaults to process.env["POSTGRES_CONNECTION_STRING"], if not set, it defaults to "memory://" (in-memory pglite)
+ *
+ * @param options - Configuration options for the database connection
+ * @returns A configured Drizzle database instance
+ * @throws {Error} If an invalid database type is specified
+ */
+export function drizzle<
+  TSchema extends Record<string, unknown>,
+  TOptions extends DrizzleOptions,
+>(
+  options?: TOptions & {
+    /**
+     * Schema to use for the database
+     * @default {}
+     */
+    schema?: TSchema;
+  },
+): Database<TOptions, TSchema> {
+  const {
+    connectionString = process.env["POSTGRES_CONNECTION_STRING"] ?? "memory://",
+    schema,
+    type = "pglite",
+    config,
+    poolConfig,
+  } = options ?? {};
+
+  if (
+    isPostgresConnectionString(connectionString) ||
+    type === "node-postgres"
+  ) {
+    const pool = new pg.Pool({
+      connectionString,
+      ...(poolConfig || {}),
+    });
+    return drizzleNode(pool, { schema, ...(config || {}) }) as Database<
+      TOptions,
+      TSchema
+    >;
+  }
+
+  if (type === "pglite") {
+    return drizzlePGLite({
+      schema: schema as TSchema,
+      connection: {
+        dataDir: connectionString || "memory://pglite",
+      },
+      ...(config || {}),
+    }) as Database<TOptions, TSchema>;
+  }
+
+  throw new Error("Invalid database type");
+}
+
+/**
+ * Options for database migration
+ */
+export type MigrateOptions = MigrationConfig;
+
+/**
+ * Performs database migration based on the provided configuration
+ * @param db - The database instance to migrate
+ * @param options - Migration configuration options
+ *
+ * @important This function runs migrations on the database instance provided to the `drizzleStorage` plugin.
+ * It automatically detects the type of database and runs the appropriate migrate function
+ * (PGLite or Node-Postgres).
+ *
+ * @example
+ * ```ts
+ * await migrate(db, { migrationsFolder: "./drizzle" });
+ * ```
+ */
+export async function migrate<TSchema extends Record<string, unknown>>(
+  db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>,
+  options: MigrateOptions,
+) {
+  const isPglite = !!("$client" in db && db.$client instanceof PGlite);
+
+  try {
+    if (isPglite) {
+      await migratePGLite(db as PgliteDatabase<TSchema>, options);
+    } else {
+      await migrateNode(db as NodePgDatabase<TSchema>, options);
+    }
+  } catch (error) {
+    throw new DrizzleStorageError(
+      "Failed to apply migrations! Please check if you have generated migrations using drizzle:generate",
+      {
+        cause: error,
+      },
+    );
+  }
+}
+
+function isPostgresConnectionString(conn: string) {
+  return conn.startsWith("postgres://") || conn.startsWith("postgresql://");
+}

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,13 @@ 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";
+    let prevFinality: DataFinality | undefined;
 
     try {
       tableNames = Object.values((schema as TSchema) ?? db._.schema ?? {}).map(
@@ -113,15 +141,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 +190,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",
               {
@@ -238,7 +300,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,6 +317,11 @@ export function drizzleStorage<
               TSchema
             >;
 
+            if (prevFinality === "pending") {
+              // invalidate if previous block's finality was "pending"
+              await invalidate(tx, cursor, idColumn, indexerId);
+            }
+
             if (finality !== "finalized") {
               await registerTriggers(
                 tx,
@@ -267,13 +335,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") {

package/src/persistence.ts

@@ -1,24 +1,29 @@
 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, text } 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 +38,8 @@ export const filters = pgTable(
   ],
 );
 
-export const schemaVersion = pgTable(SCHEMA_VERSION_TABLE_NAME, {
+/** 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 +59,22 @@ 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
     );
-  `);
+  `),
+  );
 
   // Get current schema version
   const versionRows = await tx
@@ -80,23 +95,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({
@@ -155,7 +174,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,
           },
         });
 
@@ -297,3 +318,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

@@ -11,19 +11,25 @@ import {
   char,
   integer,
   jsonb,
-  pgTable,
+  pgSchema,
   serial,
   text,
 } from "drizzle-orm/pg-core";
+import { SCHEMA_NAME } from "./constants";
 import { DrizzleStorageError } from "./utils";
 
+const ROLLBACK_TABLE_NAME = "reorg_rollback";
+
+const schema = pgSchema(SCHEMA_NAME);
+
 function getReorgTriggerName(table: string, indexerId: string) {
   return `${table}_reorg_${indexerId}`;
 }
 
 export type ReorgOperation = "I" | "U" | "D";
 
-export const reorgRollbackTable = pgTable("__reorg_rollback", {
+/** This table is not used for migrations, its only used for ease of internal operations with drizzle. */
+export const reorgRollbackTable = schema.table(ROLLBACK_TABLE_NAME, {
   n: serial("n").primaryKey(),
   op: char("op", { length: 1 }).$type<ReorgOperation>().notNull(),
   table_name: text("table_name").notNull(),
@@ -42,10 +48,14 @@ export async function initializeReorgRollbackTable<
     TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
 >(tx: PgTransaction<TQueryResult, TFullSchema, TSchema>, indexerId: string) {
   try {
+    // Create schema if it doesn't exist
+    await tx.execute(`
+    CREATE SCHEMA IF NOT EXISTS ${SCHEMA_NAME};
+    `);
     // Create the audit log table
     await tx.execute(
       sql.raw(`
-        CREATE TABLE IF NOT EXISTS __reorg_rollback(
+        CREATE TABLE IF NOT EXISTS ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(
           n SERIAL PRIMARY KEY,
           op CHAR(1) NOT NULL,
           table_name TEXT NOT NULL,
@@ -59,7 +69,7 @@ export async function initializeReorgRollbackTable<
 
     await tx.execute(
       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 ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(indexer_id, cursor);
       `),
     );
   } catch (error) {
@@ -72,7 +82,7 @@ export async function initializeReorgRollbackTable<
     // Create the trigger function
     await tx.execute(
       sql.raw(`
-      CREATE OR REPLACE FUNCTION reorg_checkpoint()
+      CREATE OR REPLACE FUNCTION ${SCHEMA_NAME}.reorg_checkpoint()
       RETURNS TRIGGER AS $$
       DECLARE
         id_col TEXT := TG_ARGV[0]::TEXT;
@@ -82,13 +92,13 @@ export async function initializeReorgRollbackTable<
         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)
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(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;
         ELSIF (TG_OP = 'UPDATE') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(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;
         ELSIF (TG_OP = 'INSERT') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
             SELECT 'I', TG_TABLE_NAME, order_key, new_id_value, null, indexer_id;
         END IF;
         RETURN NULL;
@@ -130,7 +140,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 reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
+          FOR EACH ROW EXECUTE FUNCTION ${SCHEMA_NAME}.reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
         `),
       );
     }
@@ -181,7 +191,7 @@ export async function invalidate<
   const { rows: result } = (await tx.execute(
     sql.raw(`
       WITH deleted AS (
-        DELETE FROM __reorg_rollback
+        DELETE FROM ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}
         WHERE cursor > ${Number(cursor.orderKey)}
         AND indexer_id = '${indexerId}'
         RETURNING *
@@ -305,7 +315,7 @@ export async function finalize<
   try {
     await tx.execute(
       sql.raw(`
-      DELETE FROM __reorg_rollback
+      DELETE FROM ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}
       WHERE cursor <= ${Number(cursor.orderKey)}
       AND indexer_id = '${indexerId}'
     `),
@@ -316,3 +326,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 ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}
+        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>;
+}