@apibara/indexer 2.0.0-beta.0 → 2.0.0-beta.4

package/src/sinks/csv.ts

@@ -1,7 +1,7 @@
 import fs from "node:fs";
 import type { Cursor } from "@apibara/protocol";
 import { type Options, type Stringifier, stringify } from "csv-stringify";
-import { Sink, type SinkData, type SinkWriteArgs } from "../sink";
+import { Sink, type SinkCursorParams, type SinkData } from "../sink";
 
 export type CsvArgs = {
   /**
@@ -23,6 +23,50 @@ export type CsvSinkOptions = {
   cursorColumn?: string;
 };
 
+type TxnContext = {
+  buffer: SinkData[];
+};
+
+type TxnParams = {
+  writer: {
+    insert: (data: SinkData[]) => void;
+  };
+};
+
+const transactionHelper = (context: TxnContext) => {
+  return {
+    insert: (data: SinkData[]) => {
+      context.buffer.push(...data);
+    },
+  };
+};
+
+/**
+ * A sink that writes data to a CSV file.
+ *
+ * @example
+ *
+ * ```ts
+ * const sink = csv({
+ *   filepath: "./data.csv",
+ *   csvOptions: {
+ *     header: true,
+ *   },
+ * });
+ *
+ * ...
+ * async transform({context, endCursor}){
+ *  const { writer } = useSink(context);
+ *  const insertHelper = writer(endCursor);
+ *
+ *  insertHelper.insert([
+ *    { id: 1, name: "John" },
+ *    { id: 2, name: "Jane" },
+ *  ]);
+ * }
+ *
+ * ```
+ */
 export class CsvSink extends Sink {
   constructor(
     private _stringifier: Stringifier,
@@ -31,14 +75,33 @@ export class CsvSink extends Sink {
     super();
   }
 
-  async write({ data, endCursor, finality }: SinkWriteArgs) {
-    await this.callHook("write", { data });
+  private async write({
+    data,
+    endCursor,
+  }: { data: SinkData[]; endCursor?: Cursor }) {
     // adds a "_cursor" property if "cursorColumn" is not specified by user
     data = this.processCursorColumn(data, endCursor);
     // Insert the data into csv
     await this.insertToCSV(data);
+  }
+
+  async transaction(
+    { cursor, endCursor, finality }: SinkCursorParams,
+    cb: (params: TxnParams) => Promise<void>,
+  ) {
+    const context: TxnContext = {
+      buffer: [],
+    };
+
+    const writer = transactionHelper(context);
+
+    await cb({ writer });
+    await this.write({ data: context.buffer, endCursor });
+  }
 
-    await this.callHook("flush", { endCursor, finality });
+  async invalidate(cursor?: Cursor) {
+    // TODO: Implement
+    throw new Error("Not implemented");
   }
 
   private async insertToCSV(data: SinkData[]) {
@@ -86,9 +149,7 @@ export class CsvSink extends Sink {
   }
 }
 
-export const csv = <TData extends Record<string, unknown>>(
-  args: CsvArgs & CsvSinkOptions,
-) => {
+export const csv = (args: CsvArgs & CsvSinkOptions) => {
   const { csvOptions, filepath, ...sinkOptions } = args;
   const stringifier = stringify({ ...csvOptions });
 

package/src/sinks/drizzle/Int8Range.ts

@@ -0,0 +1,52 @@
+import { customType } from "drizzle-orm/pg-core";
+import type { Range } from "postgres-range";
+import {
+  parse as rangeParse,
+  serialize as rangeSerialize,
+} from "postgres-range";
+
+type Comparable = string | number;
+
+type RangeBound<T extends Comparable> =
+  | T
+  | {
+      value: T;
+      inclusive: boolean;
+    };
+
+export class Int8Range {
+  constructor(public readonly range: Range<number>) {}
+
+  get start(): RangeBound<number> | null {
+    return this.range.lower != null
+      ? {
+          value: this.range.lower,
+          inclusive: this.range.isLowerBoundClosed(),
+        }
+      : null;
+  }
+
+  get end(): RangeBound<number> | null {
+    return this.range.upper != null
+      ? {
+          value: this.range.upper,
+          inclusive: this.range.isUpperBoundClosed(),
+        }
+      : null;
+  }
+}
+
+export const int8range = customType<{
+  data: Int8Range;
+}>({
+  dataType: () => "int8range",
+  fromDriver: (value: unknown): Int8Range => {
+    if (typeof value !== "string") {
+      throw new Error("Expected string");
+    }
+
+    const parsed = rangeParse(value, (val) => Number.parseInt(val, 10));
+    return new Int8Range(parsed);
+  },
+  toDriver: (value: Int8Range): string => rangeSerialize(value.range),
+});

package/src/sinks/drizzle/delete.ts

@@ -0,0 +1,42 @@
+import type { Cursor } from "@apibara/protocol";
+import {
+  type ExtractTablesWithRelations,
+  type SQL,
+  type TablesRelationalConfig,
+  sql,
+} from "drizzle-orm";
+import type {
+  PgQueryResultHKT,
+  PgTable,
+  PgTransaction,
+  PgUpdateBase,
+} from "drizzle-orm/pg-core";
+
+export class DrizzleSinkDelete<
+  TTable extends PgTable,
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> {
+  constructor(
+    private db: PgTransaction<TQueryResult, TFullSchema, TSchema>,
+    private table: TTable,
+    private endCursor?: Cursor,
+  ) {}
+
+  //@ts-ignore
+  where(where: SQL): Omit<
+    // for type safety used PgUpdateBase instead of PgDeleteBase
+    PgUpdateBase<TTable, TQueryResult, undefined, false, "where">,
+    "where"
+  > {
+    return this.db
+      .update(this.table)
+      .set({
+        // @ts-ignore
+        _cursor: sql`int8range(lower(_cursor), ${Number(this.endCursor?.orderKey!)}, '[)')`,
+      })
+      .where(where);
+  }
+}

package/src/sinks/drizzle/drizzle.test.ts

@@ -0,0 +1,239 @@
+import type { Cursor } from "@apibara/protocol";
+import {
+  type MockBlock,
+  MockClient,
+  type MockFilter,
+} from "@apibara/protocol/testing";
+import { asc, eq, sql } from "drizzle-orm";
+import { drizzle } from "drizzle-orm/node-postgres";
+import { serial, text } from "drizzle-orm/pg-core";
+import { Client } from "pg";
+import { beforeAll, beforeEach, describe, expect, it } from "vitest";
+import {
+  type Int8Range,
+  drizzle as drizzleSink,
+  getDrizzleCursor,
+  pgTable,
+} from ".";
+import { useSink } from "../../hooks";
+import { run } from "../../indexer";
+import { generateMockMessages } from "../../testing";
+import { getMockIndexer } from "../../testing/indexer";
+
+const testTable = pgTable("test_table", {
+  id: serial("id").primaryKey(),
+  data: text("data"),
+});
+
+const client = new Client({
+  connectionString: "postgres://postgres:postgres@localhost:5432/postgres",
+});
+
+await client.connect();
+
+const db = drizzle(client);
+
+describe("Drizzle Test", () => {
+  beforeAll(async () => {
+    // drop test_table if exists
+    await db.execute(sql`DROP TABLE IF EXISTS test_table`);
+    // create test_table with db
+    await db.execute(
+      sql`CREATE TABLE test_table (id SERIAL PRIMARY KEY, data TEXT, _cursor INT8RANGE)`,
+    );
+  });
+
+  beforeEach(async () => {
+    await db.delete(testTable).execute();
+  });
+
+  it("should insert data", async () => {
+    const client = new MockClient<MockFilter, MockBlock>((request, options) => {
+      return generateMockMessages(5);
+    });
+
+    const sink = drizzleSink({ database: db, tables: [testTable] });
+
+    const indexer = getMockIndexer({
+      sink,
+      override: {
+        transform: async ({ context, endCursor, block: { data } }) => {
+          const { db } = useSink({ context });
+          // Insert a new row into the test_table
+          // The id is set to the current cursor's orderKey
+          // The data is set to the block data
+          await db
+            .insert(testTable)
+            .values([{ id: Number(endCursor?.orderKey), data }]);
+        },
+      },
+    });
+
+    await run(client, indexer);
+
+    const result = await db.select().from(testTable).orderBy(asc(testTable.id));
+
+    expect(result).toHaveLength(5);
+    expect(result[0].data).toBe("5000000");
+    expect(result[2].data).toBe("5000002");
+  });
+
+  it("should update data", async () => {
+    const client = new MockClient<MockFilter, MockBlock>((request, options) => {
+      return generateMockMessages(5);
+    });
+
+    const sink = drizzleSink({ database: db, tables: [testTable] });
+
+    const indexer = getMockIndexer({
+      sink,
+      override: {
+        transform: async ({ context, endCursor, block: { data } }) => {
+          const { db } = useSink({ context });
+
+          // insert data for each message in db
+          await db
+            .insert(testTable)
+            .values([{ id: Number(endCursor?.orderKey), data }]);
+
+          // update data for id 5000002 when orderKey is 5000004
+          // this is to test if the update query is working
+          if (endCursor?.orderKey === 5000004n) {
+            await db
+              .update(testTable)
+              .set({ data: "0000000" })
+              .where(eq(testTable.id, 5000002));
+          }
+        },
+      },
+    });
+
+    await run(client, indexer);
+
+    const result = await db.select().from(testTable).orderBy(asc(testTable.id));
+
+    expect(result).toHaveLength(5);
+    expect(result[2].data).toBe("0000000");
+  });
+
+  it("should delete data", async () => {
+    const client = new MockClient<MockFilter, MockBlock>((request, options) => {
+      return generateMockMessages(5);
+    });
+
+    const sink = drizzleSink({ database: db, tables: [testTable] });
+
+    const indexer = getMockIndexer({
+      sink,
+      override: {
+        transform: async ({ context, endCursor, block: { data } }) => {
+          const { db } = useSink({ context });
+
+          // insert data for each message in db
+          await db
+            .insert(testTable)
+            .values([{ id: Number(endCursor?.orderKey), data }]);
+
+          // delete data for id 5000002 when orderKey is 5000004
+          // this is to test if the delete query is working
+          if (endCursor?.orderKey === 5000004n) {
+            await db.delete(testTable).where(eq(testTable.id, 5000002));
+          }
+        },
+      },
+    });
+
+    await run(client, indexer);
+
+    const result = await db.select().from(testTable).orderBy(asc(testTable.id));
+
+    expect(result).toHaveLength(5);
+
+    // as when you run delete query on a data, it isnt literally deleted from the db,
+    // instead, we just update the upper bound of that row to the current cursor
+    // check if the cursor upper bound has been set correctly
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    expect(((result[2] as any)._cursor as Int8Range).range.upper).toBe(5000004);
+  });
+
+  it("should select data", async () => {
+    const client = new MockClient<MockFilter, MockBlock>((request, options) => {
+      return generateMockMessages(5);
+    });
+
+    const sink = drizzleSink({ database: db, tables: [testTable] });
+
+    let result: (typeof testTable.$inferSelect)[] = [];
+
+    const indexer = getMockIndexer({
+      sink,
+      override: {
+        transform: async ({ context, endCursor, block: { data } }) => {
+          const { db } = useSink({ context });
+
+          // insert data for each message in db
+          await db
+            .insert(testTable)
+            .values([{ id: Number(endCursor?.orderKey), data }]);
+
+          // delete data for id 5000002 when orderKey is 5000004
+          // this will update the upper bound of the row with id 5000002 from infinity to 5000004
+          // so when we select all rows, row with id 5000002 will not be included
+          // as when we run select query it should only return rows with upper bound infinity
+          if (endCursor?.orderKey === 5000003n) {
+            await db.delete(testTable).where(eq(testTable.id, 5000002));
+          }
+
+          // when on last message of mock stream, select all rows from db
+          if (endCursor?.orderKey === 5000004n) {
+            result = await db
+              .select()
+              .from(testTable)
+              .orderBy(asc(testTable.id));
+          }
+        },
+      },
+    });
+
+    await run(client, indexer);
+
+    expect(result).toHaveLength(4);
+    expect(result.find((r) => r.id === 5000002)).toBeUndefined();
+    // check if all rows are still in db
+    const allRows = await db.select().from(testTable);
+    expect(allRows).toHaveLength(5);
+  });
+
+  it("should invalidate data correctly", async () => {
+    const sink = drizzleSink({ database: db, tables: [testTable] });
+
+    // Insert some test data
+    await db.insert(testTable).values(
+      // @ts-ignore
+      [
+        { id: 1, data: "data1", _cursor: getDrizzleCursor([1n, 5n]) },
+        { id: 2, data: "data2", _cursor: getDrizzleCursor([2n, 5n]) },
+        { id: 3, data: "data3", _cursor: getDrizzleCursor(3n) },
+        { id: 4, data: "data4", _cursor: getDrizzleCursor(4n) },
+        { id: 5, data: "data5", _cursor: getDrizzleCursor(5n) },
+      ],
+    );
+
+    // Create a cursor at position 3
+    const cursor: Cursor = { orderKey: 3n };
+
+    // Invalidate data
+    await sink.invalidate(cursor);
+
+    // Check the results
+    const result = await db.select().from(testTable).orderBy(asc(testTable.id));
+
+    expect(result).toHaveLength(3);
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    expect(((result[0] as any)._cursor as Int8Range).range.upper).toBe(null);
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    expect(((result[1] as any)._cursor as Int8Range).range.upper).toBe(null);
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    expect(((result[2] as any)._cursor as Int8Range).range.upper).toBe(null);
+  });
+});

package/src/sinks/drizzle/drizzle.ts

@@ -0,0 +1,115 @@
+import type { Cursor } from "@apibara/protocol";
+import {
+  type ExtractTablesWithRelations,
+  type TablesRelationalConfig,
+  gt,
+  sql,
+} from "drizzle-orm";
+import type {
+  AnyPgTable,
+  PgDatabase,
+  PgQueryResultHKT,
+  PgTableWithColumns,
+  TableConfig,
+} from "drizzle-orm/pg-core";
+import { Sink, type SinkCursorParams } from "../../sink";
+import { DrizzleSinkTransaction } from "./transaction";
+
+export type DrizzleSinkTables<
+  TTableConfig extends Record<string, TableConfig>,
+> = {
+  [K in keyof TTableConfig]: PgTableWithColumns<TTableConfig[K]>;
+};
+
+export type DrizzleSinkOptions<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> = {
+  /**
+   * Database instance of drizzle-orm
+   */
+  database: PgDatabase<TQueryResult, TFullSchema, TSchema>;
+  tables: AnyPgTable[];
+};
+
+/**
+ * A sink that writes data to a PostgreSQL database using Drizzle ORM.
+ *
+ * @example
+ *
+ * ```ts
+ * const sink = drizzle({
+ *   database: db,
+ * });
+ *
+ * ...
+ * async transform({context, endCursor}){
+ *  const { transaction } = useSink(context);
+ *  const db = transaction(endCursor);
+ *
+ *  db.insert(users).values([
+ *    { id: 1, name: "John" },
+ *    { id: 2, name: "Jane" },
+ *  ]);
+ * }
+ *
+ * ```
+ */
+export class DrizzleSink<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> extends Sink {
+  private _db: PgDatabase<TQueryResult, TFullSchema, TSchema>;
+  private _tables: AnyPgTable[];
+  constructor(options: DrizzleSinkOptions<TQueryResult, TFullSchema, TSchema>) {
+    super();
+    const { database, tables } = options;
+    this._db = database;
+    this._tables = tables;
+  }
+
+  async transaction(
+    { cursor, endCursor, finality }: SinkCursorParams,
+    cb: (params: {
+      db: DrizzleSinkTransaction<TQueryResult, TFullSchema, TSchema>;
+    }) => Promise<void>,
+  ): Promise<void> {
+    await this._db.transaction(async (db) => {
+      await cb({ db: new DrizzleSinkTransaction(db, endCursor) });
+    });
+  }
+
+  async invalidate(cursor?: Cursor) {
+    await this._db.transaction(async (db) => {
+      for (const table of this._tables) {
+        // delete all rows whose lowerbound of "_cursor" (int8range) column is greater than the invalidate cursor
+        await db
+          .delete(table)
+          .where(gt(sql`lower(_cursor)`, sql`${Number(cursor?.orderKey)}`))
+          .returning();
+        // and for rows whose upperbound of "_cursor" (int8range) column is greater than the invalidate cursor, set the upperbound to infinity
+        await db
+          .update(table)
+          .set({
+            _cursor: sql`int8range(lower(_cursor), NULL, '[)')`,
+          })
+          .where(gt(sql`upper(_cursor)`, sql`${Number(cursor?.orderKey)}`));
+      }
+    });
+  }
+}
+
+export const drizzle = <
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+>(
+  args: DrizzleSinkOptions<TQueryResult, TFullSchema, TSchema>,
+) => {
+  return new DrizzleSink(args);
+};

package/src/sinks/drizzle/index.ts

@@ -0,0 +1,6 @@
+export * from "./drizzle";
+export * from "./Int8Range";
+export * from "./utils";
+export * from "./transaction";
+export * from "./update";
+export * from "./delete";

package/src/sinks/drizzle/insert.ts

@@ -0,0 +1,39 @@
+import type { Cursor } from "@apibara/protocol";
+import type {
+  ExtractTablesWithRelations,
+  TablesRelationalConfig,
+} from "drizzle-orm";
+import type {
+  PgInsertValue,
+  PgQueryResultHKT,
+  PgTable,
+  PgTransaction,
+} from "drizzle-orm/pg-core";
+import { getDrizzleCursor } from "./utils";
+
+export class DrizzleSinkInsert<
+  TTable extends PgTable,
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> {
+  constructor(
+    private db: PgTransaction<TQueryResult, TFullSchema, TSchema>,
+    private table: TTable,
+    private endCursor?: Cursor,
+  ) {}
+
+  values(values: PgInsertValue<TTable> | PgInsertValue<TTable>[]) {
+    const originalInsert = this.db.insert(this.table);
+    const cursoredValues = (Array.isArray(values) ? values : [values]).map(
+      (v) => {
+        return {
+          ...v,
+          _cursor: getDrizzleCursor(this.endCursor?.orderKey),
+        };
+      },
+    );
+    return originalInsert.values(cursoredValues as PgInsertValue<TTable>[]);
+  }
+}

package/src/sinks/drizzle/select.ts

@@ -0,0 +1,44 @@
+import type { Cursor } from "@apibara/protocol";
+import {
+  type ExtractTablesWithRelations,
+  type SQL,
+  type Subquery,
+  type TablesRelationalConfig,
+  sql,
+} from "drizzle-orm";
+import type {
+  PgQueryResultHKT,
+  PgTable,
+  PgTransaction,
+  SelectedFields,
+} from "drizzle-orm/pg-core";
+import type { PgViewBase } from "drizzle-orm/pg-core/view-base";
+
+export class DrizzleSinkSelect<
+  TSelection extends SelectedFields,
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> {
+  constructor(
+    private db: PgTransaction<TQueryResult, TFullSchema, TSchema>,
+    private fields?: TSelection,
+    private endCursor?: Cursor,
+  ) {}
+
+  from<TFrom extends PgTable | Subquery | PgViewBase | SQL>(source: TFrom) {
+    if (this.fields) {
+      const originalFrom = this.db.select(this.fields).from(source);
+      return {
+        ...originalFrom,
+        where: (where?: SQL) => {
+          const combinedWhere = sql`${where ? sql`${where} AND ` : sql``}upper_inf(_cursor)`;
+          return originalFrom.where(combinedWhere);
+        },
+      };
+    }
+
+    return this.db.select().from(source).where(sql`upper_inf(_cursor)`);
+  }
+}

package/src/sinks/drizzle/transaction.ts

@@ -0,0 +1,49 @@
+import type { Cursor } from "@apibara/protocol";
+import type {
+  ExtractTablesWithRelations,
+  TablesRelationalConfig,
+} from "drizzle-orm";
+import type {
+  PgQueryResultHKT,
+  PgSelectBuilder,
+  PgTable,
+  PgTransaction,
+  SelectedFields,
+} from "drizzle-orm/pg-core";
+import { DrizzleSinkDelete } from "./delete";
+import { DrizzleSinkInsert } from "./insert";
+import { DrizzleSinkSelect } from "./select";
+import { DrizzleSinkUpdate } from "./update";
+
+export class DrizzleSinkTransaction<
+  TQueryResult extends PgQueryResultHKT,
+  TFullSchema extends Record<string, unknown> = Record<string, never>,
+  TSchema extends
+    TablesRelationalConfig = ExtractTablesWithRelations<TFullSchema>,
+> {
+  constructor(
+    private db: PgTransaction<TQueryResult, TFullSchema, TSchema>,
+    private endCursor?: Cursor,
+  ) {}
+
+  insert<TTable extends PgTable>(table: TTable) {
+    return new DrizzleSinkInsert(this.db, table, this.endCursor);
+  }
+
+  update<TTable extends PgTable>(table: TTable) {
+    return new DrizzleSinkUpdate(this.db, table, this.endCursor);
+  }
+
+  delete<TTable extends PgTable>(table: TTable) {
+    return new DrizzleSinkDelete(this.db, table, this.endCursor);
+  }
+
+  // @ts-ignore
+  select(): PgSelectBuilder<undefined>;
+  select<TSelection extends SelectedFields>(
+    fields: TSelection,
+  ): PgSelectBuilder<TSelection>;
+  select(fields?: SelectedFields) {
+    return new DrizzleSinkSelect(this.db, fields, this.endCursor);
+  }
+}