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            // Views are read-only mappings; seeding them is not meaningful.
203            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    /// Applies all pending migrations from `dir` (EF Core `Database.Migrate()`).
217    ///
218    /// Loads `up.sql`/`down.sql` pairs from the migrations directory, applies
219    /// every migration not recorded in `__ef_migrations_history`, and returns
220    /// the number of applied migrations. This is the production upgrade path;
221    /// [`Self::ensure_created`] only creates missing tables and never alters an
222    /// existing schema.
223    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    /// Drops all tables for registered entity types.
233    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    /// Actionable error for a context that discovered no entities.
248    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
263// ---------------------------------------------------------------------------
264// DbContext inherent methods (formerly on IDbContext / IDbContextExt traits)
265// ---------------------------------------------------------------------------
266
267impl DbContext {
268    /// Returns the database provider.
269    pub fn provider(&self) -> &dyn IDatabaseProvider {
270        &*self.provider
271    }
272
273    /// Returns a read-only reference to the change tracker.
274    pub fn change_tracker(&self) -> &ChangeTracker {
275        &self.change_tracker
276    }
277
278    /// Returns a mutable reference to the change tracker.
279    pub fn change_tracker_mut(&mut self) -> &mut ChangeTracker {
280        &mut self.change_tracker
281    }
282
283    /// Detaches every tracked entity and clears the change tracker.
284    ///
285    /// Call this after a rollback when the handler is not going to retry, so
286    /// the rolled-back entities cannot leak into a later `save_changes`.
287    /// Entity instances themselves are not modified — only the context's
288    /// tracking state is dropped.
289    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    /// Executes a raw SQL query and materializes the result rows into entities.
304    ///
305    /// This is the escape hatch for complex queries (multi-table JOINs, CTEs,
306    /// window functions) that are hard to express via LINQ. The caller is
307    /// responsible for SQL correctness and parameterization.
308    ///
309    /// # Example
310    /// ```rust,ignore
311    /// let blogs: Vec<Blog> = ctx
312    ///     .sql_query("SELECT * FROM blogs WHERE id = ?", &[DbValue::I32(1)])
313    ///     .await?;
314    /// ```
315    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    /// Begins a context-scoped transaction and returns its handle.
326    ///
327    /// This is the EF Core-style API: the transaction is registered on the
328    /// context, so queries, navigation loading, bulk DML, raw SQL and
329    /// `save_changes()` all use the transaction connection. Nothing is
330    /// committed or rolled back until the returned handle is consumed.
331    ///
332    /// ```rust,ignore
333    /// let txn = ctx.begin_transaction().await?;
334    /// ctx.add::<Account>(account);
335    /// ctx.save_changes().await?; // joins, does not commit
336    /// txn.commit().await?;       // or txn.rollback().await?
337    /// ```
338    ///
339    /// Dropping the handle without committing rolls back (best effort).
340    /// Nested transactions are not supported.
341    pub async fn begin_transaction(&mut self) -> EFResult<Box<dyn ITransaction>> {
342        self.begin_transaction_with_opt(None).await
343    }
344
345    /// Like [`DbContext::begin_transaction`], with an isolation level applied
346    /// at transaction start.
347    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    /// Executes a closure within a context-scoped transaction.
368    ///
369    /// Registers the transaction on the context session for the duration of
370    /// `f`, so every statement issued through the context — queries, navigation
371    /// loading, bulk DML, raw SQL and `save_changes()` — reuses the same
372    /// transaction connection. Commits on `Ok`, rolls back on `Err`.
373    ///
374    /// For a business rejection that must roll back and still return a value,
375    /// use the manual [`DbContext::begin_transaction`] handle and call
376    /// `rollback()` explicitly.
377    ///
378    /// The closure receives `&mut DbContext` and must return a pinned boxed
379    /// future. This signature works around Rust's async borrow checker by
380    /// letting the closure capture `ctx` by mutable reference while still
381    /// producing a `Send` future.
382    ///
383    /// # Example
384    ///
385    /// ```rust,ignore
386    /// ctx.use_transaction(|ctx| Box::pin(async move {
387    ///     ctx.add::<Blog>(blog);
388    ///     ctx.save_changes().await?;
389    ///     Ok(())
390    /// })).await?;
391    /// ```
392    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    /// Like [`DbContext::use_transaction`], with an explicit isolation level
401    /// applied at transaction start.
402    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        // `txn` is a local guard: if `f` returns Err, panics, or is cancelled,
436        // its Drop discards the connection and rolls back.
437        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
451// ---------------------------------------------------------------------------
452// Entity mutation methods — coordinate DbSet + ChangeTracker
453// ---------------------------------------------------------------------------
454
455impl DbContext {
456    /// Ensures the entity type is registered (saver, metadata, model builder).
457    /// Returns `(type_id, type_name)`.
458    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    /// Adds a new entity in `Added` state. The entity will be INSERTed during
484    /// `save_changes()`.
485    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    /// Attaches an existing entity in `Unchanged` state with a snapshot for
504    /// future change detection.
505    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    /// Marks an entity as `Modified` — all fields will be UPDATEd during
529    /// `save_changes()`.
530    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    /// Marks an entity for upsert (INSERT ... ON CONFLICT DO UPDATE).
549    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    /// Marks the entity at the given index as `Deleted`.
568    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    /// Marks all entities of type `T` as `Deleted`.
589    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    /// Loads all rows from the database into the DbSet as `Unchanged`.
608    /// Clears existing entries and tracker state for type `T` first.
609    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}