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