ponder-with-flashblocks 0.16.3

package/src/indexing-store/index.ts

@@ -0,0 +1,628 @@
+import type { QB } from "@/database/queryBuilder.js";
+import { onchain } from "@/drizzle/onchain.js";
+import type { Common } from "@/internal/common.js";
+import {
+  DbConnectionError,
+  InvalidStoreAccessError,
+  InvalidStoreMethodError,
+  NonRetryableUserError,
+  RawSqlError,
+  RecordNotFoundError,
+  RetryableError,
+  UndefinedTableError,
+  UniqueConstraintError,
+} from "@/internal/errors.js";
+import type { Schema } from "@/internal/types.js";
+import type { IndexingErrorHandler, SchemaBuild } from "@/internal/types.js";
+import type { Db } from "@/types/db.js";
+import { copy, copyOnWrite } from "@/utils/copy.js";
+import { createLock } from "@/utils/mutex.js";
+import { prettyPrint } from "@/utils/print.js";
+import { getSQLQueryRelations, isReadonlySQLQuery } from "@/utils/sql-parse.js";
+import { startClock } from "@/utils/timer.js";
+import {
+  type Column,
+  type QueryWithTypings,
+  type Table,
+  getTableName,
+  getViewName,
+  isTable,
+  isView,
+} from "drizzle-orm";
+import { getTableConfig } from "drizzle-orm/pg-core";
+import { drizzle } from "drizzle-orm/pg-proxy";
+import type { IndexingCache, Row } from "./cache.js";
+import { getPrimaryKeyCache } from "./utils.js";
+
+export type IndexingStore = {
+  db: Db<Schema>;
+  qb: QB;
+  isProcessingEvents: boolean;
+};
+
+export const validateUpdateSet = (
+  table: Table,
+  set: Object,
+  prev: Row,
+  cache: Map<Table, [string, Column][]>,
+) => {
+  const primaryKeys = cache.get(table)!;
+
+  for (let i = 0; i < primaryKeys.length; i++) {
+    const [js] = primaryKeys[i]!;
+    if (js in set) {
+      // Note: Noop on the primary keys if they are identical, otherwise throw an error.
+      if ((set as Row)[js] !== prev[js]) {
+        throw new NonRetryableUserError(
+          `Primary key column '${js}' cannot be updated`,
+        );
+      }
+    }
+  }
+};
+
+/** Throw an error if `table` is not an `onchainTable`. */
+export const checkOnchainTable = (
+  table: Table,
+  method: "find" | "insert" | "update" | "delete",
+) => {
+  if (table === undefined)
+    throw new UndefinedTableError(
+      `Table object passed to db.${method}() is undefined`,
+    );
+
+  if (onchain in table) return;
+
+  throw new InvalidStoreMethodError(
+    method === "find"
+      ? `db.find() can only be used with onchain tables, and '${getTableConfig(table).name}' is an offchain table or a view.`
+      : `Indexing functions can only write to onchain tables, and '${getTableConfig(table).name}' is an offchain table or a view.`,
+  );
+};
+
+export const checkTableAccess = (
+  table: Table,
+  method: "find" | "insert" | "update" | "delete",
+  key: object,
+  chainId?: number,
+) => {
+  if (chainId === undefined) return;
+  if ("chainId" in key && String(key.chainId) === String(chainId)) return;
+  throw new InvalidStoreAccessError(
+    "chainId" in key
+      ? `db.${method}(${getTableConfig(table).name}) cannot access rows on different chains when ordering is 'isolated'.`
+      : `db.${method}(${getTableConfig(table).name}) must specify 'chainId' when ordering is 'isolated'.`,
+  );
+};
+
+export const createIndexingStore = ({
+  common,
+  schemaBuild: { schema },
+  indexingCache,
+  indexingErrorHandler,
+  chainId,
+}: {
+  common: Common;
+  schemaBuild: Pick<SchemaBuild, "schema">;
+  indexingCache: IndexingCache;
+  indexingErrorHandler: IndexingErrorHandler;
+  chainId?: number;
+}): IndexingStore => {
+  let qb: QB = undefined!;
+  let isProcessingEvents = true;
+
+  const tables = Object.values(schema).filter(isTable);
+  const views = Object.values(schema).filter(isView);
+  const primaryKeyCache = getPrimaryKeyCache(tables);
+
+  const lock = createLock();
+
+  const storeMethodWrapper = (fn: (...args: any[]) => Promise<any>) => {
+    return async (...args: any[]) => {
+      try {
+        if (isProcessingEvents === false) {
+          throw new NonRetryableUserError(
+            "A store API method (find, update, insert, delete) was called after the indexing function returned. Hint: Did you forget to await the store API method call (an unawaited promise)?",
+          );
+        }
+        await lock.lock();
+        const result = await fn(...args);
+        // @ts-expect-error typescript bug lol
+        if (isProcessingEvents === false) {
+          throw new NonRetryableUserError(
+            "A store API method (find, update, insert, delete) was called after the indexing function returned. Hint: Did you forget to await the store API method call (an unawaited promise)?",
+          );
+        }
+        return result;
+      } catch (error) {
+        if (isProcessingEvents === false) {
+          throw new NonRetryableUserError(
+            "A store API method (find, update, insert, delete) was called after the indexing function returned. Hint: Did you forget to await the store API method call (an unawaited promise)?",
+          );
+        }
+
+        if (error instanceof RetryableError) {
+          indexingErrorHandler.setRetryableError(error);
+        }
+
+        throw error;
+      } finally {
+        lock.unlock();
+      }
+    };
+  };
+
+  // Note: the naming convention is used to separate user space from ponder space.
+  // "user" prefix is user space, "ponder" prefix is ponder space.
+
+  return {
+    db: {
+      // @ts-ignore
+      find: storeMethodWrapper(async (table: Table, key) => {
+        common.metrics.ponder_indexing_store_queries_total.inc({
+          table: getTableName(table),
+          method: "find",
+        });
+        checkOnchainTable(table, "find");
+        checkTableAccess(table, "find", key, chainId);
+        const ponderRow = await indexingCache.get({ table, key });
+        const userRow = ponderRow === null ? null : copyOnWrite(ponderRow);
+        return userRow;
+      }),
+      // @ts-ignore
+      insert(table: Table) {
+        return {
+          values: (userValues: any) => {
+            // @ts-ignore
+            const inner = {
+              onConflictDoNothing: storeMethodWrapper(async () => {
+                common.metrics.ponder_indexing_store_queries_total.inc({
+                  table: getTableName(table),
+                  method: "insert",
+                });
+                checkOnchainTable(table, "insert");
+
+                const ponderValues = copy(userValues);
+
+                if (Array.isArray(ponderValues)) {
+                  const ponderRows = [];
+                  for (const value of ponderValues) {
+                    checkTableAccess(table, "insert", value, chainId);
+                    const row = await indexingCache.get({ table, key: value });
+
+                    if (row) {
+                      ponderRows.push(null);
+                    } else {
+                      ponderRows.push(
+                        indexingCache.set({
+                          table,
+                          key: value,
+                          row: value,
+                          isUpdate: false,
+                        }),
+                      );
+                    }
+                  }
+                  const userRows = ponderRows.map((row) =>
+                    row === null ? row : copyOnWrite(row),
+                  );
+                  return userRows;
+                } else {
+                  checkTableAccess(table, "insert", ponderValues, chainId);
+                  const row = await indexingCache.get({
+                    table,
+                    key: ponderValues,
+                  });
+
+                  if (row) {
+                    return null;
+                  }
+
+                  const ponderRow = indexingCache.set({
+                    table,
+                    key: ponderValues,
+                    row: ponderValues,
+                    isUpdate: false,
+                  });
+
+                  const userRow = copyOnWrite(ponderRow);
+                  return userRow;
+                }
+              }),
+              onConflictDoUpdate: storeMethodWrapper(
+                async (userUpdateValues: any) => {
+                  common.metrics.ponder_indexing_store_queries_total.inc({
+                    table: getTableName(table),
+                    method: "insert",
+                  });
+                  checkOnchainTable(table, "insert");
+
+                  if (Array.isArray(userValues)) {
+                    const ponderRows: Row[] = [];
+                    for (const value of userValues) {
+                      checkTableAccess(table, "insert", value, chainId);
+                      const ponderRowUpdate = await indexingCache.get({
+                        table,
+                        key: value,
+                      });
+
+                      if (ponderRowUpdate) {
+                        if (typeof userUpdateValues === "function") {
+                          const userRowUpdate = copyOnWrite(ponderRowUpdate);
+                          const userSet = userUpdateValues(userRowUpdate);
+                          const ponderSet = copy(userSet);
+                          validateUpdateSet(
+                            table,
+                            ponderSet,
+                            ponderRowUpdate,
+                            primaryKeyCache,
+                          );
+                          for (const [key, value] of Object.entries(
+                            ponderSet,
+                          )) {
+                            if (value === undefined) continue;
+                            ponderRowUpdate[key] = value;
+                          }
+                        } else {
+                          const userSet = userUpdateValues;
+                          const ponderSet = copy(userSet);
+                          validateUpdateSet(
+                            table,
+                            ponderSet,
+                            ponderRowUpdate,
+                            primaryKeyCache,
+                          );
+                          for (const [key, value] of Object.entries(
+                            ponderSet,
+                          )) {
+                            if (value === undefined) continue;
+                            ponderRowUpdate[key] = value;
+                          }
+                        }
+                        ponderRows.push(
+                          indexingCache.set({
+                            table,
+                            key: ponderRowUpdate,
+                            row: ponderRowUpdate,
+                            isUpdate: true,
+                          }),
+                        );
+                      } else {
+                        const ponderValue = copy(value);
+                        ponderRows.push(
+                          indexingCache.set({
+                            table,
+                            key: ponderValue,
+                            row: ponderValue,
+                            isUpdate: false,
+                          }),
+                        );
+                      }
+                    }
+                    const userRows = ponderRows.map((row) =>
+                      row === null ? row : copyOnWrite(row),
+                    );
+                    return userRows;
+                  } else {
+                    checkTableAccess(table, "insert", userValues, chainId);
+                    const ponderRowUpdate = await indexingCache.get({
+                      table,
+                      key: userValues,
+                    });
+
+                    if (ponderRowUpdate) {
+                      if (typeof userUpdateValues === "function") {
+                        const userRowUpdate = copyOnWrite(ponderRowUpdate);
+                        const userSet = userUpdateValues(userRowUpdate);
+                        const ponderSet = copy(userSet);
+                        validateUpdateSet(
+                          table,
+                          ponderSet,
+                          ponderRowUpdate,
+                          primaryKeyCache,
+                        );
+                        for (const [key, value] of Object.entries(ponderSet)) {
+                          if (value === undefined) continue;
+                          ponderRowUpdate[key] = value;
+                        }
+                      } else {
+                        const userSet = userUpdateValues;
+                        const ponderSet = copy(userSet);
+                        validateUpdateSet(
+                          table,
+                          ponderSet,
+                          ponderRowUpdate,
+                          primaryKeyCache,
+                        );
+                        for (const [key, value] of Object.entries(ponderSet)) {
+                          if (value === undefined) continue;
+                          ponderRowUpdate[key] = value;
+                        }
+                      }
+                      const ponderRow = indexingCache.set({
+                        table,
+                        key: ponderRowUpdate,
+                        row: ponderRowUpdate,
+                        isUpdate: true,
+                      });
+                      const userRow = copyOnWrite(ponderRow);
+                      return userRow;
+                    }
+
+                    const ponderValues = copy(userValues);
+
+                    const ponderRowInsert = indexingCache.set({
+                      table,
+                      key: ponderValues,
+                      row: ponderValues,
+                      isUpdate: false,
+                    });
+
+                    const userRow = copyOnWrite(ponderRowInsert);
+                    return userRow;
+                  }
+                },
+              ),
+              // biome-ignore lint/suspicious/noThenProperty: <explanation>
+              then: (onFulfilled, onRejected) =>
+                storeMethodWrapper(async () => {
+                  common.metrics.ponder_indexing_store_queries_total.inc({
+                    table: getTableName(table),
+                    method: "insert",
+                  });
+                  checkOnchainTable(table, "insert");
+                  const ponderValues = copy(userValues);
+
+                  if (Array.isArray(ponderValues)) {
+                    const ponderRows = [];
+                    for (const value of ponderValues) {
+                      checkTableAccess(table, "insert", value, chainId);
+
+                      if (qb.$dialect === "pglite") {
+                        const row = await indexingCache.get({
+                          table,
+                          key: value,
+                        });
+
+                        if (row) {
+                          throw new UniqueConstraintError(
+                            `Primary key conflict in table '${getTableName(table)}'`,
+                          );
+                        } else {
+                          ponderRows.push(
+                            indexingCache.set({
+                              table,
+                              key: value,
+                              row: value,
+                              isUpdate: false,
+                            }),
+                          );
+                        }
+                      } else {
+                        // Note: optimistic assumption that no conflict exists
+                        // because error is recovered at flush time
+
+                        ponderRows.push(
+                          indexingCache.set({
+                            table,
+                            key: value,
+                            row: value,
+                            isUpdate: false,
+                          }),
+                        );
+                      }
+                    }
+                    const userRows = ponderRows.map((row) =>
+                      row === null ? row : copyOnWrite(row),
+                    );
+                    return Promise.resolve(userRows).then(
+                      onFulfilled,
+                      onRejected,
+                    );
+                  } else {
+                    checkTableAccess(table, "insert", ponderValues, chainId);
+
+                    let ponderRow: Row;
+                    if (qb.$dialect === "pglite") {
+                      const row = await indexingCache.get({
+                        table,
+                        key: ponderValues,
+                      });
+
+                      if (row) {
+                        throw new UniqueConstraintError(
+                          `Primary key conflict in table '${getTableName(table)}'`,
+                        );
+                      } else {
+                        ponderRow = indexingCache.set({
+                          table,
+                          key: ponderValues,
+                          row: ponderValues,
+                          isUpdate: false,
+                        });
+                      }
+                    } else {
+                      // Note: optimistic assumption that no conflict exists
+                      // because error is recovered at flush time
+
+                      ponderRow = indexingCache.set({
+                        table,
+                        key: ponderValues,
+                        row: ponderValues,
+                        isUpdate: false,
+                      });
+                    }
+                    const userRow = copyOnWrite(ponderRow);
+                    return Promise.resolve(userRow).then(
+                      onFulfilled,
+                      onRejected,
+                    );
+                  }
+                })().then(onFulfilled, onRejected),
+              catch: (onRejected) => inner.then(undefined, onRejected),
+              finally: (onFinally) =>
+                inner.then(
+                  (value: any) => {
+                    onFinally?.();
+                    return value;
+                  },
+                  (reason: any) => {
+                    onFinally?.();
+                    throw reason;
+                  },
+                ),
+              // @ts-ignore
+            } satisfies ReturnType<
+              ReturnType<IndexingStore["db"]["insert"]>["values"]
+            >;
+
+            return inner;
+          },
+        };
+      },
+      // @ts-ignore
+      update(table: Table, key) {
+        return {
+          set: storeMethodWrapper(async (userValues: any) => {
+            common.metrics.ponder_indexing_store_queries_total.inc({
+              table: getTableName(table),
+              method: "update",
+            });
+            checkOnchainTable(table, "update");
+            checkTableAccess(table, "update", key, chainId);
+
+            const ponderRowUpdate = await indexingCache.get({ table, key });
+
+            if (ponderRowUpdate === null) {
+              const error = new RecordNotFoundError(
+                `No existing record found in table '${getTableName(table)}'`,
+              );
+              error.meta.push(`db.update arguments:\n${prettyPrint(key)}`);
+              throw error;
+            }
+
+            if (typeof userValues === "function") {
+              const userRow = copyOnWrite(ponderRowUpdate);
+              const userSet = userValues(userRow);
+              const ponderSet = copy(userSet);
+              validateUpdateSet(
+                table,
+                ponderSet,
+                ponderRowUpdate,
+                primaryKeyCache,
+              );
+              for (const [key, value] of Object.entries(ponderSet)) {
+                if (value === undefined) continue;
+                ponderRowUpdate[key] = value;
+              }
+            } else {
+              const userSet = userValues;
+              const ponderSet = copy(userSet);
+              validateUpdateSet(
+                table,
+                ponderSet,
+                ponderRowUpdate,
+                primaryKeyCache,
+              );
+              for (const [key, value] of Object.entries(ponderSet)) {
+                if (value === undefined) continue;
+                ponderRowUpdate[key] = value;
+              }
+            }
+
+            const ponderRow = indexingCache.set({
+              table,
+              key: ponderRowUpdate,
+              row: ponderRowUpdate,
+              isUpdate: true,
+            });
+            const userRow = copyOnWrite(ponderRow);
+            return userRow;
+          }),
+        };
+      },
+      // @ts-ignore
+      delete: storeMethodWrapper(async (table: Table, key) => {
+        common.metrics.ponder_indexing_store_queries_total.inc({
+          table: getTableName(table),
+          method: "delete",
+        });
+        checkOnchainTable(table, "delete");
+        checkTableAccess(table, "delete", key, chainId);
+        return indexingCache.delete({ table, key });
+      }),
+      // @ts-ignore
+      sql: drizzle(
+        storeMethodWrapper(async (_sql, params, method, typings) => {
+          const isSelectOnly = await isReadonlySQLQuery(_sql);
+
+          if (isSelectOnly === false) {
+            await indexingCache.flush();
+            indexingCache.invalidate();
+            indexingCache.clear();
+          } else {
+            // Note: Not all nodes are implemented in the parser,
+            // so we need to try/catch to avoid throwing an error.
+            let relations: Set<string> | undefined;
+            try {
+              relations = await getSQLQueryRelations(_sql);
+            } catch {}
+
+            if (
+              Array.from(relations ?? []).some((refName) =>
+                views.some((view) => getViewName(view) === refName),
+              )
+            ) {
+              await indexingCache.flush();
+            } else {
+              await indexingCache.flush({ tableNames: relations });
+            }
+          }
+
+          const query: QueryWithTypings = { sql: _sql, params, typings };
+          const endClock = startClock();
+
+          try {
+            // Note: Use transaction so that user-land queries don't affect the
+            // in-progress transaction.
+            return await qb.transaction(async (tx) => {
+              const result = await tx.wrap((tx) =>
+                tx._.session
+                  .prepareQuery(query, undefined, undefined, method === "all")
+                  .execute(),
+              );
+
+              if (method === "all") {
+                return {
+                  // @ts-ignore
+                  ...result,
+                  // @ts-ignore
+                  rows: result.rows.map((row) => Object.values(row)),
+                };
+              }
+
+              return result;
+            });
+          } catch (error) {
+            if (error instanceof DbConnectionError) {
+              throw error;
+            }
+
+            throw new RawSqlError((error as Error).message);
+          } finally {
+            common.metrics.ponder_indexing_store_raw_sql_duration.observe(
+              endClock(),
+            );
+          }
+        }),
+        { schema, casing: "snake_case" },
+      ),
+    },
+    set qb(_qb: QB) {
+      qb = _qb;
+    },
+    set isProcessingEvents(_isProcessingEvents: boolean) {
+      isProcessingEvents = _isProcessingEvents;
+    },
+  };
+};