envio 3.0.2 → 3.1.0-rc.1

package/src/Hasura.res

@@ -110,7 +110,13 @@ let clearHasuraMetadata = async (~endpoint, ~auth) => {
   }
 }
 
-let trackTables = async (~endpoint, ~auth, ~pgSchema, ~tableNames: array<string>) => {
+type trackTableConfig = {
+  tableName: string,
+  description: option<string>,
+  columnDescriptions: dict<string>,
+}
+
+let trackTables = async (~endpoint, ~auth, ~pgSchema, ~tableConfigs: array<trackTableConfig>) => {
   try {
     let result = await trackTablesRoute->Rest.fetch(
       {
@@ -118,17 +124,33 @@ let trackTables = async (~endpoint, ~auth, ~pgSchema, ~tableNames: array<string>
         "args": {
           // If set to false, any warnings will cause the API call to fail and no new tables to be tracked. Otherwise tables that fail to track will be raised as warnings. (default: true)
           "allow_warnings": false,
-          "tables": tableNames->Array.map(tableName =>
+          "tables": tableConfigs->Array.map(({tableName, description, columnDescriptions}) => {
+            let configuration = dict{
+              "custom_name": tableName->(Utils.magic: string => JSON.t),
+            }
+            switch description {
+            | Some(d) => configuration->Dict.set("comment", d->(Utils.magic: string => JSON.t))
+            | None => ()
+            }
+            let columnConfigEntries = columnDescriptions->Dict.toArray
+            if columnConfigEntries->Array.length > 0 {
+              let columnConfig = dict{}
+              columnConfigEntries->Array.forEach(((column, comment)) =>
+                columnConfig->Dict.set(column, {"comment": comment}->(Utils.magic: {..} => JSON.t))
+              )
+              configuration->Dict.set(
+                "column_config",
+                columnConfig->(Utils.magic: dict<JSON.t> => JSON.t),
+              )
+            }
             {
               "table": {
                 "name": tableName,
                 "schema": pgSchema,
               },
-              "configuration": {
-                "custom_name": tableName,
-              },
+              "configuration": configuration,
             }
-          ),
+          }),
         }->(Utils.magic: 'a => JSON.t),
       },
       ~client=Rest.client(endpoint),
@@ -139,13 +161,13 @@ let trackTables = async (~endpoint, ~auth, ~pgSchema, ~tableNames: array<string>
     }
     Logging.trace({
       "msg": msg,
-      "tableNames": tableNames,
+      "tableNames": tableConfigs->Array.map(c => c.tableName),
     })
   } catch {
   | exn =>
     Logging.error({
       "msg": `There was an issue tracking tables in hasura - indexing may still work - but you may have issues querying the data in hasura.`,
-      "tableNames": tableNames,
+      "tableNames": tableConfigs->Array.map(c => c.tableName),
       "err": exn->Utils.prettifyExn,
     })
   }
@@ -192,31 +214,41 @@ let createEntityRelationship = async (
   ~objectName: string,
   ~mappedEntity: string,
   ~isDerivedFrom: bool,
+  ~comment: option<string>=?,
 ) => {
   let derivedFromTo = isDerivedFrom ? `"id": "${relationalKey}"` : `"${relationalKey}_id" : "id"`
 
+  let tableJson = {
+    "schema": pgSchema,
+    "name": tableName,
+  }->(Utils.magic: {..} => JSON.t)
+  let usingJson = {
+    "manual_configuration": {
+      "remote_table": {
+        "schema": pgSchema,
+        "name": mappedEntity,
+      },
+      "column_mapping": JSON.parseOrThrow(`{${derivedFromTo}}`),
+    },
+  }->(Utils.magic: {..} => JSON.t)
+
+  let args = dict{
+    "table": tableJson,
+    "name": objectName->(Utils.magic: string => JSON.t),
+    "source": "default"->(Utils.magic: string => JSON.t),
+    "using": usingJson,
+  }
+  switch comment {
+  | Some(c) => args->Dict.set("comment", c->(Utils.magic: string => JSON.t))
+  | None => ()
+  }
+
   await sendOperation(
     ~endpoint,
     ~auth,
     ~operation={
       "type": `pg_create_${relationshipType}_relationship`,
-      "args": {
-        "table": {
-          "schema": pgSchema,
-          "name": tableName,
-        },
-        "name": objectName,
-        "source": "default",
-        "using": {
-          "manual_configuration": {
-            "remote_table": {
-              "schema": pgSchema,
-              "name": mappedEntity,
-            },
-            "column_mapping": JSON.parseOrThrow(`{${derivedFromTo}}`),
-          },
-        },
-      },
+      "args": args,
     }->(Utils.magic: 'a => JSON.t),
   )
 }
@@ -230,19 +262,49 @@ let trackDatabase = async (
   ~responseLimit,
   ~schema,
 ) => {
-  let exposedInternalTableNames = [
-    InternalTable.RawEvents.table.tableName,
-    InternalTable.Views.metaViewName,
-    InternalTable.Views.chainMetadataViewName,
+  let exposedInternalTableConfigs = [
+    {
+      tableName: InternalTable.RawEvents.table.tableName,
+      description: None,
+      columnDescriptions: dict{},
+    },
+    {
+      tableName: InternalTable.Views.metaViewName,
+      description: None,
+      columnDescriptions: dict{},
+    },
+    {
+      tableName: InternalTable.Views.chainMetadataViewName,
+      description: None,
+      columnDescriptions: dict{},
+    },
   ]
-  let userTableNames = userEntities->Array.map(entity => entity.table.tableName)
-  let tableNames = [exposedInternalTableNames, userTableNames]->Belt.Array.concatMany
+  let userTableConfigs = userEntities->Array.map(entity => {
+    let columnDescriptions = dict{}
+    entity.table.fields->Array.forEach(fieldOrDerived =>
+      switch fieldOrDerived {
+      | Table.Field(field) =>
+        switch field.description {
+        | Some(d) => columnDescriptions->Dict.set(field->Table.getDbFieldName, d)
+        | None => ()
+        }
+      | Table.DerivedFrom(_) => ()
+      }
+    )
+    {
+      tableName: entity.table.tableName,
+      description: entity.table.description,
+      columnDescriptions,
+    }
+  })
+  let tableConfigs = [exposedInternalTableConfigs, userTableConfigs]->Belt.Array.concatMany
+  let tableNames = tableConfigs->Array.map(c => c.tableName)
 
   Logging.info("Tracking tables in Hasura")
 
   let _ = await clearHasuraMetadata(~endpoint, ~auth)
 
-  await trackTables(~endpoint, ~auth, ~pgSchema, ~tableNames)
+  await trackTables(~endpoint, ~auth, ~pgSchema, ~tableConfigs)
 
   for i in 0 to tableNames->Array.length - 1 {
     let tableName = tableNames->Array.getUnsafe(i)
@@ -278,6 +340,7 @@ let trackDatabase = async (
         ~objectName=derivedFromField.fieldName,
         ~relationalKey=relationalFieldName,
         ~mappedEntity=derivedFromField.derivedFromEntity,
+        ~comment=?derivedFromField.description,
       )
     }
 
@@ -295,6 +358,7 @@ let trackDatabase = async (
         ~objectName=field.fieldName,
         ~relationalKey=field.fieldName,
         ~mappedEntity=linkedEntityName,
+        ~comment=?field.description,
       )
     }
   }

package/src/Hasura.res.mjs

@@ -112,34 +112,52 @@ async function clearHasuraMetadata(endpoint, auth) {
   }
 }
 
-async function trackTables(endpoint, auth, pgSchema, tableNames) {
+async function trackTables(endpoint, auth, pgSchema, tableConfigs) {
   try {
     let result = await Rest.fetch(trackTablesRoute, {
       auth: auth,
       args: {
         allow_warnings: false,
-        tables: tableNames.map(tableName => ({
-          table: {
-            name: tableName,
-            schema: pgSchema
-          },
-          configuration: {
+        tables: tableConfigs.map(param => {
+          let description = param.description;
+          let tableName = param.tableName;
+          let configuration = {
             custom_name: tableName
+          };
+          if (description !== undefined) {
+            configuration["comment"] = description;
           }
-        }))
+          let columnConfigEntries = Object.entries(param.columnDescriptions);
+          if (columnConfigEntries.length !== 0) {
+            let columnConfig = {};
+            columnConfigEntries.forEach(param => {
+              columnConfig[param[0]] = {
+                comment: param[1]
+              };
+            });
+            configuration["column_config"] = columnConfig;
+          }
+          return {
+            table: {
+              name: tableName,
+              schema: pgSchema
+            },
+            configuration: configuration
+          };
+        })
       }
     }, Rest.client(endpoint, undefined));
     let msg;
     msg = result === "QuerySucceeded" ? "Hasura finished tracking tables" : "Hasura tables already tracked";
     return Logging.trace({
       msg: msg,
-      tableNames: tableNames
+      tableNames: tableConfigs.map(c => c.tableName)
     });
   } catch (raw_exn) {
     let exn = Primitive_exceptions.internalToException(raw_exn);
     return Logging.error({
       msg: `There was an issue tracking tables in hasura - indexing may still work - but you may have issues querying the data in hasura.`,
-      tableNames: tableNames,
+      tableNames: tableConfigs.map(c => c.tableName),
       err: Utils.prettifyExn(exn)
     });
   }
@@ -165,44 +183,81 @@ async function createSelectPermission(endpoint, auth, tableName, pgSchema, respo
   });
 }
 
-async function createEntityRelationship(endpoint, auth, pgSchema, tableName, relationshipType, relationalKey, objectName, mappedEntity, isDerivedFrom) {
+async function createEntityRelationship(endpoint, auth, pgSchema, tableName, relationshipType, relationalKey, objectName, mappedEntity, isDerivedFrom, comment) {
   let derivedFromTo = isDerivedFrom ? `"id": "` + relationalKey + `"` : `"` + relationalKey + `_id" : "id"`;
-  return await sendOperation(endpoint, auth, {
-    type: `pg_create_` + relationshipType + `_relationship`,
-    args: {
-      table: {
+  let tableJson = {
+    schema: pgSchema,
+    name: tableName
+  };
+  let usingJson = {
+    manual_configuration: {
+      remote_table: {
         schema: pgSchema,
-        name: tableName
+        name: mappedEntity
       },
-      name: objectName,
-      source: "default",
-      using: {
-        manual_configuration: {
-          remote_table: {
-            schema: pgSchema,
-            name: mappedEntity
-          },
-          column_mapping: JSON.parse(`{` + derivedFromTo + `}`)
-        }
-      }
+      column_mapping: JSON.parse(`{` + derivedFromTo + `}`)
     }
+  };
+  let args = {
+    table: tableJson,
+    name: objectName,
+    source: "default",
+    using: usingJson
+  };
+  if (comment !== undefined) {
+    args["comment"] = comment;
+  }
+  return await sendOperation(endpoint, auth, {
+    type: `pg_create_` + relationshipType + `_relationship`,
+    args: args
   });
 }
 
 async function trackDatabase(endpoint, auth, pgSchema, userEntities, aggregateEntities, responseLimit, schema) {
-  let exposedInternalTableNames = [
-    InternalTable.RawEvents.table.tableName,
-    InternalTable.Views.metaViewName,
-    InternalTable.Views.chainMetadataViewName
+  let exposedInternalTableConfigs = [
+    {
+      tableName: InternalTable.RawEvents.table.tableName,
+      description: undefined,
+      columnDescriptions: {}
+    },
+    {
+      tableName: InternalTable.Views.metaViewName,
+      description: undefined,
+      columnDescriptions: {}
+    },
+    {
+      tableName: InternalTable.Views.chainMetadataViewName,
+      description: undefined,
+      columnDescriptions: {}
+    }
   ];
-  let userTableNames = userEntities.map(entity => entity.table.tableName);
-  let tableNames = Belt_Array.concatMany([
-    exposedInternalTableNames,
-    userTableNames
+  let userTableConfigs = userEntities.map(entity => {
+    let columnDescriptions = {};
+    entity.table.fields.forEach(fieldOrDerived => {
+      if (fieldOrDerived.TAG !== "Field") {
+        return;
+      }
+      let field = fieldOrDerived._0;
+      let d = field.description;
+      if (d !== undefined) {
+        columnDescriptions[Table.getDbFieldName(field)] = d;
+        return;
+      }
+    });
+    return {
+      tableName: entity.table.tableName,
+      description: entity.table.description,
+      columnDescriptions: columnDescriptions
+    };
+  });
+  let tableConfigs = Belt_Array.concatMany([
+    exposedInternalTableConfigs,
+    userTableConfigs
   ]);
+  let tableNames = tableConfigs.map(c => c.tableName);
   Logging.info("Tracking tables in Hasura");
   await clearHasuraMetadata(endpoint, auth);
-  await trackTables(endpoint, auth, pgSchema, tableNames);
+  await trackTables(endpoint, auth, pgSchema, tableConfigs);
   for (let i = 0, i_finish = tableNames.length; i < i_finish; ++i) {
     let tableName = tableNames[i];
     await createSelectPermission(endpoint, auth, tableName, pgSchema, responseLimit, aggregateEntities);
@@ -215,13 +270,13 @@ async function trackDatabase(endpoint, auth, pgSchema, userEntities, aggregateEn
     for (let j = 0, j_finish = derivedFromFields.length; j < j_finish; ++j) {
       let derivedFromField = derivedFromFields[j];
       let relationalFieldName = Utils.unwrapResultExn(Schema.getDerivedFromFieldName(schema, derivedFromField));
-      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "array", relationalFieldName, derivedFromField.fieldName, derivedFromField.derivedFromEntity, true);
+      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "array", relationalFieldName, derivedFromField.fieldName, derivedFromField.derivedFromEntity, true, derivedFromField.description);
     }
     let linkedEntityFields = Table.getLinkedEntityFields(entityConfig.table);
     for (let j$1 = 0, j_finish$1 = linkedEntityFields.length; j$1 < j_finish$1; ++j$1) {
       let match$1 = linkedEntityFields[j$1];
       let field = match$1[0];
-      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "object", field.fieldName, field.fieldName, match$1[1], false);
+      await createEntityRelationship(endpoint, auth, pgSchema, tableName$1, "object", field.fieldName, field.fieldName, match$1[1], false, field.description);
     }
   }
   return Logging.info("Hasura configuration completed");

package/src/InMemoryStore.res

@@ -1,42 +1,23 @@
-type rawEventsKey = {
-  chainId: int,
-  eventId: string,
-}
-
-let hashRawEventsKey = (key: rawEventsKey) =>
-  EventUtils.getEventIdKeyString(~chainId=key.chainId, ~eventId=key.eventId)
-
 module EntityTables = {
-  type t = dict<InMemoryTable.Entity.t<Internal.entity>>
+  type t = dict<InMemoryTable.Entity.t>
   exception UndefinedEntity({entityName: string})
   let make = (entities: array<Internal.entityConfig>): t => {
     let init = Dict.make()
-    entities->Belt.Array.forEach(entityConfig => {
+    entities->Array.forEach(entityConfig => {
       init->Dict.set((entityConfig.name :> string), InMemoryTable.Entity.make())
     })
     init
   }
 
-  let get = (type entity, self: t, ~entityName: string) => {
+  let get = (self: t, ~entityName: string) => {
     switch self->Utils.Dict.dangerouslyGetNonOption(entityName) {
-    | Some(table) =>
-      table->(
-        Utils.magic: InMemoryTable.Entity.t<Internal.entity> => InMemoryTable.Entity.t<entity>
-      )
-
+    | Some(table) => table
     | None =>
       UndefinedEntity({entityName: entityName})->ErrorHandling.mkLogAndRaise(
         ~msg="Unexpected, entity InMemoryTable is undefined",
       )
     }
   }
-
-  let clone = (self: t) => {
-    self
-    ->Dict.toArray
-    ->Belt.Array.map(((k, v)) => (k, v->InMemoryTable.Entity.clone))
-    ->Dict.fromArray
-  }
 }
 
 type effectCacheInMemTable = {
@@ -47,30 +28,29 @@ type effectCacheInMemTable = {
 }
 
 type t = {
-  rawEvents: InMemoryTable.t<rawEventsKey, InternalTable.RawEvents.t>,
-  entities: dict<InMemoryTable.Entity.t<Internal.entity>>,
-  effects: dict<effectCacheInMemTable>,
-  rollbackTargetCheckpointId: option<Internal.checkpointId>,
+  allEntities: array<Internal.entityConfig>,
+  mutable rawEvents: array<InternalTable.RawEvents.t>,
+  mutable entities: dict<InMemoryTable.Entity.t>,
+  mutable effects: dict<effectCacheInMemTable>,
+  mutable rollback: option<Persistence.rollback>,
+  mutable committedCheckpointId: Internal.checkpointId,
 }
 
-let make = (~entities: array<Internal.entityConfig>, ~rollbackTargetCheckpointId=?): t => {
-  rawEvents: InMemoryTable.make(~hash=hashRawEventsKey),
+let make = (
+  ~entities: array<Internal.entityConfig>,
+  ~committedCheckpointId=Internal.initialCheckpointId,
+): t => {
+  allEntities: entities,
+  rawEvents: [],
   entities: EntityTables.make(entities),
   effects: Dict.make(),
-  rollbackTargetCheckpointId,
+  rollback: None,
+  committedCheckpointId,
 }
 
-let clone = (self: t) => {
-  rawEvents: self.rawEvents->InMemoryTable.clone,
-  entities: self.entities->EntityTables.clone,
-  effects: Dict.mapValues(self.effects, table => {
-    idsToStore: table.idsToStore->Array.copy,
-    invalidationsCount: table.invalidationsCount,
-    dict: table.dict->Utils.Dict.shallowCopy,
-    effect: table.effect,
-  }),
-  rollbackTargetCheckpointId: self.rollbackTargetCheckpointId,
-}
+// Once the store holds this many entities across all tables, we drop them
+// after a batch write so it doesn't grow unbounded on long running indexers.
+let keepLatestChangesLimit = 50_000.
 
 let getEffectInMemTable = (inMemoryStore: t, ~effect: Internal.effect) => {
   let key = effect.name
@@ -91,13 +71,193 @@ let getEffectInMemTable = (inMemoryStore: t, ~effect: Internal.effect) => {
 let getInMemTable = (
   inMemoryStore: t,
   ~entityConfig: Internal.entityConfig,
-): InMemoryTable.Entity.t<Internal.entity> => {
+): InMemoryTable.Entity.t => {
   inMemoryStore.entities->EntityTables.get(~entityName=entityConfig.name)
 }
 
-let isRollingBack = (inMemoryStore: t) => inMemoryStore.rollbackTargetCheckpointId !== None
+let isRollingBack = (inMemoryStore: t) => inMemoryStore.rollback !== None
+
+let writeBatch = async (
+  inMemoryStore: t,
+  ~persistence: Persistence.t,
+  ~batch,
+  ~config,
+  ~isInReorgThreshold,
+) =>
+  switch persistence.storageStatus {
+  | Unknown
+  | Initializing(_) =>
+    JsError.throwWithMessage(`Failed to access the indexer storage. The Persistence layer is not initialized.`)
+  | Ready({cache}) =>
+    let committedCheckpointId = inMemoryStore.committedCheckpointId
+    // Decide before the keepMap below trims prevEntityChanges from changesCount,
+    // so the signal still reflects every change currently held in memory.
+    let keepLatestChanges = {
+      let totalChanges = ref(0.)
+      persistence.allEntities->Array.forEach(entityConfig => {
+        totalChanges :=
+          totalChanges.contents +. (inMemoryStore->getInMemTable(~entityConfig)).changesCount
+      })
+      totalChanges.contents < keepLatestChangesLimit
+    }
+    let updatedEntities = persistence.allEntities->Array.filterMap(entityConfig => {
+      let table = inMemoryStore->getInMemTable(~entityConfig)
+
+      // The reset below drops prevEntityChanges and we reuse the array as the
+      // write buffer here, so drop it from the count before appending to it.
+      table.changesCount = table.changesCount -. table.prevEntityChanges->Array.length->Int.toFloat
+      let changes = table.prevEntityChanges
+      table.latestEntityChangeById->Utils.Dict.forEach(change =>
+        if change->Change.getCheckpointId > committedCheckpointId {
+          changes->Array.push(change)
+        }
+      )
+      if changes->Utils.Array.isEmpty {
+        None
+      } else {
+        Some(({entityConfig, changes}: Persistence.updatedEntity))
+      }
+    })
+    await persistence.storage.writeBatch(
+      ~batch,
+      ~rawEvents=inMemoryStore.rawEvents,
+      ~rollback=inMemoryStore.rollback,
+      ~isInReorgThreshold,
+      ~config,
+      ~allEntities=persistence.allEntities,
+      ~updatedEntities,
+      ~updatedEffectsCache={
+        let acc = []
+        inMemoryStore.effects->Utils.Dict.forEach(inMemTable => {
+          let {idsToStore, dict, effect, invalidationsCount} = inMemTable
+          switch idsToStore {
+          | [] => ()
+          | ids =>
+            let items = ids->Array.map((id): Internal.effectCacheItem => {
+              id,
+              output: dict->Dict.getUnsafe(id),
+            })
+            let effectName = effect.name
+            let effectCacheRecord = switch cache->Utils.Dict.dangerouslyGetNonOption(effectName) {
+            | Some(c) => c
+            | None =>
+              let c: Persistence.effectCacheRecord = {effectName, count: 0}
+              cache->Dict.set(effectName, c)
+              c
+            }
+            let shouldInitialize = effectCacheRecord.count === 0
+            effectCacheRecord.count =
+              effectCacheRecord.count + items->Array.length - invalidationsCount
+            Prometheus.EffectCacheCount.set(~count=effectCacheRecord.count, ~effectName)
+            acc
+            ->Array.push(({effect, items, shouldInitialize}: Persistence.updatedEffectCache))
+            ->ignore
+          }
+        })
+        acc
+      },
+    )
+
+    inMemoryStore.rawEvents = []
+    inMemoryStore.effects = Dict.make()
+    inMemoryStore.rollback = None
+    inMemoryStore.committedCheckpointId = switch batch.checkpointIds->Utils.Array.last {
+    | Some(checkpointId) => checkpointId
+    | None => committedCheckpointId
+    }
+    if keepLatestChanges {
+      persistence.allEntities->Array.forEach(entityConfig => {
+        let table = inMemoryStore->getInMemTable(~entityConfig)
+        inMemoryStore.entities->Dict.set(
+          (entityConfig.name :> string),
+          table->InMemoryTable.Entity.resetButKeepLatestChanges,
+        )
+      })
+    } else {
+      // Over the limit: drop everything written in a batch and keep only the
+      // entities loaded from the db, so the next batch can still read them
+      // without hitting the database.
+      let loadedFromDbCount = ref(0.)
+      let resetTables = persistence.allEntities->Array.map(entityConfig => {
+        let resetTable =
+          inMemoryStore
+          ->getInMemTable(~entityConfig)
+          ->InMemoryTable.Entity.resetButKeepLoadedFromDbChanges
+        loadedFromDbCount := loadedFromDbCount.contents +. resetTable.changesCount
+        resetTable
+      })
+      // Even the loaded-from-db entities alone exceed the limit, so there's no
+      // point keeping them around - drop everything.
+      let dropEverything = loadedFromDbCount.contents >= keepLatestChangesLimit
+      persistence.allEntities->Array.forEachWithIndex((entityConfig, idx) => {
+        inMemoryStore.entities->Dict.set(
+          (entityConfig.name :> string),
+          dropEverything ? InMemoryTable.Entity.make() : resetTables->Array.getUnsafe(idx),
+        )
+      })
+    }
+  }
+
+let prepareRollbackDiff = async (
+  inMemoryStore: t,
+  ~persistence: Persistence.t,
+  ~rollbackTargetCheckpointId,
+  ~rollbackDiffCheckpointId,
+) => {
+  inMemoryStore.rawEvents = []
+  inMemoryStore.entities = EntityTables.make(inMemoryStore.allEntities)
+  inMemoryStore.effects = Dict.make()
+  inMemoryStore.rollback = Some({
+    targetCheckpointId: rollbackTargetCheckpointId,
+    diffCheckpointId: rollbackDiffCheckpointId,
+  })
+
+  let deletedEntities = Dict.make()
+  let setEntities = Dict.make()
+
+  let _ = await persistence.allEntities
+  ->Array.map(async entityConfig => {
+    let entityTable = inMemoryStore->getInMemTable(~entityConfig)
+
+    let (removedIdsResult, restoredEntitiesResult) = await persistence.storage.getRollbackData(
+      ~entityConfig,
+      ~rollbackTargetCheckpointId,
+    )
+
+    removedIdsResult->Array.forEach(data => {
+      deletedEntities->Utils.Dict.push(entityConfig.name, data["id"])
+      entityTable->InMemoryTable.Entity.set(
+        ~committedCheckpointId=inMemoryStore.committedCheckpointId,
+        Delete({
+          entityId: data["id"],
+          checkpointId: rollbackDiffCheckpointId,
+        }),
+      )
+    })
+
+    let restoredEntities = restoredEntitiesResult->S.parseOrThrow(entityConfig.rowsSchema)
+
+    restoredEntities->Array.forEach((entity: Internal.entity) => {
+      setEntities->Utils.Dict.push(entityConfig.name, entity.id)
+      entityTable->InMemoryTable.Entity.set(
+        ~committedCheckpointId=inMemoryStore.committedCheckpointId,
+        Set({
+          entityId: entity.id,
+          checkpointId: rollbackDiffCheckpointId,
+          entity,
+        }),
+      )
+    })
+  })
+  ->Promise.all
+
+  {
+    "deletedEntities": deletedEntities,
+    "setEntities": setEntities,
+  }
+}
 
-let setBatchDcs = (inMemoryStore: t, ~batch: Batch.t, ~shouldSaveHistory) => {
+let setBatchDcs = (inMemoryStore: t, ~batch: Batch.t) => {
   let inMemTable =
     inMemoryStore->getInMemTable(~entityConfig=InternalTable.EnvioAddresses.entityConfig)
 
@@ -126,12 +286,12 @@ let setBatchDcs = (inMemoryStore: t, ~batch: Batch.t, ~shouldSaveHistory) => {
           }
 
           inMemTable->InMemoryTable.Entity.set(
+            ~committedCheckpointId=inMemoryStore.committedCheckpointId,
             Set({
               entityId: entity.id,
               checkpointId,
               entity: entity->InternalTable.EnvioAddresses.castToInternal,
             }),
-            ~shouldSaveHistory,
           )
         }
       }