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
232pub fn read_bigint(entity: &Option<Entity>, field: &str) -> BigInt {
233    entity
234        .as_ref()
235        .and_then(|e| e.get(field))
236        .map(|v| match v {
237            Value::BigInt(b) => b.clone(),
238            _ => BigInt::zero(),
239        })
240        .unwrap_or_else(BigInt::zero)
241}
242
243pub fn read_bigdecimal(entity: &Option<Entity>, field: &str) -> BigDecimal {
244    entity
245        .as_ref()
246        .and_then(|e| e.get(field))
247        .map(|v| match v {
248            Value::BigDecimal(b) => b.clone(),
249            _ => BigDecimal::zero(),
250        })
251        .unwrap_or_else(BigDecimal::zero)
252}
253
254pub fn read_string(entity: &Option<Entity>, field: &str) -> String {
255    entity
256        .as_ref()
257        .and_then(|e| e.get(field))
258        .map(|v| match v {
259            Value::String(s) => s.clone(),
260            _ => String::new(),
261        })
262        .unwrap_or_default()
263}
264
265pub fn read_bytes(entity: &Option<Entity>, field: &str) -> Vec<u8> {
266    entity
267        .as_ref()
268        .and_then(|e| e.get(field))
269        .map(|v| match v {
270            Value::Bytes(b) => b.clone(),
271            _ => Vec::new(),
272        })
273        .unwrap_or_default()
274}
275
276pub fn read_bool(entity: &Option<Entity>, field: &str) -> bool {
277    entity
278        .as_ref()
279        .and_then(|e| e.get(field))
280        .map(|v| match v {
281            Value::Bool(b) => *b,
282            _ => false,
283        })
284        .unwrap_or(false)
285}
286
287pub fn read_i32(entity: &Option<Entity>, field: &str) -> i32 {
288    entity
289        .as_ref()
290        .and_then(|e| e.get(field))
291        .map(|v| match v {
292            Value::Int(i) => *i,
293            _ => 0,
294        })
295        .unwrap_or(0)
296}
297
298pub fn read_value(entity: &Option<Entity>, field: &str) -> Value {
299    entity
300        .as_ref()
301        .and_then(|e| e.get(field))
302        .cloned()
303        .unwrap_or(Value::Null)
304}
305
306#[derive(Debug, Clone)]
307pub struct Entity {
308    entity_type: String,
309    id: String,
310    fields: HashMap<String, Value>,
311}
312
313impl Entity {
314    pub fn new(entity_type: &str, id: &str) -> Self {
315        Self::new_with_id(entity_type, id, EntityIdValue::String(id.to_string()))
316    }
317
318    pub fn new_with_id(entity_type: &str, key: &str, id: EntityIdValue) -> Self {
319        let mut fields = HashMap::new();
320        fields.insert("id".to_string(), id.into_value());
321        Self {
322            entity_type: entity_type.to_string(),
323            id: key.to_string(),
324            fields,
325        }
326    }
327
328    pub fn from_fields(entity_type: &str, id: &str, fields: HashMap<String, Value>) -> Self {
329        Self {
330            entity_type: entity_type.to_string(),
331            id: id.to_string(),
332            fields,
333        }
334    }
335
336    pub fn id(&self) -> &str {
337        &self.id
338    }
339
340    pub fn entity_type(&self) -> &str {
341        &self.entity_type
342    }
343
344    pub fn set<V: IntoEntityValue>(&mut self, field: &str, value: V) {
345        self.fields.insert(field.to_string(), value.into_value());
346    }
347
348    pub fn get(&self, field: &str) -> Option<&Value> {
349        self.fields.get(field)
350    }
351
352    pub fn get_bigint(&self, field: &str) -> BigInt {
353        match self.fields.get(field) {
354            Some(Value::BigInt(v)) => v.clone(),
355            _ => BigInt::zero(),
356        }
357    }
358
359    pub fn get_bigint_opt(&self, field: &str) -> Option<BigInt> {
360        match self.fields.get(field) {
361            Some(Value::BigInt(v)) => Some(v.clone()),
362            _ => None,
363        }
364    }
365
366    pub fn get_bigdecimal(&self, field: &str) -> BigDecimal {
367        match self.fields.get(field) {
368            Some(Value::BigDecimal(v)) => v.clone(),
369            _ => BigDecimal::zero(),
370        }
371    }
372
373    pub fn get_bigdecimal_opt(&self, field: &str) -> Option<BigDecimal> {
374        match self.fields.get(field) {
375            Some(Value::BigDecimal(v)) => Some(v.clone()),
376            _ => None,
377        }
378    }
379
380    pub fn get_string(&self, field: &str) -> String {
381        match self.fields.get(field) {
382            Some(Value::String(v)) => v.clone(),
383            _ => String::new(),
384        }
385    }
386
387    pub fn get_string_opt(&self, field: &str) -> Option<String> {
388        match self.fields.get(field) {
389            Some(Value::String(v)) => Some(v.clone()),
390            _ => None,
391        }
392    }
393
394    pub fn get_bytes(&self, field: &str) -> Vec<u8> {
395        match self.fields.get(field) {
396            Some(Value::Bytes(v)) => v.clone(),
397            _ => Vec::new(),
398        }
399    }
400
401    pub fn get_bytes_opt(&self, field: &str) -> Option<Vec<u8>> {
402        match self.fields.get(field) {
403            Some(Value::Bytes(v)) => Some(v.clone()),
404            _ => None,
405        }
406    }
407
408    pub fn get_bool(&self, field: &str) -> bool {
409        match self.fields.get(field) {
410            Some(Value::Bool(v)) => *v,
411            _ => false,
412        }
413    }
414
415    pub fn get_bool_opt(&self, field: &str) -> Option<bool> {
416        match self.fields.get(field) {
417            Some(Value::Bool(v)) => Some(*v),
418            _ => None,
419        }
420    }
421
422    pub fn get_i32(&self, field: &str) -> i32 {
423        match self.fields.get(field) {
424            Some(Value::Int(v)) => *v,
425            _ => 0,
426        }
427    }
428
429    pub fn get_i32_opt(&self, field: &str) -> Option<i32> {
430        match self.fields.get(field) {
431            Some(Value::Int(v)) => Some(*v),
432            _ => None,
433        }
434    }
435
436    pub fn fields(&self) -> &HashMap<String, Value> {
437        &self.fields
438    }
439}
440
441pub struct EntityCache {
442    entities: HashMap<String, HashMap<String, Entity>>,
443    removals: Vec<(String, String)>,
444}
445
446impl EntityCache {
447    pub fn new() -> Self {
448        Self {
449            entities: HashMap::new(),
450            removals: Vec::new(),
451        }
452    }
453
454    pub fn get(&self, entity_type: &str, id: &str) -> Option<Entity> {
455        self.entities
456            .get(entity_type)
457            .and_then(|m| m.get(id))
458            .cloned()
459    }
460
461    pub fn set(&mut self, entity_type: &str, id: &str, entity: Entity) {
462        self.entities
463            .entry(entity_type.to_string())
464            .or_insert_with(HashMap::new)
465            .insert(id.to_string(), entity);
466    }
467
468    pub fn remove(&mut self, entity_type: &str, id: &str) {
469        if let Some(m) = self.entities.get_mut(entity_type) {
470            m.remove(id);
471        }
472        let key = (entity_type.to_string(), id.to_string());
473        if !self.removals.contains(&key) {
474            self.removals.push(key);
475        }
476    }
477
478    pub fn get_all(&self, entity_type: &str) -> Vec<&Entity> {
479        self.entities
480            .get(entity_type)
481            .map(|m| m.values().collect())
482            .unwrap_or_default()
483    }
484}
485
486/// DataSourceContext — per-instance configuration passed via createWithContext.
487/// Mirrors graph-ts DataSourceContext: typed key-value store.
488#[derive(Debug, Clone)]
489pub struct DataSourceContext {
490    fields: HashMap<String, Value>,
491}
492
493impl DataSourceContext {
494    /// Create a new empty context (for write-side construction).
495    pub fn empty() -> Self {
496        Self { fields: HashMap::new() }
497    }
498
499    /// Create a context from an existing field map (for runtime population).
500    pub fn new(fields: HashMap<String, Value>) -> Self {
501        Self { fields }
502    }
503
504    /// Raw getter — returns Option for presence checks.
505    pub fn get(&self, key: &str) -> Option<&Value> {
506        self.fields.get(key)
507    }
508
509    /// Check if a key is present in the context.
510    pub fn is_set(&self, key: &str) -> bool {
511        self.fields.contains_key(key)
512    }
513
514    /// Typed getter — panics if key absent (matches graph-ts behavior).
515    pub fn get_string(&self, key: &str) -> String {
516        match self.fields.get(key) {
517            Some(Value::String(s)) => s.clone(),
518            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
519        }
520    }
521
522    pub fn get_bytes(&self, key: &str) -> Vec<u8> {
523        match self.fields.get(key) {
524            Some(Value::Bytes(b)) => b.clone(),
525            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
526        }
527    }
528
529    pub fn get_big_int(&self, key: &str) -> BigInt {
530        match self.fields.get(key) {
531            Some(Value::BigInt(n)) => n.clone(),
532            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
533        }
534    }
535
536    pub fn get_i32(&self, key: &str) -> i32 {
537        match self.fields.get(key) {
538            Some(&Value::Int(n)) => n,
539            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
540        }
541    }
542
543    pub fn get_boolean(&self, key: &str) -> bool {
544        match self.fields.get(key) {
545            Some(&Value::Bool(b)) => b,
546            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
547        }
548    }
549
550    /// Typed setters — for constructing context before createWithContext.
551    pub fn set_string(&mut self, key: &str, value: impl Into<String>) {
552        self.fields.insert(key.to_string(), Value::String(value.into()));
553    }
554
555    pub fn set_bytes(&mut self, key: &str, value: Vec<u8>) {
556        self.fields.insert(key.to_string(), Value::Bytes(value));
557    }
558
559    pub fn set_big_int(&mut self, key: &str, value: BigInt) {
560        self.fields.insert(key.to_string(), Value::BigInt(value));
561    }
562
563    pub fn set_i32(&mut self, key: &str, value: i32) {
564        self.fields.insert(key.to_string(), Value::Int(value));
565    }
566
567    pub fn set_boolean(&mut self, key: &str, value: bool) {
568        self.fields.insert(key.to_string(), Value::Bool(value));
569    }
570
571    /// Access the underlying fields (for runtime conversion).
572    pub fn into_fields(self) -> HashMap<String, Value> {
573        self.fields
574    }
575}
576
577#[async_trait]
578pub trait HandlerContext: Send {
579    async fn load_entity(&mut self, entity_type: &str, id: &str) -> Option<Entity>;
580
581    async fn save_entity(&mut self, entity_type: &str, entity: &Entity) -> Result<(), Error>;
582
583    fn entity_remove(&mut self, entity_type: &str, id: &str) -> Result<(), Error>;
584
585    /// Returns the current data source's network name.
586    fn data_source_network(&self) -> String;
587
588    /// Returns the current data source's context (set via createWithContext).
589    /// Returns an empty context if none was provided.
590    fn data_source_context(&self) -> DataSourceContext;
591
592    /// Mirrors graph-ts `dataSource.stringParam()`.
593    /// Returns the UTF-8 string interpretation of the datasource address/source bytes.
594    /// For file/offchain datasources, this is the source identifier (e.g., content hash).
595    /// For onchain templates, the result depends on how graph-node populates
596    /// `dataSource.address()` — typically the hex-decoded contract address bytes
597    /// interpreted as UTF-8, which may not produce a meaningful string.
598    fn data_source_string_param(&self) -> String;
599
600    /// Create a dynamic data source from a template.
601    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
602    fn data_source_create(&mut self, template: &str, params: &[String]) -> Result<(), Error>;
603
604    /// Create a dynamic data source from a template with context.
605    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
606    fn data_source_create_with_context(
607        &mut self,
608        template: &str,
609        params: &[String],
610        context: DataSourceContext,
611    ) -> Result<(), Error>;
612}