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

package/dist/index.mjs

@@ -1,12 +1,542 @@
+import { useIndexerContext } from '@apibara/indexer';
+import { defineIndexerPlugin } from '@apibara/indexer/plugins';
+import { and, eq, isNull, gt, lt, sql } from 'drizzle-orm';
+import { pgTable, text, integer, primaryKey, serial, char, jsonb } from 'drizzle-orm/pg-core';
+
 class DrizzleStorageError extends Error {
-  constructor(message) {
-    super(message);
+  constructor(message, options) {
+    super(message, options);
     this.name = "DrizzleStorageError";
   }
 }
+async function withTransaction(db, cb) {
+  return await db.transaction(async (txnDb) => {
+    return await cb(txnDb);
+  });
+}
+function deserialize(str) {
+  return JSON.parse(
+    str,
+    (_, value) => typeof value === "string" && value.match(/^\d+n$/) ? BigInt(value.slice(0, -1)) : value
+  );
+}
+function serialize(obj) {
+  return JSON.stringify(
+    obj,
+    (_, value) => typeof value === "bigint" ? `${value.toString()}n` : value,
+    "	"
+  );
+}
+
+const CHECKPOINTS_TABLE_NAME = "__indexer_checkpoints";
+const FILTERS_TABLE_NAME = "__indexer_filters";
+const SCHEMA_VERSION_TABLE_NAME = "__indexer_schema_version";
+const checkpoints = pgTable(CHECKPOINTS_TABLE_NAME, {
+  id: text("id").notNull().primaryKey(),
+  orderKey: integer("order_key").notNull(),
+  uniqueKey: text("unique_key").$type().notNull().default(void 0)
+});
+const filters = pgTable(
+  FILTERS_TABLE_NAME,
+  {
+    id: text("id").notNull(),
+    filter: text("filter").notNull(),
+    fromBlock: integer("from_block").notNull(),
+    toBlock: integer("to_block").$type().default(null)
+  },
+  (table) => [
+    {
+      pk: primaryKey({ columns: [table.id, table.fromBlock] })
+    }
+  ]
+);
+const schemaVersion = pgTable(SCHEMA_VERSION_TABLE_NAME, {
+  k: integer("k").notNull().primaryKey(),
+  version: integer("version").notNull()
+});
+const CURRENT_SCHEMA_VERSION = 0;
+const MIGRATIONS = [
+  // migrations[0]: v0 -> v1 (for future use)
+  []
+  // Add more migration arrays for future versions
+];
+async function initializePersistentState(tx) {
+  await tx.execute(`
+    CREATE TABLE IF NOT EXISTS ${SCHEMA_VERSION_TABLE_NAME} (
+      k INTEGER PRIMARY KEY,
+      version INTEGER NOT NULL
+    );
+  `);
+  const versionRows = await tx.select().from(schemaVersion).where(eq(schemaVersion.k, 0));
+  const storedVersion = versionRows[0]?.version ?? -1;
+  if (storedVersion > CURRENT_SCHEMA_VERSION) {
+    throw new DrizzleStorageError(
+      `Database Persistence schema version v${storedVersion} is newer than supported version v${CURRENT_SCHEMA_VERSION}`
+    );
+  }
+  try {
+    if (storedVersion === -1) {
+      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 ''
+        );
+      `);
+      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)
+        );
+      `);
+      await tx.insert(schemaVersion).values({
+        k: 0,
+        version: CURRENT_SCHEMA_VERSION
+      });
+    } else {
+      let currentVersion = storedVersion;
+      while (currentVersion < CURRENT_SCHEMA_VERSION) {
+        const migrationStatements = MIGRATIONS[currentVersion];
+        for (const statement of migrationStatements) {
+          await tx.execute(statement);
+        }
+        currentVersion++;
+      }
+      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 }
+    );
+  }
+}
+async function persistState(props) {
+  const { tx, endCursor, filter, indexerName } = props;
+  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 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
+    });
+  }
+}
+async function getState(props) {
+  const { tx, indexerName } = props;
+  try {
+    const checkpointRows = await tx.select().from(checkpoints).where(eq(checkpoints.id, indexerName));
+    const cursor = checkpointRows[0] ? {
+      orderKey: BigInt(checkpointRows[0].orderKey),
+      uniqueKey: checkpointRows[0].uniqueKey
+    } : void 0;
+    const filterRows = await tx.select().from(filters).where(and(eq(filters.id, indexerName), isNull(filters.toBlock)));
+    const filter = filterRows[0] ? deserialize(filterRows[0].filter) : void 0;
+    return { cursor, filter };
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to get persistent state", {
+      cause: error
+    });
+  }
+}
+async function invalidateState(props) {
+  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: null }).where(
+      and(
+        eq(filters.id, indexerName),
+        gt(filters.toBlock, Number(cursor.orderKey))
+      )
+    );
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to invalidate state", {
+      cause: error
+    });
+  }
+}
+async function finalizeState(props) {
+  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
+    });
+  }
+}
 
-function drizzleStorage() {
-  throw new DrizzleStorageError("Not implemented");
+pgTable("__reorg_rollback", {
+  n: serial("n").primaryKey(),
+  op: char("op", { length: 1 }).$type().notNull(),
+  table_name: text("table_name").notNull(),
+  cursor: integer("cursor").notNull(),
+  row_id: text("row_id"),
+  row_value: jsonb("row_value")
+});
+async function initializeReorgRollbackTable(tx) {
+  try {
+    await tx.execute(
+      sql.raw(`
+        CREATE TABLE IF NOT EXISTS __reorg_rollback(
+          n SERIAL PRIMARY KEY,
+          op CHAR(1) NOT NULL,
+          table_name TEXT NOT NULL,
+          cursor INTEGER NOT NULL,
+          row_id TEXT,
+          row_value JSONB
+        );
+      `)
+    );
+    await tx.execute(
+      sql.raw(`
+        CREATE OR REPLACE FUNCTION reorg_checkpoint()
+        RETURNS TRIGGER AS $$
+        DECLARE
+          id_col TEXT := TG_ARGV[0]::TEXT;
+          order_key INTEGER := TG_ARGV[1]::INTEGER;
+          new_id_value TEXT := row_to_json(NEW.*)->>id_col;
+          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)
+              SELECT 'D', TG_TABLE_NAME, order_key, old_id_value, row_to_json(OLD.*);
+          ELSIF (TG_OP = 'UPDATE') THEN
+            INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value)
+              SELECT 'U', TG_TABLE_NAME, order_key, new_id_value, row_to_json(OLD.*);
+          ELSIF (TG_OP = 'INSERT') THEN
+            INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value)
+              SELECT 'I', TG_TABLE_NAME, order_key, new_id_value, null;
+          END IF;
+          RETURN NULL;
+        END;
+        $$ LANGUAGE plpgsql;
+      `)
+    );
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to initialize reorg rollback table", {
+      cause: error
+    });
+  }
+}
+async function registerTriggers(tx, tables, endCursor, idColumn) {
+  try {
+    for (const table of tables) {
+      await tx.execute(
+        sql.raw(`DROP TRIGGER IF EXISTS ${table}_reorg ON ${table};`)
+      );
+      await tx.execute(
+        sql.raw(`
+          CREATE CONSTRAINT TRIGGER ${table}_reorg
+          AFTER INSERT OR UPDATE OR DELETE ON ${table}
+          DEFERRABLE INITIALLY DEFERRED
+          FOR EACH ROW EXECUTE FUNCTION reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`});
+        `)
+      );
+    }
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to register triggers", {
+      cause: error
+    });
+  }
+}
+async function removeTriggers(db, tables) {
+  try {
+    for (const table of tables) {
+      await db.execute(
+        sql.raw(`DROP TRIGGER IF EXISTS ${table}_reorg ON ${table};`)
+      );
+    }
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to remove triggers", {
+      cause: error
+    });
+  }
+}
+async function invalidate(tx, cursor, idColumn) {
+  const { rows: result } = await tx.execute(
+    sql.raw(`
+      WITH deleted AS (
+        DELETE FROM __reorg_rollback
+        WHERE cursor > ${Number(cursor.orderKey)}
+        RETURNING *
+      )
+      SELECT * FROM deleted ORDER BY n DESC;
+    `)
+  );
+  if (!Array.isArray(result)) {
+    throw new DrizzleStorageError(
+      "Invalid result format from reorg_rollback query"
+    );
+  }
+  for (const op of result) {
+    switch (op.op) {
+      case "I":
+        try {
+          if (!op.row_id) {
+            throw new DrizzleStorageError("Insert operation has no row_id");
+          }
+          await tx.execute(
+            sql.raw(`
+                DELETE FROM ${op.table_name}
+                WHERE ${idColumn} = '${op.row_id}'
+              `)
+          );
+        } catch (error) {
+          throw new DrizzleStorageError(
+            "Failed to invalidate | Operation - I",
+            {
+              cause: error
+            }
+          );
+        }
+        break;
+      case "D":
+        try {
+          if (!op.row_value) {
+            throw new DrizzleStorageError("Delete operation has no row_value");
+          }
+          await tx.execute(
+            sql.raw(`
+              INSERT INTO ${op.table_name}
+              SELECT * FROM json_populate_record(null::${op.table_name}, '${JSON.stringify(op.row_value)}'::json)
+            `)
+          );
+        } catch (error) {
+          throw new DrizzleStorageError(
+            "Failed to invalidate | Operation - D",
+            {
+              cause: error
+            }
+          );
+        }
+        break;
+      case "U":
+        try {
+          if (!op.row_value || !op.row_id) {
+            throw new DrizzleStorageError(
+              "Update operation has no row_value or row_id"
+            );
+          }
+          const rowValue = typeof op.row_value === "string" ? JSON.parse(op.row_value) : op.row_value;
+          const nonIdKeys = Object.keys(rowValue).filter((k) => k !== idColumn);
+          const fields = nonIdKeys.map((c) => `${c} = prev.${c}`).join(", ");
+          const query = sql.raw(`
+              UPDATE ${op.table_name}
+              SET ${fields}
+              FROM (
+                SELECT * FROM json_populate_record(null::${op.table_name}, '${JSON.stringify(op.row_value)}'::json)
+              ) as prev
+              WHERE ${op.table_name}.${idColumn} = '${op.row_id}'
+              `);
+          await tx.execute(query);
+        } catch (error) {
+          throw new DrizzleStorageError(
+            "Failed to invalidate | Operation - U",
+            {
+              cause: error
+            }
+          );
+        }
+        break;
+      default: {
+        throw new DrizzleStorageError(`Unknown operation: ${op.op}`);
+      }
+    }
+  }
+}
+async function finalize(tx, cursor) {
+  try {
+    await tx.execute(
+      sql.raw(`
+      DELETE FROM __reorg_rollback
+      WHERE cursor <= ${Number(cursor.orderKey)}
+    `)
+    );
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to finalize", {
+      cause: error
+    });
+  }
+}
+
+const DRIZZLE_PROPERTY = "_drizzle";
+function useDrizzleStorage(_db) {
+  const context = useIndexerContext();
+  if (!context[DRIZZLE_PROPERTY]) {
+    throw new DrizzleStorageError(
+      "drizzle storage is not available. Did you register the plugin?"
+    );
+  }
+  return context[DRIZZLE_PROPERTY];
+}
+function drizzleStorage({
+  db,
+  persistState: enablePersistence = true,
+  indexerName = "default",
+  schema,
+  idColumn = "id"
+}) {
+  return defineIndexerPlugin((indexer) => {
+    let tableNames = [];
+    try {
+      tableNames = Object.values(schema ?? 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({
+          tx,
+          indexerName
+        });
+        if (cursor) {
+          request.startingCursor = cursor;
+        }
+        if (filter) {
+          request.filter[1] = filter;
+        }
+      });
+    });
+    indexer.hooks.hook("connect:after", async ({ request }) => {
+      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;
+      }
+      const { db: tx } = useDrizzleStorage();
+      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;
+          if (!endCursor) {
+            throw new DrizzleStorageError("End Cursor is undefined");
+          }
+          await withTransaction(db, async (tx) => {
+            context[DRIZZLE_PROPERTY] = { db: tx };
+            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") {
+            await removeTriggers(db, tableNames);
+          }
+        } catch (error) {
+          await removeTriggers(db, tableNames);
+          throw new DrizzleStorageError("Failed to run handler:middleware", {
+            cause: error
+          });
+        }
+      });
+    });
+  });
 }
 
-export { drizzleStorage };
+export { drizzleStorage, useDrizzleStorage };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@apibara/plugin-drizzle",
-  "version": "2.0.0-beta.27",
+  "version": "2.0.0-beta.29",
   "type": "module",
   "files": [
     "dist",
@@ -25,17 +25,22 @@
     "test": "vitest",
     "test:ci": "vitest run"
   },
+  "peerDependencies": {
+    "drizzle-orm": "^0.37.0",
+    "pg": "^8.13.1"
+  },
   "devDependencies": {
     "@electric-sql/pglite": "^0.2.14",
     "@types/node": "^20.14.0",
+    "@types/pg": "^8.11.10",
+    "drizzle-orm": "^0.37.0",
+    "pg": "^8.13.1",
     "unbuild": "^2.0.0",
     "vitest": "^1.6.0"
   },
   "dependencies": {
-    "@apibara/indexer": "2.0.0-beta.28",
-    "@apibara/protocol": "2.0.0-beta.28",
-    "drizzle-orm": "^0.37.0",
-    "pg": "^8.13.1",
+    "@apibara/indexer": "2.0.0-beta.30",
+    "@apibara/protocol": "2.0.0-beta.30",
     "postgres-range": "^1.1.4"
   }
 }