envio 2.32.2 → 3.0.0-alpha-main-clickhouse-sink

package/src/PgStorage.res

@@ -32,12 +32,13 @@ let makeCreateTableQuery = (table: Table.table, ~pgSchema, ~isNumericArrayAsText
       let fieldName = field->Table.getDbFieldName
 
       {
-        `"${fieldName}" ${switch fieldType {
-          | Custom(name) if !(name->Js.String2.startsWith("NUMERIC(")) => `"${pgSchema}".${name}`
-          // Workaround for Hasura bug https://github.com/enviodev/hyperindex/issues/788
-          | Numeric if isArray && isNumericArrayAsText => (Table.Text :> string)
-          | _ => (fieldType :> string)
-          }}${isArray ? "[]" : ""}${switch defaultValue {
+        `"${fieldName}" ${Table.getPgFieldType(
+            ~fieldType,
+            ~pgSchema,
+            ~isArray,
+            ~isNullable,
+            ~isNumericArrayAsText,
+          )}${switch defaultValue {
           | Some(defaultValue) => ` DEFAULT ${defaultValue}`
           | None => isNullable ? `` : ` NOT NULL`
           }}`
@@ -56,6 +57,73 @@ let makeCreateTableQuery = (table: Table.table, ~pgSchema, ~isNumericArrayAsText
       : ""});`
 }
 
+let getEntityHistory = (~entityConfig: Internal.entityConfig): EntityHistory.pgEntityHistory<
+  'entity,
+> => {
+  switch entityConfig.pgEntityHistoryCache {
+  | Some(cache) => cache
+  | None =>
+    let cache = {
+      let id = "id"
+
+      let dataFields = entityConfig.table.fields->Belt.Array.keepMap(field =>
+        switch field {
+        | Field(field) =>
+          switch field.fieldName {
+          //id is not nullable and should be part of the pk
+          | "id" => {...field, fieldName: id, isPrimaryKey: true}->Table.Field->Some
+          | _ =>
+            {
+              ...field,
+              isNullable: true, //All entity fields are nullable in the case
+              isIndex: false, //No need to index any additional entity data fields in entity history
+            }
+            ->Field
+            ->Some
+          }
+
+        | DerivedFrom(_) => None
+        }
+      )
+
+      let actionField = Table.mkField(
+        EntityHistory.changeFieldName,
+        EntityHistory.changeFieldType,
+        ~fieldSchema=S.never,
+      )
+
+      let checkpointIdField = Table.mkField(
+        EntityHistory.checkpointIdFieldName,
+        EntityHistory.checkpointIdFieldType,
+        ~fieldSchema=EntityHistory.unsafeCheckpointIdSchema,
+        ~isPrimaryKey=true,
+      )
+
+      let entityTableName = entityConfig.table.tableName
+      let historyTableName = EntityHistory.historyTableName(
+        ~entityName=entityTableName,
+        ~entityIndex=entityConfig.index,
+      )
+      //ignore composite indices
+      let table = Table.mkTable(
+        historyTableName,
+        ~fields=dataFields->Belt.Array.concat([checkpointIdField, actionField]),
+      )
+
+      let setChangeSchema = EntityHistory.makeSetUpdateSchema(entityConfig.schema)
+
+      {
+        EntityHistory.table,
+        setChangeSchema,
+        setChangeSchemaRows: S.array(setChangeSchema),
+      }
+    }
+
+    entityConfig.pgEntityHistoryCache = Some(cache)
+    cache
+  }
+}
+
 let makeInitializeTransaction = (
   ~pgSchema,
   ~pgUser,
@@ -74,10 +142,10 @@ let makeInitializeTransaction = (
 
   let allTables = generalTables->Array.copy
   let allEntityTables = []
-  entities->Js.Array2.forEach((entity: Internal.entityConfig) => {
-    allEntityTables->Js.Array2.push(entity.table)->ignore
-    allTables->Js.Array2.push(entity.table)->ignore
-    allTables->Js.Array2.push(entity.entityHistory.table)->ignore
+  entities->Js.Array2.forEach((entityConfig: Internal.entityConfig) => {
+    allEntityTables->Js.Array2.push(entityConfig.table)->ignore
+    allTables->Js.Array2.push(entityConfig.table)->ignore
+    allTables->Js.Array2.push(getEntityHistory(~entityConfig).table)->ignore
   })
   let derivedSchema = Schema.make(allEntityTables)
 
@@ -97,10 +165,10 @@ GRANT ALL ON SCHEMA "${pgSchema}" TO public;`,
   )
 
   // Optimized enum creation - direct when cleanRun, conditional otherwise
-  enums->Js.Array2.forEach((enumConfig: Internal.enumConfig<Internal.enum>) => {
+  enums->Js.Array2.forEach((enumConfig: Table.enumConfig<Table.enum>) => {
     // Create base enum creation query once
     let enumCreateQuery = `CREATE TYPE "${pgSchema}".${enumConfig.name} AS ENUM(${enumConfig.variants
-      ->Js.Array2.map(v => `'${v->(Utils.magic: Internal.enum => string)}'`)
+      ->Js.Array2.map(v => `'${v->(Utils.magic: Table.enum => string)}'`)
       ->Js.Array2.joinWith(", ")});`
 
     query := query.contents ++ "\n" ++ enumCreateQuery
@@ -176,6 +244,14 @@ let makeLoadByIdsQuery = (~pgSchema, ~tableName) => {
   `SELECT * FROM "${pgSchema}"."${tableName}" WHERE id = ANY($1::text[]);`
 }
 
+let makeDeleteByIdQuery = (~pgSchema, ~tableName) => {
+  `DELETE FROM "${pgSchema}"."${tableName}" WHERE id = $1;`
+}
+
+let makeDeleteByIdsQuery = (~pgSchema, ~tableName) => {
+  `DELETE FROM "${pgSchema}"."${tableName}" WHERE id = ANY($1::text[]);`
+}
+
 let makeLoadAllQuery = (~pgSchema, ~tableName) => {
   `SELECT * FROM "${pgSchema}"."${tableName}";`
 }
@@ -505,6 +581,463 @@ let getConnectedPsqlExec = {
   }
 }
 
+let deleteByIdsOrThrow = async (sql, ~pgSchema, ~ids, ~table: Table.table) => {
+  switch await (
+    switch ids {
+    | [_] =>
+      sql->Postgres.preparedUnsafe(
+        makeDeleteByIdQuery(~pgSchema, ~tableName=table.tableName),
+        ids->Obj.magic,
+      )
+    | _ =>
+      sql->Postgres.preparedUnsafe(
+        makeDeleteByIdsQuery(~pgSchema, ~tableName=table.tableName),
+        [ids]->Obj.magic,
+      )
+    }
+  ) {
+  | exception exn =>
+    raise(
+      Persistence.StorageError({
+        message: `Failed deleting "${table.tableName}" from storage by ids`,
+        reason: exn,
+      }),
+    )
+  | _ => ()
+  }
+}
+
+let makeInsertDeleteUpdatesQuery = (~entityConfig: Internal.entityConfig, ~pgSchema) => {
+  let historyTableName = EntityHistory.historyTableName(
+    ~entityName=entityConfig.name,
+    ~entityIndex=entityConfig.index,
+  )
+
+  // Get all field names for the INSERT statement
+  let allHistoryFieldNames = entityConfig.table.fields->Belt.Array.keepMap(fieldOrDerived =>
+    switch fieldOrDerived {
+    | Field(field) => field->Table.getDbFieldName->Some
+    | DerivedFrom(_) => None
+    }
+  )
+  allHistoryFieldNames->Js.Array2.push(EntityHistory.checkpointIdFieldName)->ignore
+  allHistoryFieldNames->Js.Array2.push(EntityHistory.changeFieldName)->ignore
+
+  let allHistoryFieldNamesStr =
+    allHistoryFieldNames->Belt.Array.map(name => `"${name}"`)->Js.Array2.joinWith(", ")
+
+  // Build the SELECT part: id from unnest, envio_checkpoint_id from unnest, 'DELETE' for action, NULL for all other fields
+  let selectParts = allHistoryFieldNames->Belt.Array.map(fieldName => {
+    switch fieldName {
+    | field if field == Table.idFieldName => `u.${Table.idFieldName}`
+    | field if field == EntityHistory.checkpointIdFieldName =>
+      `u.${EntityHistory.checkpointIdFieldName}`
+    | field if field == EntityHistory.changeFieldName =>
+      `'${(EntityHistory.RowAction.DELETE :> string)}'`
+    | _ => "NULL"
+    }
+  })
+  let selectPartsStr = selectParts->Js.Array2.joinWith(", ")
+
+  // Get the PostgreSQL type for the checkpoint ID field
+  let checkpointIdPgType = Table.getPgFieldType(
+    ~fieldType=EntityHistory.checkpointIdFieldType,
+    ~pgSchema,
+    ~isArray=false,
+    ~isNumericArrayAsText=false,
+    ~isNullable=false,
+  )
+
+  `INSERT INTO "${pgSchema}"."${historyTableName}" (${allHistoryFieldNamesStr})
+SELECT ${selectPartsStr}
+FROM UNNEST($1::text[], $2::${checkpointIdPgType}[]) AS u(${Table.idFieldName}, ${EntityHistory.checkpointIdFieldName})`
+}
+
+let executeSet = (
+  sql: Postgres.sql,
+  ~items: array<'a>,
+  ~dbFunction: (Postgres.sql, array<'a>) => promise<unit>,
+) => {
+  if items->Array.length > 0 {
+    sql->dbFunction(items)
+  } else {
+    Promise.resolve()
+  }
+}
+
+let rec writeBatch = async (
+  sql,
+  ~batch: Batch.t,
+  ~rawEvents,
+  ~pgSchema,
+  ~rollbackTargetCheckpointId,
+  ~isInReorgThreshold,
+  ~config: Config.t,
+  ~allEntities: array<Internal.entityConfig>,
+  ~setEffectCacheOrThrow,
+  ~updatedEffectsCache,
+  ~updatedEntities: array<Persistence.updatedEntity>,
+  ~sinkPromise: option<promise<option<exn>>>,
+  ~escapeTables=?,
+) => {
+  try {
+    let shouldSaveHistory = config->Config.shouldSaveHistory(~isInReorgThreshold)
+
+    let specificError = ref(None)
+
+    let setRawEvents = executeSet(
+      _,
+      ~dbFunction=(sql, items) => {
+        sql->setOrThrow(
+          ~items,
+          ~table=InternalTable.RawEvents.table,
+          ~itemSchema=InternalTable.RawEvents.schema,
+          ~pgSchema,
+        )
+      },
+      ~items=rawEvents,
+    )
+
+    let setEntities = updatedEntities->Belt.Array.map(({entityConfig, updates}) => {
+      let entitiesToSet = []
+      let idsToDelete = []
+
+      updates->Js.Array2.forEach(row => {
+        switch row {
+        | {latestChange: Set({entity})} => entitiesToSet->Belt.Array.push(entity)
+        | {latestChange: Delete({entityId})} => idsToDelete->Belt.Array.push(entityId)
+        }
+      })
+
+      let shouldRemoveInvalidUtf8 = switch escapeTables {
+      | Some(tables) if tables->Utils.Set.has(entityConfig.table) => true
+      | _ => false
+      }
+
+      async sql => {
+        try {
+          let promises = []
+
+          if shouldSaveHistory {
+            let backfillHistoryIds = Utils.Set.make()
+            let batchSetUpdates = []
+            // Use unnest approach
+            let batchDeleteCheckpointIds = []
+            let batchDeleteEntityIds = []
+
+            updates->Js.Array2.forEach(update => {
+              switch update {
+              | {history, containsRollbackDiffChange} =>
+                history->Js.Array2.forEach(
+                  (change: Change.t<'a>) => {
+                    if !containsRollbackDiffChange {
+                      // For every update we want to make sure that there's an existing history item
+                      // with the current entity state. So we backfill history with checkpoint id 0,
+                      // before writing updates. Don't do this if the update has a rollback diff change.
+                      backfillHistoryIds->Utils.Set.add(change->Change.getEntityId)->ignore
+                    }
+                    switch change {
+                    | Delete({entityId}) => {
+                        batchDeleteEntityIds->Belt.Array.push(entityId)->ignore
+                        batchDeleteCheckpointIds
+                        ->Belt.Array.push(change->Change.getCheckpointId)
+                        ->ignore
+                      }
+                    | Set(_) => batchSetUpdates->Js.Array2.push(change)->ignore
+                    }
+                  },
+                )
+              }
+            })
+
+            if backfillHistoryIds->Utils.Set.size !== 0 {
+              // This must run before updating entity or entity history tables
+              await EntityHistory.backfillHistory(
+                sql,
+                ~pgSchema,
+                ~entityName=entityConfig.name,
+                ~entityIndex=entityConfig.index,
+                ~ids=backfillHistoryIds->Utils.Set.toArray,
+              )
+            }
+
+            if batchDeleteCheckpointIds->Utils.Array.notEmpty {
+              promises->Belt.Array.push(
+                sql
+                ->Postgres.preparedUnsafe(
+                  makeInsertDeleteUpdatesQuery(~entityConfig, ~pgSchema),
+                  (batchDeleteEntityIds, batchDeleteCheckpointIds)->Obj.magic,
+                )
+                ->Promise.ignoreValue,
+              )
+            }
+
+            if batchSetUpdates->Utils.Array.notEmpty {
+              if shouldRemoveInvalidUtf8 {
+                let entities = batchSetUpdates->Js.Array2.map(batchSetUpdate => {
+                  switch batchSetUpdate {
+                  | Set({entity}) => entity
+                  | _ => Js.Exn.raiseError("Expected Set action")
+                  }
+                })
+                entities->removeInvalidUtf8InPlace
+              }
+
+              let entityHistory = getEntityHistory(~entityConfig)
+
+              promises
+              ->Js.Array2.push(
+                sql->setOrThrow(
+                  ~items=batchSetUpdates,
+                  ~itemSchema=entityHistory.setChangeSchema,
+                  ~table=entityHistory.table,
+                  ~pgSchema,
+                ),
+              )
+              ->ignore
+            }
+          }
+
+          if entitiesToSet->Utils.Array.notEmpty {
+            if shouldRemoveInvalidUtf8 {
+              entitiesToSet->removeInvalidUtf8InPlace
+            }
+            promises->Belt.Array.push(
+              sql->setOrThrow(
+                ~items=entitiesToSet,
+                ~table=entityConfig.table,
+                ~itemSchema=entityConfig.schema,
+                ~pgSchema,
+              ),
+            )
+          }
+          if idsToDelete->Utils.Array.notEmpty {
+            promises->Belt.Array.push(
+              sql->deleteByIdsOrThrow(~pgSchema, ~ids=idsToDelete, ~table=entityConfig.table),
+            )
+          }
+
+          let _ = await promises->Promise.all
+        } catch {
+        // There's a race condition that sql->Postgres.beginSql
+        // might throw PG error, earlier, than the handled error
+        // from setOrThrow will be passed through.
+        // This is needed for the utf8 encoding fix.
+        | exn => {
+            /* Note: Entity History doesn't return StorageError yet, and directly throws JsError */
+            let normalizedExn = switch exn {
+            | JsError(_) => exn
+            | Persistence.StorageError({reason: exn}) => exn
+            | _ => exn
+            }->Js.Exn.anyToExnInternal
+
+            switch normalizedExn {
+            | JsError(error) =>
+              // Workaround for https://github.com/enviodev/hyperindex/issues/446
+              // We do escaping only when we actually got an error writing for the first time.
+              // This is not perfect, but an optimization to avoid escaping for every single item.
+
+              switch error->S.parseOrThrow(pgErrorMessageSchema) {
+              | `current transaction is aborted, commands ignored until end of transaction block` => ()
+              | `invalid byte sequence for encoding "UTF8": 0x00` =>
+                // Since the transaction is aborted at this point,
+                // we can't simply retry the function with escaped items,
+                // so propagate the error, to restart the whole batch write.
+                // Also, pass the failing table, to escape only its items.
+                // TODO: Ideally all this should be done in the file,
+                // so it'll be easier to work on PG specific logic.
+                specificError.contents = Some(PgEncodingError({table: entityConfig.table}))
+              | _ => specificError.contents = Some(exn->Utils.prettifyExn)
+              | exception _ => ()
+              }
+            | S.Raised(_) => raise(normalizedExn) // But rethrow this one, since it's not a PG error
+            | _ => ()
+            }
+
+            // Improtant: Don't rethrow here, since it'll result in
+            // an unhandled rejected promise error.
+            // That's fine not to throw, since sql->Postgres.beginSql
+            // will fail anyways.
+          }
+        }
+      }
+    })
+
+    //In the event of a rollback, rollback all meta tables based on the given
+    //valid event identifier, where all rows created after this eventIdentifier should
+    //be deleted
+    let rollbackTables = switch rollbackTargetCheckpointId {
+    | Some(rollbackTargetCheckpointId) =>
+      Some(
+        sql => {
+          let promises = allEntities->Js.Array2.map(entityConfig => {
+            sql->EntityHistory.rollback(
+              ~pgSchema,
+              ~entityName=entityConfig.name,
+              ~entityIndex=entityConfig.index,
+              ~rollbackTargetCheckpointId,
+            )
+          })
+          promises
+          ->Js.Array2.push(
+            sql->InternalTable.Checkpoints.rollback(~pgSchema, ~rollbackTargetCheckpointId),
+          )
+          ->ignore
+          Promise.all(promises)
+        },
+      )
+    | None => None
+    }
+
+    try {
+      let _ = await Promise.all2((
+        sql->Postgres.beginSql(async sql => {
+          //Rollback tables need to happen first in the traction
+          switch rollbackTables {
+          | Some(rollbackTables) =>
+            let _ = await rollbackTables(sql)
+          | None => ()
+          }
+
+          let setOperations = [
+            sql =>
+              sql->InternalTable.Chains.setProgressedChains(
+                ~pgSchema,
+                ~progressedChains=batch.progressedChainsById->Utils.Dict.mapValuesToArray((
+                  chainAfterBatch
+                ): InternalTable.Chains.progressedChain => {
+                  chainId: chainAfterBatch.fetchState.chainId,
+                  progressBlockNumber: chainAfterBatch.progressBlockNumber,
+                  totalEventsProcessed: chainAfterBatch.totalEventsProcessed,
+                }),
+              ),
+            setRawEvents,
+          ]->Belt.Array.concat(setEntities)
+
+          if shouldSaveHistory {
+            setOperations->Belt.Array.push(sql =>
+              sql->InternalTable.Checkpoints.insert(
+                ~pgSchema,
+                ~checkpointIds=batch.checkpointIds,
+                ~checkpointChainIds=batch.checkpointChainIds,
+                ~checkpointBlockNumbers=batch.checkpointBlockNumbers,
+                ~checkpointBlockHashes=batch.checkpointBlockHashes,
+                ~checkpointEventsProcessed=batch.checkpointEventsProcessed,
+              )
+            )
+          }
+
+          await setOperations
+          ->Belt.Array.map(dbFunc => sql->dbFunc)
+          ->Promise.all
+          ->Promise.ignoreValue
+
+          switch sinkPromise {
+          | Some(sinkPromise) =>
+            switch await sinkPromise {
+            | Some(exn) => raise(exn)
+            | None => ()
+            }
+          | None => ()
+          }
+        }),
+        // Since effect cache currently doesn't support rollback,
+        // we can run it outside of the transaction for simplicity.
+        updatedEffectsCache
+        ->Belt.Array.map(({effect, items, shouldInitialize}: Persistence.updatedEffectCache) => {
+          setEffectCacheOrThrow(~effect, ~items, ~initialize=shouldInitialize)
+        })
+        ->Promise.all,
+      ))
+
+      // Just in case, if there's a not PG-specific error.
+      switch specificError.contents {
+      | Some(specificError) => raise(specificError)
+      | None => ()
+      }
+    } catch {
+    | exn =>
+      raise(
+        switch specificError.contents {
+        | Some(specificError) => specificError
+        | None => exn
+        },
+      )
+    }
+  } catch {
+  | PgEncodingError({table}) =>
+    let escapeTables = switch escapeTables {
+    | Some(set) => set
+    | None => Utils.Set.make()
+    }
+    let _ = escapeTables->Utils.Set.add(table)
+    // Retry with specifying which tables to escape.
+    await writeBatch(
+      sql,
+      ~escapeTables,
+      ~rawEvents,
+      ~batch,
+      ~pgSchema,
+      ~rollbackTargetCheckpointId,
+      ~isInReorgThreshold,
+      ~config,
+      ~setEffectCacheOrThrow,
+      ~updatedEffectsCache,
+      ~allEntities,
+      ~updatedEntities,
+      ~sinkPromise,
+    )
+  }
+}
+
+// Returns the most recent entity state for IDs that need to be restored during rollback.
+// For each ID modified after the rollback target, retrieves its latest state at or before the target.
+let makeGetRollbackRestoredEntitiesQuery = (~entityConfig: Internal.entityConfig, ~pgSchema) => {
+  let dataFieldNames = entityConfig.table.fields->Belt.Array.keepMap(fieldOrDerived =>
+    switch fieldOrDerived {
+    | Field(field) => field->Table.getDbFieldName->Some
+    | DerivedFrom(_) => None
+    }
+  )
+
+  let dataFieldsCommaSeparated =
+    dataFieldNames->Belt.Array.map(name => `"${name}"`)->Js.Array2.joinWith(", ")
+
+  let historyTableName = EntityHistory.historyTableName(
+    ~entityName=entityConfig.name,
+    ~entityIndex=entityConfig.index,
+  )
+
+  `SELECT DISTINCT ON (${Table.idFieldName}) ${dataFieldsCommaSeparated}
+FROM "${pgSchema}"."${historyTableName}"
+WHERE "${EntityHistory.checkpointIdFieldName}" <= $1
+  AND EXISTS (
+    SELECT 1
+    FROM "${pgSchema}"."${historyTableName}" h
+    WHERE h.${Table.idFieldName} = "${historyTableName}".${Table.idFieldName}
+      AND h."${EntityHistory.checkpointIdFieldName}" > $1
+  )
+ORDER BY ${Table.idFieldName}, "${EntityHistory.checkpointIdFieldName}" DESC`
+}
+
+// Returns entity IDs that were created after the rollback target and have no history before it.
+// These entities should be deleted during rollback.
+let makeGetRollbackRemovedIdsQuery = (~entityConfig: Internal.entityConfig, ~pgSchema) => {
+  let historyTableName = EntityHistory.historyTableName(
+    ~entityName=entityConfig.name,
+    ~entityIndex=entityConfig.index,
+  )
+  `SELECT DISTINCT ${Table.idFieldName}
+FROM "${pgSchema}"."${historyTableName}"
+WHERE "${EntityHistory.checkpointIdFieldName}" > $1
+AND NOT EXISTS (
+  SELECT 1
+  FROM "${pgSchema}"."${historyTableName}" h
+  WHERE h.${Table.idFieldName} = "${historyTableName}".${Table.idFieldName}
+    AND h."${EntityHistory.checkpointIdFieldName}" <= $1
+)`
+}
+
 let make = (
   ~sql: Postgres.sql,
   ~pgHost,
@@ -514,6 +1047,7 @@ let make = (
   ~pgDatabase,
   ~pgPassword,
   ~isHasuraEnabled,
+  ~sink: option<Sink.t>=?,
   ~onInitialize=?,
   ~onNewTables=?,
 ): Persistence.storage => {
@@ -645,6 +1179,12 @@ let make = (
       )
     }
 
+    // Call sink.initialize before executing PG queries
+    switch sink {
+    | Some(sink) => await sink.initialize(~chainConfigs, ~entities, ~enums)
+    | None => ()
+    }
+
     let queries = makeInitializeTransaction(
       ~pgSchema,
       ~pgUser,
@@ -889,19 +1429,156 @@ let make = (
       }),
       sql
       ->Postgres.unsafe(InternalTable.Checkpoints.makeCommitedCheckpointIdQuery(~pgSchema))
-      ->(Utils.magic: promise<array<unknown>> => promise<array<{"id": int}>>),
+      ->(Utils.magic: promise<array<unknown>> => promise<array<{"id": float}>>),
       sql
       ->Postgres.unsafe(InternalTable.Checkpoints.makeGetReorgCheckpointsQuery(~pgSchema))
       ->(Utils.magic: promise<array<unknown>> => promise<array<Internal.reorgCheckpoint>>),
     ))
 
+    let checkpointId = (checkpointIdResult->Belt.Array.getUnsafe(0))["id"]
+
+    // Resume sink if present - needed to rollback any reorg changes
+    switch sink {
+    | Some(sink) => await sink.resume(~checkpointId)
+    | None => ()
+    }
+
     {
       cleanRun: false,
       reorgCheckpoints,
       cache,
       chains,
-      checkpointId: (checkpointIdResult->Belt.Array.getUnsafe(0))["id"],
+      checkpointId,
+    }
+  }
+
+  let executeUnsafe = query => sql->Postgres.unsafe(query)
+
+  let hasEntityHistoryRows = async () => {
+    // Query for all entity history tables (they have the prefix "envio_history_")
+    let historyTables = await sql->Postgres.unsafe(
+      `SELECT table_name FROM information_schema.tables 
+       WHERE table_schema = '${pgSchema}' 
+       AND table_name LIKE 'envio_history_%';`,
+    )
+
+    if historyTables->Utils.Array.isEmpty {
+      false
+    } else {
+      // Check if any of these tables have rows
+      let checks =
+        await historyTables
+        ->Belt.Array.map(async (table: {"table_name": string}) => {
+          try {
+            let query = `SELECT EXISTS(SELECT 1 FROM "${pgSchema}"."${table["table_name"]}" LIMIT 1);`
+            let result: array<{"exists": bool}> = (await sql->Postgres.unsafe(query))->Utils.magic
+            switch result {
+            | [row] => row["exists"]
+            | _ => false
+            }
+          } catch {
+          | _ => false
+          }
+        })
+        ->Promise.all
+      checks->Belt.Array.some(v => v)
+    }
+  }
+
+  let setChainMeta = chainsData =>
+    InternalTable.Chains.setMeta(sql, ~pgSchema, ~chainsData)->Promise.thenResolve(_ =>
+      %raw(`undefined`)
+    )
+
+  let pruneStaleCheckpoints = (~safeCheckpointId) =>
+    InternalTable.Checkpoints.pruneStaleCheckpoints(sql, ~pgSchema, ~safeCheckpointId)
+
+  let pruneStaleEntityHistory = (~entityName, ~entityIndex, ~safeCheckpointId) =>
+    EntityHistory.pruneStaleEntityHistory(
+      sql,
+      ~pgSchema,
+      ~entityName,
+      ~entityIndex,
+      ~safeCheckpointId,
+    )
+
+  let getRollbackTargetCheckpoint = (~reorgChainId, ~lastKnownValidBlockNumber) =>
+    InternalTable.Checkpoints.getRollbackTargetCheckpoint(
+      sql,
+      ~pgSchema,
+      ~reorgChainId,
+      ~lastKnownValidBlockNumber,
+    )
+
+  let getRollbackProgressDiff = (~rollbackTargetCheckpointId) =>
+    InternalTable.Checkpoints.getRollbackProgressDiff(sql, ~pgSchema, ~rollbackTargetCheckpointId)
+
+  let getRollbackData = async (
+    ~entityConfig: Internal.entityConfig,
+    ~rollbackTargetCheckpointId,
+  ) => {
+    await Promise.all2((
+      // Get IDs of entities that should be deleted (created after rollback target with no prior history)
+      sql
+      ->Postgres.preparedUnsafe(
+        makeGetRollbackRemovedIdsQuery(~entityConfig, ~pgSchema),
+        [rollbackTargetCheckpointId]->Utils.magic,
+      )
+      ->(Utils.magic: promise<unknown> => promise<array<{"id": string}>>),
+      // Get entities that should be restored to their state at or before rollback target
+      sql
+      ->Postgres.preparedUnsafe(
+        makeGetRollbackRestoredEntitiesQuery(~entityConfig, ~pgSchema),
+        [rollbackTargetCheckpointId]->Utils.magic,
+      )
+      ->(Utils.magic: promise<unknown> => promise<array<unknown>>),
+    ))
+  }
+
+  let writeBatchMethod = async (
+    ~batch,
+    ~rawEvents,
+    ~rollbackTargetCheckpointId,
+    ~isInReorgThreshold,
+    ~config,
+    ~allEntities,
+    ~updatedEffectsCache,
+    ~updatedEntities,
+  ) => {
+    // Initialize sink if configured
+    let sinkPromise = switch sink {
+    | Some(sink) => {
+        let timerRef = Hrtime.makeTimer()
+        Some(
+          sink.writeBatch(~batch, ~updatedEntities)
+          ->Promise.thenResolve(_ => {
+            Prometheus.SinkWrite.increment(
+              ~sinkName=sink.name,
+              ~timeMillis=timerRef->Hrtime.timeSince->Hrtime.toMillis->Hrtime.intFromMillis,
+            )
+            None
+          })
+          // Otherwise it fails with unhandled exception
+          ->Promise.catchResolve(exn => Some(exn)),
+        )
+      }
+    | None => None
     }
+
+    await writeBatch(
+      sql,
+      ~batch,
+      ~rawEvents,
+      ~pgSchema,
+      ~rollbackTargetCheckpointId,
+      ~isInReorgThreshold,
+      ~config,
+      ~allEntities,
+      ~setEffectCacheOrThrow,
+      ~updatedEffectsCache,
+      ~updatedEntities,
+      ~sinkPromise,
+    )
   }
 
   {
@@ -913,5 +1590,14 @@ let make = (
     setOrThrow,
     setEffectCacheOrThrow,
     dumpEffectCache,
+    executeUnsafe,
+    hasEntityHistoryRows,
+    setChainMeta,
+    pruneStaleCheckpoints,
+    pruneStaleEntityHistory,
+    getRollbackTargetCheckpoint,
+    getRollbackProgressDiff,
+    getRollbackData,
+    writeBatch: writeBatchMethod,
   }
 }