@apibara/plugin-drizzle 2.1.0-beta.3 → 2.1.0-beta.30

package/src/storage.ts

@@ -11,11 +11,20 @@ import {
   char,
   integer,
   jsonb,
-  pgTable,
+  pgSchema,
   serial,
   text,
 } from "drizzle-orm/pg-core";
-import { DrizzleStorageError } from "./utils";
+import { SCHEMA_NAME } from "./constants";
+import {
+  DrizzleStorageError,
+  type IdColumnMap,
+  getIdColumnForTable,
+} from "./utils";
+
+const ROLLBACK_TABLE_NAME = "reorg_rollback";
+
+const schema = pgSchema(SCHEMA_NAME);
 
 function getReorgTriggerName(table: string, indexerId: string) {
   return `${table}_reorg_${indexerId}`;
@@ -23,7 +32,8 @@ function getReorgTriggerName(table: string, indexerId: string) {
 
 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 +52,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 +73,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,24 +86,25 @@ 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;
-        order_key INTEGER := TG_ARGV[1]::INTEGER;
-        indexer_id TEXT := TG_ARGV[2]::TEXT;
+        table_name TEXT := TG_ARGV[0]::TEXT;
+        id_col TEXT := TG_ARGV[1]::TEXT;
+        order_key INTEGER := TG_ARGV[2]::INTEGER;
+        indexer_id TEXT := TG_ARGV[3]::TEXT;
         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, indexer_id)
-            SELECT 'D', TG_TABLE_NAME, order_key, old_id_value, row_to_json(OLD.*), indexer_id;
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
+            SELECT 'D', 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)
-            SELECT 'U', TG_TABLE_NAME, order_key, new_id_value, row_to_json(OLD.*), indexer_id;
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
+            SELECT 'U', 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)
-            SELECT 'I', TG_TABLE_NAME, order_key, new_id_value, null, indexer_id;
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
+            SELECT 'I', table_name, order_key, new_id_value, null, indexer_id;
         END IF;
         RETURN NULL;
       END;
@@ -115,11 +130,14 @@ export async function registerTriggers<
   tx: PgTransaction<TQueryResult, TFullSchema, TSchema>,
   tables: string[],
   endCursor: Cursor,
-  idColumn: string,
+  idColumnMap: IdColumnMap,
   indexerId: string,
 ) {
   try {
     for (const table of tables) {
+      // Determine the column ID for this specific table
+      const tableIdColumn = getIdColumnForTable(table, idColumnMap);
+
       await tx.execute(
         sql.raw(
           `DROP TRIGGER IF EXISTS ${getReorgTriggerName(table, indexerId)} ON ${table};`,
@@ -130,7 +148,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('${table}', '${tableIdColumn}', ${Number(endCursor.orderKey)}, '${indexerId}');
         `),
       );
     }
@@ -174,14 +192,14 @@ export async function invalidate<
 >(
   tx: PgTransaction<TQueryResult, TFullSchema, TSchema>,
   cursor: Cursor,
-  idColumn: string,
+  idColumnMap: IdColumnMap,
   indexerId: string,
 ) {
   // Get and delete operations after cursor in one query, ordered by newest first
   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 *
@@ -198,6 +216,9 @@ export async function invalidate<
 
   // Process each operation in reverse order
   for (const op of result) {
+    // Determine the column ID for this specific table
+    const tableIdColumn = getIdColumnForTable(op.table_name, idColumnMap);
+
     switch (op.op) {
       case "I":
         try {
@@ -208,7 +229,7 @@ export async function invalidate<
           await tx.execute(
             sql.raw(`
                 DELETE FROM ${op.table_name}
-                WHERE ${idColumn} = '${op.row_id}'
+                WHERE ${tableIdColumn} = '${op.row_id}'
               `),
           );
         } catch (error) {
@@ -261,7 +282,9 @@ export async function invalidate<
               ? JSON.parse(op.row_value)
               : op.row_value;
 
-          const nonIdKeys = Object.keys(rowValue).filter((k) => k !== idColumn);
+          const nonIdKeys = Object.keys(rowValue).filter(
+            (k) => k !== tableIdColumn,
+          );
 
           const fields = nonIdKeys.map((c) => `${c} = prev.${c}`).join(", ");
 
@@ -271,7 +294,7 @@ export async function invalidate<
               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}'
+              WHERE ${op.table_name}.${tableIdColumn} = '${op.row_id}'
               `);
 
           await tx.execute(query);
@@ -305,7 +328,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 +339,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>;
+}

package/src/utils.ts

@@ -48,3 +48,22 @@ export function serialize<T>(obj: T): string {
 export function sleep(ms: number) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
+
+export interface IdColumnMap extends Record<string, string> {
+  /**
+   * Wildcard mapping for all tables.
+   */
+  "*": string;
+}
+
+export const getIdColumnForTable = (
+  tableName: string,
+  idColumn: IdColumnMap,
+): string => {
+  // If there's a specific mapping for this table, use it
+  if (idColumn[tableName]) {
+    return idColumn[tableName];
+  }
+  // Default fallback
+  return idColumn["*"];
+};