Skip to main content

teaql_runtime/data_service/
context.rs

1use std::fmt::Write;
2use std::sync::Arc;
3use std::{collections::BTreeMap, future::Future};
4
5use teaql_core::{
6    CompactRow, DeleteCommand, Entity, InsertCommand, RecoverCommand, SelectQuery, SmartList,
7    UpdateCommand,
8};
9
10use crate::{
11    ContextError, DataServiceError, GraphMutationPlan, GraphNode, RuntimeError, UserContext,
12};
13
14use super::{
15    AggregationCacheBackend, ContextDataService, EntityDataService, InMemoryAggregationCache,
16    RuntimeDataService, UserContextMetadata, helpers::invalidate_aggregation_cache_namespace,
17};
18
19impl UserContext {
20    pub(crate) fn data_service_internal<E>(&self) -> Result<ContextDataService<'_, E>, ContextError>
21    where
22        E: teaql_data_service::QueryExecutor
23            + teaql_data_service::MutationExecutor
24            + Send
25            + Sync
26            + 'static,
27    {
28        if self.metadata.is_none() {
29            return Err(ContextError::MissingResource("metadata".to_owned()));
30        }
31
32        let executor = self.require_resource::<E>()?;
33        Ok(ContextDataService {
34            metadata: UserContextMetadata { context: self },
35            executor,
36        })
37    }
38
39    pub fn entity_data_service<E>(
40        &self,
41        entity: impl Into<String>,
42    ) -> Result<EntityDataService<'_, E>, ContextError>
43    where
44        E: teaql_data_service::QueryExecutor
45            + teaql_data_service::MutationExecutor
46            + Send
47            + Sync
48            + 'static,
49    {
50        let entity = entity.into();
51        if !self.has_entity_data_service(&entity) {
52            return Err(ContextError::MissingEntityDataService(entity));
53        }
54        Ok(EntityDataService {
55            entity,
56            data_service: self.data_service_internal::<E>()?,
57            trace_context: Vec::new(),
58        })
59    }
60
61    /// Register a data-service executor and automatically set up the
62    /// type-erased graph saver so that
63    /// [`Audited::save`](crate::AuditedSaveExt::save) works.
64    pub fn register_executor<E>(&mut self, executor: E)
65    where
66        E: teaql_data_service::QueryExecutor
67            + teaql_data_service::MutationExecutor
68            + Send
69            + Sync
70            + 'static,
71    {
72        use std::sync::Arc;
73        self.insert_resource::<Arc<dyn crate::entity_save::DynGraphSaver>>(Arc::new(
74            crate::entity_save::GraphSaverFor::<E>::new(),
75        ));
76        self.insert_resource(executor);
77    }
78}
79
80impl<'a, E> ContextDataService<'a, E>
81where
82    E: teaql_data_service::QueryExecutor
83        + teaql_data_service::MutationExecutor
84        + Send
85        + Sync
86        + 'static,
87{
88    async fn observe<T, F, N>(
89        &self,
90        family: &str,
91        name: N,
92        entity: &str,
93        work: F,
94    ) -> Result<T, DataServiceError<E::Error>>
95    where
96        F: Future<Output = Result<T, DataServiceError<E::Error>>>,
97        N: FnOnce() -> String,
98    {
99        if self.metadata.context.runtime_telemetry_is_noop() {
100            return work.await;
101        }
102        let operation = crate::RuntimeOperation::new(family, name())
103            .attribute("teaql.entity.type", entity.to_owned());
104        let scope = self.metadata.context.start_runtime_operation(operation);
105        let provider_kind = std::any::type_name::<E>().to_owned();
106        let provider_operation = family.to_owned();
107        let result = scope
108            .run(async {
109                let provider_scope = self.metadata.context.start_runtime_operation(
110                    crate::RuntimeOperation::new(
111                        "provider",
112                        format!("{provider_kind}.{provider_operation}"),
113                    )
114                    .attribute("teaql.provider.kind", provider_kind)
115                    .attribute("teaql.provider.operation", provider_operation),
116                );
117                let result = provider_scope.run(work).await;
118                match &result {
119                    Ok(_) => provider_scope.success(BTreeMap::new()),
120                    Err(_) => provider_scope.failure("data_service_error"),
121                }
122                result
123            })
124            .await;
125        match result {
126            Ok(value) => {
127                scope.success(BTreeMap::new());
128                Ok(value)
129            }
130            Err(error) => {
131                scope.failure("data_service_error");
132                Err(error)
133            }
134        }
135    }
136
137    fn data_service(&self) -> RuntimeDataService<'_, UserContextMetadata<'_>, E> {
138        RuntimeDataService::new(&self.metadata, self.executor)
139    }
140
141    pub(crate) async fn fetch_all(
142        &self,
143        mut query: SelectQuery,
144    ) -> Result<Vec<CompactRow>, DataServiceError<E::Error>> {
145        let final_comment = self.resolve_final_comment(&query.trace_chain, query.comment.clone());
146        query.comment = final_comment;
147        self.observe(
148            "query",
149            || format!("{}.list", query.entity),
150            &query.entity,
151            self.data_service().fetch_all(&query),
152        )
153        .await
154    }
155
156    pub(crate) async fn fetch_smart_list(
157        &self,
158        query: &SelectQuery,
159    ) -> Result<SmartList<CompactRow>, DataServiceError<E::Error>> {
160        self.observe(
161            "query",
162            || format!("{}.list", query.entity),
163            &query.entity,
164            self.data_service().fetch_smart_list(query),
165        )
166        .await
167    }
168
169    pub(crate) async fn fetch_entities<T>(
170        &self,
171        query: &SelectQuery,
172    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
173    where
174        T: Entity,
175    {
176        self.observe(
177            "query",
178            || format!("{}.list", query.entity),
179            &query.entity,
180            self.data_service().fetch_entities(query),
181        )
182        .await
183    }
184
185    pub(crate) async fn fetch_enhanced_entities<T>(
186        &self,
187        query: &SelectQuery,
188    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
189    where
190        T: Entity,
191    {
192        self.observe(
193            "query",
194            || format!("{}.list", query.entity),
195            &query.entity,
196            self.data_service().fetch_enhanced_entities(query),
197        )
198        .await
199    }
200
201    pub(crate) async fn insert(
202        &self,
203        command: &InsertCommand,
204    ) -> Result<u64, DataServiceError<E::Error>> {
205        let affected = self
206            .observe(
207                "mutation",
208                || format!("{}.insert", command.entity),
209                &command.entity,
210                self.data_service().insert(command),
211            )
212            .await?;
213        self.invalidate_aggregation_cache_for(&command.entity);
214        Ok(affected)
215    }
216
217    pub(crate) async fn update(
218        &self,
219        command: &UpdateCommand,
220    ) -> Result<u64, DataServiceError<E::Error>> {
221        let affected = self
222            .observe(
223                "mutation",
224                || format!("{}.update", command.entity),
225                &command.entity,
226                self.data_service().update(command),
227            )
228            .await?;
229        self.invalidate_aggregation_cache_for(&command.entity);
230        Ok(affected)
231    }
232
233    pub(crate) async fn batch_insert(
234        &self,
235        command: &teaql_core::BatchInsertCommand,
236    ) -> Result<u64, DataServiceError<E::Error>> {
237        let affected = self
238            .observe(
239                "mutation",
240                || format!("{}.batch_insert", command.entity),
241                &command.entity,
242                self.data_service().batch_insert(command),
243            )
244            .await?;
245        self.invalidate_aggregation_cache_for(&command.entity);
246        Ok(affected)
247    }
248
249    pub(crate) async fn batch_update(
250        &self,
251        command: &teaql_core::BatchUpdateCommand,
252    ) -> Result<u64, DataServiceError<E::Error>> {
253        let affected = self
254            .observe(
255                "mutation",
256                || format!("{}.batch_update", command.entity),
257                &command.entity,
258                self.data_service().batch_update(command),
259            )
260            .await?;
261        self.invalidate_aggregation_cache_for(&command.entity);
262        Ok(affected)
263    }
264
265    pub(crate) async fn delete(
266        &self,
267        command: &DeleteCommand,
268    ) -> Result<u64, DataServiceError<E::Error>> {
269        let affected = self
270            .observe(
271                "mutation",
272                || format!("{}.delete", command.entity),
273                &command.entity,
274                self.data_service().delete(command),
275            )
276            .await?;
277        self.invalidate_aggregation_cache_for(&command.entity);
278        Ok(affected)
279    }
280
281    pub(crate) async fn recover(
282        &self,
283        command: &RecoverCommand,
284    ) -> Result<u64, DataServiceError<E::Error>> {
285        let affected = self
286            .observe(
287                "mutation",
288                || format!("{}.recover", command.entity),
289                &command.entity,
290                self.data_service().recover(command),
291            )
292            .await?;
293        self.invalidate_aggregation_cache_for(&command.entity);
294        Ok(affected)
295    }
296
297    pub(super) fn invalidate_aggregation_cache_for(&self, entity: &str) {
298        if let Some(cache) = self
299            .metadata
300            .context
301            .get_resource::<Arc<dyn AggregationCacheBackend>>()
302        {
303            invalidate_aggregation_cache_namespace(cache.as_ref(), entity);
304        }
305        if let Some(cache) = self
306            .metadata
307            .context
308            .get_resource::<InMemoryAggregationCache>()
309        {
310            invalidate_aggregation_cache_namespace(cache, entity);
311        }
312    }
313
314    pub(crate) fn resolve_final_comment(
315        &self,
316        trace_chain: &[teaql_core::TraceNode],
317        comment: Option<String>,
318    ) -> Option<String> {
319        let chain_str = (!trace_chain.is_empty()).then(|| {
320            let mut chain = String::with_capacity(trace_chain.len().saturating_mul(64));
321            for (index, node) in trace_chain.iter().enumerate() {
322                if index > 0 {
323                    chain.push_str(" -> ");
324                }
325                match node.entity_id {
326                    Some(id) => {
327                        let _ = write!(chain, "{}({id}): {}", node.entity_type, node.comment);
328                    }
329                    None => {
330                        let _ = write!(chain, "{}(pending): {}", node.entity_type, node.comment);
331                    }
332                }
333            }
334            chain
335        });
336
337        let business_comment = chain_str.or(comment);
338        let user_id = self
339            .metadata
340            .context
341            .user_identifier()
342            .map(|s| s.to_owned());
343
344        match (user_id, business_comment) {
345            (Some(user), Some(bus)) if !user.is_empty() && !bus.is_empty() => {
346                Some(format!("[{user}] {bus}"))
347            }
348            (Some(user), _) if !user.is_empty() => Some(format!("[{user}]")),
349            (_, Some(bus)) if !bus.is_empty() => Some(bus),
350            _ => None,
351        }
352    }
353}