contrag_core/data_sources/
canister_state.rs

1use candid::{CandidType, Principal, encode_one};
2use crate::data_sources::DataSource;
3use crate::entity::RagEntity;
4use crate::error::{ContragError, Result};
5use crate::config::EntityConfig;
6
7/// Data source that reads from other ICP canisters via inter-canister calls
8pub struct CanisterStateSource {
9    entity_configs: std::collections::HashMap<String, EntityConfig>,
10}
11
12impl CanisterStateSource {
13    /// Create a new canister state source with entity configurations
14    pub fn new(entity_configs: Vec<EntityConfig>) -> Self {
15        let mut map = std::collections::HashMap::new();
16        for config in entity_configs {
17            map.insert(config.name.clone(), config);
18        }
19        Self {
20            entity_configs: map,
21        }
22    }
23
24    /// Get entity configuration by type
25    fn get_config(&self, entity_type: &str) -> Result<&EntityConfig> {
26        self.entity_configs
27            .get(entity_type)
28            .ok_or_else(|| {
29                ContragError::ConfigError(format!("No configuration found for entity type: {}", entity_type))
30            })
31    }
32
33    /// Make inter-canister call to fetch entity
34    /// 
35    /// Note: This is a placeholder for the actual ICP inter-canister call.
36    /// In a real canister, you would use ic_cdk::call
37    async fn call_canister<T: CandidType>(
38        &self,
39        canister_id: Principal,
40        method: &str,
41        args: Vec<u8>,
42    ) -> Result<T> {
43        // This would be implemented using ic_cdk::call in actual canister code
44        // For now, this is a placeholder that will be replaced in the canister implementation
45        
46        #[cfg(target_family = "wasm")]
47        {
48            use ic_cdk::api::call::call_raw;
49            
50            let result = call_raw(canister_id, method, args, 0)
51                .await
52                .map_err(|(code, msg)| {
53                    ContragError::CanisterCallError(format!("Call failed: {:?} - {}", code, msg))
54                })?;
55            
56            decode_one(&result).map_err(|e| {
57                ContragError::CanisterCallError(format!("Failed to decode response: {}", e))
58            })
59        }
60        
61        #[cfg(not(target_family = "wasm"))]
62        {
63            Err(ContragError::CanisterCallError(
64                "Canister calls only work in WASM environment".to_string()
65            ))
66        }
67    }
68}
69
70#[async_trait::async_trait]
71impl DataSource for CanisterStateSource {
72    async fn read_entity<T: RagEntity + CandidType>(
73        &self,
74        entity_type: &str,
75        entity_id: &str,
76    ) -> Result<T> {
77        let config = self.get_config(entity_type)?;
78        
79        let canister_id = Principal::from_text(&config.canister_id)
80            .map_err(|e| ContragError::ConfigError(format!("Invalid canister ID: {}", e)))?;
81        
82        // Encode the entity_id as argument
83        let args = encode_one(&entity_id)
84            .map_err(|e| ContragError::SerializationError(format!("Failed to encode args: {}", e)))?;
85        
86        self.call_canister(canister_id, &config.fetch_method, args)
87            .await
88    }
89
90    async fn read_entities<T: RagEntity + CandidType + Send>(
91        &self,
92        entity_type: &str,
93        entity_ids: Vec<String>,
94    ) -> Result<Vec<T>> {
95        let mut entities = vec![];
96        
97        // Fetch entities one by one
98        // TODO: Optimize with batch fetch if fetch_many_method is configured
99        for id in entity_ids {
100            match self.read_entity(entity_type, &id).await {
101                Ok(entity) => entities.push(entity),
102                Err(e) => {
103                    ic_cdk::println!("Failed to fetch entity {} of type {}: {:?}", id, entity_type, e);
104                    // Continue with other entities
105                }
106            }
107        }
108        
109        Ok(entities)
110    }
111
112    async fn query_entities<T: RagEntity + CandidType>(
113        &self,
114        entity_type: &str,
115        _filter: Option<String>,
116    ) -> Result<Vec<T>> {
117        let config = self.get_config(entity_type)?;
118        
119        if let Some(fetch_many_method) = &config.fetch_many_method {
120            let canister_id = Principal::from_text(&config.canister_id)
121                .map_err(|e| ContragError::ConfigError(format!("Invalid canister ID: {}", e)))?;
122            
123            let args = encode_one(&_filter)
124                .map_err(|e| ContragError::SerializationError(format!("Failed to encode args: {}", e)))?;
125            
126            self.call_canister(canister_id, fetch_many_method, args)
127                .await
128        } else {
129            Ok(vec![])
130        }
131    }
132}
133
134/// Helper to create a canister state source from config
135pub fn create_from_config(entity_configs: Vec<EntityConfig>) -> CanisterStateSource {
136    CanisterStateSource::new(entity_configs)
137}
contrag_core/data_sources/canister_state.rs

contrag_core/data_sources/
canister_state.rs
