graph, runtime: Autogenerate ids

graph/src/components/store/entity_cache.rs

@@ -6,7 +6,7 @@ use std::sync::Arc;
 
 use crate::components::store::write::EntityModification;
 use crate::components::store::{self as s, Entity, EntityOperation};
-use crate::data::store::{EntityValidationError, IntoEntityIterator};
+use crate::data::store::{EntityValidationError, Id, IdType, IntoEntityIterator};
 use crate::prelude::ENV_VARS;
 use crate::schema::{EntityKey, InputSchema};
 use crate::util::intern::Error as InternError;
@@ -86,6 +86,11 @@ pub struct EntityCache {
     pub store: Arc<dyn s::ReadStore>,
 
     pub schema: InputSchema,
+
+    /// A sequence number for generating entity IDs. We use one number for
+    /// all id's as the id's are scoped by block and a u32 has plenty of
+    /// room for all changes in one block.
+    seq: u32,
 }
 
 impl Debug for EntityCache {
@@ -112,6 +117,7 @@ impl EntityCache {
             in_handler: false,
             schema: store.input_schema(),
             store,
+            seq: 0,
         }
     }
 
@@ -134,6 +140,7 @@ impl EntityCache {
             in_handler: false,
             schema: store.input_schema(),
             store,
+            seq: 0,
         }
     }
 
@@ -384,6 +391,13 @@ impl EntityCache {
         }
     }
 
+    /// Generate an id.
+    pub fn generate_id(&mut self, id_type: IdType, block: BlockNumber) -> anyhow::Result<Id> {
+        let id = id_type.generate_id(block, self.seq)?;
+        self.seq += 1;
+        Ok(id)
+    }
+
     /// Return the changes that have been made via `set` and `remove` as
     /// `EntityModification`, making sure to only produce one when a change
     /// to the current state is actually needed.

graph/src/data/store/id.rs

@@ -6,6 +6,8 @@ use std::convert::TryFrom;
 use std::fmt;
 
 use crate::{
+    anyhow, bail,
+    components::store::BlockNumber,
     data::graphql::{ObjectTypeExt, TypeExt},
     prelude::s,
 };
@@ -46,6 +48,31 @@ impl IdType {
             IdType::Int8 => "Int8",
         }
     }
+
+    /// Generate an entity id from the block number and a sequence number.
+    ///
+    /// * Bytes: `[block:4, seq:4]`
+    /// * Int8: `[block:4, seq:4]`
+    /// * String: Always an error; users should use `Bytes` or `Int8`
+    ///   instead
+    pub fn generate_id(&self, block: BlockNumber, seq: u32) -> anyhow::Result<Id> {
+        match self {
+            IdType::String => bail!("String does not support generating ids"),
+            IdType::Bytes => {
+                let mut bytes = [0u8; 8];
+                bytes[0..4].copy_from_slice(&block.to_be_bytes());
+                bytes[4..8].copy_from_slice(&seq.to_be_bytes());
+                let bytes = scalar::Bytes::from(bytes);
+                Ok(Id::Bytes(bytes))
+            }
+            IdType::Int8 => {
+                let mut bytes = [0u8; 8];
+                bytes[0..4].copy_from_slice(&seq.to_le_bytes());
+                bytes[4..8].copy_from_slice(&block.to_le_bytes());
+                Ok(Id::Int8(i64::from_le_bytes(bytes)))
+            }
+        }
+    }
 }
 
 impl<'a> TryFrom<&s::ObjectType> for IdType {
@@ -490,3 +517,34 @@ impl IdList {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::data::store::{Id, IdType};
+
+    #[test]
+    fn generate_id() {
+        let id = IdType::Bytes.generate_id(1, 2).unwrap();
+        let exp = IdType::Bytes.parse("0x0000000100000002".into()).unwrap();
+        assert_eq!(exp, id);
+
+        let id = IdType::Bytes.generate_id(3, 2).unwrap();
+        let exp = IdType::Bytes.parse("0x0000000300000002".into()).unwrap();
+        assert_eq!(exp, id);
+
+        let id = IdType::Int8.generate_id(3, 2).unwrap();
+        let exp = Id::Int8(0x0000_0003__0000_0002);
+        assert_eq!(exp, id);
+
+        // Should be id + 1
+        let id2 = IdType::Int8.generate_id(3, 3).unwrap();
+        let d = id2.to_string().parse::<i64>().unwrap() - id.to_string().parse::<i64>().unwrap();
+        assert_eq!(1, d);
+        // Should be id + 2^32
+        let id3 = IdType::Int8.generate_id(4, 2).unwrap();
+        let d = id3.to_string().parse::<i64>().unwrap() - id.to_string().parse::<i64>().unwrap();
+        assert_eq!(1 << 32, d);
+
+        IdType::String.generate_id(3, 2).unwrap_err();
+    }
+}

runtime/test/src/test.rs

@@ -1,5 +1,5 @@
 use graph::components::metrics::gas::GasMetrics;
-use graph::data::store::scalar;
+use graph::data::store::{scalar, Id, IdType};
 use graph::data::subgraph::*;
 use graph::data::value::Word;
 use graph::prelude::web3::types::U256;
@@ -1510,3 +1510,66 @@ async fn test_store_set_invalid_fields() {
 
     assert!(err_is_none);
 }
+
+/// Test generating ids through `store_set`
+#[tokio::test]
+async fn generate_id() {
+    const AUTO: &str = "auto";
+    const INT8: &str = "Int8";
+    const BINARY: &str = "Binary";
+
+    let schema = "type Int8 @entity(immutable: true) {
+        id: Int8!,
+        name: String,
+    }
+
+    type Binary @entity(immutable: true) {
+        id: Bytes!,
+        name: String,
+    }";
+
+    let mut host = Host::new(schema, "hostGenerateId", "boolean.wasm", None).await;
+
+    // Since these entities are immutable, storing twice would generate an
+    // error; but since the ids are autogenerated, each invocation creates a
+    // new id. Note that the types of the ids have an incorrect type, but
+    // that doesn't matter since they get overwritten.
+    host.store_set(INT8, AUTO, vec![("id", "u1"), ("name", "int1")])
+        .expect("setting auto works");
+    host.store_set(INT8, AUTO, vec![("id", "u1"), ("name", "int2")])
+        .expect("setting auto works");
+    host.store_set(BINARY, AUTO, vec![("id", "u1"), ("name", "bin1")])
+        .expect("setting auto works");
+    host.store_set(BINARY, AUTO, vec![("id", "u1"), ("name", "bin2")])
+        .expect("setting auto works");
+
+    let entity_cache = host.ctx.state.entity_cache;
+    let mods = entity_cache.as_modifications(12).unwrap().modifications;
+    let id_map: HashMap<&str, Id> = HashMap::from_iter(
+        vec![
+            (
+                "bin1",
+                IdType::Bytes.parse("0x0000000c00000002".into()).unwrap(),
+            ),
+            (
+                "bin2",
+                IdType::Bytes.parse("0x0000000c00000003".into()).unwrap(),
+            ),
+            ("int1", Id::Int8(0x0000_000c__0000_0000)),
+            ("int2", Id::Int8(0x0000_000c__0000_0001)),
+        ]
+        .into_iter(),
+    );
+    assert_eq!(4, mods.len());
+    for m in &mods {
+        match m {
+            EntityModification::Insert { data, .. } => {
+                let id = data.get("id").unwrap();
+                let name = data.get("name").unwrap().as_str().unwrap();
+                let exp = id_map.get(name).unwrap();
+                assert_eq!(exp, id, "Wrong id for entity with name `{name}`");
+            }
+            _ => panic!("expected Insert modification"),
+        }
+    }
+}

runtime/wasm/src/host_exports.rs

@@ -204,7 +204,7 @@ impl<C: Blockchain> HostExports<C> {
         &self,
         logger: &Logger,
         state: &mut BlockState<C>,
-        _block: BlockNumber,
+        block: BlockNumber,
         proof_of_indexing: &SharedProofOfIndexing,
         entity_type: String,
         entity_id: String,
@@ -213,6 +213,22 @@ impl<C: Blockchain> HostExports<C> {
         gas: &GasCounter,
     ) -> Result<(), HostExportError> {
         let entity_type = state.entity_cache.schema.entity_type(&entity_type)?;
+
+        let entity_id = if entity_id == "auto" {
+            if self.data_source_causality_region != CausalityRegion::ONCHAIN {
+                return Err(anyhow!(
+                    "Autogenerated IDs are only supported for onchain data sources"
+                )
+                .into());
+            }
+            let id_type = entity_type.id_type()?;
+            let id = state.entity_cache.generate_id(id_type, block)?;
+            data.insert(store::ID.clone(), id.clone().into());
+            id.to_string()
+        } else {
+            entity_id
+        };
+
         let key = entity_type.parse_key_in(entity_id, self.data_source_causality_region)?;
         self.check_entity_type_access(&key.entity_type)?;