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    // Array field getters — extract a Vec<T> from a Value::List, mapping
443    // each inner element through the corresponding scalar getter.
444    pub fn get_string_array(&self, field: &str) -> Vec<String> {
445        match self.fields.get(field) {
446            Some(Value::List(items)) => items
447                .iter()
448                .filter_map(|v| match v {
449                    Value::String(s) => Some(s.clone()),
450                    _ => None,
451                })
452                .collect(),
453            _ => Vec::new(),
454        }
455    }
456
457    pub fn get_string_array_opt(&self, field: &str) -> Option<Vec<String>> {
458        match self.fields.get(field) {
459            Some(Value::List(_)) => Some(self.get_string_array(field)),
460            _ => None,
461        }
462    }
463
464    pub fn get_bytes_array(&self, field: &str) -> Vec<Vec<u8>> {
465        match self.fields.get(field) {
466            Some(Value::List(items)) => items
467                .iter()
468                .filter_map(|v| match v {
469                    Value::Bytes(b) => Some(b.clone()),
470                    _ => None,
471                })
472                .collect(),
473            _ => Vec::new(),
474        }
475    }
476
477    pub fn get_bytes_array_opt(&self, field: &str) -> Option<Vec<Vec<u8>>> {
478        match self.fields.get(field) {
479            Some(Value::List(_)) => Some(self.get_bytes_array(field)),
480            _ => None,
481        }
482    }
483
484    pub fn get_bigint_array(&self, field: &str) -> Vec<BigInt> {
485        match self.fields.get(field) {
486            Some(Value::List(items)) => items
487                .iter()
488                .filter_map(|v| match v {
489                    Value::BigInt(b) => Some(b.clone()),
490                    _ => None,
491                })
492                .collect(),
493            _ => Vec::new(),
494        }
495    }
496
497    pub fn get_bigint_array_opt(&self, field: &str) -> Option<Vec<BigInt>> {
498        match self.fields.get(field) {
499            Some(Value::List(_)) => Some(self.get_bigint_array(field)),
500            _ => None,
501        }
502    }
503
504    pub fn get_bigdecimal_array(&self, field: &str) -> Vec<BigDecimal> {
505        match self.fields.get(field) {
506            Some(Value::List(items)) => items
507                .iter()
508                .filter_map(|v| match v {
509                    Value::BigDecimal(b) => Some(b.clone()),
510                    _ => None,
511                })
512                .collect(),
513            _ => Vec::new(),
514        }
515    }
516
517    pub fn get_bigdecimal_array_opt(&self, field: &str) -> Option<Vec<BigDecimal>> {
518        match self.fields.get(field) {
519            Some(Value::List(_)) => Some(self.get_bigdecimal_array(field)),
520            _ => None,
521        }
522    }
523
524    pub fn get_bool_array(&self, field: &str) -> Vec<bool> {
525        match self.fields.get(field) {
526            Some(Value::List(items)) => items
527                .iter()
528                .filter_map(|v| match v {
529                    Value::Bool(b) => Some(*b),
530                    _ => None,
531                })
532                .collect(),
533            _ => Vec::new(),
534        }
535    }
536
537    pub fn get_bool_array_opt(&self, field: &str) -> Option<Vec<bool>> {
538        match self.fields.get(field) {
539            Some(Value::List(_)) => Some(self.get_bool_array(field)),
540            _ => None,
541        }
542    }
543
544    pub fn get_i32_array(&self, field: &str) -> Vec<i32> {
545        match self.fields.get(field) {
546            Some(Value::List(items)) => items
547                .iter()
548                .filter_map(|v| match v {
549                    Value::Int(v) => Some(*v),
550                    _ => None,
551                })
552                .collect(),
553            _ => Vec::new(),
554        }
555    }
556
557    pub fn get_i32_array_opt(&self, field: &str) -> Option<Vec<i32>> {
558        match self.fields.get(field) {
559            Some(Value::List(_)) => Some(self.get_i32_array(field)),
560            _ => None,
561        }
562    }
563
564    pub fn fields(&self) -> &HashMap<String, Value> {
565        &self.fields
566    }
567}
568
569pub struct EntityCache {
570    entities: HashMap<String, HashMap<String, Entity>>,
571    removals: Vec<(String, String)>,
572}
573
574impl EntityCache {
575    pub fn new() -> Self {
576        Self {
577            entities: HashMap::new(),
578            removals: Vec::new(),
579        }
580    }
581
582    pub fn get(&self, entity_type: &str, id: &str) -> Option<Entity> {
583        self.entities
584            .get(entity_type)
585            .and_then(|m| m.get(id))
586            .cloned()
587    }
588
589    pub fn set(&mut self, entity_type: &str, id: &str, entity: Entity) {
590        self.entities
591            .entry(entity_type.to_string())
592            .or_insert_with(HashMap::new)
593            .insert(id.to_string(), entity);
594    }
595
596    pub fn remove(&mut self, entity_type: &str, id: &str) {
597        if let Some(m) = self.entities.get_mut(entity_type) {
598            m.remove(id);
599        }
600        let key = (entity_type.to_string(), id.to_string());
601        if !self.removals.contains(&key) {
602            self.removals.push(key);
603        }
604    }
605
606    pub fn get_all(&self, entity_type: &str) -> Vec<&Entity> {
607        self.entities
608            .get(entity_type)
609            .map(|m| m.values().collect())
610            .unwrap_or_default()
611    }
612}
613
614/// DataSourceContext — per-instance configuration passed via createWithContext.
615/// Mirrors graph-ts DataSourceContext: typed key-value store.
616#[derive(Debug, Clone)]
617pub struct DataSourceContext {
618    fields: HashMap<String, Value>,
619}
620
621impl DataSourceContext {
622    /// Create a new empty context (for write-side construction).
623    pub fn empty() -> Self {
624        Self { fields: HashMap::new() }
625    }
626
627    /// Create a context from an existing field map (for runtime population).
628    pub fn new(fields: HashMap<String, Value>) -> Self {
629        Self { fields }
630    }
631
632    /// Raw getter — returns Option for presence checks.
633    pub fn get(&self, key: &str) -> Option<&Value> {
634        self.fields.get(key)
635    }
636
637    /// Check if a key is present in the context.
638    pub fn is_set(&self, key: &str) -> bool {
639        self.fields.contains_key(key)
640    }
641
642    /// Typed getter — panics if key absent (matches graph-ts behavior).
643    pub fn get_string(&self, key: &str) -> String {
644        match self.fields.get(key) {
645            Some(Value::String(s)) => s.clone(),
646            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
647        }
648    }
649
650    pub fn get_bytes(&self, key: &str) -> Vec<u8> {
651        match self.fields.get(key) {
652            Some(Value::Bytes(b)) => b.clone(),
653            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
654        }
655    }
656
657    pub fn get_big_int(&self, key: &str) -> BigInt {
658        match self.fields.get(key) {
659            Some(Value::BigInt(n)) => n.clone(),
660            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
661        }
662    }
663
664    pub fn get_i32(&self, key: &str) -> i32 {
665        match self.fields.get(key) {
666            Some(&Value::Int(n)) => n,
667            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
668        }
669    }
670
671    pub fn get_boolean(&self, key: &str) -> bool {
672        match self.fields.get(key) {
673            Some(&Value::Bool(b)) => b,
674            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
675        }
676    }
677
678    /// Typed setters — for constructing context before createWithContext.
679    pub fn set_string(&mut self, key: &str, value: impl Into<String>) {
680        self.fields.insert(key.to_string(), Value::String(value.into()));
681    }
682
683    pub fn set_bytes(&mut self, key: &str, value: Vec<u8>) {
684        self.fields.insert(key.to_string(), Value::Bytes(value));
685    }
686
687    pub fn set_big_int(&mut self, key: &str, value: BigInt) {
688        self.fields.insert(key.to_string(), Value::BigInt(value));
689    }
690
691    pub fn set_i32(&mut self, key: &str, value: i32) {
692        self.fields.insert(key.to_string(), Value::Int(value));
693    }
694
695    pub fn set_boolean(&mut self, key: &str, value: bool) {
696        self.fields.insert(key.to_string(), Value::Bool(value));
697    }
698
699    /// Access the underlying fields (for runtime conversion).
700    pub fn into_fields(self) -> HashMap<String, Value> {
701        self.fields
702    }
703}
704
705#[async_trait]
706pub trait HandlerContext: Send {
707    async fn load_entity(&mut self, entity_type: &str, id: &str) -> Option<Entity>;
708
709    async fn save_entity(&mut self, entity_type: &str, entity: &Entity) -> Result<(), Error>;
710
711    fn entity_remove(&mut self, entity_type: &str, id: &str) -> Result<(), Error>;
712
713    /// Returns the current data source's network name.
714    fn data_source_network(&self) -> String;
715
716    /// Returns the current data source's context (set via createWithContext).
717    /// Returns an empty context if none was provided.
718    fn data_source_context(&self) -> DataSourceContext;
719
720    /// Mirrors graph-ts `dataSource.stringParam()`.
721    /// Returns the UTF-8 string interpretation of the datasource address/source bytes.
722    /// For file/offchain datasources, this is the source identifier (e.g., content hash).
723    /// For onchain templates, the result depends on how graph-node populates
724    /// `dataSource.address()` — typically the hex-decoded contract address bytes
725    /// interpreted as UTF-8, which may not produce a meaningful string.
726    fn data_source_string_param(&self) -> String;
727
728    /// Create a dynamic data source from a template.
729    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
730    fn data_source_create(&mut self, template: &str, params: &[String]) -> Result<(), Error>;
731
732    /// Create a dynamic data source from a template with context.
733    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
734    fn data_source_create_with_context(
735        &mut self,
736        template: &str,
737        params: &[String],
738        context: DataSourceContext,
739    ) -> Result<(), Error>;
740}
741
742#[cfg(test)]
743mod tests {
744    use super::*;
745
746    #[test]
747    fn option_some_into_value_delegates_to_inner() {
748        let v: Value = Some(BigInt::from(1)).into_value();
749        match v {
750            Value::BigInt(b) => assert_eq!(b, BigInt::from(1)),
751            other => panic!("expected Value::BigInt, got {:?}", other),
752        }
753    }
754
755    #[test]
756    fn option_none_into_value_is_null() {
757        let v: Value = None::<BigInt>.into_value();
758        assert!(matches!(v, Value::Null), "expected Value::Null, got {:?}", v);
759    }
760
761    #[test]
762    fn vec_of_t_into_value_maps_to_list() {
763        let xs: Vec<BigInt> = vec![BigInt::from(1), BigInt::from(2), BigInt::from(3)];
764        let v: Value = xs.into_value();
765        match v {
766            Value::List(items) => {
767                assert_eq!(items.len(), 3);
768                assert!(matches!(items[0], Value::BigInt(_)));
769            }
770            other => panic!("expected Value::List, got {:?}", other),
771        }
772    }
773
774    #[test]
775    fn vec_of_bytes_into_value_maps_to_list_of_bytes() {
776        // Verifies the blanket Vec<T> impl does not collide with the
777        // concrete Vec<u8> impl: Vec<Vec<u8>> goes through the blanket
778        // (each inner Vec<u8> uses the concrete impl), producing a list
779        // of byte values rather than a single byte blob.
780        let xs: Vec<Vec<u8>> = vec![vec![0x01, 0x02], vec![0x03]];
781        let v: Value = xs.into_value();
782        match v {
783            Value::List(items) => {
784                assert_eq!(items.len(), 2);
785                assert!(matches!(items[0], Value::Bytes(_)));
786            }
787            other => panic!("expected Value::List, got {:?}", other),
788        }
789    }
790
791    #[test]
792    fn concrete_vec_u8_still_routes_to_bytes_variant() {
793        let v: Value = vec![0x01u8, 0x02, 0x03].into_value();
794        assert!(matches!(v, Value::Bytes(_)), "Vec<u8> must keep using the concrete Bytes impl, got {:?}", v);
795    }
796}