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 let rows = self.model_builder.seed_rows_for(&meta.type_id);
203 if !rows.is_empty() {
204 MigrationEngine::new(dialect)
205 .apply_seed_data(&*self.provider, meta, rows)
206 .await?;
207 }
208 }
209 Ok(())
210 }
211
212 pub async fn ensure_deleted(&self) -> EFResult<()> {
214 let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
215 if metas.is_empty() {
216 metas = self.entity_metas.values().cloned().collect();
217 }
218 if metas.is_empty() {
219 return Err(self.empty_model_error("ensure_deleted"));
220 }
221 let dialect = self.provider.migration_dialect();
222 MigrationEngine::new(dialect)
223 .ensure_deleted(&*self.provider, &metas)
224 .await
225 }
226
227 fn empty_model_error(&self, operation: &str) -> EFError {
229 let mut message = format!(
230 "No entity types registered. Call ctx.set::<T>() or ctx.discover_entities() before {operation}()."
231 );
232 if let Some(key) = &self.context_key {
233 message.push_str(&format!(
234 " This is keyed context '{key}', which only discovers entities tagged #[context(\"{key}\")]. \
235 For a read/write replica sharing the default schema, call \
236 DbContextOptionsBuilder::use_default_entity_model()."
237 ));
238 }
239 EFError::configuration(message)
240 }
241}
242
243impl DbContext {
248 pub fn provider(&self) -> &dyn IDatabaseProvider {
250 &*self.provider
251 }
252
253 pub fn change_tracker(&self) -> &ChangeTracker {
255 &self.change_tracker
256 }
257
258 pub fn change_tracker_mut(&mut self) -> &mut ChangeTracker {
260 &mut self.change_tracker
261 }
262
263 pub fn clear_tracking(&mut self) {
270 let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
271 for type_id in type_ids {
272 let Some(saver) = self.savers.get(&type_id) else {
273 continue;
274 };
275 let Some(set) = self.sets.get_mut(&type_id) else {
276 continue;
277 };
278 saver.clear_entries(set.as_mut());
279 }
280 self.change_tracker.clear();
281 }
282
283 pub async fn sql_query<T: IFromRow + IEntityType>(
296 &self,
297 sql: &str,
298 params: &[DbValue],
299 ) -> EFResult<Vec<T>> {
300 let mut conn = self.provider.get_connection().await?;
301 let rows = conn.query(sql, params).await?;
302 crate::entity::materialize_entities(&rows)
303 }
304
305 pub async fn begin_transaction(&mut self) -> EFResult<Box<dyn ITransaction>> {
322 self.begin_transaction_with_opt(None).await
323 }
324
325 pub async fn begin_transaction_with(
328 &mut self,
329 isolation: IsolationLevel,
330 ) -> EFResult<Box<dyn ITransaction>> {
331 self.begin_transaction_with_opt(Some(isolation)).await
332 }
333
334 async fn begin_transaction_with_opt(
335 &mut self,
336 isolation: Option<IsolationLevel>,
337 ) -> EFResult<Box<dyn ITransaction>> {
338 let shared = self.session.begin_with(isolation).await?;
339 let facade = SessionConnection::new(Arc::clone(&shared));
340 Ok(Box::new(DbTransaction::new_shared(
341 shared,
342 facade,
343 Arc::clone(&self.session),
344 )))
345 }
346
347 pub async fn use_transaction<F, R>(&mut self, f: F) -> EFResult<R>
373 where
374 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
375 R: Send + 'static,
376 {
377 self.use_transaction_with_opt(None, f).await
378 }
379
380 pub async fn use_transaction_with<F, R>(
383 &mut self,
384 isolation: IsolationLevel,
385 f: F,
386 ) -> EFResult<R>
387 where
388 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
389 R: Send + 'static,
390 {
391 self.use_transaction_with_opt(Some(isolation), f).await
392 }
393
394 async fn use_transaction_with_opt<F, R>(
395 &mut self,
396 isolation: Option<IsolationLevel>,
397 f: F,
398 ) -> EFResult<R>
399 where
400 for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
401 R: Send + 'static,
402 {
403 if self.session.is_active().await {
404 return Err(EFError::transaction(
405 "a transaction is already active on this DbContext",
406 ));
407 }
408 let shared = self.session.begin_with(isolation).await?;
409 let facade = SessionConnection::new(Arc::clone(&shared));
410 let txn: Box<dyn ITransaction> = Box::new(DbTransaction::new_shared(
411 shared,
412 facade,
413 Arc::clone(&self.session),
414 ));
415 let result = f(self).await;
418 match result {
419 Ok(r) => {
420 txn.commit().await?;
421 Ok(r)
422 }
423 Err(e) => {
424 let _ = txn.rollback().await;
425 Err(e)
426 }
427 }
428 }
429}
430
431impl DbContext {
436 fn ensure_registered<T>(&mut self) -> (TypeId, String)
439 where
440 T: IEntityType
441 + IEntitySnapshot
442 + IGetKeyValues
443 + IFromRow
444 + INavigationSetter
445 + Send
446 + Sync
447 + 'static,
448 {
449 let type_id = TypeId::of::<T>();
450 self.savers
451 .entry(type_id)
452 .or_insert_with(|| Box::new(SetOps::<T>::new()));
453 self.entity_metas
454 .entry(type_id)
455 .or_insert_with(T::entity_meta);
456 if !self.model_builder.has_entity(type_id) {
457 self.model_builder.register_entity_meta(T::entity_meta());
458 }
459 let type_name = self.entity_metas[&type_id].type_name.to_string();
460 (type_id, type_name)
461 }
462
463 pub fn add<T>(&mut self, entity: T)
466 where
467 T: IEntityType
468 + IEntitySnapshot
469 + IGetKeyValues
470 + IFromRow
471 + INavigationSetter
472 + Send
473 + Sync
474 + 'static,
475 {
476 let (type_id, type_name) = self.ensure_registered::<T>();
477 let entry_id =
478 self.change_tracker
479 .track(type_id, &type_name, EntityState::Added, None, false);
480 self.set::<T>().push_entry(entity, entry_id);
481 }
482
483 pub fn attach<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 snapshot = entity.snapshot();
498 let entry_id = self.change_tracker.track(
499 type_id,
500 &type_name,
501 EntityState::Unchanged,
502 Some(snapshot),
503 false,
504 );
505 self.set::<T>().push_entry(entity, entry_id);
506 }
507
508 pub fn update<T>(&mut self, entity: T)
511 where
512 T: IEntityType
513 + IEntitySnapshot
514 + IGetKeyValues
515 + IFromRow
516 + INavigationSetter
517 + Send
518 + Sync
519 + 'static,
520 {
521 let (type_id, type_name) = self.ensure_registered::<T>();
522 let entry_id =
523 self.change_tracker
524 .track(type_id, &type_name, EntityState::Modified, None, false);
525 self.set::<T>().push_entry(entity, entry_id);
526 }
527
528 pub fn upsert<T>(&mut self, entity: T)
530 where
531 T: IEntityType
532 + IEntitySnapshot
533 + IGetKeyValues
534 + IFromRow
535 + INavigationSetter
536 + Send
537 + Sync
538 + 'static,
539 {
540 let (type_id, type_name) = self.ensure_registered::<T>();
541 let entry_id =
542 self.change_tracker
543 .track(type_id, &type_name, EntityState::Added, None, true);
544 self.set::<T>().push_entry(entity, entry_id);
545 }
546
547 pub fn remove_at<T>(&mut self, index: usize) -> EFResult<()>
549 where
550 T: IEntityType
551 + IEntitySnapshot
552 + IGetKeyValues
553 + IFromRow
554 + INavigationSetter
555 + Send
556 + Sync
557 + 'static,
558 {
559 let entry_id = self
560 .set::<T>()
561 .entry_id_at(index)
562 .ok_or_else(|| EFError::not_found("Entity not found at the given index".to_string()))?;
563 self.change_tracker
564 .set_state(entry_id, EntityState::Deleted);
565 Ok(())
566 }
567
568 pub fn remove_all<T>(&mut self)
570 where
571 T: IEntityType
572 + IEntitySnapshot
573 + IGetKeyValues
574 + IFromRow
575 + INavigationSetter
576 + Send
577 + Sync
578 + 'static,
579 {
580 let entry_ids: Vec<u64> = self.set::<T>().entries.iter().map(|e| e.entry_id).collect();
581 for entry_id in entry_ids {
582 self.change_tracker
583 .set_state(entry_id, EntityState::Deleted);
584 }
585 }
586
587 pub async fn load_all<T>(&mut self) -> EFResult<()>
590 where
591 T: IEntityType
592 + IEntitySnapshot
593 + IGetKeyValues
594 + IFromRow
595 + INavigationSetter
596 + crate::entity::ILazyInit
597 + Send
598 + Sync
599 + 'static,
600 {
601 let type_id = TypeId::of::<T>();
602 let items = self.set::<T>().query().to_list().await?;
603 self.set::<T>().clear_entries();
604 self.change_tracker.clear_by_type(type_id);
605 for item in items {
606 self.attach::<T>(item);
607 }
608 Ok(())
609 }
610}