@apibara/plugin-drizzle 2.1.0-beta.6 → 2.1.0-beta.7

package/dist/index.cjs

@@ -4,10 +4,21 @@ 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 constants = require('./shared/plugin-drizzle.e884ca32.cjs');
+const pglite$1 = require('@electric-sql/pglite');
+const nodePostgres = require('drizzle-orm/node-postgres');
+const migrator$1 = require('drizzle-orm/node-postgres/migrator');
+const pglite = require('drizzle-orm/pglite');
+const migrator = require('drizzle-orm/pglite/migrator');
+const pg = require('pg');
 const protocol = require('@apibara/protocol');
 const drizzleOrm = require('drizzle-orm');
 const pgCore = require('drizzle-orm/pg-core');
 
+function _interopDefaultCompat (e) { return e && typeof e === 'object' && 'default' in e ? e.default : e; }
+
+const pg__default = /*#__PURE__*/_interopDefaultCompat(pg);
+
 class DrizzleStorageError extends Error {
   constructor(message, options) {
     super(message, options);
@@ -36,6 +47,50 @@ function sleep(ms) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+function drizzle(options) {
+  const {
+    connectionString = "memory://",
+    schema,
+    type = "pglite",
+    config,
+    poolConfig
+  } = options ?? {};
+  if (connectionString.startsWith("postgres://") || type === "node-postgres") {
+    const pool = new pg__default.Pool({
+      connectionString,
+      ...poolConfig || {}
+    });
+    return nodePostgres.drizzle(pool, { schema, ...config || {} });
+  }
+  if (type === "pglite") {
+    return pglite.drizzle({
+      schema,
+      connection: {
+        dataDir: connectionString || "memory://pglite"
+      },
+      ...config || {}
+    });
+  }
+  throw new Error("Invalid database type");
+}
+async function migrate(db, options) {
+  const isPglite = !!("$client" in db && db.$client instanceof pglite$1.PGlite);
+  try {
+    if (isPglite) {
+      await migrator.migrate(db, options);
+    } else {
+      await migrator$1.migrate(db, options);
+    }
+  } catch (error) {
+    throw new DrizzleStorageError(
+      "Failed to apply migrations! Please check if you have generated migrations using drizzle:generate",
+      {
+        cause: error
+      }
+    );
+  }
+}
+
 const CHECKPOINTS_TABLE_NAME = "__indexer_checkpoints";
 const FILTERS_TABLE_NAME = "__indexer_filters";
 const SCHEMA_VERSION_TABLE_NAME = "__indexer_schema_version";
@@ -434,25 +489,25 @@ async function finalize(tx, cursor, indexerId) {
   }
 }
 
-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 drizzleStorage({
   db,
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
   schema,
-  idColumn = "id"
+  idColumn = "id",
+  migrate: migrateOptions
 }) {
-  return plugins.defineIndexerPlugin((indexer) => {
+  return plugins.defineIndexerPlugin((indexer$1) => {
     let tableNames = [];
     let indexerId = "";
     try {
@@ -464,12 +519,22 @@ function drizzleStorage({
         cause: error
       });
     }
-    indexer.hooks.hook("run:before", async () => {
-      const { indexerName: indexerFileName, availableIndexers } = plugins$1.useInternalContext();
+    indexer$1.hooks.hook("run:before", 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;
       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) {
@@ -479,6 +544,9 @@ function drizzleStorage({
           break;
         } catch (error) {
           if (retries === MAX_RETRIES) {
+            if (error instanceof DrizzleStorageError) {
+              throw error;
+            }
             throw new DrizzleStorageError(
               "Initialization failed after 5 retries",
               {
@@ -491,7 +559,7 @@ function drizzleStorage({
         }
       }
     });
-    indexer.hooks.hook("connect:before", async ({ request }) => {
+    indexer$1.hooks.hook("connect:before", async ({ request }) => {
       if (!enablePersistence) {
         return;
       }
@@ -508,7 +576,7 @@ function drizzleStorage({
         }
       });
     });
-    indexer.hooks.hook("connect:after", async ({ request }) => {
+    indexer$1.hooks.hook("connect:after", async ({ request }) => {
       const cursor = request.startingCursor;
       if (!cursor) {
         return;
@@ -520,7 +588,7 @@ function drizzleStorage({
         }
       });
     });
-    indexer.hooks.hook("connect:factory", async ({ request, endCursor }) => {
+    indexer$1.hooks.hook("connect:factory", async ({ request, endCursor }) => {
       if (!enablePersistence) {
         return;
       }
@@ -534,7 +602,7 @@ function drizzleStorage({
         });
       }
     });
-    indexer.hooks.hook("message:finalize", async ({ message }) => {
+    indexer$1.hooks.hook("message:finalize", async ({ message }) => {
       const { cursor } = message.finalize;
       if (!cursor) {
         throw new DrizzleStorageError("Finalized Cursor is undefined");
@@ -546,7 +614,7 @@ function drizzleStorage({
         }
       });
     });
-    indexer.hooks.hook("message:invalidate", async ({ message }) => {
+    indexer$1.hooks.hook("message:invalidate", async ({ message }) => {
       const { cursor } = message.invalidate;
       if (!cursor) {
         throw new DrizzleStorageError("Invalidate Cursor is undefined");
@@ -558,7 +626,7 @@ function drizzleStorage({
         }
       });
     });
-    indexer.hooks.hook("handler:middleware", async ({ use }) => {
+    indexer$1.hooks.hook("handler:middleware", async ({ use }) => {
       use(async (context, next) => {
         try {
           const { endCursor, finality } = context;
@@ -566,7 +634,7 @@ function drizzleStorage({
             throw new DrizzleStorageError("End Cursor is undefined");
           }
           await withTransaction(db, async (tx) => {
-            context[DRIZZLE_PROPERTY] = { db: tx };
+            context[constants.DRIZZLE_PROPERTY] = { db: tx };
             if (finality !== "finalized") {
               await registerTriggers(
                 tx,
@@ -577,7 +645,7 @@ function drizzleStorage({
               );
             }
             await next();
-            delete context[DRIZZLE_PROPERTY];
+            delete context[constants.DRIZZLE_PROPERTY];
             if (enablePersistence) {
               await persistState({
                 tx,
@@ -600,5 +668,7 @@ function drizzleStorage({
   });
 }
 
+exports.drizzle = drizzle;
 exports.drizzleStorage = drizzleStorage;
+exports.migrate = migrate;
 exports.useDrizzleStorage = useDrizzleStorage;

package/dist/index.d.cts

@@ -1,17 +1,144 @@
 import * as _apibara_indexer_plugins from '@apibara/indexer/plugins';
-import { TablesRelationalConfig, ExtractTablesWithRelations } from 'drizzle-orm';
+import { DrizzleConfig, TablesRelationalConfig, ExtractTablesWithRelations } from 'drizzle-orm';
 import { PgQueryResultHKT, PgTransaction, PgDatabase } from 'drizzle-orm/pg-core';
+import { PGliteOptions, PGlite } from '@electric-sql/pglite';
+import { MigrationConfig } from 'drizzle-orm/migrator';
+import { NodePgDatabase as NodePgDatabase$1 } from 'drizzle-orm/node-postgres';
+import { PgliteDatabase as PgliteDatabase$1 } from 'drizzle-orm/pglite';
+import pg from 'pg';
+
+/**
+ * Union type of all possible drizzle database options
+ */
+type DrizzleOptions = PgliteDrizzleOptions | NodePgDrizzleOptions;
+/**
+ * Configuration options for Node-Postgres database connection
+ */
+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
+ */
+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 "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
+ */
+type PgliteDatabase<TSchema extends Record<string, unknown>> = PgliteDatabase$1<TSchema> & {
+    $client: PGlite;
+};
+/**
+ * Extended Node-Postgres database type with client information
+ */
+type NodePgDatabase<TSchema extends Record<string, unknown>> = NodePgDatabase$1<TSchema> & {
+    $client: pg.Pool;
+};
+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
+ * @param options - Configuration options for the database connection
+ * @returns A configured Drizzle database instance
+ * @throws {Error} If an invalid database type is specified
+ */
+declare function drizzle<TSchema extends Record<string, unknown>, TOptions extends DrizzleOptions>(options?: TOptions & {
+    /**
+     * Schema to use for the database
+     * @default {}
+     */
+    schema?: TSchema;
+}): Database<TOptions, TSchema>;
+/**
+ * Options for database migration
+ */
+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" });
+ * ```
+ */
+declare function migrate<TSchema extends Record<string, unknown>>(db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>, options: MigrateOptions): Promise<void>;
 
 type DrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>> = {
     db: PgTransaction<TQueryResult, TFullSchema, TSchema>;
 };
 declare function useDrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>(_db?: PgDatabase<TQueryResult, TFullSchema, TSchema>): DrizzleStorage<TQueryResult, TFullSchema, TSchema>;
 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>;
+    /**
+     * 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;
 }
 /**
  * Creates a plugin that uses Drizzle as the storage layer.
@@ -22,7 +149,8 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
  * @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.
  */
-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, }: 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, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
 
-export { type DrizzleStorage, type DrizzleStorageOptions, drizzleStorage, useDrizzleStorage };
+export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/index.d.mts

@@ -1,17 +1,144 @@
 import * as _apibara_indexer_plugins from '@apibara/indexer/plugins';
-import { TablesRelationalConfig, ExtractTablesWithRelations } from 'drizzle-orm';
+import { DrizzleConfig, TablesRelationalConfig, ExtractTablesWithRelations } from 'drizzle-orm';
 import { PgQueryResultHKT, PgTransaction, PgDatabase } from 'drizzle-orm/pg-core';
+import { PGliteOptions, PGlite } from '@electric-sql/pglite';
+import { MigrationConfig } from 'drizzle-orm/migrator';
+import { NodePgDatabase as NodePgDatabase$1 } from 'drizzle-orm/node-postgres';
+import { PgliteDatabase as PgliteDatabase$1 } from 'drizzle-orm/pglite';
+import pg from 'pg';
+
+/**
+ * Union type of all possible drizzle database options
+ */
+type DrizzleOptions = PgliteDrizzleOptions | NodePgDrizzleOptions;
+/**
+ * Configuration options for Node-Postgres database connection
+ */
+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
+ */
+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 "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
+ */
+type PgliteDatabase<TSchema extends Record<string, unknown>> = PgliteDatabase$1<TSchema> & {
+    $client: PGlite;
+};
+/**
+ * Extended Node-Postgres database type with client information
+ */
+type NodePgDatabase<TSchema extends Record<string, unknown>> = NodePgDatabase$1<TSchema> & {
+    $client: pg.Pool;
+};
+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
+ * @param options - Configuration options for the database connection
+ * @returns A configured Drizzle database instance
+ * @throws {Error} If an invalid database type is specified
+ */
+declare function drizzle<TSchema extends Record<string, unknown>, TOptions extends DrizzleOptions>(options?: TOptions & {
+    /**
+     * Schema to use for the database
+     * @default {}
+     */
+    schema?: TSchema;
+}): Database<TOptions, TSchema>;
+/**
+ * Options for database migration
+ */
+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" });
+ * ```
+ */
+declare function migrate<TSchema extends Record<string, unknown>>(db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>, options: MigrateOptions): Promise<void>;
 
 type DrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>> = {
     db: PgTransaction<TQueryResult, TFullSchema, TSchema>;
 };
 declare function useDrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>(_db?: PgDatabase<TQueryResult, TFullSchema, TSchema>): DrizzleStorage<TQueryResult, TFullSchema, TSchema>;
 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>;
+    /**
+     * 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;
 }
 /**
  * Creates a plugin that uses Drizzle as the storage layer.
@@ -22,7 +149,8 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
  * @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.
  */
-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, }: 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, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
 
-export { type DrizzleStorage, type DrizzleStorageOptions, drizzleStorage, useDrizzleStorage };
+export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/index.d.ts

@@ -1,17 +1,144 @@
 import * as _apibara_indexer_plugins from '@apibara/indexer/plugins';
-import { TablesRelationalConfig, ExtractTablesWithRelations } from 'drizzle-orm';
+import { DrizzleConfig, TablesRelationalConfig, ExtractTablesWithRelations } from 'drizzle-orm';
 import { PgQueryResultHKT, PgTransaction, PgDatabase } from 'drizzle-orm/pg-core';
+import { PGliteOptions, PGlite } from '@electric-sql/pglite';
+import { MigrationConfig } from 'drizzle-orm/migrator';
+import { NodePgDatabase as NodePgDatabase$1 } from 'drizzle-orm/node-postgres';
+import { PgliteDatabase as PgliteDatabase$1 } from 'drizzle-orm/pglite';
+import pg from 'pg';
+
+/**
+ * Union type of all possible drizzle database options
+ */
+type DrizzleOptions = PgliteDrizzleOptions | NodePgDrizzleOptions;
+/**
+ * Configuration options for Node-Postgres database connection
+ */
+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
+ */
+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 "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
+ */
+type PgliteDatabase<TSchema extends Record<string, unknown>> = PgliteDatabase$1<TSchema> & {
+    $client: PGlite;
+};
+/**
+ * Extended Node-Postgres database type with client information
+ */
+type NodePgDatabase<TSchema extends Record<string, unknown>> = NodePgDatabase$1<TSchema> & {
+    $client: pg.Pool;
+};
+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
+ * @param options - Configuration options for the database connection
+ * @returns A configured Drizzle database instance
+ * @throws {Error} If an invalid database type is specified
+ */
+declare function drizzle<TSchema extends Record<string, unknown>, TOptions extends DrizzleOptions>(options?: TOptions & {
+    /**
+     * Schema to use for the database
+     * @default {}
+     */
+    schema?: TSchema;
+}): Database<TOptions, TSchema>;
+/**
+ * Options for database migration
+ */
+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" });
+ * ```
+ */
+declare function migrate<TSchema extends Record<string, unknown>>(db: PgliteDatabase<TSchema> | NodePgDatabase<TSchema>, options: MigrateOptions): Promise<void>;
 
 type DrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>> = {
     db: PgTransaction<TQueryResult, TFullSchema, TSchema>;
 };
 declare function useDrizzleStorage<TQueryResult extends PgQueryResultHKT, TFullSchema extends Record<string, unknown> = Record<string, never>, TSchema extends TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>>(_db?: PgDatabase<TQueryResult, TFullSchema, TSchema>): DrizzleStorage<TQueryResult, TFullSchema, TSchema>;
 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>;
+    /**
+     * 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;
 }
 /**
  * Creates a plugin that uses Drizzle as the storage layer.
@@ -22,7 +149,8 @@ interface DrizzleStorageOptions<TQueryResult extends PgQueryResultHKT, TFullSche
  * @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.
  */
-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, }: 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, idColumn, migrate: migrateOptions, }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>): _apibara_indexer_plugins.IndexerPlugin<TFilter, TBlock>;
 
-export { type DrizzleStorage, type DrizzleStorageOptions, drizzleStorage, useDrizzleStorage };
+export { type Database, type DrizzleOptions, type DrizzleStorage, type DrizzleStorageOptions, type MigrateOptions, type NodePgDatabase, type NodePgDrizzleOptions, type PgliteDatabase, type PgliteDrizzleOptions, drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/index.mjs

@@ -1,7 +1,14 @@
 import { useIndexerContext } from '@apibara/indexer';
-import { defineIndexerPlugin } from '@apibara/indexer/plugins';
+import { defineIndexerPlugin, useLogger } from '@apibara/indexer/plugins';
 import { generateIndexerId } from '@apibara/indexer/internal';
 import { useInternalContext } from '@apibara/indexer/internal/plugins';
+import { D as DRIZZLE_PROPERTY, a as DRIZZLE_STORAGE_DB_PROPERTY } from './shared/plugin-drizzle.f8d1b186.mjs';
+import { PGlite } from '@electric-sql/pglite';
+import { drizzle as drizzle$1 } from 'drizzle-orm/node-postgres';
+import { migrate as migrate$2 } from 'drizzle-orm/node-postgres/migrator';
+import { drizzle as drizzle$2 } from 'drizzle-orm/pglite';
+import { migrate as migrate$1 } from 'drizzle-orm/pglite/migrator';
+import pg from 'pg';
 import { normalizeCursor } from '@apibara/protocol';
 import { eq, and, isNull, gt, lt, sql } from 'drizzle-orm';
 import { pgTable, text, integer, primaryKey, serial, char, jsonb } from 'drizzle-orm/pg-core';
@@ -34,6 +41,50 @@ function sleep(ms) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+function drizzle(options) {
+  const {
+    connectionString = "memory://",
+    schema,
+    type = "pglite",
+    config,
+    poolConfig
+  } = options ?? {};
+  if (connectionString.startsWith("postgres://") || type === "node-postgres") {
+    const pool = new pg.Pool({
+      connectionString,
+      ...poolConfig || {}
+    });
+    return drizzle$1(pool, { schema, ...config || {} });
+  }
+  if (type === "pglite") {
+    return drizzle$2({
+      schema,
+      connection: {
+        dataDir: connectionString || "memory://pglite"
+      },
+      ...config || {}
+    });
+  }
+  throw new Error("Invalid database type");
+}
+async function migrate(db, options) {
+  const isPglite = !!("$client" in db && db.$client instanceof PGlite);
+  try {
+    if (isPglite) {
+      await migrate$1(db, options);
+    } else {
+      await migrate$2(db, options);
+    }
+  } catch (error) {
+    throw new DrizzleStorageError(
+      "Failed to apply migrations! Please check if you have generated migrations using drizzle:generate",
+      {
+        cause: error
+      }
+    );
+  }
+}
+
 const CHECKPOINTS_TABLE_NAME = "__indexer_checkpoints";
 const FILTERS_TABLE_NAME = "__indexer_filters";
 const SCHEMA_VERSION_TABLE_NAME = "__indexer_schema_version";
@@ -432,7 +483,6 @@ async function finalize(tx, cursor, indexerId) {
   }
 }
 
-const DRIZZLE_PROPERTY = "_drizzle";
 const MAX_RETRIES = 5;
 function useDrizzleStorage(_db) {
   const context = useIndexerContext();
@@ -448,7 +498,8 @@ function drizzleStorage({
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
   schema,
-  idColumn = "id"
+  idColumn = "id",
+  migrate: migrateOptions
 }) {
   return defineIndexerPlugin((indexer) => {
     let tableNames = [];
@@ -463,11 +514,21 @@ function drizzleStorage({
       });
     }
     indexer.hooks.hook("run:before", async () => {
-      const { indexerName: indexerFileName, availableIndexers } = useInternalContext();
+      const internalContext = useInternalContext();
+      const context = useIndexerContext();
+      const logger = useLogger();
+      context[DRIZZLE_STORAGE_DB_PROPERTY] = db;
+      const { indexerName: indexerFileName, availableIndexers } = internalContext;
       indexerId = generateIndexerId(indexerFileName, identifier);
       let retries = 0;
+      let migrationsApplied = 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) {
@@ -477,6 +538,9 @@ function drizzleStorage({
           break;
         } catch (error) {
           if (retries === MAX_RETRIES) {
+            if (error instanceof DrizzleStorageError) {
+              throw error;
+            }
             throw new DrizzleStorageError(
               "Initialization failed after 5 retries",
               {
@@ -598,4 +662,4 @@ function drizzleStorage({
   });
 }
 
-export { drizzleStorage, useDrizzleStorage };
+export { drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/shared/plugin-drizzle.e884ca32.cjs

@@ -0,0 +1,7 @@
+'use strict';
+
+const DRIZZLE_PROPERTY = "_drizzle";
+const DRIZZLE_STORAGE_DB_PROPERTY = "_drizzleStorageDB";
+
+exports.DRIZZLE_PROPERTY = DRIZZLE_PROPERTY;
+exports.DRIZZLE_STORAGE_DB_PROPERTY = DRIZZLE_STORAGE_DB_PROPERTY;

package/dist/shared/plugin-drizzle.f8d1b186.mjs

@@ -0,0 +1,4 @@
+const DRIZZLE_PROPERTY = "_drizzle";
+const DRIZZLE_STORAGE_DB_PROPERTY = "_drizzleStorageDB";
+
+export { DRIZZLE_PROPERTY as D, DRIZZLE_STORAGE_DB_PROPERTY as a };

package/dist/testing.cjs

@@ -0,0 +1,13 @@
+'use strict';
+
+const constants = require('./shared/plugin-drizzle.e884ca32.cjs');
+
+function getTestDatabase(context) {
+  const db = context[constants.DRIZZLE_STORAGE_DB_PROPERTY];
+  if (!db) {
+    throw new Error("Drizzle database not found in context");
+  }
+  return db;
+}
+
+exports.getTestDatabase = getTestDatabase;

package/dist/testing.d.cts

@@ -0,0 +1,6 @@
+import { VcrResult } from '@apibara/indexer/testing';
+import { PgDatabase, PgQueryResultHKT } from 'drizzle-orm/pg-core';
+
+declare function getTestDatabase(context: VcrResult): PgDatabase<PgQueryResultHKT>;
+
+export { getTestDatabase };

package/dist/testing.d.mts

@@ -0,0 +1,6 @@
+import { VcrResult } from '@apibara/indexer/testing';
+import { PgDatabase, PgQueryResultHKT } from 'drizzle-orm/pg-core';
+
+declare function getTestDatabase(context: VcrResult): PgDatabase<PgQueryResultHKT>;
+
+export { getTestDatabase };

package/dist/testing.d.ts

@@ -0,0 +1,6 @@
+import { VcrResult } from '@apibara/indexer/testing';
+import { PgDatabase, PgQueryResultHKT } from 'drizzle-orm/pg-core';
+
+declare function getTestDatabase(context: VcrResult): PgDatabase<PgQueryResultHKT>;
+
+export { getTestDatabase };

package/dist/testing.mjs

@@ -0,0 +1,11 @@
+import { a as DRIZZLE_STORAGE_DB_PROPERTY } from './shared/plugin-drizzle.f8d1b186.mjs';
+
+function getTestDatabase(context) {
+  const db = context[DRIZZLE_STORAGE_DB_PROPERTY];
+  if (!db) {
+    throw new Error("Drizzle database not found in context");
+  }
+  return db;
+}
+
+export { getTestDatabase };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@apibara/plugin-drizzle",
-  "version": "2.1.0-beta.6",
+  "version": "2.1.0-beta.7",
   "type": "module",
   "files": [
     "dist",
@@ -15,6 +15,12 @@
       "import": "./dist/index.mjs",
       "require": "./dist/index.cjs",
       "default": "./dist/index.mjs"
+    },
+    "./testing": {
+      "types": "./dist/testing.d.ts",
+      "import": "./dist/testing.mjs",
+      "require": "./dist/testing.cjs",
+      "default": "./dist/testing.mjs"
     }
   },
   "scripts": {
@@ -39,8 +45,8 @@
     "vitest": "^1.6.0"
   },
   "dependencies": {
-    "@apibara/indexer": "2.1.0-beta.6",
-    "@apibara/protocol": "2.1.0-beta.6",
+    "@apibara/indexer": "2.1.0-beta.7",
+    "@apibara/protocol": "2.1.0-beta.7",
     "postgres-range": "^1.1.4"
   }
 }

package/src/constants.ts

@@ -0,0 +1,2 @@
+export const DRIZZLE_PROPERTY = "_drizzle";
+export const DRIZZLE_STORAGE_DB_PROPERTY = "_drizzleStorageDB";

package/src/helper.ts

@@ -0,0 +1,192 @@
+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 "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
+ * @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 = "memory://",
+    schema,
+    type = "pglite",
+    config,
+    poolConfig,
+  } = options ?? {};
+
+  if (connectionString.startsWith("postgres://") || 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,
+      },
+    );
+  }
+}

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,6 +14,8 @@ 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,
@@ -30,7 +32,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 +70,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 +105,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,6 +120,7 @@ export function drizzleStorage<
   indexerName: identifier = "default",
   schema,
   idColumn = "id",
+  migrate: migrateOptions,
 }: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>) {
   return defineIndexerPlugin<TFilter, TBlock>((indexer) => {
     let tableNames: string[] = [];
@@ -113,15 +137,31 @@ 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);
 
       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 +171,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/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>;
+}