@apibara/plugin-drizzle 2.0.0-beta.28 → 2.0.0-beta.29

package/src/index.ts

@@ -1,5 +1,260 @@
-import { DrizzleStorageError } from "./utils";
+import { useIndexerContext } from "@apibara/indexer";
+import { defineIndexerPlugin } from "@apibara/indexer/plugins";
 
-export function drizzleStorage<TFilter, TBlock>() {
-  throw new DrizzleStorageError("Not implemented");
+import type {
+  ExtractTablesWithRelations,
+  TablesRelationalConfig,
+} from "drizzle-orm";
+
+import type { Cursor, DataFinality } from "@apibara/protocol";
+import type {
+  PgDatabase,
+  PgQueryResultHKT,
+  PgTransaction,
+} from "drizzle-orm/pg-core";
+import {
+  finalizeState,
+  getState,
+  initializePersistentState,
+  invalidateState,
+  persistState,
+} from "./persistence";
+import {
+  finalize,
+  initializeReorgRollbackTable,
+  invalidate,
+  registerTriggers,
+  removeTriggers,
+} from "./storage";
+import { DrizzleStorageError, withTransaction } from "./utils";
+
+const DRIZZLE_PROPERTY = "_drizzle";
+
+export type DrizzleStorage<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> = {
+  db: PgTransaction<TQueryResult, TFullSchema, TSchema>;
+};
+
+export 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> {
+  const context = useIndexerContext();
+
+  if (!context[DRIZZLE_PROPERTY]) {
+    throw new DrizzleStorageError(
+      "drizzle storage is not available. Did you register the plugin?",
+    );
+  }
+
+  return context[DRIZZLE_PROPERTY];
+}
+
+export interface DrizzleStorageOptions<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> {
+  db: PgDatabase<TQueryResult, TFullSchema, TSchema>;
+  persistState?: boolean;
+  indexerName?: string;
+  schema?: Record<string, unknown>;
+  idColumn?: string;
+}
+
+/**
+ * Creates a plugin that uses Drizzle as the storage layer.
+ *
+ * Supports storing the indexer's state and provides a simple Key-Value store.
+ * @param options.db - The Drizzle database instance.
+ * @param options.persistState - Whether to persist the indexer's state. Defaults to true.
+ * @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'.
+ */
+export function drizzleStorage<
+  TFilter,
+  TBlock,
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>({
+  db,
+  persistState: enablePersistence = true,
+  indexerName = "default",
+  schema,
+  idColumn = "id",
+}: DrizzleStorageOptions<TQueryResult, TFullSchema, TSchema>) {
+  return defineIndexerPlugin<TFilter, TBlock>((indexer) => {
+    let tableNames: string[] = [];
+
+    try {
+      tableNames = Object.values((schema as TSchema) ?? db._.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 () => {
+      await withTransaction(db, async (tx) => {
+        await initializeReorgRollbackTable(tx);
+        if (enablePersistence) {
+          await initializePersistentState(tx);
+        }
+      });
+    });
+
+    indexer.hooks.hook("connect:before", async ({ request }) => {
+      if (!enablePersistence) {
+        return;
+      }
+
+      await withTransaction(db, async (tx) => {
+        const { cursor, filter } = await getState<
+          TFilter,
+          TQueryResult,
+          TFullSchema,
+          TSchema
+        >({
+          tx,
+          indexerName,
+        });
+        if (cursor) {
+          request.startingCursor = cursor;
+        }
+        if (filter) {
+          request.filter[1] = filter;
+        }
+      });
+    });
+
+    indexer.hooks.hook("connect:after", async ({ request }) => {
+      // On restart, we need to invalidate data for blocks that were processed but not persisted.
+      const cursor = request.startingCursor;
+
+      if (!cursor) {
+        return;
+      }
+
+      await withTransaction(db, async (tx) => {
+        await invalidate(tx, cursor, idColumn);
+
+        if (enablePersistence) {
+          await invalidateState({ tx, cursor, indexerName });
+        }
+      });
+    });
+
+    indexer.hooks.hook("connect:factory", async ({ request, endCursor }) => {
+      if (!enablePersistence) {
+        return;
+      }
+      // We can call this hook because this hook is called inside the transaction of handler:middleware
+      // so we have access to the transaction from the context
+      const { db: tx } = useDrizzleStorage(db);
+
+      if (endCursor && request.filter[1]) {
+        await persistState({
+          tx,
+          endCursor,
+          filter: request.filter[1],
+          indexerName,
+        });
+      }
+    });
+
+    indexer.hooks.hook("message:finalize", async ({ message }) => {
+      const { cursor } = message.finalize;
+
+      if (!cursor) {
+        throw new DrizzleStorageError("Finalized Cursor is undefined");
+      }
+
+      await withTransaction(db, async (tx) => {
+        await finalize(tx, cursor);
+
+        if (enablePersistence) {
+          await finalizeState({ tx, cursor, indexerName });
+        }
+      });
+    });
+
+    indexer.hooks.hook("message:invalidate", async ({ message }) => {
+      const { cursor } = message.invalidate;
+
+      if (!cursor) {
+        throw new DrizzleStorageError("Invalidate Cursor is undefined");
+      }
+
+      await withTransaction(db, async (tx) => {
+        await invalidate(tx, cursor, idColumn);
+
+        if (enablePersistence) {
+          await invalidateState({ tx, cursor, indexerName });
+        }
+      });
+    });
+
+    indexer.hooks.hook("handler:middleware", async ({ use }) => {
+      use(async (context, next) => {
+        try {
+          const { endCursor, finality } = context as {
+            endCursor: Cursor;
+            finality: DataFinality;
+          };
+
+          if (!endCursor) {
+            throw new DrizzleStorageError("End Cursor is undefined");
+          }
+
+          await withTransaction(db, async (tx) => {
+            context[DRIZZLE_PROPERTY] = { db: tx } as DrizzleStorage<
+              TQueryResult,
+              TFullSchema,
+              TSchema
+            >;
+
+            if (finality !== "finalized") {
+              await registerTriggers(tx, tableNames, endCursor, idColumn);
+            }
+
+            await next();
+            delete context[DRIZZLE_PROPERTY];
+
+            if (enablePersistence) {
+              await persistState({
+                tx,
+                endCursor,
+                indexerName,
+              });
+            }
+          });
+
+          if (finality !== "finalized") {
+            // remove trigger outside of the transaction or it won't be triggered.
+            await removeTriggers(db, tableNames);
+          }
+        } catch (error) {
+          await removeTriggers(db, tableNames);
+
+          throw new DrizzleStorageError("Failed to run handler:middleware", {
+            cause: error,
+          });
+        }
+      });
+    });
+  });
 }

package/src/persistence.ts

@@ -1,24 +1,18 @@
-/*
 import type { Cursor } from "@apibara/protocol";
-import {
-  type ExtractTablesWithRelations,
-  type TablesRelationalConfig,
-  and,
-  eq,
-  isNull,
+import { and, eq, gt, isNull, lt } from "drizzle-orm";
+import type {
+  ExtractTablesWithRelations,
+  TablesRelationalConfig,
 } from "drizzle-orm";
-import {
-  type PgDatabase,
-  type PgQueryResultHKT,
-  integer,
-  pgTable,
-  primaryKey,
-  text,
-} from "drizzle-orm/pg-core";
-import { deserialize, serialize } from "../vcr";
-import { defineIndexerPlugin } from "./config";
-
-export const checkpoints = pgTable("checkpoints", {
+import type { PgQueryResultHKT, PgTransaction } from "drizzle-orm/pg-core";
+import { integer, pgTable, primaryKey, text } from "drizzle-orm/pg-core";
+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";
+
+export const checkpoints = pgTable(CHECKPOINTS_TABLE_NAME, {
   id: text("id").notNull().primaryKey(),
   orderKey: integer("order_key").notNull(),
   uniqueKey: text("unique_key")
@@ -28,167 +22,281 @@ export const checkpoints = pgTable("checkpoints", {
 });
 
 export const filters = pgTable(
-  "filters",
+  FILTERS_TABLE_NAME,
   {
     id: text("id").notNull(),
     filter: text("filter").notNull(),
     fromBlock: integer("from_block").notNull(),
-    toBlock: integer("to_block"),
+    toBlock: integer("to_block").$type<number | null>().default(null),
   },
-  (table) => ({
-    pk: primaryKey({ columns: [table.id, table.fromBlock] }),
-  }),
+  (table) => [
+    {
+      pk: primaryKey({ columns: [table.id, table.fromBlock] }),
+    },
+  ],
 );
 
-export function drizzlePersistence<
-  TFilter,
-  TBlock,
-  TTxnParams,
+export const schemaVersion = pgTable(SCHEMA_VERSION_TABLE_NAME, {
+  k: integer("k").notNull().primaryKey(),
+  version: integer("version").notNull(),
+});
+
+export const CURRENT_SCHEMA_VERSION = 0;
+
+// migrations for future schema updates
+const MIGRATIONS: string[][] = [
+  // migrations[0]: v0 -> v1 (for future use)
+  [],
+  // Add more migration arrays for future versions
+];
+
+export async function initializePersistentState<
   TQueryResult extends PgQueryResultHKT,
   TFullSchema extends Record<string, unknown> = Record<string, never>,
   TSchema extends
     TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
->({
-  database,
-  indexerName = "default",
-}: {
-  database: PgDatabase<TQueryResult, TFullSchema, TSchema>;
-  indexerName?: string;
-}) {
-  return defineIndexerPlugin<TFilter, TBlock, TTxnParams>((indexer) => {
-    let store: DrizzlePersistence<TFilter, TQueryResult, TFullSchema, TSchema>;
+>(tx: PgTransaction<TQueryResult, TFullSchema, TSchema>) {
+  // Create schema version table
+  await tx.execute(`
+    CREATE TABLE IF NOT EXISTS ${SCHEMA_VERSION_TABLE_NAME} (
+      k INTEGER PRIMARY KEY,
+      version INTEGER NOT NULL
+    );
+  `);
 
-    indexer.hooks.hook("run:before", async () => {
-      store = new DrizzlePersistence(database, indexerName);
-      // Tables are created by user via migrations in Drizzle
-    });
+  // Get current schema version
+  const versionRows = await tx
+    .select()
+    .from(schemaVersion)
+    .where(eq(schemaVersion.k, 0));
 
-    indexer.hooks.hook("connect:before", async ({ request }) => {
-      const { cursor, filter } = await store.get();
+  const storedVersion = versionRows[0]?.version ?? -1;
 
-      if (cursor) {
-        request.startingCursor = cursor;
-      }
+  // Check for incompatible version
+  if (storedVersion > CURRENT_SCHEMA_VERSION) {
+    throw new DrizzleStorageError(
+      `Database Persistence schema version v${storedVersion} is newer than supported version v${CURRENT_SCHEMA_VERSION}`,
+    );
+  }
 
-      if (filter) {
-        request.filter[1] = filter;
-      }
-    });
+  // Begin schema updates
+  try {
+    if (storedVersion === -1) {
+      // First time initialization
+      await tx.execute(`
+        CREATE TABLE IF NOT EXISTS ${CHECKPOINTS_TABLE_NAME} (
+          id TEXT PRIMARY KEY,
+          order_key INTEGER NOT NULL,
+          unique_key TEXT NOT NULL DEFAULT ''
+        );
+      `);
 
-    indexer.hooks.hook("transaction:commit", async ({ endCursor }) => {
-      if (endCursor) {
-        await store.put({ cursor: endCursor });
-      }
-    });
+      await tx.execute(`
+        CREATE TABLE IF NOT EXISTS ${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)
+        );
+      `);
 
-    indexer.hooks.hook("connect:factory", async ({ request, endCursor }) => {
-      if (request.filter[1]) {
-        await store.put({ cursor: endCursor, filter: request.filter[1] });
+      // Set initial schema version
+      await tx.insert(schemaVersion).values({
+        k: 0,
+        version: CURRENT_SCHEMA_VERSION,
+      });
+    } else {
+      // Run any necessary migrations
+      let currentVersion = storedVersion;
+      while (currentVersion < CURRENT_SCHEMA_VERSION) {
+        const migrationStatements = MIGRATIONS[currentVersion];
+        for (const statement of migrationStatements) {
+          await tx.execute(statement);
+        }
+        currentVersion++;
       }
-    });
-  });
+
+      // Update schema version
+      await tx
+        .update(schemaVersion)
+        .set({ version: CURRENT_SCHEMA_VERSION })
+        .where(eq(schemaVersion.k, 0));
+    }
+  } catch (error) {
+    throw new DrizzleStorageError(
+      "Failed to initialize or migrate database schema",
+      { cause: error },
+    );
+  }
 }
 
-export class DrizzlePersistence<
+export async function persistState<
   TFilter,
   TQueryResult extends PgQueryResultHKT,
   TFullSchema extends Record<string, unknown> = Record<string, never>,
   TSchema extends
     TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
-> {
-  constructor(
-    private _db: PgDatabase<TQueryResult, TFullSchema, TSchema>,
-    private _indexerName: string,
-  ) {}
-
-  public async get(): Promise<{ cursor?: Cursor; filter?: TFilter }> {
-    const cursor = await this._getCheckpoint();
-    const filter = await this._getFilter();
-
-    return { cursor, filter };
-  }
+>(props: {
+  tx: PgTransaction<TQueryResult, TFullSchema, TSchema>;
+  endCursor: Cursor;
+  filter?: TFilter;
+  indexerName: string;
+}) {
+  const { tx, endCursor, filter, indexerName } = props;
 
-  public async put({ cursor, filter }: { cursor?: Cursor; filter?: TFilter }) {
-    if (cursor) {
-      await this._putCheckpoint(cursor);
+  try {
+    if (endCursor) {
+      await tx
+        .insert(checkpoints)
+        .values({
+          id: indexerName,
+          orderKey: Number(endCursor.orderKey),
+          uniqueKey: endCursor.uniqueKey,
+        })
+        .onConflictDoUpdate({
+          target: checkpoints.id,
+          set: {
+            orderKey: Number(endCursor.orderKey),
+            uniqueKey: endCursor.uniqueKey,
+          },
+        });
 
       if (filter) {
-        await this._putFilter(filter, cursor);
+        await tx
+          .update(filters)
+          .set({ toBlock: Number(endCursor.orderKey) })
+          .where(and(eq(filters.id, indexerName), isNull(filters.toBlock)));
+
+        await tx
+          .insert(filters)
+          .values({
+            id: indexerName,
+            filter: serialize(filter),
+            fromBlock: Number(endCursor.orderKey),
+            toBlock: null,
+          })
+          .onConflictDoUpdate({
+            target: [filters.id, filters.fromBlock],
+            set: {
+              filter: serialize(filter),
+              fromBlock: Number(endCursor.orderKey),
+              toBlock: null,
+            },
+          });
       }
     }
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to persist state", {
+      cause: error,
+    });
   }
+}
 
-  // --- CHECKPOINTS TABLE METHODS ---
+export async function getState<
+  TFilter,
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(props: {
+  tx: PgTransaction<TQueryResult, TFullSchema, TSchema>;
+  indexerName: string;
+}): Promise<{ cursor?: Cursor; filter?: TFilter }> {
+  const { tx, indexerName } = props;
 
-  private async _getCheckpoint(): Promise<Cursor | undefined> {
-    const rows = await this._db
+  try {
+    const checkpointRows = await tx
       .select()
       .from(checkpoints)
-      .where(eq(checkpoints.id, this._indexerName));
-
-    const row = rows[0];
-    if (!row) return undefined;
-
-    return {
-      orderKey: BigInt(row.orderKey),
-      uniqueKey: row.uniqueKey,
-    };
-  }
-
-  private async _putCheckpoint(cursor: Cursor) {
-    await this._db
-      .insert(checkpoints)
-      .values({
-        id: this._indexerName,
-        orderKey: Number(cursor.orderKey),
-        uniqueKey: cursor.uniqueKey,
-      })
-      .onConflictDoUpdate({
-        target: checkpoints.id,
-        set: {
-          orderKey: Number(cursor.orderKey),
-          uniqueKey: cursor.uniqueKey,
-        },
-      });
-  }
+      .where(eq(checkpoints.id, indexerName));
 
-  // --- FILTERS TABLE METHODS ---
+    const cursor = checkpointRows[0]
+      ? {
+          orderKey: BigInt(checkpointRows[0].orderKey),
+          uniqueKey: checkpointRows[0].uniqueKey,
+        }
+      : undefined;
 
-  private async _getFilter(): Promise<TFilter | undefined> {
-    const rows = await this._db
+    const filterRows = await tx
       .select()
       .from(filters)
-      .where(and(eq(filters.id, this._indexerName), isNull(filters.toBlock)));
+      .where(and(eq(filters.id, indexerName), isNull(filters.toBlock)));
 
-    const row = rows[0];
+    const filter = filterRows[0]
+      ? deserialize<TFilter>(filterRows[0].filter)
+      : undefined;
 
-    if (!row) return undefined;
-
-    return deserialize(row.filter) as TFilter;
+    return { cursor, filter };
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to get persistent state", {
+      cause: error,
+    });
   }
+}
 
-  private async _putFilter(filter: TFilter, endCursor: Cursor) {
-    // Update existing filter's to_block
-    await this._db
+export async function invalidateState<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(props: {
+  tx: PgTransaction<TQueryResult, TFullSchema, TSchema>;
+  cursor: Cursor;
+  indexerName: string;
+}) {
+  const { tx, cursor, indexerName } = props;
+
+  try {
+    await tx
+      .delete(filters)
+      .where(
+        and(
+          eq(filters.id, indexerName),
+          gt(filters.fromBlock, Number(cursor.orderKey)),
+        ),
+      );
+
+    await tx
       .update(filters)
-      .set({ toBlock: Number(endCursor.orderKey) })
-      .where(and(eq(filters.id, this._indexerName), isNull(filters.toBlock)));
-
-    // Insert new filter
-    await this._db
-      .insert(filters)
-      .values({
-        id: this._indexerName,
-        filter: serialize(filter as Record<string, unknown>),
-        fromBlock: Number(endCursor.orderKey),
-      })
-      .onConflictDoUpdate({
-        target: [filters.id, filters.fromBlock],
-        set: {
-          filter: serialize(filter as Record<string, unknown>),
-          fromBlock: Number(endCursor.orderKey),
-        },
-      });
+      .set({ toBlock: null })
+      .where(
+        and(
+          eq(filters.id, indexerName),
+          gt(filters.toBlock, Number(cursor.orderKey)),
+        ),
+      );
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to invalidate state", {
+      cause: error,
+    });
+  }
+}
+
+export async function finalizeState<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(props: {
+  tx: PgTransaction<TQueryResult, TFullSchema, TSchema>;
+  cursor: Cursor;
+  indexerName: string;
+}) {
+  const { tx, cursor, indexerName } = props;
+
+  try {
+    await tx
+      .delete(filters)
+      .where(
+        and(
+          eq(filters.id, indexerName),
+          lt(filters.toBlock, Number(cursor.orderKey)),
+        ),
+      );
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to finalize state", {
+      cause: error,
+    });
   }
 }
-*/