use crate::db_context::{DbContext, DbContextOptions, DbContextOptionsBuilder};
use rust_dix::{IServiceResolver, RdiError, Scope};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
pub trait IDbContextKey: Send + Sync {
fn name(&self) -> &str;
}
impl IDbContextKey for &str {
fn name(&self) -> &str {
self
}
}
impl IDbContextKey for String {
fn name(&self) -> &str {
self.as_str()
}
}
impl IDbContextKey for &String {
fn name(&self) -> &str {
self.as_str()
}
}
pub trait DbContextResolver {
fn get_owned_context(
&self,
key: impl IDbContextKey + Into<ServiceKey>,
) -> Result<DbContext, RdiError>;
fn get_context(
&self,
key: impl IDbContextKey + Into<ServiceKey>,
) -> Result<Arc<DbContext>, RdiError>;
}
impl DbContextResolver for ServiceProvider {
fn get_owned_context(
&self,
key: impl IDbContextKey + Into<ServiceKey>,
) -> Result<DbContext, RdiError> {
self.get_keyed_owned(key)
}
fn get_context(
&self,
key: impl IDbContextKey + Into<ServiceKey>,
) -> Result<Arc<DbContext>, RdiError> {
self.get_keyed(key)
}
}
impl DbContextResolver for Scope {
fn get_owned_context(
&self,
key: impl IDbContextKey + Into<ServiceKey>,
) -> Result<DbContext, RdiError> {
self.get_keyed_owned(key)
}
fn get_context(
&self,
key: impl IDbContextKey + Into<ServiceKey>,
) -> Result<Arc<DbContext>, RdiError> {
self.get_keyed(key)
}
}
struct AsyncDbContextOptions(Arc<DbContextOptions>);
fn resolve_arc<T: Send + Sync + 'static>(resolver: &dyn IServiceResolver) -> Option<Arc<T>> {
let any = resolver.get_any(std::any::type_name::<T>())?;
any.downcast::<Arc<T>>().ok().map(|arc| Arc::clone(&*arc))
}
fn resolve_keyed_arc_by_key<T: Send + Sync + 'static>(
resolver: &dyn IServiceResolver,
key: &ServiceKey,
) -> Option<Arc<T>> {
let any = resolver.get_keyed_any_key(std::any::type_name::<T>(), key)?;
any.downcast::<Arc<T>>().ok().map(|arc| Arc::clone(&*arc))
}
fn build_options(configure: impl FnOnce(&mut DbContextOptionsBuilder)) -> Arc<DbContextOptions> {
let mut builder = DbContextOptionsBuilder::new();
configure(&mut builder);
Arc::new(builder.build())
}
pub const DEFAULT_ANYKEY_CAPACITY: usize = 256;
struct AnyKeyOptionsCache {
capacity: usize,
clock: u64,
entries: HashMap<ServiceKey, (Arc<DbContextOptions>, u64)>,
}
impl AnyKeyOptionsCache {
fn new(capacity: usize) -> Self {
Self {
capacity: capacity.max(1),
clock: 0,
entries: HashMap::new(),
}
}
fn get_or_build(
&mut self,
key: &ServiceKey,
build: impl FnOnce() -> Arc<DbContextOptions>,
) -> Arc<DbContextOptions> {
self.clock += 1;
let tick = self.clock;
if let Some((options, last_used)) = self.entries.get_mut(key) {
*last_used = tick;
return Arc::clone(options);
}
let options = build();
if self.entries.len() >= self.capacity {
if let Some(oldest) = self
.entries
.iter()
.min_by_key(|(_, (_, used))| *used)
.map(|(key, _)| key.clone())
{
self.entries.remove(&oldest);
}
}
self.entries
.insert(key.clone(), (Arc::clone(&options), tick));
options
}
}
fn async_options_factory(
options: Arc<DbContextOptions>,
) -> impl Fn(
Arc<rust_dix::ServiceProvider>,
)
-> std::pin::Pin<Box<dyn std::future::Future<Output = Arc<AsyncDbContextOptions>> + Send>>
+ Send
+ Sync
+ 'static {
move |_provider| {
let options = Arc::clone(&options);
Box::pin(async move {
options
.create_provider()
.expect("DbContext provider initialization failed during build_async");
if let Some(preflight) = options.async_preflight.clone() {
if let Err(err) = preflight(Arc::clone(&options)).await {
panic!("DbContext async preflight failed: {err}");
}
}
Arc::new(AsyncDbContextOptions(options))
})
}
}
fn context_from_async_options(runtime: Arc<AsyncDbContextOptions>) -> Arc<DbContext> {
let ctx = DbContext::from_options(&runtime.0).expect("Failed to create DbContext");
Arc::new(ctx)
}
pub trait DbContextServiceCollectionExt {
fn add_dbcontext(
self,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self;
fn add_dbcontext_keyed(
self,
key: impl IDbContextKey + Into<ServiceKey>,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self;
fn add_dbcontext_anykey(
self,
configure: impl Fn(&ServiceKey, &mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self;
fn add_dbcontext_anykey_with_capacity(
self,
capacity: usize,
configure: impl Fn(&ServiceKey, &mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self;
fn add_dbcontext_async(
self,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self;
fn add_dbcontext_keyed_async(
self,
key: impl IDbContextKey + Into<ServiceKey>,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self;
}
impl DbContextServiceCollectionExt for ::rust_dix::ServiceCollection {
fn add_dbcontext(
self,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self {
let options = build_options(configure);
options
.create_provider()
.expect("DbContext provider initialization failed at startup");
self.scoped(move |_| {
let ctx = DbContext::from_options(&options).expect("Failed to create DbContext");
Arc::new(ctx) as Arc<DbContext>
})
}
fn add_dbcontext_keyed(
self,
key: impl IDbContextKey + Into<ServiceKey>,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self {
let key_name = key.name().to_string();
let service_key: ServiceKey = key.into();
let mut builder = DbContextOptionsBuilder::new();
configure(&mut builder);
if !builder.uses_default_entity_model() {
builder.context_key(key_name);
}
let options = Arc::new(builder.build());
options
.create_provider()
.expect("DbContext provider initialization failed at startup");
self.keyed_scoped(service_key, move |_| {
let ctx = DbContext::from_options(&options).expect("Failed to create DbContext");
Arc::new(ctx) as Arc<DbContext>
})
}
fn add_dbcontext_anykey(
self,
configure: impl Fn(&ServiceKey, &mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self {
self.add_dbcontext_anykey_with_capacity(DEFAULT_ANYKEY_CAPACITY, configure)
}
fn add_dbcontext_anykey_with_capacity(
self,
capacity: usize,
configure: impl Fn(&ServiceKey, &mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self {
let options_cache = Arc::new(Mutex::new(AnyKeyOptionsCache::new(capacity)));
self.keyed_any_scoped::<DbContext>(move |_resolver, key| {
let options = {
let mut cache = options_cache
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
cache.get_or_build(key, || {
let mut builder = DbContextOptionsBuilder::new();
configure(key, &mut builder);
builder.use_default_entity_model();
let options = Arc::new(builder.build());
options
.create_provider()
.expect("AnyKey DbContext provider initialization failed");
options
})
};
let ctx = DbContext::from_options(&options).expect("Failed to create DbContext");
Arc::new(ctx) as Arc<DbContext>
})
}
fn add_dbcontext_async(
self,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self {
let options = build_options(configure);
self.async_singleton::<AsyncDbContextOptions>(async_options_factory(options))
.scoped(move |resolver| {
let runtime = resolve_arc::<AsyncDbContextOptions>(resolver).expect(
"async DbContext options are not initialized; call ServiceCollection::build_async()",
);
context_from_async_options(runtime)
})
}
fn add_dbcontext_keyed_async(
self,
key: impl IDbContextKey + Into<ServiceKey>,
configure: impl FnOnce(&mut DbContextOptionsBuilder) + Send + Sync + 'static,
) -> Self {
let key_name = key.name().to_string();
let service_key: ServiceKey = key.into();
let mut builder = DbContextOptionsBuilder::new();
configure(&mut builder);
if !builder.uses_default_entity_model() {
builder.context_key(key_name);
}
let options = Arc::new(builder.build());
let scoped_key = service_key.clone();
self.async_keyed_singleton::<AsyncDbContextOptions>(
service_key,
async_options_factory(options),
)
.keyed_scoped(scoped_key.clone(), move |resolver| {
let runtime = resolve_keyed_arc_by_key::<AsyncDbContextOptions>(resolver, &scoped_key)
.expect("async keyed DbContext options are not initialized; call build_async()");
context_from_async_options(runtime)
})
}
}
pub use rust_dix::{ServiceCollection, ServiceKey, ServiceProvider, ServiceProviderOptions};