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_bigdecimal(&self, field: &str) -> BigDecimal {
351        match self.fields.get(field) {
352            Some(Value::BigDecimal(v)) => v.clone(),
353            _ => BigDecimal::zero(),
354        }
355    }
356
357    pub fn get_string(&self, field: &str) -> String {
358        match self.fields.get(field) {
359            Some(Value::String(v)) => v.clone(),
360            _ => String::new(),
361        }
362    }
363
364    pub fn get_bytes(&self, field: &str) -> Vec<u8> {
365        match self.fields.get(field) {
366            Some(Value::Bytes(v)) => v.clone(),
367            _ => Vec::new(),
368        }
369    }
370
371    pub fn get_bool(&self, field: &str) -> bool {
372        match self.fields.get(field) {
373            Some(Value::Bool(v)) => *v,
374            _ => false,
375        }
376    }
377
378    pub fn get_i32(&self, field: &str) -> i32 {
379        match self.fields.get(field) {
380            Some(Value::Int(v)) => *v,
381            _ => 0,
382        }
383    }
384
385    pub fn fields(&self) -> &HashMap<String, Value> {
386        &self.fields
387    }
388}
389
390pub struct EntityCache {
391    entities: HashMap<String, HashMap<String, Entity>>,
392    removals: Vec<(String, String)>,
393}
394
395impl EntityCache {
396    pub fn new() -> Self {
397        Self {
398            entities: HashMap::new(),
399            removals: Vec::new(),
400        }
401    }
402
403    pub fn get(&self, entity_type: &str, id: &str) -> Option<Entity> {
404        self.entities
405            .get(entity_type)
406            .and_then(|m| m.get(id))
407            .cloned()
408    }
409
410    pub fn set(&mut self, entity_type: &str, id: &str, entity: Entity) {
411        self.entities
412            .entry(entity_type.to_string())
413            .or_insert_with(HashMap::new)
414            .insert(id.to_string(), entity);
415    }
416
417    pub fn remove(&mut self, entity_type: &str, id: &str) {
418        if let Some(m) = self.entities.get_mut(entity_type) {
419            m.remove(id);
420        }
421        let key = (entity_type.to_string(), id.to_string());
422        if !self.removals.contains(&key) {
423            self.removals.push(key);
424        }
425    }
426
427    pub fn get_all(&self, entity_type: &str) -> Vec<&Entity> {
428        self.entities
429            .get(entity_type)
430            .map(|m| m.values().collect())
431            .unwrap_or_default()
432    }
433}
434
435/// DataSourceContext — per-instance configuration passed via createWithContext.
436/// Mirrors graph-ts DataSourceContext: typed key-value store.
437#[derive(Debug, Clone)]
438pub struct DataSourceContext {
439    fields: HashMap<String, Value>,
440}
441
442impl DataSourceContext {
443    /// Create a new empty context (for write-side construction).
444    pub fn empty() -> Self {
445        Self { fields: HashMap::new() }
446    }
447
448    /// Create a context from an existing field map (for runtime population).
449    pub fn new(fields: HashMap<String, Value>) -> Self {
450        Self { fields }
451    }
452
453    /// Raw getter — returns Option for presence checks.
454    pub fn get(&self, key: &str) -> Option<&Value> {
455        self.fields.get(key)
456    }
457
458    /// Check if a key is present in the context.
459    pub fn is_set(&self, key: &str) -> bool {
460        self.fields.contains_key(key)
461    }
462
463    /// Typed getter — panics if key absent (matches graph-ts behavior).
464    pub fn get_string(&self, key: &str) -> String {
465        match self.fields.get(key) {
466            Some(Value::String(s)) => s.clone(),
467            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
468        }
469    }
470
471    pub fn get_bytes(&self, key: &str) -> Vec<u8> {
472        match self.fields.get(key) {
473            Some(Value::Bytes(b)) => b.clone(),
474            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
475        }
476    }
477
478    pub fn get_big_int(&self, key: &str) -> BigInt {
479        match self.fields.get(key) {
480            Some(Value::BigInt(n)) => n.clone(),
481            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
482        }
483    }
484
485    pub fn get_i32(&self, key: &str) -> i32 {
486        match self.fields.get(key) {
487            Some(&Value::Int(n)) => n,
488            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
489        }
490    }
491
492    pub fn get_boolean(&self, key: &str) -> bool {
493        match self.fields.get(key) {
494            Some(&Value::Bool(b)) => b,
495            _ => panic!("DataSourceContext: missing or wrong type for key '{}'", key),
496        }
497    }
498
499    /// Typed setters — for constructing context before createWithContext.
500    pub fn set_string(&mut self, key: &str, value: impl Into<String>) {
501        self.fields.insert(key.to_string(), Value::String(value.into()));
502    }
503
504    pub fn set_bytes(&mut self, key: &str, value: Vec<u8>) {
505        self.fields.insert(key.to_string(), Value::Bytes(value));
506    }
507
508    pub fn set_big_int(&mut self, key: &str, value: BigInt) {
509        self.fields.insert(key.to_string(), Value::BigInt(value));
510    }
511
512    pub fn set_i32(&mut self, key: &str, value: i32) {
513        self.fields.insert(key.to_string(), Value::Int(value));
514    }
515
516    pub fn set_boolean(&mut self, key: &str, value: bool) {
517        self.fields.insert(key.to_string(), Value::Bool(value));
518    }
519
520    /// Access the underlying fields (for runtime conversion).
521    pub fn into_fields(self) -> HashMap<String, Value> {
522        self.fields
523    }
524}
525
526#[async_trait]
527pub trait HandlerContext: Send {
528    async fn load_entity(&mut self, entity_type: &str, id: &str) -> Option<Entity>;
529
530    async fn save_entity(&mut self, entity_type: &str, entity: &Entity) -> Result<(), Error>;
531
532    fn entity_remove(&mut self, entity_type: &str, id: &str) -> Result<(), Error>;
533
534    /// Returns the current data source's network name.
535    fn data_source_network(&self) -> String;
536
537    /// Returns the current data source's context (set via createWithContext).
538    /// Returns an empty context if none was provided.
539    fn data_source_context(&self) -> DataSourceContext;
540
541    /// Mirrors graph-ts `dataSource.stringParam()`.
542    /// Returns the UTF-8 string interpretation of the datasource address/source bytes.
543    /// For file/offchain datasources, this is the source identifier (e.g., content hash).
544    /// For onchain templates, the result depends on how graph-node populates
545    /// `dataSource.address()` — typically the hex-decoded contract address bytes
546    /// interpreted as UTF-8, which may not produce a meaningful string.
547    fn data_source_string_param(&self) -> String;
548
549    /// Create a dynamic data source from a template.
550    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
551    fn data_source_create(&mut self, template: &str, params: &[String]) -> Result<(), Error>;
552
553    /// Create a dynamic data source from a template with context.
554    /// `params` are template-specific: for onchain templates, typically `[address_hex]`.
555    fn data_source_create_with_context(
556        &mut self,
557        template: &str,
558        params: &[String],
559        context: DataSourceContext,
560    ) -> Result<(), Error>;
561}