Skip to main content

graph_native/
entity.rs

1use crate::types::{Address, BigDecimal, BigInt};
2use anyhow::Error;
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8pub enum Value {
9    String(String),
10    Int(i32),
11    Int8(i64),
12    Timestamp(i64),
13    BigDecimal(BigDecimal),
14    Bool(bool),
15    List(Vec<Value>),
16    Null,
17    Bytes(Vec<u8>),
18    BigInt(BigInt),
19}
20
21impl From<&str> for Value {
22    fn from(value: &str) -> Self {
23        Self::String(value.to_string())
24    }
25}
26
27impl From<String> for Value {
28    fn from(value: String) -> Self {
29        Self::String(value)
30    }
31}
32
33impl From<i32> for Value {
34    fn from(value: i32) -> Self {
35        Self::Int(value)
36    }
37}
38
39impl From<i64> for Value {
40    fn from(value: i64) -> Self {
41        Self::Int8(value)
42    }
43}
44
45impl From<u64> for Value {
46    fn from(value: u64) -> Self {
47        Self::BigInt(value.into())
48    }
49}
50
51impl From<bool> for Value {
52    fn from(value: bool) -> Self {
53        Self::Bool(value)
54    }
55}
56
57impl From<BigInt> for Value {
58    fn from(value: BigInt) -> Self {
59        Self::BigInt(value)
60    }
61}
62
63impl From<BigDecimal> for Value {
64    fn from(value: BigDecimal) -> Self {
65        Self::BigDecimal(value)
66    }
67}
68
69impl From<Address> for Value {
70    fn from(value: Address) -> Self {
71        Self::Bytes(value.to_bytes())
72    }
73}
74
75pub fn bytes_to_value(v: &[u8]) -> Value {
76    Value::Bytes(v.to_vec())
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
80pub enum EntityIdValue {
81    String(String),
82    Bytes(Vec<u8>),
83}
84
85impl EntityIdValue {
86    fn into_value(self) -> Value {
87        match self {
88            Self::String(value) => Value::String(value),
89            Self::Bytes(value) => Value::Bytes(value),
90        }
91    }
92}
93
94pub fn bytes_to_hex_id(bytes: &[u8]) -> String {
95    let mut id = String::with_capacity(2 + bytes.len() * 2);
96    id.push_str("0x");
97    for byte in bytes {
98        id.push(nibble_to_hex(byte >> 4));
99        id.push(nibble_to_hex(byte & 0x0f));
100    }
101    id
102}
103
104pub fn hex_id_to_bytes(id: &str) -> Vec<u8> {
105    let hex = id.strip_prefix("0x").unwrap_or(id);
106    if hex.len() % 2 != 0 {
107        return Vec::new();
108    }
109
110    let mut bytes = Vec::with_capacity(hex.len() / 2);
111    let chars: Vec<char> = hex.chars().collect();
112    let mut i = 0;
113    while i < chars.len() {
114        let high = match hex_to_nibble(chars[i]) {
115            Some(value) => value,
116            None => return Vec::new(),
117        };
118        let low = match hex_to_nibble(chars[i + 1]) {
119            Some(value) => value,
120            None => return Vec::new(),
121        };
122        bytes.push((high << 4) | low);
123        i += 2;
124    }
125
126    bytes
127}
128
129fn hex_to_nibble(ch: char) -> Option<u8> {
130    ch.to_digit(16).map(|value| value as u8)
131}
132
133fn nibble_to_hex(nibble: u8) -> char {
134    match nibble {
135        0..=9 => (b'0' + nibble) as char,
136        10..=15 => (b'a' + nibble - 10) as char,
137        _ => '0',
138    }
139}
140
141pub trait IntoEntityValue {
142    fn into_value(self) -> Value;
143}
144
145impl IntoEntityValue for String {
146    fn into_value(self) -> Value {
147        Value::String(self)
148    }
149}
150
151impl IntoEntityValue for &str {
152    fn into_value(self) -> Value {
153        Value::String(self.to_string())
154    }
155}
156
157impl IntoEntityValue for i32 {
158    fn into_value(self) -> Value {
159        Value::Int(self)
160    }
161}
162
163impl IntoEntityValue for u64 {
164    fn into_value(self) -> Value {
165        Value::from(self)
166    }
167}
168
169impl IntoEntityValue for bool {
170    fn into_value(self) -> Value {
171        Value::Bool(self)
172    }
173}
174
175impl IntoEntityValue for BigInt {
176    fn into_value(self) -> Value {
177        Value::BigInt(self)
178    }
179}
180
181impl IntoEntityValue for BigDecimal {
182    fn into_value(self) -> Value {
183        Value::BigDecimal(self)
184    }
185}
186
187impl IntoEntityValue for Address {
188    fn into_value(self) -> Value {
189        Value::Bytes(self.to_bytes())
190    }
191}
192
193impl IntoEntityValue for Vec<u8> {
194    fn into_value(self) -> Value {
195        Value::Bytes(self)
196    }
197}
198
199impl IntoEntityValue for &BigInt {
200    fn into_value(self) -> Value {
201        Value::BigInt(self.clone())
202    }
203}
204
205impl IntoEntityValue for &BigDecimal {
206    fn into_value(self) -> Value {
207        Value::BigDecimal(self.clone())
208    }
209}
210
211impl IntoEntityValue for &Address {
212    fn into_value(self) -> Value {
213        Value::Bytes(self.to_bytes())
214    }
215}
216
217impl IntoEntityValue for &Vec<u8> {
218    fn into_value(self) -> Value {
219        Value::Bytes(self.clone())
220    }
221}
222
223impl<T: IntoEntityValue> IntoEntityValue for Option<T> {
224    fn into_value(self) -> Value {
225        match self {
226            Some(v) => v.into_value(),
227            None => Value::Null,
228        }
229    }
230}
231
232impl<T: IntoEntityValue> IntoEntityValue for Vec<T> {
233    fn into_value(self) -> Value {
234        Value::List(self.into_iter().map(|v| v.into_value()).collect())
235    }
236}
237
238pub fn read_bigint(entity: &Option<Entity>, field: &str) -> BigInt {
239    entity
240        .as_ref()
241        .and_then(|e| e.get(field))
242        .map(|v| match v {
243            Value::BigInt(b) => b.clone(),
244            _ => BigInt::zero(),
245        })
246        .unwrap_or_else(BigInt::zero)
247}
248
249pub fn read_bigdecimal(entity: &Option<Entity>, field: &str) -> BigDecimal {
250    entity
251        .as_ref()
252        .and_then(|e| e.get(field))
253        .map(|v| match v {
254            Value::BigDecimal(b) => b.clone(),
255            _ => BigDecimal::zero(),
256        })
257        .unwrap_or_else(BigDecimal::zero)
258}
259
260pub fn read_string(entity: &Option<Entity>, field: &str) -> String {
261    entity
262        .as_ref()
263        .and_then(|e| e.get(field))
264        .map(|v| match v {
265            Value::String(s) => s.clone(),
266            _ => String::new(),
267        })
268        .unwrap_or_default()
269}
270
271pub fn read_bytes(entity: &Option<Entity>, field: &str) -> Vec<u8> {
272    entity
273        .as_ref()
274        .and_then(|e| e.get(field))
275        .map(|v| match v {
276            Value::Bytes(b) => b.clone(),
277            _ => Vec::new(),
278        })
279        .unwrap_or_default()
280}
281
282pub fn read_bool(entity: &Option<Entity>, field: &str) -> bool {
283    entity
284        .as_ref()
285        .and_then(|e| e.get(field))
286        .map(|v| match v {
287            Value::Bool(b) => *b,
288            _ => false,
289        })
290        .unwrap_or(false)
291}
292
293pub fn read_i32(entity: &Option<Entity>, field: &str) -> i32 {
294    entity
295        .as_ref()
296        .and_then(|e| e.get(field))
297        .map(|v| match v {
298            Value::Int(i) => *i,
299            _ => 0,
300        })
301        .unwrap_or(0)
302}
303
304pub fn read_value(entity: &Option<Entity>, field: &str) -> Value {
305    entity
306        .as_ref()
307        .and_then(|e| e.get(field))
308        .cloned()
309        .unwrap_or(Value::Null)
310}
311
312#[derive(Debug, Clone)]
313pub struct Entity {
314    entity_type: String,
315    id: String,
316    fields: HashMap<String, Value>,
317}
318
319impl Entity {
320    pub fn new(entity_type: &str, id: &str) -> Self {
321        Self::new_with_id(entity_type, id, EntityIdValue::String(id.to_string()))
322    }
323
324    pub fn new_with_id(entity_type: &str, key: &str, id: EntityIdValue) -> Self {
325        let mut fields = HashMap::new();
326        fields.insert("id".to_string(), id.into_value());
327        Self {
328            entity_type: entity_type.to_string(),
329            id: key.to_string(),
330            fields,
331        }
332    }
333
334    pub fn from_fields(entity_type: &str, id: &str, fields: HashMap<String, Value>) -> Self {
335        Self {
336            entity_type: entity_type.to_string(),
337            id: id.to_string(),
338            fields,
339        }
340    }
341
342    pub fn id(&self) -> &str {
343        &self.id
344    }
345
346    pub fn entity_type(&self) -> &str {
347        &self.entity_type
348    }
349
350    pub fn set<V: IntoEntityValue>(&mut self, field: &str, value: V) {
351        self.fields.insert(field.to_string(), value.into_value());
352    }
353
354    pub fn get(&self, field: &str) -> Option<&Value> {
355        self.fields.get(field)
356    }
357
358    pub fn get_bigint(&self, field: &str) -> BigInt {
359        match self.fields.get(field) {
360            Some(Value::BigInt(v)) => v.clone(),
361            _ => BigInt::zero(),
362        }
363    }
364
365    pub fn get_bigint_opt(&self, field: &str) -> Option<BigInt> {
366        match self.fields.get(field) {
367            Some(Value::BigInt(v)) => Some(v.clone()),
368            _ => None,
369        }
370    }
371
372    pub fn get_bigdecimal(&self, field: &str) -> BigDecimal {
373        match self.fields.get(field) {
374            Some(Value::BigDecimal(v)) => v.clone(),
375            _ => BigDecimal::zero(),
376        }
377    }
378
379    pub fn get_bigdecimal_opt(&self, field: &str) -> Option<BigDecimal> {
380        match self.fields.get(field) {
381            Some(Value::BigDecimal(v)) => Some(v.clone()),
382            _ => None,
383        }
384    }
385
386    pub fn get_string(&self, field: &str) -> String {
387        match self.fields.get(field) {
388            Some(Value::String(v)) => v.clone(),
389            _ => String::new(),
390        }
391    }
392
393    pub fn get_string_opt(&self, field: &str) -> Option<String> {
394        match self.fields.get(field) {
395            Some(Value::String(v)) => Some(v.clone()),
396            _ => None,
397        }
398    }
399
400    pub fn get_bytes(&self, field: &str) -> Vec<u8> {
401        match self.fields.get(field) {
402            Some(Value::Bytes(v)) => v.clone(),
403            _ => Vec::new(),
404        }
405    }
406
407    pub fn get_bytes_opt(&self, field: &str) -> Option<Vec<u8>> {
408        match self.fields.get(field) {
409            Some(Value::Bytes(v)) => Some(v.clone()),
410            _ => None,
411        }
412    }
413
414    pub fn get_bool(&self, field: &str) -> bool {
415        match self.fields.get(field) {
416            Some(Value::Bool(v)) => *v,
417            _ => false,
418        }
419    }
420
421    pub fn get_bool_opt(&self, field: &str) -> Option<bool> {
422        match self.fields.get(field) {
423            Some(Value::Bool(v)) => Some(*v),
424            _ => None,
425        }
426    }
427
428    pub fn get_i32(&self, field: &str) -> i32 {
429        match self.fields.get(field) {
430            Some(Value::Int(v)) => *v,
431            _ => 0,
432        }
433    }
434
435    pub fn get_i32_opt(&self, field: &str) -> Option<i32> {
436        match self.fields.get(field) {
437            Some(Value::Int(v)) => Some(*v),
438            _ => None,
439        }
440    }
441
442    pub fn fields(&self) -> &HashMap<String, Value> {
443        &self.fields
444    }
445}
446
447pub struct EntityCache {
448    entities: HashMap<String, HashMap<String, Entity>>,
449    removals: Vec<(String, String)>,
450}
451
452impl EntityCache {
453    pub fn new() -> Self {
454        Self {
455            entities: HashMap::new(),
456            removals: Vec::new(),
457        }
458    }
459
460    pub fn get(&self, entity_type: &str, id: &str) -> Option<Entity> {
461        self.entities
462            .get(entity_type)
463            .and_then(|m| m.get(id))
464            .cloned()
465    }
466
467    pub fn set(&mut self, entity_type: &str, id: &str, entity: Entity) {
468        self.entities
469            .entry(entity_type.to_string())
470            .or_insert_with(HashMap::new)
471            .insert(id.to_string(), entity);
472    }
473
474    pub fn remove(&mut self, entity_type: &str, id: &str) {
475        if let Some(m) = self.entities.get_mut(entity_type) {
476            m.remove(id);
477        }
478        let key = (entity_type.to_string(), id.to_string());
479        if !self.removals.contains(&key) {
480            self.removals.push(key);
481        }
482    }
483
484    pub fn get_all(&self, entity_type: &str) -> Vec<&Entity> {
485        self.entities
486            .get(entity_type)
487            .map(|m| m.values().collect())
488            .unwrap_or_default()
489    }
490}
491
492/// DataSourceContext — per-instance configuration passed via createWithContext.
493/// Mirrors graph-ts DataSourceContext: typed key-value store.
494#[derive(Debug, Clone)]
495pub struct DataSourceContext {
496    fields: HashMap<String, Value>,
497}
498
499impl DataSourceContext {
500    /// Create a new empty context (for write-side construction).
501    pub fn empty() -> Self {
502        Self { fields: HashMap::new() }
503    }
504
505    /// Create a context from an existing field map (for runtime population).
506    pub fn new(fields: HashMap<String, Value>) -> Self {
507        Self { fields }
508    }
509
510    /// Raw getter — returns Option for presence checks.
511    pub fn get(&self, key: &str) -> Option<&Value> {
512        self.fields.get(key)
513    }
514
515    /// Check if a key is present in the context.
516    pub fn is_set(&self, key: &str) -> bool {
517        self.fields.contains_key(key)
518    }
519
520    /// Typed getter — panics if key absent (matches graph-ts behavior).
521    pub fn get_string(&self, key: &str) -> String {
522        match self.fields.get(key) {
523            Some(Value::String(s)) => s.clone(),
524            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
525        }
526    }
527
528    pub fn get_bytes(&self, key: &str) -> Vec<u8> {
529        match self.fields.get(key) {
530            Some(Value::Bytes(b)) => b.clone(),
531            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
532        }
533    }
534
535    pub fn get_big_int(&self, key: &str) -> BigInt {
536        match self.fields.get(key) {
537            Some(Value::BigInt(n)) => n.clone(),
538            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
539        }
540    }
541
542    pub fn get_i32(&self, key: &str) -> i32 {
543        match self.fields.get(key) {
544            Some(&Value::Int(n)) => n,
545            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
546        }
547    }
548
549    pub fn get_boolean(&self, key: &str) -> bool {
550        match self.fields.get(key) {
551            Some(&Value::Bool(b)) => b,
552            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
553        }
554    }
555
556    /// Typed setters — for constructing context before createWithContext.
557    pub fn set_string(&mut self, key: &str, value: impl Into<String>) {
558        self.fields.insert(key.to_string(), Value::String(value.into()));
559    }
560
561    pub fn set_bytes(&mut self, key: &str, value: Vec<u8>) {
562        self.fields.insert(key.to_string(), Value::Bytes(value));
563    }
564
565    pub fn set_big_int(&mut self, key: &str, value: BigInt) {
566        self.fields.insert(key.to_string(), Value::BigInt(value));
567    }
568
569    pub fn set_i32(&mut self, key: &str, value: i32) {
570        self.fields.insert(key.to_string(), Value::Int(value));
571    }
572
573    pub fn set_boolean(&mut self, key: &str, value: bool) {
574        self.fields.insert(key.to_string(), Value::Bool(value));
575    }
576
577    /// Access the underlying fields (for runtime conversion).
578    pub fn into_fields(self) -> HashMap<String, Value> {
579        self.fields
580    }
581}
582
583#[async_trait]
584pub trait HandlerContext: Send {
585    async fn load_entity(&mut self, entity_type: &str, id: &str) -> Option<Entity>;
586
587    async fn save_entity(&mut self, entity_type: &str, entity: &Entity) -> Result<(), Error>;
588
589    fn entity_remove(&mut self, entity_type: &str, id: &str) -> Result<(), Error>;
590
591    /// Returns the current data source's network name.
592    fn data_source_network(&self) -> String;
593
594    /// Returns the current data source's context (set via createWithContext).
595    /// Returns an empty context if none was provided.
596    fn data_source_context(&self) -> DataSourceContext;
597
598    /// Mirrors graph-ts `dataSource.stringParam()`.
599    /// Returns the UTF-8 string interpretation of the datasource address/source bytes.
600    /// For file/offchain datasources, this is the source identifier (e.g., content hash).
601    /// For onchain templates, the result depends on how graph-node populates
602    /// `dataSource.address()` — typically the hex-decoded contract address bytes
603    /// interpreted as UTF-8, which may not produce a meaningful string.
604    fn data_source_string_param(&self) -> String;
605
606    /// Create a dynamic data source from a template.
607    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
608    fn data_source_create(&mut self, template: &str, params: &[String]) -> Result<(), Error>;
609
610    /// Create a dynamic data source from a template with context.
611    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
612    fn data_source_create_with_context(
613        &mut self,
614        template: &str,
615        params: &[String],
616        context: DataSourceContext,
617    ) -> Result<(), Error>;
618}
619
620#[cfg(test)]
621mod tests {
622    use super::*;
623
624    #[test]
625    fn option_some_into_value_delegates_to_inner() {
626        let v: Value = Some(BigInt::from(1)).into_value();
627        match v {
628            Value::BigInt(b) => assert_eq!(b, BigInt::from(1)),
629            other => panic!("expected Value::BigInt, got {:?}", other),
630        }
631    }
632
633    #[test]
634    fn option_none_into_value_is_null() {
635        let v: Value = None::<BigInt>.into_value();
636        assert!(matches!(v, Value::Null), "expected Value::Null, got {:?}", v);
637    }
638
639    #[test]
640    fn vec_of_t_into_value_maps_to_list() {
641        let xs: Vec<BigInt> = vec![BigInt::from(1), BigInt::from(2), BigInt::from(3)];
642        let v: Value = xs.into_value();
643        match v {
644            Value::List(items) => {
645                assert_eq!(items.len(), 3);
646                assert!(matches!(items[0], Value::BigInt(_)));
647            }
648            other => panic!("expected Value::List, got {:?}", other),
649        }
650    }
651
652    #[test]
653    fn vec_of_bytes_into_value_maps_to_list_of_bytes() {
654        // Verifies the blanket Vec<T> impl does not collide with the
655        // concrete Vec<u8> impl: Vec<Vec<u8>> goes through the blanket
656        // (each inner Vec<u8> uses the concrete impl), producing a list
657        // of byte values rather than a single byte blob.
658        let xs: Vec<Vec<u8>> = vec![vec![0x01, 0x02], vec![0x03]];
659        let v: Value = xs.into_value();
660        match v {
661            Value::List(items) => {
662                assert_eq!(items.len(), 2);
663                assert!(matches!(items[0], Value::Bytes(_)));
664            }
665            other => panic!("expected Value::List, got {:?}", other),
666        }
667    }
668
669    #[test]
670    fn concrete_vec_u8_still_routes_to_bytes_variant() {
671        let v: Value = vec![0x01u8, 0x02, 0x03].into_value();
672        assert!(matches!(v, Value::Bytes(_)), "Vec<u8> must keep using the concrete Bytes impl, got {:?}", v);
673    }
674}