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
79pub trait IntoEntityValue {
80    fn into_value(self) -> Value;
81}
82
83impl IntoEntityValue for String {
84    fn into_value(self) -> Value {
85        Value::String(self)
86    }
87}
88
89impl IntoEntityValue for &str {
90    fn into_value(self) -> Value {
91        Value::String(self.to_string())
92    }
93}
94
95impl IntoEntityValue for i32 {
96    fn into_value(self) -> Value {
97        Value::Int(self)
98    }
99}
100
101impl IntoEntityValue for u64 {
102    fn into_value(self) -> Value {
103        Value::from(self)
104    }
105}
106
107impl IntoEntityValue for bool {
108    fn into_value(self) -> Value {
109        Value::Bool(self)
110    }
111}
112
113impl IntoEntityValue for BigInt {
114    fn into_value(self) -> Value {
115        Value::BigInt(self)
116    }
117}
118
119impl IntoEntityValue for BigDecimal {
120    fn into_value(self) -> Value {
121        Value::BigDecimal(self)
122    }
123}
124
125impl IntoEntityValue for Address {
126    fn into_value(self) -> Value {
127        Value::Bytes(self.to_bytes())
128    }
129}
130
131impl IntoEntityValue for Vec<u8> {
132    fn into_value(self) -> Value {
133        Value::Bytes(self)
134    }
135}
136
137impl IntoEntityValue for &BigInt {
138    fn into_value(self) -> Value {
139        Value::BigInt(self.clone())
140    }
141}
142
143impl IntoEntityValue for &BigDecimal {
144    fn into_value(self) -> Value {
145        Value::BigDecimal(self.clone())
146    }
147}
148
149impl IntoEntityValue for &Address {
150    fn into_value(self) -> Value {
151        Value::Bytes(self.to_bytes())
152    }
153}
154
155impl IntoEntityValue for &Vec<u8> {
156    fn into_value(self) -> Value {
157        Value::Bytes(self.clone())
158    }
159}
160
161pub fn read_bigint(entity: &Option<Entity>, field: &str) -> BigInt {
162    entity
163        .as_ref()
164        .and_then(|e| e.get(field))
165        .map(|v| match v {
166            Value::BigInt(b) => b.clone(),
167            _ => BigInt::zero(),
168        })
169        .unwrap_or_else(BigInt::zero)
170}
171
172pub fn read_bigdecimal(entity: &Option<Entity>, field: &str) -> BigDecimal {
173    entity
174        .as_ref()
175        .and_then(|e| e.get(field))
176        .map(|v| match v {
177            Value::BigDecimal(b) => b.clone(),
178            _ => BigDecimal::zero(),
179        })
180        .unwrap_or_else(BigDecimal::zero)
181}
182
183pub fn read_string(entity: &Option<Entity>, field: &str) -> String {
184    entity
185        .as_ref()
186        .and_then(|e| e.get(field))
187        .map(|v| match v {
188            Value::String(s) => s.clone(),
189            _ => String::new(),
190        })
191        .unwrap_or_default()
192}
193
194pub fn read_bytes(entity: &Option<Entity>, field: &str) -> Vec<u8> {
195    entity
196        .as_ref()
197        .and_then(|e| e.get(field))
198        .map(|v| match v {
199            Value::Bytes(b) => b.clone(),
200            _ => Vec::new(),
201        })
202        .unwrap_or_default()
203}
204
205pub fn read_bool(entity: &Option<Entity>, field: &str) -> bool {
206    entity
207        .as_ref()
208        .and_then(|e| e.get(field))
209        .map(|v| match v {
210            Value::Bool(b) => *b,
211            _ => false,
212        })
213        .unwrap_or(false)
214}
215
216pub fn read_i32(entity: &Option<Entity>, field: &str) -> i32 {
217    entity
218        .as_ref()
219        .and_then(|e| e.get(field))
220        .map(|v| match v {
221            Value::Int(i) => *i,
222            _ => 0,
223        })
224        .unwrap_or(0)
225}
226
227pub fn read_value(entity: &Option<Entity>, field: &str) -> Value {
228    entity
229        .as_ref()
230        .and_then(|e| e.get(field))
231        .cloned()
232        .unwrap_or(Value::Null)
233}
234
235#[derive(Debug, Clone)]
236pub struct Entity {
237    entity_type: String,
238    id: String,
239    fields: HashMap<String, Value>,
240}
241
242impl Entity {
243    pub fn new(entity_type: &str, id: &str) -> Self {
244        let mut fields = HashMap::new();
245        fields.insert("id".to_string(), Value::String(id.to_string()));
246        Self {
247            entity_type: entity_type.to_string(),
248            id: id.to_string(),
249            fields,
250        }
251    }
252
253    pub fn from_fields(entity_type: &str, id: &str, fields: HashMap<String, Value>) -> Self {
254        Self {
255            entity_type: entity_type.to_string(),
256            id: id.to_string(),
257            fields,
258        }
259    }
260
261    pub fn id(&self) -> &str {
262        &self.id
263    }
264
265    pub fn entity_type(&self) -> &str {
266        &self.entity_type
267    }
268
269    pub fn set<V: IntoEntityValue>(&mut self, field: &str, value: V) {
270        self.fields.insert(field.to_string(), value.into_value());
271    }
272
273    pub fn get(&self, field: &str) -> Option<&Value> {
274        self.fields.get(field)
275    }
276
277    pub fn get_bigint(&self, field: &str) -> BigInt {
278        match self.fields.get(field) {
279            Some(Value::BigInt(v)) => v.clone(),
280            _ => BigInt::zero(),
281        }
282    }
283
284    pub fn get_bigdecimal(&self, field: &str) -> BigDecimal {
285        match self.fields.get(field) {
286            Some(Value::BigDecimal(v)) => v.clone(),
287            _ => BigDecimal::zero(),
288        }
289    }
290
291    pub fn get_string(&self, field: &str) -> String {
292        match self.fields.get(field) {
293            Some(Value::String(v)) => v.clone(),
294            _ => String::new(),
295        }
296    }
297
298    pub fn get_bytes(&self, field: &str) -> Vec<u8> {
299        match self.fields.get(field) {
300            Some(Value::Bytes(v)) => v.clone(),
301            _ => Vec::new(),
302        }
303    }
304
305    pub fn get_bool(&self, field: &str) -> bool {
306        match self.fields.get(field) {
307            Some(Value::Bool(v)) => *v,
308            _ => false,
309        }
310    }
311
312    pub fn get_i32(&self, field: &str) -> i32 {
313        match self.fields.get(field) {
314            Some(Value::Int(v)) => *v,
315            _ => 0,
316        }
317    }
318
319    pub fn fields(&self) -> &HashMap<String, Value> {
320        &self.fields
321    }
322}
323
324pub struct EntityCache {
325    entities: HashMap<String, HashMap<String, Entity>>,
326    removals: Vec<(String, String)>,
327}
328
329impl EntityCache {
330    pub fn new() -> Self {
331        Self {
332            entities: HashMap::new(),
333            removals: Vec::new(),
334        }
335    }
336
337    pub fn get(&self, entity_type: &str, id: &str) -> Option<Entity> {
338        self.entities
339            .get(entity_type)
340            .and_then(|m| m.get(id))
341            .cloned()
342    }
343
344    pub fn set(&mut self, entity_type: &str, id: &str, entity: Entity) {
345        self.entities
346            .entry(entity_type.to_string())
347            .or_insert_with(HashMap::new)
348            .insert(id.to_string(), entity);
349    }
350
351    pub fn remove(&mut self, entity_type: &str, id: &str) {
352        if let Some(m) = self.entities.get_mut(entity_type) {
353            m.remove(id);
354        }
355        let key = (entity_type.to_string(), id.to_string());
356        if !self.removals.contains(&key) {
357            self.removals.push(key);
358        }
359    }
360
361    pub fn get_all(&self, entity_type: &str) -> Vec<&Entity> {
362        self.entities
363            .get(entity_type)
364            .map(|m| m.values().collect())
365            .unwrap_or_default()
366    }
367}
368
369#[async_trait]
370pub trait HandlerContext: Send {
371    async fn load_entity(&mut self, entity_type: &str, id: &str) -> Option<Entity>;
372
373    async fn save_entity(&mut self, entity_type: &str, entity: &Entity) -> Result<(), Error>;
374
375    fn entity_remove(&mut self, entity_type: &str, id: &str) -> Result<(), Error>;
376}