graph, store: Include vid in entity modifications use in ClampRangeQuery

lutter · lutter · commit 4396f8a41b61 · 2021-12-01T15:11:58.000-08:00
diff --git a/graph/src/components/store.rs b/graph/src/components/store.rs
@@ -553,7 +553,7 @@ impl<'a> FromIterator<&'a EntityModification> for StoreEvent {
             .map(|op| {
                 use self::EntityModification::*;
                 match op {
-                    Insert { key, .. } | Overwrite { key, .. } | Remove { key } => {
+                    Insert { key, .. } | Overwrite { key, .. } | Remove { key, .. } => {
                         EntityChange::for_data(key.clone())
                     }
                 }
@@ -1451,16 +1451,20 @@ pub enum EntityModification {
     /// Insert the entity
     Insert { key: EntityKey, data: Entity },
     /// Update the entity by overwriting it
-    Overwrite { key: EntityKey, data: Entity },
+    Overwrite {
+        key: EntityKey,
+        data: Entity,
+        prev_vid: Vid,
+    },
     /// Remove the entity
-    Remove { key: EntityKey },
+    Remove { key: EntityKey, prev_vid: Vid },
 }
 
 impl EntityModification {
     pub fn entity_key(&self) -> &EntityKey {
         use EntityModification::*;
         match self {
-            Insert { key, .. } | Overwrite { key, .. } | Remove { key } => key,
+            Insert { key, .. } | Overwrite { key, .. } | Remove { key, .. } => key,
         }
     }
 
@@ -1731,7 +1735,11 @@ impl EntityCache {
                     let ev = EntityVersion::new(data.clone(), current.vid);
                     self.current.insert(key.clone(), Some(ev));
                     if current.data != data {
-                        Some(Overwrite { key, data })
+                        Some(Overwrite {
+                            key,
+                            data,
+                            prev_vid: current.vid,
+                        })
                     } else {
                         None
                     }
@@ -1741,15 +1749,22 @@ impl EntityCache {
                     let ev = EntityVersion::new(data.clone(), current.vid);
                     self.current.insert(key.clone(), Some(ev));
                     if current.data != data {
-                        Some(Overwrite { key, data })
+                        Some(Overwrite {
+                            key,
+                            data,
+                            prev_vid: current.vid,
+                        })
                     } else {
                         None
                     }
                 }
                 // Existing entity was deleted
-                (Some(_), EntityOp::Remove) => {
+                (Some(current), EntityOp::Remove) => {
                     self.current.insert(key.clone(), None);
-                    Some(Remove { key })
+                    Some(Remove {
+                        key,
+                        prev_vid: current.vid,
+                    })
                 }
                 // Entity was deleted, but it doesn't exist in the store
                 (None, EntityOp::Remove) => None,
diff --git a/graph/tests/entity_cache.rs b/graph/tests/entity_cache.rs
@@ -144,18 +144,22 @@ fn overwrite_modifications() {
     cache.set(sigurros_key.clone(), sigurros_data.clone());
 
     let result = cache.as_modifications();
+    let one = NonZeroU64::new(1);
+    let two = NonZeroU64::new(2);
     assert_eq!(
         sort_by_entity_key(result.unwrap().modifications),
-        sort_by_entity_key(vec![
+        vec![
             EntityModification::Overwrite {
                 key: mogwai_key,
                 data: mogwai_data,
+                prev_vid: one,
             },
             EntityModification::Overwrite {
                 key: sigurros_key,
                 data: sigurros_data,
+                prev_vid: two,
             }
-        ])
+        ]
     );
 }
 
@@ -205,15 +209,17 @@ fn consecutive_modifications() {
     // We expect a single overwrite modification for the above that leaves "id"
     // and "name" untouched, sets "founded" and removes the "label" field.
     let result = cache.as_modifications();
+    let one = NonZeroU64::new(1);
     assert_eq!(
-        sort_by_entity_key(result.unwrap().modifications),
-        sort_by_entity_key(vec![EntityModification::Overwrite {
+        result.unwrap().modifications,
+        vec![EntityModification::Overwrite {
             key: update_key,
             data: Entity::from(vec![
                 ("id", "mogwai".into()),
                 ("name", "Mogwai".into()),
                 ("founded", 1995.into()),
             ]),
-        },])
+            prev_vid: one
+        },]
     );
 }
diff --git a/store/postgres/src/deployment_store.rs b/store/postgres/src/deployment_store.rs
@@ -4,7 +4,7 @@ use diesel::pg::PgConnection;
 use diesel::prelude::*;
 use diesel::r2d2::{ConnectionManager, PooledConnection};
 use graph::components::store::{EntityType, StoredDynamicDataSource};
-use graph::data::store::EntityVersion;
+use graph::data::store::{EntityVersion, Vid};
 use graph::data::subgraph::status;
 use graph::prelude::{
     tokio, CancelHandle, CancelToken, CancelableError, PoolWaitStats, SubgraphDeploymentEntity,
@@ -327,17 +327,23 @@ impl DeploymentStore {
                         .or_insert_with(Vec::new)
                         .push((key, Cow::from(data)));
                 }
-                Overwrite { key, data } => {
-                    overwrites
+                Overwrite {
+                    key,
+                    data,
+                    prev_vid,
+                } => {
+                    let (entities, vids) = overwrites
                         .entry(key.entity_type.clone())
-                        .or_insert_with(Vec::new)
-                        .push((key, Cow::from(data)));
+                        .or_insert_with(|| (Vec::new(), Vec::new()));
+
+                    entities.push((key, Cow::from(data)));
+                    vids.push(*prev_vid);
                 }
-                Remove { key } => {
+                Remove { key, prev_vid } => {
                     removals
                         .entry(key.entity_type.clone())
                         .or_insert_with(Vec::new)
-                        .push(key.entity_id.as_str());
+                        .push(*prev_vid);
                 }
             }
         }
@@ -351,21 +357,24 @@ impl DeploymentStore {
         }
 
         // Overwrites:
-        for (entity_type, mut entities) in overwrites.into_iter() {
+        for (entity_type, (mut entities, vids)) in overwrites.into_iter() {
             // we do not update the count since the number of entities remains the same
-            self.overwrite_entities(&entity_type, &mut entities, conn, layout, ptr, &stopwatch)?;
-        }
-
-        // Removals
-        for (entity_type, entity_keys) in removals.into_iter() {
-            count -= self.remove_entities(
+            self.overwrite_entities(
                 &entity_type,
-                entity_keys.as_slice(),
+                &mut entities,
+                vids.as_slice(),
                 conn,
                 layout,
                 ptr,
                 &stopwatch,
-            )? as i32;
+            )?;
+        }
+
+        // Removals
+        for (entity_type, vids) in removals.into_iter() {
+            count -=
+                self.remove_entities(&entity_type, vids.as_slice(), conn, layout, ptr, &stopwatch)?
+                    as i32;
         }
         Ok(count)
     }
@@ -394,6 +403,7 @@ impl DeploymentStore {
         &'a self,
         entity_type: &'a EntityType,
         data: &'a mut [(&'a EntityKey, Cow<'a, Entity>)],
+        vids: &'a [Vid],
         conn: &PgConnection,
         layout: &'a Layout,
         ptr: &BlockPtr,
@@ -407,28 +417,29 @@ impl DeploymentStore {
         section.end();
 
         let _section = stopwatch.start_section("apply_entity_modifications_update");
-        layout.update(conn, &entity_type, data, block_number(ptr), stopwatch)
+        layout.update(conn, &entity_type, data, vids, block_number(ptr), stopwatch)
     }
 
     fn remove_entities(
         &self,
         entity_type: &EntityType,
-        entity_keys: &[&str],
+        vids: &[Vid],
         conn: &PgConnection,
         layout: &Layout,
         ptr: &BlockPtr,
         stopwatch: &StopwatchMetrics,
     ) -> Result<usize, StoreError> {
         let _section = stopwatch.start_section("apply_entity_modifications_delete");
         layout
-            .delete(
-                conn,
-                entity_type,
-                &entity_keys,
-                block_number(ptr),
-                stopwatch,
-            )
-            .map_err(|_error| anyhow!("Failed to remove entities: {:?}", entity_keys).into())
+            .delete(conn, entity_type, vids, block_number(ptr), stopwatch)
+            .map_err(|_error| {
+                anyhow!(
+                    "Failed to remove entities for type {}: {:?}",
+                    entity_type,
+                    vids
+                )
+                .into()
+            })
     }
 
     /// Execute a closure with a connection to the database.
diff --git a/store/postgres/src/relational.rs b/store/postgres/src/relational.rs
@@ -32,7 +32,7 @@ use crate::{
 use graph::components::store::EntityType;
 use graph::data::graphql::ext::{DirectiveFinder, DocumentExt, ObjectTypeExt};
 use graph::data::schema::{FulltextConfig, FulltextDefinition, Schema, SCHEMA_TYPE_NAME};
-use graph::data::store::{EntityVersion, BYTES_SCALAR};
+use graph::data::store::{EntityVersion, Vid, BYTES_SCALAR};
 use graph::data::subgraph::schema::{POI_OBJECT, POI_TABLE};
 use graph::prelude::{
     anyhow, info, BlockNumber, DeploymentHash, Entity, EntityChange, EntityCollection,
@@ -681,17 +681,14 @@ impl Layout {
         conn: &PgConnection,
         entity_type: &'a EntityType,
         entities: &'a mut [(&'a EntityKey, Cow<'a, Entity>)],
+        vids: &'a [Vid],
         block: BlockNumber,
         stopwatch: &StopwatchMetrics,
     ) -> Result<usize, StoreError> {
         let table = self.table_for_entity(&entity_type)?;
-        let entity_keys: Vec<&str> = entities
-            .iter()
-            .map(|(key, _)| key.entity_id.as_str())
-            .collect();
 
         let section = stopwatch.start_section("update_modification_clamp_range_query");
-        ClampRangeQuery::new(table, &entity_type, &entity_keys, block).execute(conn)?;
+        ClampRangeQuery::new(table, &vids, block)?.execute(conn)?;
         section.end();
 
         let _section = stopwatch.start_section("update_modification_insert_query");
@@ -712,15 +709,15 @@ impl Layout {
         &self,
         conn: &PgConnection,
         entity_type: &EntityType,
-        entity_ids: &[&str],
+        vids: &[Vid],
         block: BlockNumber,
         stopwatch: &StopwatchMetrics,
     ) -> Result<usize, StoreError> {
         let table = self.table_for_entity(&entity_type)?;
         let _section = stopwatch.start_section("delete_modification_clamp_range_query");
         let mut count = 0;
-        for chunk in entity_ids.chunks(DELETE_OPERATION_CHUNK_SIZE) {
-            count += ClampRangeQuery::new(table, &entity_type, chunk, block).execute(conn)?
+        for chunk in vids.chunks(DELETE_OPERATION_CHUNK_SIZE) {
+            count += ClampRangeQuery::new(table, chunk, block)?.execute(conn)?
         }
         Ok(count)
     }
diff --git a/store/postgres/src/relational_queries.rs b/store/postgres/src/relational_queries.rs
@@ -2598,23 +2598,36 @@ impl<'a, Conn> RunQueryDsl<Conn> for FilterQuery<'a> {}
 
 /// Reduce the upper bound of the current entry's block range to `block` as
 /// long as that does not result in an empty block range
-#[derive(Debug, Clone, Constructor)]
-pub struct ClampRangeQuery<'a, S> {
+#[derive(Debug, Clone)]
+pub struct ClampRangeQuery<'a> {
     table: &'a Table,
-    entity_type: &'a EntityType,
-    entity_ids: &'a [S],
+    vids: Vec<i64>,
     block: BlockNumber,
 }
 
-impl<'a, S> QueryFragment<Pg> for ClampRangeQuery<'a, S>
-where
-    S: AsRef<str> + diesel::serialize::ToSql<Text, Pg>,
-{
+fn i64_from_vid(vid: &Vid) -> Result<i64, StoreError> {
+    match vid {
+        Some(vid) => i64::try_from(vid.get())
+            .map_err(|_| StoreError::ConstraintViolation(format!("vid out of bounds: {}", vid))),
+        None => Err(StoreError::ConstraintViolation("missing vid".to_string())),
+    }
+}
+
+impl<'a> ClampRangeQuery<'a> {
+    pub fn new(table: &'a Table, vids: &'a [Vid], block: BlockNumber) -> Result<Self, StoreError> {
+        let vids = vids
+            .into_iter()
+            .map(i64_from_vid)
+            .collect::<Result<Vec<_>, _>>()?;
+        Ok(Self { table, vids, block })
+    }
+}
+
+impl<'a> QueryFragment<Pg> for ClampRangeQuery<'a> {
     fn walk_ast(&self, mut out: AstPass<Pg>) -> QueryResult<()> {
         // update table
         //    set block_range = int4range(lower(block_range), $block)
-        //  where id in (id1, id2, ..., idN)
-        //    and block_range @> INTMAX
+        //  where vid = any([vid1, vid2, ..., vidN])
         out.unsafe_to_cache_prepared();
         out.push_sql("update ");
         out.push_sql(self.table.qualified_name.as_str());
@@ -2624,27 +2637,21 @@ where
         out.push_identifier(BLOCK_RANGE_COLUMN)?;
         out.push_sql("), ");
         out.push_bind_param::<Integer, _>(&self.block)?;
-        out.push_sql(")\n where ");
-
-        self.table.primary_key().is_in(self.entity_ids, &mut out)?;
-        out.push_sql(" and (");
-        out.push_sql(BLOCK_RANGE_CURRENT);
+        out.push_sql(")\n where vid = any(");
+        out.push_bind_param::<Array<BigInt>, _>(&self.vids)?;
         out.push_sql(")");
 
         Ok(())
     }
 }
 
-impl<'a, S> QueryId for ClampRangeQuery<'a, S>
-where
-    S: AsRef<str> + diesel::serialize::ToSql<Text, Pg>,
-{
+impl<'a> QueryId for ClampRangeQuery<'a> {
     type QueryId = ();
 
     const HAS_STATIC_QUERY_ID: bool = false;
 }
 
-impl<'a, S, Conn> RunQueryDsl<Conn> for ClampRangeQuery<'a, S> {}
+impl<'a, Conn> RunQueryDsl<Conn> for ClampRangeQuery<'a> {}
 
 /// Helper struct for returning the id's touched by the RevertRemove and
 /// RevertExtend queries
diff --git a/store/postgres/tests/relational.rs b/store/postgres/tests/relational.rs
diff --git a/store/postgres/tests/relational_bytes.rs b/store/postgres/tests/relational_bytes.rs
