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};
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}
62
63impl DbContext {
64 pub fn from_options(options: &DbContextOptions) -> EFResult<Self> {
72 let write_provider = options.create_provider()?;
73 let session = Arc::new(Session::new(write_provider));
74 let provider: Arc<dyn IDatabaseProvider> =
75 Arc::new(SessionProvider::new(Arc::clone(&session)));
76 let built = options
77 .metadata_cache
78 .get_or_build(options.context_key.as_deref());
79 let ctx = Self {
80 sets: HashMap::new(),
81 savers: HashMap::new(),
82 entity_metas: built.entity_metas.clone(),
83 model_builder: ModelBuilder::from_built(&built),
84 change_tracker: ChangeTracker::new(),
85 provider,
86 context_key: options.context_key.clone(),
87 session,
88 interceptor_pipeline: InterceptorPipeline::new(options.interceptors.clone()),
89 lazy_loading_enabled: options.lazy_loading_enabled,
90 };
91 Ok(ctx)
92 }
93
94 pub fn set<T>(&mut self) -> &mut DbSet<T>
95 where
96 T: IEntityType
97 + IEntitySnapshot
98 + IGetKeyValues
99 + IFromRow
100 + INavigationSetter
101 + Send
102 + Sync
103 + 'static,
104 {
105 let type_id = TypeId::of::<T>();
106 self.savers
107 .entry(type_id)
108 .or_insert_with(|| Box::new(SetOps::<T>::new()));
109 self.entity_metas
110 .entry(type_id)
111 .or_insert_with(T::entity_meta);
112 if !self.model_builder.has_entity(type_id) {
113 self.model_builder.register_entity_meta(T::entity_meta());
114 }
115 self.sets.entry(type_id).or_insert_with(|| {
116 let table_name = self
117 .model_builder
118 .build()
119 .into_iter()
120 .find(|m| m.type_id == type_id)
121 .map(|m| m.table_name.to_string())
122 .unwrap_or_else(|| T::entity_meta().table_name.to_string());
123 let mut db_set = DbSet::<T>::with_provider(table_name, Arc::clone(&self.provider));
124 if let Some(filter) = self.model_builder.get_query_filter(&type_id) {
125 db_set.set_query_filter(filter.clone());
126 }
127 db_set.set_filter_map(self.model_builder.filters_by_table());
128 db_set.set_lazy_loading_enabled(self.lazy_loading_enabled);
129 Box::new(db_set)
130 });
131 self.sets
132 .get_mut(&type_id)
133 .and_then(|b| b.downcast_mut::<DbSet<T>>())
134 .expect("DbSet type mismatch")
135 }
136
137 pub fn model(&mut self) -> &mut ModelBuilder {
139 &mut self.model_builder
140 }
141
142 pub fn model_builder(&self) -> &ModelBuilder {
144 &self.model_builder
145 }
146
147 pub fn entity_metas_contains<T: IEntityType>(&self) -> bool {
150 self.entity_metas.contains_key(&TypeId::of::<T>())
151 }
152
153 pub fn discover_entities(&mut self) -> EFResult<()> {
165 Ok(())
166 }
167
168 pub fn detect_changes(&mut self) {
170 let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
171 for type_id in type_ids {
172 let Some(set) = self.sets.get_mut(&type_id) else {
173 continue;
174 };
175 let Some(saver) = self.savers.get(&type_id) else {
176 continue;
177 };
178 saver.detect_changes(set.as_mut(), &mut self.change_tracker);
179 }
180 }
181
182 pub async fn ensure_created(&self) -> EFResult<()> {
189 let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
190 if metas.is_empty() {
191 metas = self.entity_metas.values().cloned().collect();
192 }
193 if metas.is_empty() {
194 return Err(self.empty_model_error("ensure_created"));
195 }
196 let dialect = self.provider.migration_dialect();
197 MigrationEngine::new(dialect)
198 .ensure_created(&*self.provider, &metas)
199 .await?;
200
201 for meta in &metas {
202 if meta.is_view {
204 continue;
205 }
206 let rows = self.model_builder.seed_rows_for(&meta.type_id);
207 if !rows.is_empty() {
208 MigrationEngine::new(dialect)
209 .apply_seed_data(&*self.provider, meta, rows)
210 .await?;
211 }
212 }
213 Ok(())
214 }
215
216 pub async fn migrate(&self, dir: impl AsRef<std::path::Path>) -> EFResult<usize> {
224 let store = crate::migration::MigrationStore::new(dir.as_ref());
225 let migrations = store.load_all()?;
226 let dialect = self.provider.migration_dialect();
227 MigrationEngine::new(dialect)
228 .apply_pending(&*self.provider, &migrations)
229 .await
230 }
231
232 pub async fn ensure_deleted(&self) -> EFResult<()> {
234 let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
235 if metas.is_empty() {
236 metas = self.entity_metas.values().cloned().collect();
237 }
238 if metas.is_empty() {
239 return Err(self.empty_model_error("ensure_deleted"));
240 }
241 let dialect = self.provider.migration_dialect();
242 MigrationEngine::new(dialect)
243 .ensure_deleted(&*self.provider, &metas)
244 .await
245 }
246
247 fn empty_model_error(&self, operation: &str) -> EFError {
249 let mut message = format!(
250 "No entity types registered. Call ctx.set::<T>() or ctx.discover_entities() before {operation}()."
251 );
252 if let Some(key) = &self.context_key {
253 message.push_str(&format!(
254 " This is keyed context '{key}', which only discovers entities tagged #[context(\"{key}\")]. \
255 For a read/write replica sharing the default schema, call \
256 DbContextOptionsBuilder::use_default_entity_model()."
257 ));
258 }
259 EFError::configuration(message)
260 }
261}
262
263impl DbContext {
268 pub fn provider(&self) -> &dyn IDatabaseProvider {
270 &*self.provider
271 }
272
273 pub fn change_tracker(&self) -> &ChangeTracker {
275 &self.change_tracker
276 }
277
278 pub fn change_tracker_mut(&mut self) -> &mut ChangeTracker {
280 &mut self.change_tracker
281 }
282
283 pub fn clear_tracking(&mut self) {
290 let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
291 for type_id in type_ids {
292 let Some(saver) = self.savers.get(&type_id) else {
293 continue;
294 };
295 let Some(set) = self.sets.get_mut(&type_id) else {
296 continue;
297 };
298 saver.clear_entries(set.as_mut());
299 }
300 self.change_tracker.clear();
301 }
302
303 pub async fn sql_query<T: IFromRow + IEntityType>(
316 &self,
317 sql: &str,
318 params: &[DbValue],
319 ) -> EFResult<Vec<T>> {
320 let mut conn = self.provider.get_connection().await?;
321 let rows = conn.query(sql, params).await?;
322 crate::entity::materialize_entities(&rows)
323 }
324
325 pub async fn begin_transaction(&mut self) -> EFResult<Box<dyn ITransaction>> {
342 self.begin_transaction_with_opt(None).await
343 }
344
345 pub async fn begin_transaction_with(
348 &mut self,
349 isolation: IsolationLevel,
350 ) -> EFResult<Box<dyn ITransaction>> {
351 self.begin_transaction_with_opt(Some(isolation)).await
352 }
353
354 async fn begin_transaction_with_opt(
355 &mut self,
356 isolation: Option<IsolationLevel>,
357 ) -> EFResult<Box<dyn ITransaction>> {
358 let shared = self.session.begin_with(isolation).await?;
359 let facade = SessionConnection::new(Arc::clone(&shared));
360 Ok(Box::new(DbTransaction::new_shared(
361 shared,
362 facade,
363 Arc::clone(&self.session),
364 )))
365 }
366
367 pub async fn use_transaction<F, R>(&mut self, f: F) -> EFResult<R>
393 where
394 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
395 R: Send + 'static,
396 {
397 self.use_transaction_with_opt(None, f).await
398 }
399
400 pub async fn use_transaction_with<F, R>(
403 &mut self,
404 isolation: IsolationLevel,
405 f: F,
406 ) -> EFResult<R>
407 where
408 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
409 R: Send + 'static,
410 {
411 self.use_transaction_with_opt(Some(isolation), f).await
412 }
413
414 async fn use_transaction_with_opt<F, R>(
415 &mut self,
416 isolation: Option<IsolationLevel>,
417 f: F,
418 ) -> EFResult<R>
419 where
420 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
421 R: Send + 'static,
422 {
423 if self.session.is_active().await {
424 return Err(EFError::transaction(
425 "a transaction is already active on this DbContext",
426 ));
427 }
428 let shared = self.session.begin_with(isolation).await?;
429 let facade = SessionConnection::new(Arc::clone(&shared));
430 let txn: Box<dyn ITransaction> = Box::new(DbTransaction::new_shared(
431 shared,
432 facade,
433 Arc::clone(&self.session),
434 ));
435 let result = f(self).await;
438 match result {
439 Ok(r) => {
440 txn.commit().await?;
441 Ok(r)
442 }
443 Err(e) => {
444 let _ = txn.rollback().await;
445 Err(e)
446 }
447 }
448 }
449}
450
451impl DbContext {
456 fn ensure_registered<T>(&mut self) -> (TypeId, String)
459 where
460 T: IEntityType
461 + IEntitySnapshot
462 + IGetKeyValues
463 + IFromRow
464 + INavigationSetter
465 + Send
466 + Sync
467 + 'static,
468 {
469 let type_id = TypeId::of::<T>();
470 self.savers
471 .entry(type_id)
472 .or_insert_with(|| Box::new(SetOps::<T>::new()));
473 self.entity_metas
474 .entry(type_id)
475 .or_insert_with(T::entity_meta);
476 if !self.model_builder.has_entity(type_id) {
477 self.model_builder.register_entity_meta(T::entity_meta());
478 }
479 let type_name = self.entity_metas[&type_id].type_name.to_string();
480 (type_id, type_name)
481 }
482
483 pub fn add<T>(&mut self, entity: T)
486 where
487 T: IEntityType
488 + IEntitySnapshot
489 + IGetKeyValues
490 + IFromRow
491 + INavigationSetter
492 + Send
493 + Sync
494 + 'static,
495 {
496 let (type_id, type_name) = self.ensure_registered::<T>();
497 let entry_id =
498 self.change_tracker
499 .track(type_id, &type_name, EntityState::Added, None, false);
500 self.set::<T>().push_entry(entity, entry_id);
501 }
502
503 pub fn attach<T>(&mut self, entity: T)
506 where
507 T: IEntityType
508 + IEntitySnapshot
509 + IGetKeyValues
510 + IFromRow
511 + INavigationSetter
512 + Send
513 + Sync
514 + 'static,
515 {
516 let (type_id, type_name) = self.ensure_registered::<T>();
517 let snapshot = entity.snapshot();
518 let entry_id = self.change_tracker.track(
519 type_id,
520 &type_name,
521 EntityState::Unchanged,
522 Some(snapshot),
523 false,
524 );
525 self.set::<T>().push_entry(entity, entry_id);
526 }
527
528 pub fn update<T>(&mut self, entity: T)
531 where
532 T: IEntityType
533 + IEntitySnapshot
534 + IGetKeyValues
535 + IFromRow
536 + INavigationSetter
537 + Send
538 + Sync
539 + 'static,
540 {
541 let (type_id, type_name) = self.ensure_registered::<T>();
542 let entry_id =
543 self.change_tracker
544 .track(type_id, &type_name, EntityState::Modified, None, false);
545 self.set::<T>().push_entry(entity, entry_id);
546 }
547
548 pub fn upsert<T>(&mut self, entity: T)
550 where
551 T: IEntityType
552 + IEntitySnapshot
553 + IGetKeyValues
554 + IFromRow
555 + INavigationSetter
556 + Send
557 + Sync
558 + 'static,
559 {
560 let (type_id, type_name) = self.ensure_registered::<T>();
561 let entry_id =
562 self.change_tracker
563 .track(type_id, &type_name, EntityState::Added, None, true);
564 self.set::<T>().push_entry(entity, entry_id);
565 }
566
567 pub fn remove_at<T>(&mut self, index: usize) -> EFResult<()>
569 where
570 T: IEntityType
571 + IEntitySnapshot
572 + IGetKeyValues
573 + IFromRow
574 + INavigationSetter
575 + Send
576 + Sync
577 + 'static,
578 {
579 let entry_id = self
580 .set::<T>()
581 .entry_id_at(index)
582 .ok_or_else(|| EFError::not_found("Entity not found at the given index".to_string()))?;
583 self.change_tracker
584 .set_state(entry_id, EntityState::Deleted);
585 Ok(())
586 }
587
588 pub fn remove_all<T>(&mut self)
590 where
591 T: IEntityType
592 + IEntitySnapshot
593 + IGetKeyValues
594 + IFromRow
595 + INavigationSetter
596 + Send
597 + Sync
598 + 'static,
599 {
600 let entry_ids: Vec<u64> = self.set::<T>().entries.iter().map(|e| e.entry_id).collect();
601 for entry_id in entry_ids {
602 self.change_tracker
603 .set_state(entry_id, EntityState::Deleted);
604 }
605 }
606
607 pub async fn load_all<T>(&mut self) -> EFResult<()>
610 where
611 T: IEntityType
612 + IEntitySnapshot
613 + IGetKeyValues
614 + IFromRow
615 + INavigationSetter
616 + crate::entity::ILazyInit
617 + Send
618 + Sync
619 + 'static,
620 {
621 let type_id = TypeId::of::<T>();
622 let items = self.set::<T>().query().to_list().await?;
623 self.set::<T>().clear_entries();
624 self.change_tracker.clear_by_type(type_id);
625 for item in items {
626 self.attach::<T>(item);
627 }
628 Ok(())
629 }
630}