graph, store: Do not panic in Entity.insert when given an uninterned key

Author: lutter

graph/src/components/store/entity_cache.rs

@@ -9,7 +9,7 @@ use crate::prelude::ENV_VARS;
 use crate::schema::InputSchema;
 use crate::util::lfu_cache::LfuCache;
 
-use super::{DerivedEntityQuery, EntityType, LoadRelatedRequest};
+use super::{DerivedEntityQuery, EntityType, LoadRelatedRequest, StoreError};
 
 /// The scope in which the `EntityCache` should perform a `get` operation
 pub enum GetScope {
@@ -113,11 +113,7 @@ impl EntityCache {
         self.handler_updates.clear();
     }
 
-    pub fn get(
-        &mut self,
-        key: &EntityKey,
-        scope: GetScope,
-    ) -> Result<Option<Entity>, s::QueryExecutionError> {
+    pub fn get(&mut self, key: &EntityKey, scope: GetScope) -> Result<Option<Entity>, StoreError> {
         // Get the current entity, apply any updates from `updates`, then
         // from `handler_updates`.
         let mut entity = match scope {
@@ -133,10 +129,10 @@ impl EntityCache {
         });
 
         if let Some(op) = self.updates.get(key).cloned() {
-            entity = op.apply_to(entity)
+            entity = op.apply_to(entity).map_err(|e| key.unknown_attribute(e))?;
         }
         if let Some(op) = self.handler_updates.get(key).cloned() {
-            entity = op.apply_to(entity)
+            entity = op.apply_to(entity).map_err(|e| key.unknown_attribute(e))?;
         }
         Ok(entity)
     }
@@ -266,7 +262,7 @@ impl EntityCache {
     /// to the current state is actually needed.
     ///
     /// Also returns the updated `LfuCache`.
-    pub fn as_modifications(mut self) -> Result<ModificationsAndCache, s::QueryExecutionError> {
+    pub fn as_modifications(mut self) -> Result<ModificationsAndCache, StoreError> {
         assert!(!self.in_handler);
 
         // The first step is to make sure all entities being set are in `self.current`.
@@ -305,7 +301,8 @@ impl EntityCache {
                 // Entity may have been changed
                 (Some(current), EntityOp::Update(updates)) => {
                     let mut data = current.clone();
-                    data.merge_remove_null_fields(updates);
+                    data.merge_remove_null_fields(updates)
+                        .map_err(|e| key.unknown_attribute(e))?;
                     self.current.insert(key.clone(), Some(data.clone()));
                     if current != data {
                         Some(Overwrite { key, data })

graph/src/components/store/err.rs

@@ -18,6 +18,8 @@ pub enum StoreError {
     UnknownField(String),
     #[error("unknown table '{0}'")]
     UnknownTable(String),
+    #[error("entity type '{0}' does not have an attribute '{0}'")]
+    UnknownAttribute(String, String),
     #[error("malformed directive '{0}'")]
     MalformedDirective(String),
     #[error("query execution failed: {0}")]

graph/src/components/store/mod.rs

@@ -29,6 +29,7 @@ use crate::data::store::*;
 use crate::data::value::Word;
 use crate::data_source::CausalityRegion;
 use crate::schema::InputSchema;
+use crate::util::intern::{self, Error as InternError};
 use crate::{constraint_violation, prelude::*};
 
 /// The type name of an entity. This is the string that is used in the
@@ -139,6 +140,12 @@ pub struct EntityKey {
     pub causality_region: CausalityRegion,
 }
 
+impl EntityKey {
+    pub fn unknown_attribute(&self, err: intern::Error) -> StoreError {
+        StoreError::UnknownAttribute(self.entity_type.to_string(), err.not_interned())
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct LoadRelatedRequest {
     /// Name of the entity type.
@@ -1008,14 +1015,14 @@ enum EntityOp {
 }
 
 impl EntityOp {
-    fn apply_to(self, entity: Option<Entity>) -> Option<Entity> {
+    fn apply_to(self, entity: Option<Entity>) -> Result<Option<Entity>, InternError> {
         use EntityOp::*;
         match (self, entity) {
-            (Remove, _) => None,
-            (Overwrite(new), _) | (Update(new), None) => Some(new),
+            (Remove, _) => Ok(None),
+            (Overwrite(new), _) | (Update(new), None) => Ok(Some(new)),
             (Update(updates), Some(mut entity)) => {
-                entity.merge_remove_null_fields(updates);
-                Some(entity)
+                entity.merge_remove_null_fields(updates)?;
+                Ok(Some(entity))
             }
         }
     }

graph/src/data/store/mod.rs

@@ -5,7 +5,7 @@ use crate::{
     runtime::gas::{Gas, GasSizeOf},
     schema::InputSchema,
     util::intern::AtomPool,
-    util::intern::{NullValue, Object},
+    util::intern::{Error as InternError, NullValue, Object},
 };
 use crate::{data::subgraph::DeploymentHash, prelude::EntityChange};
 use anyhow::{anyhow, Error};
@@ -744,8 +744,8 @@ impl Entity {
         self.0.get(key)
     }
 
-    pub fn insert(&mut self, key: &str, value: Value) -> Option<Value> {
-        self.0.insert(key, value).expect("key is in AtomPool")
+    pub fn insert(&mut self, key: &str, value: Value) -> Result<Option<Value>, InternError> {
+        self.0.insert(key, value)
     }
 
     pub fn remove(&mut self, key: &str) -> Option<Value> {
@@ -797,13 +797,14 @@ impl Entity {
     /// If a key exists in both entities, the value from `update` is chosen.
     /// If a key only exists on one entity, the value from that entity is chosen.
     /// If a key is set to `Value::Null` in `update`, the key/value pair is removed.
-    pub fn merge_remove_null_fields(&mut self, update: Entity) {
+    pub fn merge_remove_null_fields(&mut self, update: Entity) -> Result<(), InternError> {
         for (key, value) in update.0.into_iter() {
             match value {
                 Value::Null => self.remove(&key),
-                _ => self.insert(&key, value),
+                _ => self.insert(&key, value)?,
             };
         }
+        Ok(())
     }
 
     /// Remove all entries with value `Value::Null` from `self`

graph/src/util/intern.rs

@@ -52,7 +52,7 @@ pub enum Error {
 }
 
 impl Error {
-    pub fn not_interned(&self) -> &str {
+    pub fn not_interned(self) -> String {
         match self {
             Error::NotInterned(s) => s,
         }

store/postgres/src/relational_queries.rs

@@ -1762,7 +1762,8 @@ impl<'a> InsertQuery<'a> {
                     if !fulltext_field_values.is_empty() {
                         entity
                             .to_mut()
-                            .insert(&column.field, Value::List(fulltext_field_values));
+                            .insert(&column.field, Value::List(fulltext_field_values))
+                            .map_err(|e| entity_key.unknown_attribute(e))?;
                     }
                 }
                 if !column.is_nullable() && !entity.contains_key(&column.field) {

store/test-store/tests/postgres/relational.rs

@@ -341,7 +341,7 @@ fn make_user(
     }
     .unwrap();
     if let Some(drinks) = drinks {
-        user.insert("drinks", drinks.into());
+        user.insert("drinks", drinks.into()).unwrap();
     }
     user
 }

store/test-store/tests/postgres/store.rs

@@ -446,7 +446,7 @@ fn update_existing() {
             _ => unreachable!(),
         };
 
-        new_data.insert("bin_name", Value::Bytes(bin_name));
+        new_data.insert("bin_name", Value::Bytes(bin_name)).unwrap();
         assert_eq!(writable.get(&entity_key).unwrap(), Some(new_data));
     })
 }