use crate::change_executor::ChangeExecutor;
use crate::db_set::DbSet;
use crate::entity::{IEntitySnapshot, IEntityType, IFromRow, IGetKeyValues, INavigationSetter};
use crate::error::{EFError, EFResult};
use crate::interceptor::{InterceptorPipeline, SaveChangesContext, SaveChangesResultContext};
use crate::metadata::EntityTypeMeta;
use crate::migration::MigrationEngine;
use crate::model_builder::ModelBuilder;
use crate::provider::{DbValue, IAsyncConnection, IDatabaseProvider};
use crate::tracking::{ChangeTracker, EntityEntryView};
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;
#[derive(Clone)]
pub struct DbContextOptions {
pub(crate) connection_string: String,
pub(crate) provider_tag: Option<String>,
#[allow(clippy::type_complexity)]
pub(crate) provider_factory:
Option<Arc<dyn Fn(&str) -> EFResult<Arc<dyn IDatabaseProvider>> + Send + Sync>>,
pub(crate) provider_cache: Arc<std::sync::Mutex<Option<Arc<dyn IDatabaseProvider>>>>,
pub(crate) interceptors: Vec<Arc<dyn crate::interceptor::ISaveChangesInterceptor>>,
pub(crate) lazy_loading_enabled: bool,
pub(crate) context_key: Option<String>,
pub(crate) metadata_cache: Arc<crate::metadata_cache::MetadataCache>,
#[cfg(feature = "tracing")]
pub(crate) slow_query_threshold: Option<std::time::Duration>,
}
impl std::fmt::Debug for DbContextOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DbContextOptions")
.field(
"connection_string",
&redact_connection_string(&self.connection_string),
)
.field("provider_tag", &self.provider_tag)
.finish()
}
}
fn redact_connection_string(cs: &str) -> String {
if let Some(scheme_end) = cs.find("://") {
let (scheme, rest) = cs.split_at(scheme_end + 3);
if let Some(at) = rest.find('@') {
let (userinfo, host_and_rest) = rest.split_at(at);
let redacted_user = match userinfo.find(':') {
Some(colon) => &userinfo[..colon],
None => userinfo,
};
return format!("{}{}***@{}", scheme, redacted_user, &host_and_rest[1..]);
}
return cs.to_string();
}
if cs.contains('=') {
return cs
.split(';')
.map(|pair| {
let eq = match pair.find('=') {
Some(e) => e,
None => return pair.to_string(),
};
let key = pair[..eq].trim().to_lowercase();
if key.contains("password") || key.contains("pwd") {
format!("{}=***", &pair[..eq])
} else {
pair.to_string()
}
})
.collect::<Vec<_>>()
.join(";");
}
cs.to_string()
}
impl DbContextOptions {
pub fn connection_string(&self) -> &str {
&self.connection_string
}
pub fn provider_tag(&self) -> Option<&str> {
self.provider_tag.as_deref()
}
pub fn lazy_loading_enabled(&self) -> bool {
self.lazy_loading_enabled
}
pub fn context_key(&self) -> Option<&str> {
self.context_key.as_deref()
}
pub fn create_provider(&self) -> EFResult<Arc<dyn IDatabaseProvider>> {
let mut guard = self
.provider_cache
.lock()
.unwrap_or_else(|p| p.into_inner());
if let Some(provider) = guard.as_ref() {
return Ok(Arc::clone(provider));
}
let factory = self.provider_factory.as_ref().ok_or_else(|| {
crate::error::EFError::configuration(
"No provider configured. Call use_sqlite / use_postgres / use_mysql first.",
)
})?;
let provider = factory(self.connection_string())?;
#[cfg(feature = "tracing")]
if let Some(threshold) = self.slow_query_threshold {
provider.set_slow_query_threshold(threshold);
}
*guard = Some(Arc::clone(&provider));
Ok(provider)
}
}
#[allow(clippy::derivable_impls)]
impl Default for DbContextOptions {
fn default() -> Self {
Self {
connection_string: String::new(),
provider_tag: None,
provider_factory: None,
provider_cache: Arc::new(std::sync::Mutex::new(None)),
interceptors: Vec::new(),
lazy_loading_enabled: false,
context_key: None,
metadata_cache: Arc::new(crate::metadata_cache::MetadataCache::new()),
#[cfg(feature = "tracing")]
slow_query_threshold: None,
}
}
}
pub struct DbContextOptionsBuilder {
inner: DbContextOptions,
}
impl DbContextOptionsBuilder {
pub fn new() -> Self {
Self {
inner: DbContextOptions::default(),
}
}
pub fn connection_string(&mut self, cs: impl Into<String>) -> &mut Self {
self.inner.connection_string = cs.into();
self
}
pub fn set_provider(&mut self, tag: &str, cs: impl Into<String>) -> &mut Self {
self.inner.provider_tag = Some(tag.to_string());
self.inner.connection_string = cs.into();
self
}
#[allow(clippy::type_complexity)]
pub fn set_provider_factory(
&mut self,
tag: &str,
cs: impl Into<String>,
factory: Arc<dyn Fn(&str) -> EFResult<Arc<dyn IDatabaseProvider>> + Send + Sync>,
) -> &mut Self {
self.inner.provider_tag = Some(tag.to_string());
self.inner.connection_string = cs.into();
self.inner.provider_factory = Some(factory);
self
}
pub fn add_interceptor(
&mut self,
interceptor: impl crate::interceptor::ISaveChangesInterceptor + 'static,
) -> &mut Self {
self.inner.interceptors.push(Arc::new(interceptor));
self
}
pub fn use_lazy_loading(&mut self, enabled: bool) -> &mut Self {
self.inner.lazy_loading_enabled = enabled;
self
}
pub fn context_key(&mut self, key: impl Into<String>) -> &mut Self {
self.inner.context_key = Some(key.into());
self
}
#[cfg(feature = "tracing")]
pub fn slow_query_threshold(&mut self, threshold: std::time::Duration) -> &mut Self {
self.inner.slow_query_threshold = Some(threshold);
self
}
pub fn build(self) -> DbContextOptions {
self.inner
}
}
impl Default for DbContextOptionsBuilder {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
trait ErasedSetOps: Send + Sync {
async fn save(
&self,
conn: &mut (dyn IAsyncConnection + Send),
provider: &dyn IDatabaseProvider,
raw_set: &mut (dyn Any + Send + Sync),
meta: &EntityTypeMeta,
) -> EFResult<(usize, usize, usize)>;
fn detect_changes(&self, raw_set: &mut (dyn Any + Send + Sync));
fn clear(&self, raw_set: &mut (dyn Any + Send + Sync + 'static));
fn collect_entries(&self, raw_set: &(dyn Any + Send + Sync)) -> Vec<EntityEntryView>;
}
struct SetOps<E> {
_phantom: std::marker::PhantomData<E>,
}
impl<E> SetOps<E> {
fn new() -> Self {
Self {
_phantom: std::marker::PhantomData,
}
}
}
#[async_trait::async_trait]
impl<E> ErasedSetOps for SetOps<E>
where
E: IEntityType
+ IEntitySnapshot
+ IGetKeyValues
+ IFromRow
+ INavigationSetter
+ Send
+ Sync
+ 'static,
{
async fn save(
&self,
conn: &mut (dyn IAsyncConnection + Send),
provider: &dyn IDatabaseProvider,
raw_set: &mut (dyn Any + Send + Sync),
meta: &EntityTypeMeta,
) -> EFResult<(usize, usize, usize)> {
let db_set = raw_set
.downcast_mut::<DbSet<E>>()
.expect("SetOps type mismatch");
save_one_set(conn, provider, db_set, meta).await
}
fn detect_changes(&self, raw_set: &mut (dyn Any + Send + Sync)) {
if let Some(db_set) = raw_set.downcast_mut::<DbSet<E>>() {
db_set.detect_changes();
}
}
fn clear(&self, raw_set: &mut (dyn Any + Send + Sync + 'static)) {
if let Some(db_set) = raw_set.downcast_mut::<DbSet<E>>() {
db_set.clear_entries();
}
}
fn collect_entries(&self, raw_set: &(dyn Any + Send + Sync)) -> Vec<EntityEntryView> {
let Some(db_set) = raw_set.downcast_ref::<DbSet<E>>() else {
return Vec::new();
};
let type_name = E::entity_meta().type_name.to_string();
db_set
.entries
.iter()
.map(|e| EntityEntryView {
type_id: TypeId::of::<E>(),
type_name: type_name.clone(),
state: e.state,
})
.collect()
}
}
pub struct DbContext {
sets: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
savers: HashMap<TypeId, Box<dyn ErasedSetOps>>,
entity_metas: HashMap<TypeId, EntityTypeMeta>,
model_builder: ModelBuilder,
change_tracker: ChangeTracker,
provider: Arc<dyn IDatabaseProvider>,
interceptor_pipeline: InterceptorPipeline,
lazy_loading_enabled: bool,
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());
}
}
}
}
fn build_save_context(&self) -> SaveChangesContext {
let mut views: Vec<EntityEntryView> = Vec::new();
for (type_id, set) in &self.sets {
if let Some(saver) = self.savers.get(type_id) {
views.extend(saver.collect_entries(set.as_ref()));
}
}
SaveChangesContext::from_views(views)
}
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 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 save_changes(&mut self) -> EFResult<SaveChangesResult> {
let _save_guard = crate::observability::SaveChangesGuard::new();
let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
for type_id in &type_ids {
let set = self.sets.get_mut(type_id).unwrap();
self.savers
.get(type_id)
.unwrap()
.detect_changes(set.as_mut());
}
let configured_metas: HashMap<TypeId, EntityTypeMeta> = self
.model_builder
.build()
.into_iter()
.map(|m| (m.type_id, m))
.collect();
let save_ctx = self.build_save_context();
self.interceptor_pipeline.on_saving(&save_ctx).await?;
enum TxnSource {
Ambient(Box<dyn ITransaction>),
Managed(Box<dyn IAsyncConnection>),
}
let mut txn = match self.ambient_transaction.take() {
Some(t) => TxnSource::Ambient(t),
None => {
let mut c = self.provider.get_connection().await?;
c.begin_transaction().await?;
TxnSource::Managed(c)
}
};
let type_ids: Vec<TypeId> = self.sets.keys().copied().collect();
let mut total_added = 0usize;
let mut total_updated = 0usize;
let mut total_deleted = 0usize;
for type_id in &type_ids {
let saver = self.savers.get(type_id).expect("saver not registered");
let set = self.sets.get_mut(type_id).unwrap();
let meta = configured_metas
.get(type_id)
.or_else(|| self.entity_metas.get(type_id))
.expect("meta not found for entity type");
let conn_ref: &mut dyn IAsyncConnection = match &mut txn {
TxnSource::Ambient(t) => t.connection(),
TxnSource::Managed(c) => c.as_mut(),
};
let (a, u, d) = match saver
.save(conn_ref, &*self.provider, set.as_mut(), meta)
.await
{
Ok(r) => r,
Err(e) => {
if let TxnSource::Managed(mut conn) = txn {
let _ = conn.rollback_transaction().await;
} else if let TxnSource::Ambient(t) = txn {
self.ambient_transaction = Some(t);
}
self.interceptor_pipeline
.on_save_failed(&save_ctx, &e)
.await;
return Err(e);
}
};
total_added += a;
total_updated += u;
total_deleted += d;
}
match txn {
TxnSource::Ambient(t) => {
self.ambient_transaction = Some(t);
}
TxnSource::Managed(mut conn) => {
if let Err(e) = conn.commit_transaction().await {
self.interceptor_pipeline
.on_save_failed(&save_ctx, &e)
.await;
return Err(e);
}
}
}
self.change_tracker.accept_all_changes();
for type_id in &type_ids {
let saver = self.savers.get(type_id).unwrap();
let set = self.sets.get_mut(type_id).unwrap();
saver.clear(set.as_mut());
}
let result_ctx = SaveChangesResultContext {
added: total_added,
updated: total_updated,
deleted: total_deleted,
};
self.interceptor_pipeline
.on_saved(&save_ctx, &result_ctx)
.await?;
Ok(SaveChangesResult {
added: total_added,
updated: total_updated,
deleted: total_deleted,
})
}
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()
.expect("ambient_transaction set above");
match result {
Ok(r) => {
txn.commit().await?;
Ok(r)
}
Err(e) => {
let _ = txn.rollback().await;
Err(e)
}
}
}
}
#[allow(clippy::type_complexity)]
pub async fn save_one_set<E>(
conn: &mut dyn IAsyncConnection,
provider: &dyn IDatabaseProvider,
db_set: &mut DbSet<E>,
meta: &EntityTypeMeta,
) -> EFResult<(usize, usize, usize)>
where
E: IEntityType + IEntitySnapshot + IGetKeyValues + IFromRow,
{
let query_filter = db_set.query_filter();
let added: Vec<(&E, &EntityTypeMeta)> = db_set
.tracked_by_state(crate::entity::EntityState::Added)
.into_iter()
.map(|(e, _)| (e, meta))
.collect();
let modified: Vec<(&E, &EntityTypeMeta, Option<&HashMap<String, DbValue>>)> = db_set
.tracked_by_state(crate::entity::EntityState::Modified)
.into_iter()
.map(|(e, orig)| (e, meta, orig))
.collect();
let deleted: Vec<(&E, &EntityTypeMeta, Option<&HashMap<String, DbValue>>)> = db_set
.tracked_by_state(crate::entity::EntityState::Deleted)
.into_iter()
.map(|(e, orig)| (e, meta, orig))
.collect();
let mut ac = 0usize;
let mut uc = 0usize;
let mut dc = 0usize;
if !added.is_empty() {
ac = ChangeExecutor::execute_inserts(conn, provider, &added, |_, _| {}).await?;
}
if !modified.is_empty() {
uc = ChangeExecutor::execute_updates(conn, provider, &modified, query_filter).await?;
}
if !deleted.is_empty() {
dc = ChangeExecutor::execute_deletes(conn, provider, &deleted, query_filter).await?;
}
Ok((ac, uc, dc))
}
#[derive(Debug, Clone)]
pub struct SaveChangesResult {
pub added: usize,
pub updated: usize,
pub deleted: usize,
}
impl SaveChangesResult {
pub fn total(&self) -> usize {
self.added + self.updated + self.deleted
}
}
impl std::fmt::Display for SaveChangesResult {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{} entities modified ({} added, {} updated, {} deleted)",
self.total(),
self.added,
self.updated,
self.deleted
)
}
}