1use crate::db_set::DbSet;
18use crate::entity::{
19 EntityState, IEntitySnapshot, IEntityType, IFromRow, IGetKeyValues, INavigationSetter,
20};
21use crate::error::{EFError, EFResult};
22use crate::interceptor::InterceptorPipeline;
23use crate::metadata::EntityTypeMeta;
24use crate::migration::MigrationEngine;
25use crate::model_builder::ModelBuilder;
26use crate::provider::{DbValue, IDatabaseProvider, IsolationLevel};
27use crate::session::{Session, SessionConnection, SessionProvider};
28use crate::tracking::ChangeTracker;
29use crate::transaction::{DbTransaction, ITransaction, TxOutcome};
30use std::any::{Any, TypeId};
31use std::collections::HashMap;
32use std::future::Future;
33use std::pin::Pin;
34use std::sync::Arc;
35
36use super::{DbContextOptions, ErasedSetOps, SetOps};
37
38pub struct DbContext {
45 pub(crate) sets: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
46 pub(crate) savers: HashMap<TypeId, Box<dyn ErasedSetOps>>,
47 pub(crate) entity_metas: HashMap<TypeId, EntityTypeMeta>,
48 pub(crate) model_builder: ModelBuilder,
49 pub(crate) change_tracker: ChangeTracker,
50 pub(crate) provider: Arc<dyn IDatabaseProvider>,
51 pub(crate) context_key: Option<String>,
55 pub(crate) session: Arc<Session>,
59 pub(crate) interceptor_pipeline: InterceptorPipeline,
60 pub(crate) lazy_loading_enabled: bool,
61 pub(crate) ambient_transaction: Option<Box<dyn ITransaction>>,
66}
67
68impl DbContext {
69 pub fn from_options(options: &DbContextOptions) -> EFResult<Self> {
77 let write_provider = options.create_provider()?;
78 let session = Arc::new(Session::new(write_provider));
79 let provider: Arc<dyn IDatabaseProvider> =
80 Arc::new(SessionProvider::new(Arc::clone(&session)));
81 let built = options
82 .metadata_cache
83 .get_or_build(options.context_key.as_deref());
84 let ctx = Self {
85 sets: HashMap::new(),
86 savers: HashMap::new(),
87 entity_metas: built.entity_metas.clone(),
88 model_builder: ModelBuilder::from_built(&built),
89 change_tracker: ChangeTracker::new(),
90 provider,
91 context_key: options.context_key.clone(),
92 session,
93 interceptor_pipeline: InterceptorPipeline::new(options.interceptors.clone()),
94 lazy_loading_enabled: options.lazy_loading_enabled,
95 ambient_transaction: None,
96 };
97 Ok(ctx)
98 }
99
100 pub fn set<T>(&mut self) -> &mut DbSet<T>
101 where
102 T: IEntityType
103 + IEntitySnapshot
104 + IGetKeyValues
105 + IFromRow
106 + INavigationSetter
107 + Send
108 + Sync
109 + 'static,
110 {
111 let type_id = TypeId::of::<T>();
112 self.savers
113 .entry(type_id)
114 .or_insert_with(|| Box::new(SetOps::<T>::new()));
115 self.entity_metas
116 .entry(type_id)
117 .or_insert_with(T::entity_meta);
118 if !self.model_builder.has_entity(type_id) {
119 self.model_builder.register_entity_meta(T::entity_meta());
120 }
121 self.sets.entry(type_id).or_insert_with(|| {
122 let table_name = self
123 .model_builder
124 .build()
125 .into_iter()
126 .find(|m| m.type_id == type_id)
127 .map(|m| m.table_name.to_string())
128 .unwrap_or_else(|| T::entity_meta().table_name.to_string());
129 let mut db_set = DbSet::<T>::with_provider(table_name, Arc::clone(&self.provider));
130 if let Some(filter) = self.model_builder.get_query_filter(&type_id) {
131 db_set.set_query_filter(filter.clone());
132 }
133 db_set.set_filter_map(self.model_builder.filters_by_table());
134 db_set.set_lazy_loading_enabled(self.lazy_loading_enabled);
135 Box::new(db_set)
136 });
137 self.sets
138 .get_mut(&type_id)
139 .and_then(|b| b.downcast_mut::<DbSet<T>>())
140 .expect("DbSet type mismatch")
141 }
142
143 pub fn model(&mut self) -> &mut ModelBuilder {
145 &mut self.model_builder
146 }
147
148 pub fn model_builder(&self) -> &ModelBuilder {
150 &self.model_builder
151 }
152
153 pub fn entity_metas_contains<T: IEntityType>(&self) -> bool {
156 self.entity_metas.contains_key(&TypeId::of::<T>())
157 }
158
159 pub fn discover_entities(&mut self) -> EFResult<()> {
171 Ok(())
172 }
173
174 pub fn detect_changes(&mut self) {
176 let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
177 for type_id in type_ids {
178 let Some(set) = self.sets.get_mut(&type_id) else {
179 continue;
180 };
181 let Some(saver) = self.savers.get(&type_id) else {
182 continue;
183 };
184 saver.detect_changes(set.as_mut(), &mut self.change_tracker);
185 }
186 }
187
188 pub async fn ensure_created(&self) -> EFResult<()> {
195 let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
196 if metas.is_empty() {
197 metas = self.entity_metas.values().cloned().collect();
198 }
199 if metas.is_empty() {
200 return Err(self.empty_model_error("ensure_created"));
201 }
202 let dialect = self.provider.migration_dialect();
203 MigrationEngine::new(dialect)
204 .ensure_created(&*self.provider, &metas)
205 .await?;
206
207 for meta in &metas {
208 let rows = self.model_builder.seed_rows_for(&meta.type_id);
209 if !rows.is_empty() {
210 MigrationEngine::new(dialect)
211 .apply_seed_data(&*self.provider, meta, rows)
212 .await?;
213 }
214 }
215 Ok(())
216 }
217
218 pub async fn ensure_deleted(&self) -> EFResult<()> {
220 let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
221 if metas.is_empty() {
222 metas = self.entity_metas.values().cloned().collect();
223 }
224 if metas.is_empty() {
225 return Err(self.empty_model_error("ensure_deleted"));
226 }
227 let dialect = self.provider.migration_dialect();
228 MigrationEngine::new(dialect)
229 .ensure_deleted(&*self.provider, &metas)
230 .await
231 }
232
233 fn empty_model_error(&self, operation: &str) -> EFError {
235 let mut message = format!(
236 "No entity types registered. Call ctx.set::<T>() or ctx.discover_entities() before {operation}()."
237 );
238 if let Some(key) = &self.context_key {
239 message.push_str(&format!(
240 " This is keyed context '{key}', which only discovers entities tagged #[context(\"{key}\")]. \
241 For a read/write replica sharing the default schema, call \
242 DbContextOptionsBuilder::use_default_entity_model()."
243 ));
244 }
245 EFError::configuration(message)
246 }
247}
248
249impl DbContext {
254 pub fn provider(&self) -> &dyn IDatabaseProvider {
256 &*self.provider
257 }
258
259 pub fn change_tracker(&self) -> &ChangeTracker {
261 &self.change_tracker
262 }
263
264 pub fn change_tracker_mut(&mut self) -> &mut ChangeTracker {
266 &mut self.change_tracker
267 }
268
269 pub fn clear_tracking(&mut self) {
277 let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
278 for type_id in type_ids {
279 let Some(saver) = self.savers.get(&type_id) else {
280 continue;
281 };
282 let Some(set) = self.sets.get_mut(&type_id) else {
283 continue;
284 };
285 saver.clear_entries(set.as_mut());
286 }
287 self.change_tracker.clear();
288 }
289
290 pub async fn sql_query<T: IFromRow + IEntityType>(
303 &self,
304 sql: &str,
305 params: &[DbValue],
306 ) -> EFResult<Vec<T>> {
307 let mut conn = self.provider.get_connection().await?;
308 let rows = conn.query(sql, params).await?;
309 crate::entity::materialize_entities(&rows)
310 }
311
312 pub fn transaction_mut(&mut self) -> Option<&mut Box<dyn ITransaction>> {
321 self.ambient_transaction.as_mut()
322 }
323
324 pub async fn begin_transaction(&mut self) -> EFResult<Box<dyn ITransaction>> {
335 self.begin_transaction_with_opt(None).await
336 }
337
338 pub async fn begin_transaction_with(
341 &mut self,
342 isolation: IsolationLevel,
343 ) -> EFResult<Box<dyn ITransaction>> {
344 self.begin_transaction_with_opt(Some(isolation)).await
345 }
346
347 async fn begin_transaction_with_opt(
348 &mut self,
349 isolation: Option<IsolationLevel>,
350 ) -> EFResult<Box<dyn ITransaction>> {
351 let mut conn = self.provider.get_connection().await?;
352 match isolation {
353 Some(level) => conn.begin_transaction_with(level).await?,
354 None => conn.begin_transaction().await?,
355 }
356 Ok(Box::new(DbTransaction::new(conn)))
357 }
358
359 pub async fn use_transaction<F, R>(&mut self, f: F) -> EFResult<R>
383 where
384 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
385 R: Send + 'static,
386 {
387 self.use_transaction_with_opt(None, f).await
388 }
389
390 pub async fn use_transaction_with<F, R>(
393 &mut self,
394 isolation: IsolationLevel,
395 f: F,
396 ) -> EFResult<R>
397 where
398 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
399 R: Send + 'static,
400 {
401 self.use_transaction_with_opt(Some(isolation), f).await
402 }
403
404 async fn use_transaction_with_opt<F, R>(
405 &mut self,
406 isolation: Option<IsolationLevel>,
407 f: F,
408 ) -> EFResult<R>
409 where
410 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
411 R: Send + 'static,
412 {
413 if self.ambient_transaction.is_some() {
414 return Err(EFError::transaction(
415 "ambient transaction already active; nested use_transaction is not supported",
416 ));
417 }
418 let shared = self.session.begin_with(isolation).await?;
419 let facade = SessionConnection::new(Arc::clone(&shared));
420 self.ambient_transaction = Some(Box::new(DbTransaction::new_shared(
421 shared,
422 facade,
423 Some(Arc::clone(&self.session)),
424 )));
425 let result = f(self).await;
426 let txn = self.ambient_transaction.take().ok_or_else(|| {
427 EFError::transaction("ambient_transaction was consumed during use_transaction closure")
428 })?;
429 match result {
430 Ok(r) => {
431 txn.commit().await?;
432 Ok(r)
433 }
434 Err(e) => {
435 let _ = txn.rollback().await;
436 Err(e)
437 }
438 }
439 }
440
441 pub async fn transaction<F, R>(&mut self, f: F) -> EFResult<R>
473 where
474 for<'a> F: FnOnce(
475 &'a mut Self,
476 )
477 -> Pin<Box<dyn Future<Output = EFResult<TxOutcome<R>>> + Send + 'a>>,
478 R: Send + 'static,
479 {
480 self.transaction_with_opt(None, f).await
481 }
482
483 pub async fn transaction_with<F, R>(&mut self, isolation: IsolationLevel, f: F) -> EFResult<R>
486 where
487 for<'a> F: FnOnce(
488 &'a mut Self,
489 )
490 -> Pin<Box<dyn Future<Output = EFResult<TxOutcome<R>>> + Send + 'a>>,
491 R: Send + 'static,
492 {
493 self.transaction_with_opt(Some(isolation), f).await
494 }
495
496 async fn transaction_with_opt<F, R>(
497 &mut self,
498 isolation: Option<IsolationLevel>,
499 f: F,
500 ) -> EFResult<R>
501 where
502 for<'a> F: FnOnce(
503 &'a mut Self,
504 )
505 -> Pin<Box<dyn Future<Output = EFResult<TxOutcome<R>>> + Send + 'a>>,
506 R: Send + 'static,
507 {
508 if self.ambient_transaction.is_some() {
509 return Err(EFError::transaction(
510 "ambient transaction already active; nested transaction is not supported",
511 ));
512 }
513 let shared = self.session.begin_with(isolation).await?;
514 let facade = SessionConnection::new(Arc::clone(&shared));
515 self.ambient_transaction = Some(Box::new(DbTransaction::new_shared(
516 shared,
517 facade,
518 Some(Arc::clone(&self.session)),
519 )));
520 let result = f(self).await;
521 let txn = self.ambient_transaction.take().ok_or_else(|| {
522 EFError::transaction("ambient_transaction was consumed during transaction closure")
523 })?;
524 match result {
525 Ok(TxOutcome::Commit(value)) => {
526 txn.commit().await?;
527 Ok(value)
528 }
529 Ok(TxOutcome::Rollback(value)) => {
530 txn.rollback().await?;
531 self.clear_tracking();
534 Ok(value)
535 }
536 Err(e) => {
537 let _ = txn.rollback().await;
538 self.clear_tracking();
541 Err(e)
542 }
543 }
544 }
545}
546
547impl DbContext {
552 fn ensure_registered<T>(&mut self) -> (TypeId, String)
555 where
556 T: IEntityType
557 + IEntitySnapshot
558 + IGetKeyValues
559 + IFromRow
560 + INavigationSetter
561 + Send
562 + Sync
563 + 'static,
564 {
565 let type_id = TypeId::of::<T>();
566 self.savers
567 .entry(type_id)
568 .or_insert_with(|| Box::new(SetOps::<T>::new()));
569 self.entity_metas
570 .entry(type_id)
571 .or_insert_with(T::entity_meta);
572 if !self.model_builder.has_entity(type_id) {
573 self.model_builder.register_entity_meta(T::entity_meta());
574 }
575 let type_name = self.entity_metas[&type_id].type_name.to_string();
576 (type_id, type_name)
577 }
578
579 pub fn add<T>(&mut self, entity: T)
582 where
583 T: IEntityType
584 + IEntitySnapshot
585 + IGetKeyValues
586 + IFromRow
587 + INavigationSetter
588 + Send
589 + Sync
590 + 'static,
591 {
592 let (type_id, type_name) = self.ensure_registered::<T>();
593 let entry_id =
594 self.change_tracker
595 .track(type_id, &type_name, EntityState::Added, None, false);
596 self.set::<T>().push_entry(entity, entry_id);
597 }
598
599 pub fn attach<T>(&mut self, entity: T)
602 where
603 T: IEntityType
604 + IEntitySnapshot
605 + IGetKeyValues
606 + IFromRow
607 + INavigationSetter
608 + Send
609 + Sync
610 + 'static,
611 {
612 let (type_id, type_name) = self.ensure_registered::<T>();
613 let snapshot = entity.snapshot();
614 let entry_id = self.change_tracker.track(
615 type_id,
616 &type_name,
617 EntityState::Unchanged,
618 Some(snapshot),
619 false,
620 );
621 self.set::<T>().push_entry(entity, entry_id);
622 }
623
624 pub fn update<T>(&mut self, entity: T)
627 where
628 T: IEntityType
629 + IEntitySnapshot
630 + IGetKeyValues
631 + IFromRow
632 + INavigationSetter
633 + Send
634 + Sync
635 + 'static,
636 {
637 let (type_id, type_name) = self.ensure_registered::<T>();
638 let entry_id =
639 self.change_tracker
640 .track(type_id, &type_name, EntityState::Modified, None, false);
641 self.set::<T>().push_entry(entity, entry_id);
642 }
643
644 pub fn upsert<T>(&mut self, entity: T)
646 where
647 T: IEntityType
648 + IEntitySnapshot
649 + IGetKeyValues
650 + IFromRow
651 + INavigationSetter
652 + Send
653 + Sync
654 + 'static,
655 {
656 let (type_id, type_name) = self.ensure_registered::<T>();
657 let entry_id =
658 self.change_tracker
659 .track(type_id, &type_name, EntityState::Added, None, true);
660 self.set::<T>().push_entry(entity, entry_id);
661 }
662
663 pub fn remove_at<T>(&mut self, index: usize) -> EFResult<()>
665 where
666 T: IEntityType
667 + IEntitySnapshot
668 + IGetKeyValues
669 + IFromRow
670 + INavigationSetter
671 + Send
672 + Sync
673 + 'static,
674 {
675 let entry_id = self
676 .set::<T>()
677 .entry_id_at(index)
678 .ok_or_else(|| EFError::not_found("Entity not found at the given index".to_string()))?;
679 self.change_tracker
680 .set_state(entry_id, EntityState::Deleted);
681 Ok(())
682 }
683
684 pub fn remove_all<T>(&mut self)
686 where
687 T: IEntityType
688 + IEntitySnapshot
689 + IGetKeyValues
690 + IFromRow
691 + INavigationSetter
692 + Send
693 + Sync
694 + 'static,
695 {
696 let entry_ids: Vec<u64> = self.set::<T>().entries.iter().map(|e| e.entry_id).collect();
697 for entry_id in entry_ids {
698 self.change_tracker
699 .set_state(entry_id, EntityState::Deleted);
700 }
701 }
702
703 pub async fn load_all<T>(&mut self) -> EFResult<()>
706 where
707 T: IEntityType
708 + IEntitySnapshot
709 + IGetKeyValues
710 + IFromRow
711 + INavigationSetter
712 + crate::entity::ILazyInit
713 + Send
714 + Sync
715 + 'static,
716 {
717 let type_id = TypeId::of::<T>();
718 let items = self.set::<T>().query().to_list().await?;
719 self.set::<T>().clear_entries();
720 self.change_tracker.clear_by_type(type_id);
721 for item in items {
722 self.attach::<T>(item);
723 }
724 Ok(())
725 }
726}