Skip to main content

teaql_runtime/data_service/
resolved.rs

1use std::hash::{Hash, Hasher};
2use std::sync::Arc;
3use std::time::{Duration, SystemTime};
4
5use teaql_core::{
6    AggregationCacheOptions, DeleteCommand, Entity, Expr, InsertCommand, Record, RecoverCommand,
7    RelationAggregate, SelectQuery, SmartList, SortDirection, UpdateCommand, Value,
8};
9
10use crate::{
11    CheckObjectStatus, ContinuousPageCursor, DataServiceError, EntityDataServiceBehavior,
12    MetadataStore, PurposedSelectQuery, RawAuditEvent, RuntimeError, clear_record_status,
13    mark_record_status,
14};
15
16use super::{
17    AggregationCacheBackend, ContextDataService, EntityDataService, InMemoryAggregationCache,
18    UserContextMetadata, helpers::*,
19};
20
21#[derive(Debug, Clone)]
22struct ContinuousPageExecution {
23    query_key: String,
24    direction: SortDirection,
25    page_size: u64,
26    original_offset: u64,
27    ttl_seconds: u64,
28    optimized: bool,
29    seek_cursor_id: Option<String>,
30}
31
32impl<'a, E> EntityDataService<'a, E>
33where
34    E: teaql_data_service::QueryExecutor
35        + teaql_data_service::MutationExecutor
36        + Send
37        + Sync
38        + 'static,
39{
40    pub(super) fn query_behavior(
41        &self,
42        entity: &str,
43    ) -> Option<Arc<dyn EntityDataServiceBehavior>> {
44        self.data_service
45            .metadata
46            .context
47            .entity_data_service_behavior(entity)
48    }
49
50    pub(super) fn behavior(&self) -> Option<Arc<dyn EntityDataServiceBehavior>> {
51        self.data_service
52            .metadata
53            .context
54            .entity_data_service_behavior(&self.entity)
55    }
56
57    pub fn entity(&self) -> &str {
58        &self.entity
59    }
60
61    pub fn select(&self) -> SelectQuery {
62        SelectQuery::new(self.entity.clone())
63    }
64
65    pub fn insert_command(&self) -> InsertCommand {
66        InsertCommand::new(self.entity.clone())
67    }
68
69    fn enforce_insert_policy(&self, command: &mut InsertCommand) -> Result<(), RuntimeError> {
70        if let Some(policy) = self.data_service.metadata.context.request_policy.as_ref() {
71            policy.enforce_insert(self.data_service.metadata.context, command)?;
72        }
73        Ok(())
74    }
75
76    fn enforce_update_policy(&self, command: &mut UpdateCommand) -> Result<(), RuntimeError> {
77        if let Some(policy) = self.data_service.metadata.context.request_policy.as_ref() {
78            policy.enforce_update(self.data_service.metadata.context, command)?;
79        }
80        Ok(())
81    }
82
83    fn enforce_delete_policy(&self, command: &mut DeleteCommand) -> Result<(), RuntimeError> {
84        if let Some(policy) = self.data_service.metadata.context.request_policy.as_ref() {
85            policy.enforce_delete(self.data_service.metadata.context, command)?;
86        }
87        Ok(())
88    }
89
90    fn enforce_recover_policy(&self, command: &mut RecoverCommand) -> Result<(), RuntimeError> {
91        if let Some(policy) = self.data_service.metadata.context.request_policy.as_ref() {
92            policy.enforce_recover(self.data_service.metadata.context, command)?;
93        }
94        Ok(())
95    }
96
97    fn prepare_select_query(&self, query: &SelectQuery) -> Result<SelectQuery, RuntimeError> {
98        self.prepare_select_query_owned(query.clone())
99    }
100
101    fn prepare_select_query_owned(
102        &self,
103        mut query: SelectQuery,
104    ) -> Result<SelectQuery, RuntimeError> {
105        let mut full_trace = self.trace_context.clone();
106        full_trace.extend(query.trace_chain);
107        query.trace_chain = full_trace;
108
109        if let Some(behavior) = self.query_behavior(&query.entity) {
110            behavior.before_select(self.data_service.metadata.context, &mut query)?;
111        }
112        if let Some(policy) = self.data_service.metadata.context.request_policy.as_ref() {
113            policy.enforce_select(self.data_service.metadata.context, &mut query)?;
114        }
115        // Ensure local_key fields for relation loads are projected so that
116        // enhance_query_relations can match parent rows to child records.
117        if !query.relations.is_empty() {
118            if let Some(descriptor) = self.data_service.metadata.context.entity(&query.entity) {
119                for load in &query.relations {
120                    if let Some(relation) = descriptor.relation_by_name(&load.name) {
121                        if !query.projection.contains(&relation.local_key) {
122                            query.projection.push(relation.local_key.clone());
123                        }
124                    }
125                }
126            }
127        }
128        Ok(query)
129    }
130
131    pub fn prepare_insert_command(
132        &self,
133        command: &InsertCommand,
134    ) -> Result<InsertCommand, RuntimeError> {
135        let mut command = command.clone();
136        if let Some(behavior) = self.behavior() {
137            behavior.before_insert(self.data_service.metadata.context, &mut command)?;
138        }
139        self.enforce_insert_policy(&mut command)?;
140
141        let entity = self
142            .data_service
143            .metadata
144            .context
145            .require_entity(&command.entity)?;
146        if let Some(id_property) = entity.id_property() {
147            let needs_id = !command.values.contains_key(&id_property.name)
148                || is_unassigned_id(command.values.get(&id_property.name));
149            if needs_id {
150                let id = self
151                    .data_service
152                    .metadata
153                    .context
154                    .next_id(&command.entity)?;
155                command
156                    .values
157                    .insert(id_property.name.clone(), Value::U64(id));
158            }
159        }
160        ensure_initial_version(&mut command.values, entity);
161        mark_record_status(&mut command.values, CheckObjectStatus::Create);
162        let check_result = self
163            .data_service
164            .metadata
165            .context
166            .check_and_fix_record(&command.entity, &mut command.values);
167        clear_record_status(&mut command.values);
168        check_result?;
169
170        Ok(command)
171    }
172
173    pub fn update_command(&self, id: impl Into<Value>) -> UpdateCommand {
174        UpdateCommand::new(self.entity.clone(), id)
175    }
176
177    pub fn prepare_update_command(
178        &self,
179        command: &UpdateCommand,
180    ) -> Result<UpdateCommand, RuntimeError> {
181        let mut command = command.clone();
182        if let Some(behavior) = self.behavior() {
183            behavior.before_update(self.data_service.metadata.context, &mut command)?;
184        }
185        self.enforce_update_policy(&mut command)?;
186
187        Ok(command)
188    }
189
190    pub fn delete_command(&self, id: impl Into<Value>) -> DeleteCommand {
191        DeleteCommand::new(self.entity.clone(), id)
192    }
193
194    pub fn recover_command(&self, id: impl Into<Value>, expected_version: i64) -> RecoverCommand {
195        RecoverCommand::new(self.entity.clone(), id, expected_version)
196    }
197
198    pub(crate) async fn fetch_all_internal(
199        &self,
200        query: &SelectQuery,
201    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
202        let query = self
203            .prepare_select_query(query)
204            .map_err(DataServiceError::Runtime)?;
205        let query = query
206            .prepare_for_list()
207            .map_err(|message| DataServiceError::Runtime(RuntimeError::Graph(message)))?;
208        if query.continuous_page_fetch.is_none()
209            && query.object_group_bys.is_empty()
210            && query.child_enhancements.is_empty()
211            && query.relations.is_empty()
212        {
213            return self.fetch_prepared_query_owned(query).await;
214        }
215        self.fetch_prepared_all(&query).await
216    }
217
218    async fn fetch_all_owned_internal(
219        &self,
220        query: SelectQuery,
221    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
222        let query = self
223            .prepare_select_query_owned(query)
224            .map_err(DataServiceError::Runtime)?
225            .prepare_for_list()
226            .map_err(|message| DataServiceError::Runtime(RuntimeError::Graph(message)))?;
227        if query.continuous_page_fetch.is_none()
228            && query.object_group_bys.is_empty()
229            && query.child_enhancements.is_empty()
230            && query.relations.is_empty()
231        {
232            return self.fetch_prepared_query_owned(query).await;
233        }
234        self.fetch_prepared_all(&query).await
235    }
236
237    async fn prepare_continuous_page(
238        &self,
239        query: SelectQuery,
240    ) -> (SelectQuery, Option<ContinuousPageExecution>) {
241        let Some(options) = query.continuous_page_fetch.as_ref() else {
242            self.data_service
243                .metadata
244                .context
245                .observe_continuous_page("DISABLED", None);
246            return (query, None);
247        };
248        let Some(slice) = query.slice.as_ref() else {
249            self.data_service
250                .metadata
251                .context
252                .observe_continuous_page("OFFSET_FALLBACK:INVALID_SLICE", None);
253            return (query, None);
254        };
255        let Some(page_size) = slice.limit else {
256            self.data_service
257                .metadata
258                .context
259                .observe_continuous_page("OFFSET_FALLBACK:INVALID_SLICE", None);
260            return (query, None);
261        };
262        if query.partition_by.is_some()
263            || !query.aggregates.is_empty()
264            || !query.group_by.is_empty()
265        {
266            self.data_service
267                .metadata
268                .context
269                .observe_continuous_page("OFFSET_FALLBACK:UNSUPPORTED_QUERY_SHAPE", None);
270            return (query, None);
271        }
272        if query.order_by.len() != 1
273            || query.order_by[0].field != "id"
274            || query.order_by[0].expr.is_some()
275        {
276            self.data_service
277                .metadata
278                .context
279                .observe_continuous_page("OFFSET_FALLBACK:ORDER_NOT_SEEKABLE_ID", None);
280            return (query, None);
281        }
282        let direction = query.order_by[0].direction;
283        let query_key = self.continuous_page_query_key(&query, &options.namespace);
284        let execution = ContinuousPageExecution {
285            query_key: query_key.clone(),
286            direction,
287            page_size,
288            original_offset: slice.offset,
289            ttl_seconds: options.ttl_seconds,
290            optimized: false,
291            seek_cursor_id: None,
292        };
293        if slice.offset == 0 {
294            self.data_service
295                .metadata
296                .context
297                .observe_continuous_page("OFFSET_FALLBACK:FIRST_PAGE", None);
298            return (query, Some(execution));
299        }
300        let cursor = match self
301            .data_service
302            .metadata
303            .context
304            .continuous_page_cursor_store()
305            .get(&query_key, slice.offset)
306            .await
307        {
308            Ok(Some(cursor)) => cursor,
309            Ok(None) => {
310                self.data_service
311                    .metadata
312                    .context
313                    .observe_continuous_page("OFFSET_FALLBACK:CACHE_MISS", None);
314                return (query, Some(execution));
315            }
316            Err(_) => {
317                self.data_service
318                    .metadata
319                    .context
320                    .observe_continuous_page("OFFSET_FALLBACK:STORE_UNAVAILABLE", None);
321                return (query, Some(execution));
322            }
323        };
324        if cursor.entity != query.entity
325            || cursor.direction != direction
326            || cursor.page_size != page_size
327            || cursor.next_offset != slice.offset
328            || cursor.expires_at <= SystemTime::now()
329        {
330            self.data_service
331                .metadata
332                .context
333                .observe_continuous_page("OFFSET_FALLBACK:CURSOR_INVALID", None);
334            return (query, Some(execution));
335        }
336        let mut optimized = query;
337        optimized.slice.as_mut().expect("validated slice").offset = 0;
338        optimized = optimized.and_filter(match direction {
339            SortDirection::Asc => Expr::gt("id", cursor.boundary.clone()),
340            SortDirection::Desc => Expr::lt("id", cursor.boundary.clone()),
341        });
342        let seek_cursor_id = cursor.cursor_id;
343        self.data_service
344            .metadata
345            .context
346            .observe_continuous_page("CURSOR_SEEK", Some(seek_cursor_id.clone()));
347        (
348            optimized,
349            Some(ContinuousPageExecution {
350                optimized: true,
351                seek_cursor_id: Some(seek_cursor_id),
352                ..execution
353            }),
354        )
355    }
356
357    async fn register_continuous_page(
358        &self,
359        execution: &Option<ContinuousPageExecution>,
360        rows: &[Record],
361    ) {
362        let Some(execution) = execution else { return };
363        if rows.len() as u64 != execution.page_size {
364            return;
365        }
366        let Some(boundary) = rows.last().and_then(|row| row.get("id")).cloned() else {
367            return;
368        };
369        let cursor_id = format!(
370            "cpg_{:x}",
371            SystemTime::now()
372                .duration_since(SystemTime::UNIX_EPOCH)
373                .unwrap_or_default()
374                .as_nanos()
375        );
376        let cursor = ContinuousPageCursor {
377            cursor_id,
378            query_key: execution.query_key.clone(),
379            entity: self.entity.clone(),
380            direction: execution.direction,
381            boundary,
382            page_size: execution.page_size,
383            next_offset: execution.original_offset + rows.len() as u64,
384            expires_at: SystemTime::now() + Duration::from_secs(execution.ttl_seconds),
385        };
386        if self
387            .data_service
388            .metadata
389            .context
390            .continuous_page_cursor_store()
391            .put(cursor)
392            .await
393            .is_err()
394        {
395            self.data_service
396                .metadata
397                .context
398                .observe_continuous_page("OFFSET_FALLBACK:STORE_UNAVAILABLE", None);
399        } else if execution.optimized {
400            self.data_service
401                .metadata
402                .context
403                .observe_continuous_page("CURSOR_SEEK", execution.seek_cursor_id.clone());
404        } else {
405            self.data_service
406                .metadata
407                .context
408                .observe_continuous_page("OFFSET_FALLBACK:FIRST_PAGE", None);
409        }
410    }
411
412    fn continuous_page_query_key(&self, query: &SelectQuery, namespace: &str) -> String {
413        let mut normalized = query.clone();
414        if let Some(slice) = normalized.slice.as_mut() {
415            slice.offset = 0;
416        }
417        normalized.comment = None;
418        normalized.trace_chain.clear();
419        let mut hasher = std::collections::hash_map::DefaultHasher::new();
420        namespace.hash(&mut hasher);
421        format!("{normalized:?}").hash(&mut hasher);
422        self.data_service
423            .metadata
424            .context
425            .user_identifier()
426            .hash(&mut hasher);
427        format!("teaql:continuous-page:v1:{:016x}", hasher.finish())
428    }
429
430    /// Fetch root records from the provider cursor without materializing them.
431    /// Relation and aggregate enhancement needs a separate batched protocol and
432    /// is rejected here instead of silently returning incomplete entities.
433    pub(crate) async fn fetch_stream_internal(
434        &self,
435        query: &SelectQuery,
436    ) -> Result<
437        std::pin::Pin<
438            Box<
439                dyn futures_core::Stream<
440                        Item = Result<teaql_data_service::StreamChunk, DataServiceError<E::Error>>,
441                    > + '_,
442            >,
443        >,
444        DataServiceError<E::Error>,
445    >
446    where
447        E: teaql_data_service::StreamQueryExecutor,
448    {
449        let query = self
450            .prepare_select_query(query)
451            .map_err(DataServiceError::Runtime)?;
452        let query = query
453            .prepare_for_list()
454            .map_err(|message| DataServiceError::Runtime(RuntimeError::Graph(message)))?;
455
456        if !query.relations.is_empty()
457            || !query.child_enhancements.is_empty()
458            || !query.object_group_bys.is_empty()
459        {
460            return Err(DataServiceError::Runtime(RuntimeError::Graph(
461                "streaming relation or aggregate enhancement is not supported; stream a root query or use execute_for_list"
462                    .to_owned(),
463            )));
464        }
465
466        let chunk_size = query
467            .stream_config
468            .as_ref()
469            .map(|c| c.chunk_size)
470            .unwrap_or(1000);
471
472        let final_comment = self
473            .data_service
474            .resolve_final_comment(&query.trace_chain, query.comment.clone());
475        let mut query = query.clone();
476        query.comment = final_comment;
477
478        let request = teaql_data_service::QueryRequest {
479            query: query.clone(),
480            trace_chain: query.trace_chain.clone(),
481            comment: query.comment.clone(),
482            capture_debug_query: self.data_service.metadata.capture_query_debug(),
483        };
484
485        let chunks = self.data_service.executor.query_stream(request, chunk_size);
486        use futures_util::StreamExt;
487        Ok(Box::pin(
488            chunks.map(|item| item.map_err(DataServiceError::Executor)),
489        ))
490    }
491
492    async fn fetch_prepared_all(
493        &self,
494        query: &SelectQuery,
495    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
496        let query = query
497            .clone()
498            .prepare_for_list()
499            .map_err(|message| DataServiceError::Runtime(RuntimeError::Graph(message)))?;
500        if query.continuous_page_fetch.is_none()
501            && query.object_group_bys.is_empty()
502            && query.child_enhancements.is_empty()
503            && query.relations.is_empty()
504        {
505            return self.fetch_prepared_query(&query).await;
506        }
507        let (execution_query, continuous) = self.prepare_continuous_page(query).await;
508        let mut rows = self.fetch_prepared_query(&execution_query).await?;
509        self.enhance_object_group_bys_internal(
510            &mut rows,
511            &execution_query.object_group_bys,
512            &execution_query.trace_chain,
513        )
514        .await?;
515        self.enhance_child_queries_internal(
516            &mut rows,
517            &execution_query.child_enhancements,
518            &execution_query.trace_chain,
519        )
520        .await?;
521        self.enhance_query_relations_internal(&mut rows, &execution_query)
522            .await?;
523        self.register_continuous_page(&continuous, &rows).await;
524        Ok(rows)
525    }
526
527    async fn fetch_prepared_query(
528        &self,
529        query: &SelectQuery,
530    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
531        let final_comment = self
532            .data_service
533            .resolve_final_comment(&query.trace_chain, query.comment.clone());
534        let mut query = query.clone();
535        query.comment = final_comment;
536        if let Some(options) = query.aggregation_cache.filter(|options| options.enabled) {
537            if let Some(cache) = self
538                .data_service
539                .metadata
540                .context
541                .get_resource::<Arc<dyn AggregationCacheBackend>>()
542            {
543                return self
544                    .fetch_prepared_query_with_cache(&query, options, cache.as_ref())
545                    .await;
546            }
547            if let Some(cache) = self
548                .data_service
549                .metadata
550                .context
551                .get_resource::<InMemoryAggregationCache>()
552            {
553                return self
554                    .fetch_prepared_query_with_cache(&query, options, cache)
555                    .await;
556            }
557        }
558        let request = teaql_data_service::QueryRequest {
559            query: query.clone(),
560            trace_chain: query.trace_chain.clone(),
561            comment: query.comment.clone(),
562            capture_debug_query: self.data_service.metadata.capture_query_debug(),
563        };
564        let res = self
565            .data_service
566            .executor
567            .query(request)
568            .await
569            .map_err(DataServiceError::Executor)?;
570        self.data_service
571            .metadata
572            .context
573            .record_metadata_log(&res.metadata);
574        Ok(res.rows)
575    }
576
577    async fn fetch_prepared_query_owned(
578        &self,
579        mut query: SelectQuery,
580    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
581        if query
582            .aggregation_cache
583            .is_some_and(|options| options.enabled)
584        {
585            return self.fetch_prepared_query(&query).await;
586        }
587        query.comment = self
588            .data_service
589            .resolve_final_comment(&query.trace_chain, query.comment.take());
590        let trace_chain = std::mem::take(&mut query.trace_chain);
591        let request = teaql_data_service::QueryRequest {
592            trace_chain,
593            comment: query.comment.clone(),
594            capture_debug_query: self.data_service.metadata.capture_query_debug(),
595            query,
596        };
597        let res = self
598            .data_service
599            .executor
600            .query(request)
601            .await
602            .map_err(DataServiceError::Executor)?;
603        self.data_service
604            .metadata
605            .context
606            .record_metadata_log(&res.metadata);
607        Ok(res.rows)
608    }
609
610    async fn fetch_prepared_query_with_cache(
611        &self,
612        query: &SelectQuery,
613        options: AggregationCacheOptions,
614        cache: &dyn AggregationCacheBackend,
615    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
616        let key = aggregation_cache_key(
617            cache.namespace(),
618            &aggregation_cache_namespace(&query.entity),
619            query,
620        );
621        let scope = self.data_service.metadata.context.start_runtime_operation(
622            crate::RuntimeOperation::new("cache", format!("{}.aggregation.get", query.entity))
623                .attribute("teaql.cache.operation", "get"),
624        );
625        let result = scope
626            .run(async {
627                if let Some(rows) = cache.get(&key, options.cache_expired_millis) {
628                    return Ok((rows, "hit"));
629                }
630                let request = teaql_data_service::QueryRequest {
631                    query: query.clone(),
632                    trace_chain: query.trace_chain.clone(),
633                    comment: query.comment.clone(),
634                    capture_debug_query: self.data_service.metadata.capture_query_debug(),
635                };
636                let provider_kind = std::any::type_name::<E>().to_owned();
637                let provider_scope = self.data_service.metadata.context.start_runtime_operation(
638                    crate::RuntimeOperation::new("provider", format!("{provider_kind}.query"))
639                        .attribute("teaql.provider.kind", provider_kind)
640                        .attribute("teaql.provider.operation", "query"),
641                );
642                let provider_result = provider_scope
643                    .run(self.data_service.executor.query(request))
644                    .await;
645                let res = match provider_result {
646                    Ok(value) => {
647                        provider_scope.success(std::collections::BTreeMap::new());
648                        value
649                    }
650                    Err(error) => {
651                        provider_scope.failure("data_service_error");
652                        return Err(DataServiceError::Executor(error));
653                    }
654                };
655                self.data_service
656                    .metadata
657                    .context
658                    .record_metadata_log(&res.metadata);
659                let rows = res.rows;
660                cache.put(key, rows.clone());
661                Ok((rows, "miss"))
662            })
663            .await;
664        match result {
665            Ok((rows, cache_result)) => {
666                scope.success(std::collections::BTreeMap::from([(
667                    "teaql.cache.result".to_owned(),
668                    crate::RuntimeAttributeValue::from(cache_result),
669                )]));
670                Ok(rows)
671            }
672            Err(error) => {
673                scope.failure("cache_load_error");
674                Err(error)
675            }
676        }
677    }
678
679    pub(crate) async fn fetch_all_with_relation_aggregates_internal(
680        &self,
681        query: &SelectQuery,
682        relation_aggregates: &[RelationAggregate],
683    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
684        let query = self
685            .prepare_select_query(query)
686            .map_err(DataServiceError::Runtime)?;
687
688        let mut rows = self.fetch_prepared_all(&query).await?;
689        self.enhance_relation_aggregates_internal(
690            &mut rows,
691            relation_aggregates,
692            query.aggregation_cache,
693            &query.trace_chain,
694        )
695        .await?;
696        Ok(rows)
697    }
698
699    pub(crate) async fn fetch_smart_list_internal(
700        &self,
701        query: &SelectQuery,
702    ) -> Result<SmartList<Record>, DataServiceError<E::Error>> {
703        let query = self
704            .prepare_select_query(query)
705            .map_err(DataServiceError::Runtime)?;
706
707        self.data_service.fetch_smart_list(&query).await
708    }
709
710    pub(crate) async fn fetch_smart_list_with_relation_aggregates_internal(
711        &self,
712        query: &SelectQuery,
713        relation_aggregates: &[RelationAggregate],
714    ) -> Result<SmartList<Record>, DataServiceError<E::Error>> {
715        self.fetch_all_with_relation_aggregates_internal(query, relation_aggregates)
716            .await
717            .map(SmartList::from)
718    }
719
720    pub(crate) async fn fetch_entities_internal<T>(
721        &self,
722        query: &SelectQuery,
723    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
724    where
725        T: Entity,
726    {
727        let query = self
728            .prepare_select_query(query)
729            .map_err(DataServiceError::Runtime)?;
730
731        self.data_service.fetch_entities(&query).await
732    }
733
734    pub(crate) async fn fetch_entities_with_relation_aggregates_internal<T>(
735        &self,
736        query: &SelectQuery,
737        relation_aggregates: &[RelationAggregate],
738    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
739    where
740        T: Entity,
741    {
742        self.fetch_all_with_relation_aggregates_internal(query, relation_aggregates)
743            .await?
744            .into_iter()
745            .map(|record| {
746                let mut entity = T::from_record(record)?;
747                let root = crate::EntityRoot::default();
748                entity.on_loaded(&root as &dyn std::any::Any);
749                Ok(entity)
750            })
751            .collect::<Result<Vec<_>, _>>()
752            .map(SmartList::from)
753            .map_err(DataServiceError::Entity)
754    }
755
756    pub(crate) async fn fetch_enhanced_entities_with_relation_aggregates_internal<T>(
757        &self,
758        query: &SelectQuery,
759        relation_aggregates: &[RelationAggregate],
760    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
761    where
762        T: Entity,
763    {
764        let query = self
765            .prepare_select_query(query)
766            .map_err(DataServiceError::Runtime)?;
767        self.fetch_enhanced_entities_with_relation_aggregates_prepared(query, relation_aggregates)
768            .await
769    }
770
771    async fn fetch_enhanced_entities_with_relation_aggregates_owned_internal<T>(
772        &self,
773        query: SelectQuery,
774        relation_aggregates: &[RelationAggregate],
775    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
776    where
777        T: Entity,
778    {
779        let query = self
780            .prepare_select_query_owned(query)
781            .map_err(DataServiceError::Runtime)?;
782        self.fetch_enhanced_entities_with_relation_aggregates_prepared(query, relation_aggregates)
783            .await
784    }
785
786    async fn fetch_enhanced_entities_with_relation_aggregates_prepared<T>(
787        &self,
788        query: SelectQuery,
789        relation_aggregates: &[RelationAggregate],
790    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
791    where
792        T: Entity,
793    {
794        if relation_aggregates.is_empty()
795            && query.continuous_page_fetch.is_none()
796            && query.object_group_bys.is_empty()
797            && query.child_enhancements.is_empty()
798            && query.relations.is_empty()
799        {
800            let query = query
801                .prepare_for_list()
802                .map_err(|message| DataServiceError::Runtime(RuntimeError::Graph(message)))?;
803            return self
804                .fetch_prepared_query_owned(query)
805                .await?
806                .into_iter()
807                .map(|record| {
808                    let mut entity = T::from_record(record)?;
809                    let root = crate::EntityRoot::default();
810                    entity.on_loaded(&root as &dyn std::any::Any);
811                    Ok(entity)
812                })
813                .collect::<Result<Vec<_>, _>>()
814                .map(SmartList::from)
815                .map_err(DataServiceError::Entity);
816        }
817
818        let mut rows = self.fetch_prepared_all(&query).await?;
819        self.enhance_relation_aggregates_internal(
820            &mut rows,
821            relation_aggregates,
822            query.aggregation_cache,
823            &query.trace_chain,
824        )
825        .await?;
826        self.enhance_relations_internal(&mut rows).await?;
827        rows.into_iter()
828            .map(|record| {
829                let mut entity = T::from_record(record)?;
830                let root = crate::EntityRoot::default();
831                entity.on_loaded(&root as &dyn std::any::Any);
832                Ok(entity)
833            })
834            .collect::<Result<Vec<_>, _>>()
835            .map(SmartList::from)
836            .map_err(DataServiceError::Entity)
837    }
838
839    pub(crate) async fn fetch_enhanced_entities_internal<T>(
840        &self,
841        query: &SelectQuery,
842    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
843    where
844        T: Entity,
845    {
846        let query = self
847            .prepare_select_query(query)
848            .map_err(DataServiceError::Runtime)?;
849
850        let mut rows = self.fetch_prepared_all(&query).await?;
851        self.enhance_relations_internal(&mut rows).await?;
852        let root = self
853            .data_service
854            .metadata
855            .context
856            .get_resource::<crate::EntityRoot>()
857            .cloned();
858        rows.into_iter()
859            .map(|record| {
860                let mut entity = T::from_record(record)?;
861                if let Some(ref root) = root {
862                    entity.on_loaded(root as &dyn std::any::Any);
863                }
864                Ok(entity)
865            })
866            .collect::<Result<Vec<_>, _>>()
867            .map(SmartList::from)
868            .map_err(DataServiceError::Entity)
869    }
870
871    #[doc(hidden)]
872    pub async fn fetch_all(
873        &self,
874        query: &PurposedSelectQuery,
875    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
876        self.fetch_all_internal(query.as_query()).await
877    }
878
879    #[doc(hidden)]
880    pub async fn fetch_all_owned(
881        &self,
882        query: PurposedSelectQuery,
883    ) -> Result<Vec<Record>, DataServiceError<E::Error>> {
884        self.fetch_all_owned_internal(query.into_query()).await
885    }
886
887    #[doc(hidden)]
888    pub async fn fetch_stream(
889        &self,
890        query: &PurposedSelectQuery,
891    ) -> Result<
892        std::pin::Pin<
893            Box<
894                dyn futures_core::Stream<
895                        Item = Result<teaql_data_service::StreamChunk, DataServiceError<E::Error>>,
896                    > + '_,
897            >,
898        >,
899        DataServiceError<E::Error>,
900    >
901    where
902        E: teaql_data_service::StreamQueryExecutor,
903    {
904        self.fetch_stream_internal(query.as_query()).await
905    }
906
907    #[doc(hidden)]
908    pub async fn fetch_smart_list(
909        &self,
910        query: &PurposedSelectQuery,
911    ) -> Result<SmartList<Record>, DataServiceError<E::Error>> {
912        self.fetch_smart_list_internal(query.as_query()).await
913    }
914
915    #[doc(hidden)]
916    pub async fn fetch_smart_list_with_relation_aggregates(
917        &self,
918        query: &PurposedSelectQuery,
919        relation_aggregates: &[RelationAggregate],
920    ) -> Result<SmartList<Record>, DataServiceError<E::Error>> {
921        self.fetch_smart_list_with_relation_aggregates_internal(
922            query.as_query(),
923            relation_aggregates,
924        )
925        .await
926    }
927
928    #[doc(hidden)]
929    pub async fn fetch_entities<T>(
930        &self,
931        query: &PurposedSelectQuery,
932    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
933    where
934        T: Entity,
935    {
936        self.fetch_entities_internal(query.as_query()).await
937    }
938
939    #[doc(hidden)]
940    pub async fn fetch_enhanced_entities<T>(
941        &self,
942        query: &PurposedSelectQuery,
943    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
944    where
945        T: Entity,
946    {
947        self.fetch_enhanced_entities_internal(query.as_query())
948            .await
949    }
950
951    #[doc(hidden)]
952    pub async fn fetch_enhanced_entities_with_relation_aggregates<T>(
953        &self,
954        query: &PurposedSelectQuery,
955        relation_aggregates: &[RelationAggregate],
956    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
957    where
958        T: Entity,
959    {
960        self.fetch_enhanced_entities_with_relation_aggregates_internal(
961            query.as_query(),
962            relation_aggregates,
963        )
964        .await
965    }
966
967    #[doc(hidden)]
968    pub async fn fetch_enhanced_entities_with_relation_aggregates_owned<T>(
969        &self,
970        query: PurposedSelectQuery,
971        relation_aggregates: &[RelationAggregate],
972    ) -> Result<SmartList<T>, DataServiceError<E::Error>>
973    where
974        T: Entity,
975    {
976        self.fetch_enhanced_entities_with_relation_aggregates_owned_internal(
977            query.into_query(),
978            relation_aggregates,
979        )
980        .await
981    }
982
983    pub(crate) async fn insert_internal(
984        &self,
985        command: &InsertCommand,
986    ) -> Result<u64, DataServiceError<E::Error>> {
987        let command = self
988            .prepare_insert_command(command)
989            .map_err(DataServiceError::Runtime)?;
990        self.execute_prepared_insert_with_comment(command, self.trace_context.clone())
991            .await
992    }
993
994    pub(crate) async fn update_internal(
995        &self,
996        command: &UpdateCommand,
997    ) -> Result<u64, DataServiceError<E::Error>> {
998        let command = self
999            .prepare_update_command(command)
1000            .map_err(DataServiceError::Runtime)?;
1001        self.execute_prepared_update_with_comment(command, self.trace_context.clone())
1002            .await
1003    }
1004
1005    pub(crate) async fn delete_internal(
1006        &self,
1007        command: &DeleteCommand,
1008    ) -> Result<u64, DataServiceError<E::Error>> {
1009        self.delete_scoped_internal(command, self.trace_context.clone())
1010            .await
1011    }
1012
1013    pub(crate) async fn delete_scoped_internal(
1014        &self,
1015        command: &DeleteCommand,
1016        trace_chain: Vec<teaql_core::TraceNode>,
1017    ) -> Result<u64, DataServiceError<E::Error>> {
1018        let mut command = command.clone();
1019        command.trace_chain = trace_chain.clone();
1020        if let Some(behavior) = self.behavior() {
1021            behavior
1022                .before_delete(self.data_service.metadata.context, &mut command)
1023                .map_err(DataServiceError::Runtime)?;
1024        }
1025        self.enforce_delete_policy(&mut command)
1026            .map_err(DataServiceError::Runtime)?;
1027
1028        let old_values =
1029            self.fetch_current_event_row(&command.entity, &command.id, trace_chain.clone())?;
1030        let affected = self.data_service.delete(&command).await?;
1031
1032        let mut event = RawAuditEvent::deleted_with_old_values(
1033            command.entity,
1034            command.id,
1035            command.expected_version,
1036            old_values,
1037        );
1038        event.trace_chain = trace_chain;
1039        self.emit_event(event).map_err(DataServiceError::Runtime)?;
1040        Ok(affected)
1041    }
1042
1043    pub(crate) async fn recover_internal(
1044        &self,
1045        command: &RecoverCommand,
1046    ) -> Result<u64, DataServiceError<E::Error>> {
1047        let mut command = command.clone();
1048        command.trace_chain = self.trace_context.clone();
1049        if let Some(behavior) = self.behavior() {
1050            behavior
1051                .before_recover(self.data_service.metadata.context, &mut command)
1052                .map_err(DataServiceError::Runtime)?;
1053        }
1054        self.enforce_recover_policy(&mut command)
1055            .map_err(DataServiceError::Runtime)?;
1056        let old_values = self.fetch_current_event_row(
1057            &command.entity,
1058            &command.id,
1059            command.trace_chain.clone(),
1060        )?;
1061        let affected = self.data_service.recover(&command).await?;
1062        let event = RawAuditEvent::recovered_with_old_values(
1063            command.entity,
1064            command.id,
1065            command.expected_version,
1066            old_values,
1067        );
1068        self.emit_event(event).map_err(DataServiceError::Runtime)?;
1069        Ok(affected)
1070    }
1071
1072    fn emit_event(&self, event: RawAuditEvent) -> Result<(), RuntimeError> {
1073        self.data_service.metadata.context.send_event(event)
1074    }
1075
1076    #[allow(dead_code)]
1077    pub(super) async fn execute_prepared_insert(
1078        &self,
1079        command: InsertCommand,
1080    ) -> Result<u64, DataServiceError<E::Error>> {
1081        self.execute_prepared_insert_with_comment(command, Vec::new())
1082            .await
1083    }
1084
1085    pub(super) async fn execute_prepared_insert_with_comment(
1086        &self,
1087        mut command: InsertCommand,
1088        trace_chain: Vec<teaql_core::TraceNode>,
1089    ) -> Result<u64, DataServiceError<E::Error>> {
1090        command.trace_chain = trace_chain.clone();
1091        let affected = self.data_service.insert(&command).await?;
1092        let mut event = RawAuditEvent::created(command.entity, command.values);
1093        event.trace_chain = trace_chain;
1094        self.emit_event(event).map_err(DataServiceError::Runtime)?;
1095        Ok(affected)
1096    }
1097
1098    pub(super) async fn execute_prepared_batch_insert(
1099        &self,
1100        command: teaql_core::BatchInsertCommand,
1101    ) -> Result<u64, DataServiceError<E::Error>> {
1102        if command.batch_values.is_empty() {
1103            return Ok(0);
1104        }
1105        let affected = self.data_service.batch_insert(&command).await?;
1106
1107        let entity = command.entity.clone();
1108        for (i, values) in command.batch_values.into_iter().enumerate() {
1109            let mut event = RawAuditEvent::created(entity.clone(), values);
1110            if i < command.trace_chains.len() {
1111                event.trace_chain = command.trace_chains[i].clone();
1112            }
1113            self.emit_event(event).map_err(DataServiceError::Runtime)?;
1114        }
1115        Ok(affected)
1116    }
1117
1118    #[allow(dead_code)]
1119    pub(super) async fn execute_prepared_update(
1120        &self,
1121        command: UpdateCommand,
1122    ) -> Result<u64, DataServiceError<E::Error>> {
1123        self.execute_prepared_update_with_comment(command, Vec::new())
1124            .await
1125    }
1126
1127    pub(super) async fn execute_prepared_update_with_comment(
1128        &self,
1129        mut command: UpdateCommand,
1130        trace_chain: Vec<teaql_core::TraceNode>,
1131    ) -> Result<u64, DataServiceError<E::Error>> {
1132        command.trace_chain = trace_chain.clone();
1133
1134        let mut old_values = command.old_values.clone();
1135        let needs_fetch = match &old_values {
1136            Some(snapshot) => !command.values.keys().all(|k| snapshot.contains_key(k)),
1137            None => true,
1138        };
1139        if needs_fetch {
1140            old_values =
1141                self.fetch_current_event_row(&command.entity, &command.id, trace_chain.clone())?;
1142        }
1143
1144        let affected = self.data_service.update(&command).await?;
1145        let updated_fields = command.values.keys().cloned().collect();
1146        let mut values = command.values.clone();
1147        values.insert("id".to_owned(), command.id.clone());
1148        if let Some(version) = command.expected_version {
1149            values.insert("version".to_owned(), Value::I64(version + 1));
1150        }
1151        let mut new_values = old_values.clone().unwrap_or_default();
1152        for (field, value) in &values {
1153            new_values.insert(field.clone(), value.clone());
1154        }
1155        let mut event = RawAuditEvent::updated_with_old_values(
1156            command.entity,
1157            values,
1158            old_values,
1159            new_values,
1160            updated_fields,
1161        );
1162        event.trace_chain = trace_chain;
1163        self.emit_event(event).map_err(DataServiceError::Runtime)?;
1164        Ok(affected)
1165    }
1166
1167    pub(super) async fn execute_prepared_batch_update(
1168        &self,
1169        command: teaql_core::BatchUpdateCommand,
1170    ) -> Result<u64, DataServiceError<E::Error>> {
1171        if command.batch_values.is_empty() {
1172            return Ok(0);
1173        }
1174        let affected = self.data_service.batch_update(&command).await?;
1175
1176        let entity = command.entity.clone();
1177        for (i, values) in command.batch_values.into_iter().enumerate() {
1178            let mut full_values = values.clone();
1179            full_values.insert("id".to_owned(), command.batch_ids[i].clone());
1180            if let Some(Some(version)) = command.batch_expected_versions.get(i) {
1181                full_values.insert("version".to_owned(), teaql_core::Value::I64(*version + 1));
1182            }
1183
1184            let old_values = command.batch_old_values.get(i).cloned().unwrap_or(None);
1185            let mut new_values = old_values.clone().unwrap_or_default();
1186            for (field, value) in &full_values {
1187                new_values.insert(field.clone(), value.clone());
1188            }
1189
1190            let mut event = RawAuditEvent::updated_with_old_values(
1191                entity.clone(),
1192                full_values,
1193                old_values,
1194                new_values,
1195                command.update_fields.clone(),
1196            );
1197            if i < command.trace_chains.len() {
1198                event.trace_chain = command.trace_chains[i].clone();
1199            }
1200            self.emit_event(event).map_err(DataServiceError::Runtime)?;
1201        }
1202        Ok(affected)
1203    }
1204
1205    fn fetch_current_event_row(
1206        &self,
1207        _entity: &str,
1208        _id: &Value,
1209        _trace_chain: Vec<teaql_core::TraceNode>,
1210    ) -> Result<Option<Record>, DataServiceError<E::Error>> {
1211        // PER THE USER: "我们不需要在审计的时候去抓旧的值"
1212        // Avoid DB queries during event emission. We rely on in-memory `original_values`.
1213        Ok(None)
1214    }
1215
1216    pub(crate) fn scoped_data_service_internal(&self, entity: String) -> EntityDataService<'a, E> {
1217        EntityDataService {
1218            entity,
1219            data_service: ContextDataService {
1220                metadata: UserContextMetadata {
1221                    context: self.data_service.metadata.context,
1222                },
1223                executor: self.data_service.executor,
1224            },
1225            trace_context: Vec::new(),
1226        }
1227    }
1228}