Skip to main content

rust_ef/db_context/
context.rs

1//! `DbContext` struct and non-save methods.
2//!
3//! `DbContext` is the concrete context type. Entity sets use a type-map:
4//! `ctx.set::<Blog>()` lazy-creates `DbSet<Blog>`. `SetOps<T>` dispatchers
5//! enable `save_changes()` to iterate all entity types.
6//!
7//! ## Ownership and Mutation
8//!
9//! `DbContext` methods (`set::<T>()`, `save_changes()`, `detect_changes()`)
10//! require `&mut self` — this is idiomatic Rust, not a limitation. The DI
11//! integration (`add_dbcontext`) registers the context as **Scoped** and
12//! supports owned resolution via `get_owned::<DbContext>()`.
13//!
14//! `DbContext` is **not** thread-safe — a single instance must not be shared
15//! across threads (aligned with EFCore).
16
17use 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
38/// The session / unit-of-work entry point.
39///
40/// Entity sets are lazy-created via `ctx.set::<T>()`. Metadata is loaded from
41/// the process-level `MetadataCache` on `DbContextOptions` (built once per
42/// `context_key`, then `Arc`-shared). `save_changes()` commits all pending
43/// changes in a transaction.
44pub 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    /// Keyed context key (`None` for the default model). Used only for
52    /// actionable diagnostics when a keyed replica forgets
53    /// `use_default_entity_model()`.
54    pub(crate) context_key: Option<String>,
55    /// Per-context connection session. While an explicit transaction is open,
56    /// every statement issued through `provider` borrows the transaction's
57    /// connection instead of acquiring a new pooled connection.
58    pub(crate) session: Arc<Session>,
59    pub(crate) interceptor_pipeline: InterceptorPipeline,
60    pub(crate) lazy_loading_enabled: bool,
61}
62
63impl DbContext {
64    /// Creates the context from options (uses the provider factory stored in options).
65    ///
66    /// Entity metadata is loaded from the process-level `MetadataCache` on
67    /// `DbContextOptions` — `inventory::iter` + `IEntityTypeConfiguration::configure()`
68    /// run once per `context_key` (first call), then the result is `Arc`-shared
69    /// across all `DbContext` instances. Per-instance `ModelBuilder` mutations
70    /// (`has_query_filter`, etc.) after construction only affect this instance.
71    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    /// Returns the model builder for Fluent API configuration.
138    pub fn model(&mut self) -> &mut ModelBuilder {
139        &mut self.model_builder
140    }
141
142    /// Returns a read-only reference to the model builder.
143    pub fn model_builder(&self) -> &ModelBuilder {
144        &self.model_builder
145    }
146
147    /// Returns `true` if an entity of type `T` has been discovered and
148    /// registered in the entity metadata map.
149    pub fn entity_metas_contains<T: IEntityType>(&self) -> bool {
150        self.entity_metas.contains_key(&TypeId::of::<T>())
151    }
152
153    /// No-op since metadata caching was introduced.
154    ///
155    /// Historically, this method iterated `inventory::iter` to register entity
156    /// metadata and apply `#[entity(T)]` configurations. That work now happens
157    /// in `MetadataCache::build()`, invoked lazily by `from_options()` via
158    /// `MetadataCache::get_or_build()`. The result is `Arc`-shared across all
159    /// `DbContext` instances with the same `context_key`.
160    ///
161    /// Retained as a public method for backward compatibility — existing code
162    /// calling `ctx.discover_entities()?` after `from_options()` continues to
163    /// compile and run (as a no-op). The metadata is already populated.
164    pub fn discover_entities(&mut self) -> EFResult<()> {
165        Ok(())
166    }
167
168    /// Detects changes on all tracked DbSets by comparing property snapshots.
169    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    /// Creates all tables for registered entity types.
183    ///
184    /// Sources metas from `model_builder.build()`, which applies all Fluent
185    /// API configurations and `#[entity(T)]` overrides. Entities are
186    /// discovered automatically via `#[derive(EntityType)]`; call
187    /// `discover_entities()` first, or use `set::<T>()` to register manually.
188    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    /// Drops all tables for registered entity types.
213    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    /// Actionable error for a context that discovered no entities.
228    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
243// ---------------------------------------------------------------------------
244// DbContext inherent methods (formerly on IDbContext / IDbContextExt traits)
245// ---------------------------------------------------------------------------
246
247impl DbContext {
248    /// Returns the database provider.
249    pub fn provider(&self) -> &dyn IDatabaseProvider {
250        &*self.provider
251    }
252
253    /// Returns a read-only reference to the change tracker.
254    pub fn change_tracker(&self) -> &ChangeTracker {
255        &self.change_tracker
256    }
257
258    /// Returns a mutable reference to the change tracker.
259    pub fn change_tracker_mut(&mut self) -> &mut ChangeTracker {
260        &mut self.change_tracker
261    }
262
263    /// Detaches every tracked entity and clears the change tracker.
264    ///
265    /// Call this after a rollback when the handler is not going to retry, so
266    /// the rolled-back entities cannot leak into a later `save_changes`.
267    /// Entity instances themselves are not modified — only the context's
268    /// tracking state is dropped.
269    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    /// Executes a raw SQL query and materializes the result rows into entities.
284    ///
285    /// This is the escape hatch for complex queries (multi-table JOINs, CTEs,
286    /// window functions) that are hard to express via LINQ. The caller is
287    /// responsible for SQL correctness and parameterization.
288    ///
289    /// # Example
290    /// ```rust,ignore
291    /// let blogs: Vec<Blog> = ctx
292    ///     .sql_query("SELECT * FROM blogs WHERE id = ?", &[DbValue::I32(1)])
293    ///     .await?;
294    /// ```
295    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    /// Begins a context-scoped transaction and returns its handle.
306    ///
307    /// This is the EF Core-style API: the transaction is registered on the
308    /// context, so queries, navigation loading, bulk DML, raw SQL and
309    /// `save_changes()` all use the transaction connection. Nothing is
310    /// committed or rolled back until the returned handle is consumed.
311    ///
312    /// ```rust,ignore
313    /// let txn = ctx.begin_transaction().await?;
314    /// ctx.add::<Account>(account);
315    /// ctx.save_changes().await?; // joins, does not commit
316    /// txn.commit().await?;       // or txn.rollback().await?
317    /// ```
318    ///
319    /// Dropping the handle without committing rolls back (best effort).
320    /// Nested transactions are not supported.
321    pub async fn begin_transaction(&mut self) -> EFResult<Box<dyn ITransaction>> {
322        self.begin_transaction_with_opt(None).await
323    }
324
325    /// Like [`DbContext::begin_transaction`], with an isolation level applied
326    /// at transaction start.
327    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    /// Executes a closure within a context-scoped transaction.
348    ///
349    /// Registers the transaction on the context session for the duration of
350    /// `f`, so every statement issued through the context — queries, navigation
351    /// loading, bulk DML, raw SQL and `save_changes()` — reuses the same
352    /// transaction connection. Commits on `Ok`, rolls back on `Err`.
353    ///
354    /// For a business rejection that must roll back and still return a value,
355    /// use the manual [`DbContext::begin_transaction`] handle and call
356    /// `rollback()` explicitly.
357    ///
358    /// The closure receives `&mut DbContext` and must return a pinned boxed
359    /// future. This signature works around Rust's async borrow checker by
360    /// letting the closure capture `ctx` by mutable reference while still
361    /// producing a `Send` future.
362    ///
363    /// # Example
364    ///
365    /// ```rust,ignore
366    /// ctx.use_transaction(|ctx| Box::pin(async move {
367    ///     ctx.add::<Blog>(blog);
368    ///     ctx.save_changes().await?;
369    ///     Ok(())
370    /// })).await?;
371    /// ```
372    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    /// Like [`DbContext::use_transaction`], with an explicit isolation level
381    /// applied at transaction start.
382    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        // `txn` is a local guard: if `f` returns Err, panics, or is cancelled,
416        // its Drop discards the connection and rolls back.
417        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
431// ---------------------------------------------------------------------------
432// Entity mutation methods — coordinate DbSet + ChangeTracker
433// ---------------------------------------------------------------------------
434
435impl DbContext {
436    /// Ensures the entity type is registered (saver, metadata, model builder).
437    /// Returns `(type_id, type_name)`.
438    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    /// Adds a new entity in `Added` state. The entity will be INSERTed during
464    /// `save_changes()`.
465    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    /// Attaches an existing entity in `Unchanged` state with a snapshot for
484    /// future change detection.
485    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    /// Marks an entity as `Modified` — all fields will be UPDATEd during
509    /// `save_changes()`.
510    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    /// Marks an entity for upsert (INSERT ... ON CONFLICT DO UPDATE).
529    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    /// Marks the entity at the given index as `Deleted`.
548    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    /// Marks all entities of type `T` as `Deleted`.
569    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    /// Loads all rows from the database into the DbSet as `Unchanged`.
588    /// Clears existing entries and tracker state for type `T` first.
589    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}