Allow using Bytes as well as String/ID for the id of entities

graph/src/components/store.rs

@@ -1521,10 +1521,37 @@ impl EntityCache {
     /// with existing data. The entity will be validated against the
     /// subgraph schema, and any errors will result in an `Err` being
     /// returned.
-    pub fn set(&mut self, key: EntityKey, entity: Entity) -> Result<(), anyhow::Error> {
-        let is_valid = entity
-            .validate(&self.store.input_schema().document, &key)
-            .is_ok();
+    pub fn set(&mut self, key: EntityKey, mut entity: Entity) -> Result<(), anyhow::Error> {
+        fn check_id(key: &EntityKey, prev_id: &str) -> Result<(), anyhow::Error> {
+            if prev_id != key.entity_id {
+                return Err(anyhow!(
+                    "Value of {} attribute 'id' conflicts with ID passed to `store.set()`: \
+                {} != {}",
+                    key.entity_type,
+                    prev_id,
+                    key.entity_id,
+                ));
+            } else {
+                return Ok(());
+            }
+        }
+
+        // Set the id if there isn't one yet, and make sure that a
+        // previously set id agrees with the one in the `key`
+        match entity.get("id") {
+            Some(Value::String(s)) => check_id(&key, &s)?,
+            Some(Value::Bytes(b)) => check_id(&key, &b.to_string())?,
+            Some(_) => {
+                // The validation will catch the type mismatch
+                ()
+            }
+            None => {
+                let value = self.store.input_schema().id_value(&key)?;
+                entity.set("id", value);
+            }
+        }
+
+        let is_valid = entity.validate(&self.store.input_schema(), &key).is_ok();
 
         self.entity_op(key.clone(), EntityOp::Update(entity));
 
@@ -1539,7 +1566,7 @@ impl EntityCache {
                     key.entity_id
                 )
             })?;
-            entity.validate(&self.store.input_schema().document, &key)?;
+            entity.validate(&self.store.input_schema(), &key)?;
         }
 
         Ok(())

graph/src/data/graphql/ext.rs

@@ -4,10 +4,8 @@ use crate::prelude::s::{
     Definition, Directive, Document, EnumType, Field, InterfaceType, ObjectType, Type,
     TypeDefinition, Value,
 };
-use crate::prelude::ValueType;
 use lazy_static::lazy_static;
 use std::collections::{BTreeMap, HashMap};
-use std::str::FromStr;
 
 lazy_static! {
     static ref ALLOW_NON_DETERMINISTIC_FULLTEXT_SEARCH: bool = if cfg!(debug_assertions) {
@@ -65,8 +63,6 @@ pub trait DocumentExt {
 
     fn get_named_type(&self, name: &str) -> Option<&TypeDefinition>;
 
-    fn scalar_value_type(&self, field_type: &Type) -> ValueType;
-
     /// Return `true` if the type does not allow selection of child fields.
     ///
     /// # Panics
@@ -208,25 +204,6 @@ impl DocumentExt for Document {
             })
     }
 
-    fn scalar_value_type(&self, field_type: &Type) -> ValueType {
-        use TypeDefinition as t;
-        match field_type {
-            Type::NamedType(name) => {
-                ValueType::from_str(&name).unwrap_or_else(|_| match self.get_named_type(name) {
-                    Some(t::Object(_)) | Some(t::Interface(_)) | Some(t::Enum(_)) => {
-                        ValueType::String
-                    }
-                    Some(t::Scalar(_)) => unreachable!("user-defined scalars are not used"),
-                    Some(t::Union(_)) => unreachable!("unions are not used"),
-                    Some(t::InputObject(_)) => unreachable!("inputObjects are not used"),
-                    None => unreachable!("names of field types have been validated"),
-                })
-            }
-            Type::NonNullType(inner) => self.scalar_value_type(inner),
-            Type::ListType(inner) => self.scalar_value_type(inner),
-        }
-    }
-
     fn is_leaf_type(&self, field_type: &Type) -> bool {
         match self
             .get_named_type(field_type.get_base_type())

graph/src/data/schema.rs

@@ -1,17 +1,19 @@
 use crate::cheap_clone::CheapClone;
-use crate::components::store::{EntityType, SubgraphStore};
+use crate::components::store::{EntityKey, EntityType, SubgraphStore};
 use crate::data::graphql::ext::{DirectiveExt, DirectiveFinder, DocumentExt, TypeExt, ValueExt};
-use crate::data::store::ValueType;
+use crate::data::graphql::ObjectTypeExt;
+use crate::data::store::{self, ValueType};
 use crate::data::subgraph::{DeploymentHash, SubgraphName};
 use crate::prelude::{
-    lazy_static,
+    anyhow, lazy_static,
     q::Value,
     s::{self, Definition, InterfaceType, ObjectType, TypeDefinition, *},
 };
 
 use anyhow::{Context, Error};
 use graphql_parser::{self, Pos};
 use inflector::Inflector;
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
@@ -24,6 +26,7 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use super::graphql::ObjectOrInterface;
+use super::store::scalar;
 
 pub const SCHEMA_TYPE_NAME: &str = "_Schema_";
 
@@ -55,6 +58,8 @@ pub enum SchemaValidationError {
          the following fields: {2}"
     )]
     InterfaceFieldsMissing(String, String, Strings), // (type, interface, missing_fields)
+    #[error("Implementors of interface `{0}` use different id types `{1}`. They must all use the same type")]
+    InterfaceImplementorsMixId(String, String),
     #[error("Field `{1}` in type `{0}` has invalid @derivedFrom: {2}")]
     InvalidDerivedFrom(String, String, String), // (type, field, reason)
     #[error("The following type names are reserved: `{0}`")]
@@ -589,6 +594,39 @@ impl Schema {
         }
     }
 
+    /// Construct a value for the entity type's id attribute
+    pub fn id_value(&self, key: &EntityKey) -> Result<store::Value, Error> {
+        let base_type = self
+            .document
+            .get_object_type_definition(key.entity_type.as_str())
+            .ok_or_else(|| {
+                anyhow!(
+                    "Entity {}[{}]: unknown entity type `{}`",
+                    key.entity_type,
+                    key.entity_id,
+                    key.entity_type
+                )
+            })?
+            .field("id")
+            .unwrap()
+            .field_type
+            .get_base_type();
+
+        match base_type {
+            "ID" | "String" => Ok(store::Value::String(key.entity_id.clone())),
+            "Bytes" => Ok(store::Value::Bytes(scalar::Bytes::from_str(
+                &key.entity_id,
+            )?)),
+            s => {
+                return Err(anyhow!(
+                    "Entity type {} uses illegal type {} for id column",
+                    key.entity_type,
+                    s
+                ))
+            }
+        }
+    }
+
     pub fn resolve_schema_references<S: SubgraphStore>(
         &self,
         store: Arc<S>,
@@ -815,6 +853,7 @@ impl Schema {
             self.validate_schema_type_has_no_fields(),
             self.validate_directives_on_schema_type(),
             self.validate_reserved_types_usage(),
+            self.validate_interface_id_type(),
         ]
         .into_iter()
         .filter(Result::is_err)
@@ -1479,6 +1518,25 @@ impl Schema {
         }
     }
 
+    fn validate_interface_id_type(&self) -> Result<(), SchemaValidationError> {
+        for (intf, obj_types) in &self.types_for_interface {
+            let id_types: HashSet<&str> = HashSet::from_iter(
+                obj_types
+                    .iter()
+                    .filter_map(|obj_type| obj_type.field("id"))
+                    .map(|f| f.field_type.get_base_type())
+                    .map(|name| if name == "ID" { "String" } else { name }),
+            );
+            if id_types.len() > 1 {
+                return Err(SchemaValidationError::InterfaceImplementorsMixId(
+                    intf.to_string(),
+                    id_types.iter().join(", "),
+                ));
+            }
+        }
+        Ok(())
+    }
+
     fn subgraph_schema_object_type(&self) -> Option<&ObjectType> {
         self.document
             .get_object_type_definitions()
@@ -1547,6 +1605,39 @@ fn invalid_interface_implementation() {
     );
 }
 
+#[test]
+fn interface_implementations_id_type() {
+    fn check_schema(bar_id: &str, baz_id: &str, ok: bool) {
+        let schema = format!(
+            "interface Foo {{ x: Int }}
+             type Bar implements Foo @entity {{
+                id: {bar_id}!
+                x: Int
+             }}
+
+             type Baz implements Foo @entity {{
+                id: {baz_id}!
+                x: Int
+            }}"
+        );
+        let schema = Schema::parse(&schema, DeploymentHash::new("dummy").unwrap()).unwrap();
+        let res = schema.validate(&HashMap::new());
+        if ok {
+            assert!(matches!(res, Ok(_)));
+        } else {
+            assert!(matches!(res, Err(_)));
+            assert!(matches!(
+                res.unwrap_err()[0],
+                SchemaValidationError::InterfaceImplementorsMixId(_, _)
+            ));
+        }
+    }
+    check_schema("ID", "ID", true);
+    check_schema("ID", "String", true);
+    check_schema("ID", "Bytes", false);
+    check_schema("Bytes", "String", false);
+}
+
 #[test]
 fn test_derived_from_validation() {
     const OTHER_TYPES: &str = "

graph/src/data/store/mod.rs

@@ -1,6 +1,7 @@
 use crate::{
     components::store::{DeploymentLocator, EntityType},
-    prelude::{anyhow::Context, q, r, s, CacheWeight, EntityKey, QueryExecutionError},
+    data::graphql::ObjectTypeExt,
+    prelude::{anyhow::Context, q, r, s, CacheWeight, EntityKey, QueryExecutionError, Schema},
     runtime::gas::{Gas, GasSizeOf},
 };
 use crate::{data::subgraph::DeploymentHash, prelude::EntityChange};
@@ -563,11 +564,15 @@ impl Entity {
         v
     }
 
-    /// Try to get this entity's ID
+    /// Return the ID of this entity. If the ID is a string, return the
+    /// string. If it is `Bytes`, return it as a hex string with a `0x`
+    /// prefix. If the ID is not set or anything but a `String` or `Bytes`,
+    /// return an error
     pub fn id(&self) -> Result<String, Error> {
         match self.get("id") {
             None => Err(anyhow!("Entity is missing an `id` attribute")),
             Some(Value::String(s)) => Ok(s.to_owned()),
+            Some(Value::Bytes(b)) => Ok(b.to_string()),
             _ => Err(anyhow!("Entity has non-string `id` attribute")),
         }
     }
@@ -605,14 +610,58 @@ impl Entity {
     /// Validate that this entity matches the object type definition in the
     /// schema. An entity that passes these checks can be stored
     /// successfully in the subgraph's database schema
-    pub fn validate(&self, schema: &s::Document, key: &EntityKey) -> Result<(), anyhow::Error> {
+    pub fn validate(&self, schema: &Schema, key: &EntityKey) -> Result<(), anyhow::Error> {
+        fn scalar_value_type(schema: &Schema, field_type: &s::Type) -> ValueType {
+            use s::TypeDefinition as t;
+            match field_type {
+                s::Type::NamedType(name) => ValueType::from_str(&name).unwrap_or_else(|_| {
+                    match schema.document.get_named_type(name) {
+                        Some(t::Object(obj_type)) => {
+                            let id = obj_type.field("id").expect("all object types have an id");
+                            scalar_value_type(schema, &id.field_type)
+                        }
+                        Some(t::Interface(intf)) => {
+                            // Validation checks that all implementors of an
+                            // interface use the same type for `id`. It is
+                            // therefore enough to use the id type of one of
+                            // the implementors
+                            match schema
+                                .types_for_interface()
+                                .get(&EntityType::new(intf.name.clone()))
+                                .expect("interface type names are known")
+                                .first()
+                            {
+                                None => {
+                                    // Nothing is implementing this interface; we assume it's of type string
+                                    // see also: id-type-for-unimplemented-interfaces
+                                    ValueType::String
+                                }
+                                Some(obj_type) => {
+                                    let id =
+                                        obj_type.field("id").expect("all object types have an id");
+                                    scalar_value_type(schema, &id.field_type)
+                                }
+                            }
+                        }
+                        Some(t::Enum(_)) => ValueType::String,
+                        Some(t::Scalar(_)) => unreachable!("user-defined scalars are not used"),
+                        Some(t::Union(_)) => unreachable!("unions are not used"),
+                        Some(t::InputObject(_)) => unreachable!("inputObjects are not used"),
+                        None => unreachable!("names of field types have been validated"),
+                    }
+                }),
+                s::Type::NonNullType(inner) => scalar_value_type(schema, inner),
+                s::Type::ListType(inner) => scalar_value_type(schema, inner),
+            }
+        }
+
         if key.entity_type.is_poi() {
             // Users can't modify Poi entities, and therefore they do not
             // need to be validated. In addition, the schema has no object
             // type for them, and validation would therefore fail
             return Ok(());
         }
-        let object_type_definitions = schema.get_object_type_definitions();
+        let object_type_definitions = schema.document.get_object_type_definitions();
         let object_type = object_type_definitions
             .iter()
             .find(|object_type| key.entity_type.as_str() == &object_type.name)
@@ -627,7 +676,7 @@ impl Entity {
             let is_derived = field.is_derived();
             match (self.get(&field.name), is_derived) {
                 (Some(value), false) => {
-                    let scalar_type = schema.scalar_value_type(&field.field_type);
+                    let scalar_type = scalar_value_type(schema, &field.field_type);
                     if field.field_type.is_list() {
                         // Check for inhomgeneous lists to produce a better
                         // error message for them; other problems, like
@@ -815,7 +864,7 @@ fn entity_validation() {
             id.to_owned(),
         );
 
-        let err = thing.validate(&schema.document, &key);
+        let err = thing.validate(&schema, &key);
         if errmsg == "" {
             assert!(
                 err.is_ok(),