address CI failures

Author: MoonBoi9001

graph/src/components/store/write.rs

@@ -4,14 +4,14 @@ use std::{collections::HashSet, sync::Arc};
 use crate::{
     blockchain::{block_stream::FirehoseCursor, BlockPtr, BlockTime},
     cheap_clone::CheapClone,
+    components::metrics::block_state::BlockStateMetrics,
     components::subgraph::Entity,
     constraint_violation,
     data::{store::Id, subgraph::schema::SubgraphError},
     data_source::CausalityRegion,
     derive::CacheWeight,
     prelude::DeploymentHash,
     util::cache_weight::CacheWeight,
-    components::metrics::block_state::BlockStateMetrics,
 };
 
 use super::{BlockNumber, EntityKey, EntityType, StoreError, StoreEvent, StoredDynamicDataSource};
@@ -511,7 +511,7 @@ impl RowGroup {
             .rows
             .iter()
             .filter_map(|row| match row {
-                EntityModification::Insert { data, .. } 
+                EntityModification::Insert { data, .. }
                 | EntityModification::Overwrite { data, .. } => Some(data.as_ref().clone()),
                 EntityModification::Remove { .. } => None,
             })

runtime/wasm/src/host_exports.rs

@@ -33,6 +33,8 @@ use crate::{error::DeterminismLevel, module::IntoTrap};
 
 use super::module::WasmInstanceData;
 
+use graph::schema::EntityKey;
+
 fn write_poi_event(
     proof_of_indexing: &SharedProofOfIndexing,
     poi_event: &ProofOfIndexingEvent,
@@ -354,7 +356,7 @@ impl HostExports {
             .metrics
             .track_storage_size_change(&entity_type, &entity, false);
 
-        if !state.entity_cache.contains_key(&key) {
+        if state.entity_cache.get(&key, GetScope::Store)?.is_none() {
             state.metrics.track_entity_count_change(&entity_type, 1);
         }
 
@@ -1250,21 +1252,27 @@ impl HostExports {
             .context("Failed to decode")
     }
 
-    fn remove_entity(&mut self, key: &EntityKey) -> Result<(), HostExportError> {
+    fn remove_entity(
+        &mut self,
+        key: &EntityKey,
+        state: &mut BlockState,
+    ) -> Result<(), HostExportError> {
         let entity_type = key.entity_type.clone();
-        
-        if let Some(entity) = state.entity_cache.get(&key, GetScope::Store)? {
-            // Track removal metrics while we still have the entity data
+
+        if let Some(entity) = state.entity_cache.get(key, GetScope::Store)? {
             state
                 .metrics
                 .track_storage_size_change(&entity_type, &entity, true);
             state.metrics.track_entity_count_change(&entity_type, -1);
         }
 
-        state.modifications.push(EntityModification::remove(
-            key.clone(),
-            state.block.number,
-        ));
+        state
+            .entity_cache
+            .modifications
+            .push(EntityModification::remove(
+                key.clone(),
+                state.entity_cache.block_number,
+            ));
         Ok(())
     }
 }
@@ -1395,3 +1403,9 @@ fn bytes_to_string_is_lossy() {
         )
     )
 }
+
+impl From<StoreError> for HostExportError {
+    fn from(err: StoreError) -> Self {
+        HostExportError::Deterministic(anyhow::Error::from(err))
+    }
+}