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

package/dist/index.mjs

@@ -1,10 +1,17 @@
 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 { S as SCHEMA_NAME, D as DRIZZLE_PROPERTY, a as DRIZZLE_STORAGE_DB_PROPERTY } from './shared/plugin-drizzle.2d226351.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';
+import { sql, eq, and, isNull, gt, lt } from 'drizzle-orm';
+import { pgSchema, text, integer, primaryKey, serial, char, jsonb } from 'drizzle-orm/pg-core';
 
 class DrizzleStorageError extends Error {
   constructor(message, options) {
@@ -34,15 +41,63 @@ function sleep(ms) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
-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, {
+function drizzle(options) {
+  const {
+    connectionString = process.env["POSTGRES_CONNECTION_STRING"] ?? "memory://",
+    schema,
+    type = "pglite",
+    config,
+    poolConfig
+  } = options ?? {};
+  if (isPostgresConnectionString(connectionString) || type === "node-postgres") {
+    const pool = new pg.Pool({
+      connectionString,
+      ...poolConfig || {}
+    });
+    return 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
+      }
+    );
+  }
+}
+function isPostgresConnectionString(conn) {
+  return conn.startsWith("postgres://") || conn.startsWith("postgresql://");
+}
+
+const CHECKPOINTS_TABLE_NAME = "checkpoints";
+const FILTERS_TABLE_NAME = "filters";
+const SCHEMA_VERSION_TABLE_NAME = "schema_version";
+const schema$1 = pgSchema(SCHEMA_NAME);
+const checkpoints = schema$1.table(CHECKPOINTS_TABLE_NAME, {
   id: text("id").notNull().primaryKey(),
   orderKey: integer("order_key").notNull(),
   uniqueKey: text("unique_key")
 });
-const filters = pgTable(
+const filters = schema$1.table(
   FILTERS_TABLE_NAME,
   {
     id: text("id").notNull(),
@@ -56,7 +111,7 @@ const filters = pgTable(
     }
   ]
 );
-const schemaVersion = pgTable(SCHEMA_VERSION_TABLE_NAME, {
+const schemaVersion = schema$1.table(SCHEMA_VERSION_TABLE_NAME, {
   k: integer("k").notNull().primaryKey(),
   version: integer("version").notNull()
 });
@@ -67,12 +122,19 @@ const MIGRATIONS = [
   // Add more migration arrays for future versions
 ];
 async function initializePersistentState(tx) {
-  await tx.execute(`
-    CREATE TABLE IF NOT EXISTS ${SCHEMA_VERSION_TABLE_NAME} (
+  await tx.execute(
+    sql.raw(`
+      CREATE SCHEMA IF NOT EXISTS ${SCHEMA_NAME};
+  `)
+  );
+  await tx.execute(
+    sql.raw(`
+    CREATE TABLE IF NOT EXISTS ${SCHEMA_NAME}.${SCHEMA_VERSION_TABLE_NAME} (
       k INTEGER PRIMARY KEY,
       version INTEGER NOT NULL
     );
-  `);
+  `)
+  );
   const versionRows = await tx.select().from(schemaVersion).where(eq(schemaVersion.k, 0));
   const storedVersion = versionRows[0]?.version ?? -1;
   if (storedVersion > CURRENT_SCHEMA_VERSION) {
@@ -82,22 +144,26 @@ async function initializePersistentState(tx) {
   }
   try {
     if (storedVersion === -1) {
-      await tx.execute(`
-        CREATE TABLE IF NOT EXISTS ${CHECKPOINTS_TABLE_NAME} (
+      await tx.execute(
+        sql.raw(`
+        CREATE TABLE IF NOT EXISTS ${SCHEMA_NAME}.${CHECKPOINTS_TABLE_NAME} (
           id TEXT PRIMARY KEY,
           order_key INTEGER NOT NULL,
           unique_key TEXT
         );
-      `);
-      await tx.execute(`
-        CREATE TABLE IF NOT EXISTS ${FILTERS_TABLE_NAME} (
+      `)
+      );
+      await tx.execute(
+        sql.raw(`
+        CREATE TABLE IF NOT EXISTS ${SCHEMA_NAME}.${FILTERS_TABLE_NAME} (
           id TEXT NOT NULL,
           filter TEXT NOT NULL,
           from_block INTEGER NOT NULL,
           to_block INTEGER DEFAULT NULL,
           PRIMARY KEY (id, from_block)
         );
-      `);
+      `)
+      );
       await tx.insert(schemaVersion).values({
         k: 0,
         version: CURRENT_SCHEMA_VERSION
@@ -132,7 +198,9 @@ async function persistState(props) {
         target: checkpoints.id,
         set: {
           orderKey: Number(endCursor.orderKey),
-          uniqueKey: endCursor.uniqueKey
+          // Explicitly set the unique key to `null` to indicate that it has been deleted
+          // Otherwise drizzle will not update its value.
+          uniqueKey: endCursor.uniqueKey ? endCursor.uniqueKey : null
         }
       });
       if (filter) {
@@ -211,11 +279,24 @@ async function finalizeState(props) {
     });
   }
 }
+async function resetPersistence(props) {
+  const { tx, indexerId } = props;
+  try {
+    await tx.delete(checkpoints).where(eq(checkpoints.id, indexerId));
+    await tx.delete(filters).where(eq(filters.id, indexerId));
+  } catch (error) {
+    throw new DrizzleStorageError("Failed to reset persistence state", {
+      cause: error
+    });
+  }
+}
 
+const ROLLBACK_TABLE_NAME = "reorg_rollback";
+const schema = pgSchema(SCHEMA_NAME);
 function getReorgTriggerName(table, indexerId) {
   return `${table}_reorg_${indexerId}`;
 }
-pgTable("__reorg_rollback", {
+schema.table(ROLLBACK_TABLE_NAME, {
   n: serial("n").primaryKey(),
   op: char("op", { length: 1 }).$type().notNull(),
   table_name: text("table_name").notNull(),
@@ -226,9 +307,12 @@ pgTable("__reorg_rollback", {
 });
 async function initializeReorgRollbackTable(tx, indexerId) {
   try {
+    await tx.execute(`
+    CREATE SCHEMA IF NOT EXISTS ${SCHEMA_NAME};
+    `);
     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,
@@ -241,7 +325,7 @@ async function initializeReorgRollbackTable(tx, indexerId) {
     );
     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) {
@@ -252,7 +336,7 @@ async function initializeReorgRollbackTable(tx, indexerId) {
   try {
     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;
@@ -262,13 +346,13 @@ async function initializeReorgRollbackTable(tx, indexerId) {
         old_id_value TEXT := row_to_json(OLD.*)->>id_col;
       BEGIN
         IF (TG_OP = 'DELETE') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
             SELECT 'D', TG_TABLE_NAME, order_key, old_id_value, row_to_json(OLD.*), indexer_id;
         ELSIF (TG_OP = 'UPDATE') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
             SELECT 'U', TG_TABLE_NAME, order_key, new_id_value, row_to_json(OLD.*), indexer_id;
         ELSIF (TG_OP = 'INSERT') THEN
-          INSERT INTO __reorg_rollback(op, table_name, cursor, row_id, row_value, indexer_id)
+          INSERT INTO ${SCHEMA_NAME}.${ROLLBACK_TABLE_NAME}(op, table_name, cursor, row_id, row_value, indexer_id)
             SELECT 'I', TG_TABLE_NAME, order_key, new_id_value, null, indexer_id;
         END IF;
         RETURN NULL;
@@ -298,7 +382,7 @@ async function registerTriggers(tx, tables, endCursor, idColumn, indexerId) {
           CREATE CONSTRAINT TRIGGER ${getReorgTriggerName(table, indexerId)}
           AFTER INSERT OR UPDATE OR DELETE ON ${table}
           DEFERRABLE INITIALLY DEFERRED
-          FOR EACH ROW EXECUTE FUNCTION reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
+          FOR EACH ROW EXECUTE FUNCTION ${SCHEMA_NAME}.reorg_checkpoint('${idColumn}', ${`${Number(endCursor.orderKey)}`}, '${indexerId}');
         `)
       );
     }
@@ -327,7 +411,7 @@ async function invalidate(tx, cursor, idColumn, indexerId) {
   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 *
@@ -420,7 +504,7 @@ async function finalize(tx, cursor, indexerId) {
   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}'
     `)
@@ -431,8 +515,37 @@ async function finalize(tx, cursor, indexerId) {
     });
   }
 }
+async function cleanupStorage(tx, tables, indexerId) {
+  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
+    });
+  }
+}
 
-const DRIZZLE_PROPERTY = "_drizzle";
 const MAX_RETRIES = 5;
 function useDrizzleStorage(_db) {
   const context = useIndexerContext();
@@ -448,11 +561,14 @@ function drizzleStorage({
   persistState: enablePersistence = true,
   indexerName: identifier = "default",
   schema,
-  idColumn = "id"
+  idColumn = "id",
+  migrate: migrateOptions
 }) {
   return defineIndexerPlugin((indexer) => {
     let tableNames = [];
     let indexerId = "";
+    const alwaysReindex = process.env["APIBARA_ALWAYS_REINDEX"] === "true";
+    let prevFinality;
     try {
       tableNames = Object.values(schema ?? db._.schema ?? {}).map(
         (table) => table.dbName
@@ -463,11 +579,33 @@ 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);
+      if (alwaysReindex) {
+        logger.warn(
+          `Reindexing: Deleting all data from tables - ${tableNames.join(", ")}`
+        );
+        await withTransaction(db, async (tx) => {
+          await cleanupStorage(tx, tableNames, indexerId);
+          if (enablePersistence) {
+            await resetPersistence({ tx, indexerId });
+          }
+          logger.success("Tables have been cleaned up for reindexing");
+        });
+      }
       let retries = 0;
+      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 +615,9 @@ function drizzleStorage({
           break;
         } catch (error) {
           if (retries === MAX_RETRIES) {
+            if (error instanceof DrizzleStorageError) {
+              throw error;
+            }
             throw new DrizzleStorageError(
               "Initialization failed after 5 retries",
               {
@@ -559,12 +700,15 @@ function drizzleStorage({
     indexer.hooks.hook("handler:middleware", async ({ use }) => {
       use(async (context, next) => {
         try {
-          const { endCursor, finality } = context;
+          const { endCursor, finality, cursor } = context;
           if (!endCursor) {
             throw new DrizzleStorageError("End Cursor is undefined");
           }
           await withTransaction(db, async (tx) => {
             context[DRIZZLE_PROPERTY] = { db: tx };
+            if (prevFinality === "pending") {
+              await invalidate(tx, cursor, idColumn, indexerId);
+            }
             if (finality !== "finalized") {
               await registerTriggers(
                 tx,
@@ -576,13 +720,14 @@ function drizzleStorage({
             }
             await next();
             delete context[DRIZZLE_PROPERTY];
-            if (enablePersistence) {
+            if (enablePersistence && finality !== "pending") {
               await persistState({
                 tx,
                 endCursor,
                 indexerId
               });
             }
+            prevFinality = finality;
           });
           if (finality !== "finalized") {
             await removeTriggers(db, tableNames, indexerId);
@@ -598,4 +743,4 @@ function drizzleStorage({
   });
 }
 
-export { drizzleStorage, useDrizzleStorage };
+export { drizzle, drizzleStorage, migrate, useDrizzleStorage };

package/dist/shared/plugin-drizzle.2d226351.mjs

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

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

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

package/dist/testing.cjs

@@ -0,0 +1,13 @@
+'use strict';
+
+const constants = require('./shared/plugin-drizzle.cae20704.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.2d226351.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.2",
+  "version": "2.1.0-beta.20",
   "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": {
@@ -26,21 +32,21 @@
     "test:ci": "vitest run"
   },
   "peerDependencies": {
-    "drizzle-orm": "^0.37.0",
-    "pg": "^8.13.1"
+    "drizzle-orm": "<1",
+    "pg": ">=8"
   },
   "devDependencies": {
     "@electric-sql/pglite": "^0.2.17",
     "@types/node": "^20.14.0",
     "@types/pg": "^8.11.10",
-    "drizzle-orm": "^0.37.0",
+    "drizzle-orm": "^0.40.1",
     "pg": "^8.13.1",
     "unbuild": "^2.0.0",
     "vitest": "^1.6.0"
   },
   "dependencies": {
-    "@apibara/indexer": "2.1.0-beta.2",
-    "@apibara/protocol": "2.1.0-beta.2",
+    "@apibara/indexer": "2.1.0-beta.20",
+    "@apibara/protocol": "2.1.0-beta.20",
     "postgres-range": "^1.1.4"
   }
 }

package/src/constants.ts

@@ -0,0 +1,3 @@
+export const DRIZZLE_PROPERTY = "_drizzle";
+export const DRIZZLE_STORAGE_DB_PROPERTY = "_drizzleStorageDB";
+export const SCHEMA_NAME = "airfoil";