use crate::db_set::DbSet;
use crate::entity::{IEntitySnapshot, IEntityType, IFromRow, IGetKeyValues, INavigationSetter};
use crate::error::{EFError, EFResult};
use crate::interceptor::InterceptorPipeline;
use crate::metadata::EntityTypeMeta;
use crate::migration::MigrationEngine;
use crate::model_builder::ModelBuilder;
use crate::provider::{DbValue, IDatabaseProvider};
use crate::tracking::ChangeTracker;
use crate::transaction::{DbTransaction, ITransaction};
use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use super::{DbContextOptions, ErasedSetOps, SetOps};
pub struct DbContext {
pub(crate) sets: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
pub(crate) savers: HashMap<TypeId, Box<dyn ErasedSetOps>>,
pub(crate) entity_metas: HashMap<TypeId, EntityTypeMeta>,
pub(crate) model_builder: ModelBuilder,
pub(crate) change_tracker: ChangeTracker,
pub(crate) provider: Arc<dyn IDatabaseProvider>,
pub(crate) interceptor_pipeline: InterceptorPipeline,
pub(crate) lazy_loading_enabled: bool,
pub(crate) ambient_transaction: Option<Box<dyn ITransaction>>,
}
impl DbContext {
pub fn from_options(options: &DbContextOptions) -> EFResult<Self> {
let provider = options.create_provider()?;
let built = options
.metadata_cache
.get_or_build(options.context_key.as_deref());
let ctx = Self {
sets: HashMap::new(),
savers: HashMap::new(),
entity_metas: built.entity_metas.clone(),
model_builder: ModelBuilder::from_built(&built),
change_tracker: ChangeTracker::new(),
provider,
interceptor_pipeline: InterceptorPipeline::new(options.interceptors.clone()),
lazy_loading_enabled: options.lazy_loading_enabled,
ambient_transaction: None,
};
Ok(ctx)
}
pub fn set<T>(&mut self) -> &mut DbSet<T>
where
T: IEntityType
+ IEntitySnapshot
+ IGetKeyValues
+ IFromRow
+ INavigationSetter
+ Send
+ Sync
+ 'static,
{
let type_id = TypeId::of::<T>();
self.savers
.entry(type_id)
.or_insert_with(|| Box::new(SetOps::<T>::new()));
self.entity_metas
.entry(type_id)
.or_insert_with(T::entity_meta);
if !self.model_builder.has_entity(type_id) {
self.model_builder.register_entity_meta(T::entity_meta());
}
self.sets.entry(type_id).or_insert_with(|| {
let table_name = self
.model_builder
.build()
.into_iter()
.find(|m| m.type_id == type_id)
.map(|m| m.table_name.to_string())
.unwrap_or_else(|| T::entity_meta().table_name.to_string());
let mut db_set = DbSet::<T>::with_provider(table_name, Arc::clone(&self.provider));
if let Some(filter) = self.model_builder.get_query_filter(&type_id) {
db_set.set_query_filter(filter.clone());
}
db_set.set_filter_map(self.model_builder.filters_by_table());
db_set.set_lazy_loading_enabled(self.lazy_loading_enabled);
Box::new(db_set)
});
self.sets
.get_mut(&type_id)
.and_then(|b| b.downcast_mut::<DbSet<T>>())
.expect("DbSet type mismatch")
}
pub fn model(&mut self) -> &mut ModelBuilder {
&mut self.model_builder
}
pub fn model_builder(&self) -> &ModelBuilder {
&self.model_builder
}
pub fn entity_metas_contains<T: IEntityType>(&self) -> bool {
self.entity_metas.contains_key(&TypeId::of::<T>())
}
pub fn discover_entities(&mut self) -> EFResult<()> {
Ok(())
}
pub fn detect_changes(&mut self) {
let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
for type_id in type_ids {
if let Some(set) = self.sets.get_mut(&type_id) {
if let Some(saver) = self.savers.get(&type_id) {
saver.detect_changes(set.as_mut());
}
}
}
}
pub async fn ensure_created(&self) -> EFResult<()> {
let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
if metas.is_empty() {
metas = self.entity_metas.values().cloned().collect();
}
if metas.is_empty() {
return Err(EFError::configuration(
"No entity types registered. Call ctx.discover_entities() or ctx.set::<T>() before ensure_created().",
));
}
let dialect = self.provider.migration_dialect();
MigrationEngine::new(dialect)
.ensure_created(&*self.provider, &metas)
.await?;
for meta in &metas {
let rows = self.model_builder.seed_rows_for(&meta.type_id);
if !rows.is_empty() {
MigrationEngine::new(dialect)
.apply_seed_data(&*self.provider, meta, rows)
.await?;
}
}
Ok(())
}
pub async fn ensure_deleted(&self) -> EFResult<()> {
let mut metas: Vec<EntityTypeMeta> = self.model_builder.build();
if metas.is_empty() {
metas = self.entity_metas.values().cloned().collect();
}
if metas.is_empty() {
return Err(EFError::configuration(
"No entity types registered. Call ctx.discover_entities() or ctx.set::<T>() before ensure_deleted().",
));
}
let dialect = self.provider.migration_dialect();
MigrationEngine::new(dialect)
.ensure_deleted(&*self.provider, &metas)
.await
}
}
impl DbContext {
pub fn provider(&self) -> &dyn IDatabaseProvider {
&*self.provider
}
pub fn change_tracker(&self) -> &ChangeTracker {
&self.change_tracker
}
pub fn change_tracker_mut(&mut self) -> &mut ChangeTracker {
&mut self.change_tracker
}
pub async fn sql_query<T: IFromRow + IEntityType>(
&self,
sql: &str,
params: &[DbValue],
) -> EFResult<Vec<T>> {
let mut conn = self.provider.get_connection().await?;
let rows = conn.query(sql, params).await?;
crate::entity::materialize_entities(&rows)
}
pub fn transaction_mut(&mut self) -> Option<&mut Box<dyn ITransaction>> {
self.ambient_transaction.as_mut()
}
pub async fn begin_transaction(&mut self) -> EFResult<Box<dyn ITransaction>> {
let mut conn = self.provider.get_connection().await?;
conn.begin_transaction().await?;
Ok(Box::new(DbTransaction::new(conn)))
}
pub async fn use_transaction<F, R>(&mut self, f: F) -> EFResult<R>
where
for<'a> F: FnOnce(&'a mut Self) -> Pin<Box<dyn Future<Output = EFResult<R>> + Send + 'a>>,
R: Send + 'static,
{
if self.ambient_transaction.is_some() {
return Err(EFError::transaction(
"ambient transaction already active; nested use_transaction is not supported",
));
}
let mut conn = self.provider.get_connection().await?;
conn.begin_transaction().await?;
self.ambient_transaction = Some(Box::new(DbTransaction::new(conn)));
let result = f(self).await;
let txn = self.ambient_transaction.take().ok_or_else(|| {
EFError::transaction("ambient_transaction was consumed during use_transaction closure")
})?;
match result {
Ok(r) => {
txn.commit().await?;
Ok(r)
}
Err(e) => {
let _ = txn.rollback().await;
Err(e)
}
}
}
}