use std::any::TypeId;
use std::borrow::Cow;
use std::collections::HashMap;
use std::hash::Hash;
use std::ops::{Deref, DerefMut};
use crate::backend::Backend;
use crate::query_builder::*;
use crate::result::QueryResult;
#[allow(missing_debug_implementations, unreachable_pub)]
#[cfg_attr(
doc_cfg,
doc(cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"))
)]
pub struct StatementCache<DB: Backend, Statement> {
pub(crate) cache: HashMap<StatementCacheKey<DB>, Statement>,
}
#[derive(Debug, Clone, Copy)]
#[cfg_attr(
doc_cfg,
doc(cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"))
)]
#[allow(unreachable_pub)]
pub enum PrepareForCache {
Yes,
No,
}
#[allow(
clippy::len_without_is_empty,
clippy::new_without_default,
unreachable_pub
)]
impl<DB, Statement> StatementCache<DB, Statement>
where
DB: Backend,
DB::TypeMetadata: Clone,
DB::QueryBuilder: Default,
StatementCacheKey<DB>: Hash + Eq,
{
#[allow(unreachable_pub)]
pub fn new() -> Self {
StatementCache {
cache: HashMap::new(),
}
}
#[allow(unreachable_pub)]
#[cfg(any(
feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes",
feature = "postgres",
all(feature = "sqlite", test)
))]
#[cfg_attr(
doc_cfg,
doc(cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"))
)]
pub fn len(&self) -> usize {
self.cache.len()
}
#[allow(unreachable_pub)]
pub fn cached_statement<T, F>(
&mut self,
source: &T,
backend: &DB,
bind_types: &[DB::TypeMetadata],
prepare_fn: F,
) -> QueryResult<MaybeCached<'_, Statement>>
where
T: QueryFragment<DB> + QueryId,
F: FnOnce(&str, PrepareForCache) -> QueryResult<Statement>,
{
use std::collections::hash_map::Entry::{Occupied, Vacant};
let cache_key = StatementCacheKey::for_source(source, bind_types, backend)?;
if !source.is_safe_to_cache_prepared(backend)? {
let sql = cache_key.sql(source, backend)?;
return prepare_fn(&sql, PrepareForCache::No).map(MaybeCached::CannotCache);
}
let cached_result = match self.cache.entry(cache_key) {
Occupied(entry) => entry.into_mut(),
Vacant(entry) => {
let statement = {
let sql = entry.key().sql(source, backend)?;
prepare_fn(&sql, PrepareForCache::Yes)
};
entry.insert(statement?)
}
};
Ok(MaybeCached::Cached(cached_result))
}
}
#[allow(missing_debug_implementations, unreachable_pub)]
#[cfg_attr(
doc_cfg,
doc(cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"))
)]
#[non_exhaustive]
pub enum MaybeCached<'a, T: 'a> {
CannotCache(T),
Cached(&'a mut T),
}
impl<'a, T> Deref for MaybeCached<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
match *self {
MaybeCached::CannotCache(ref x) => x,
MaybeCached::Cached(ref x) => &**x,
}
}
}
impl<'a, T> DerefMut for MaybeCached<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
match *self {
MaybeCached::CannotCache(ref mut x) => x,
MaybeCached::Cached(ref mut x) => &mut **x,
}
}
}
#[allow(missing_debug_implementations, unreachable_pub)]
#[derive(Hash, PartialEq, Eq)]
#[cfg_attr(
doc_cfg,
doc(cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"))
)]
pub enum StatementCacheKey<DB: Backend> {
Type(TypeId),
Sql {
sql: String,
bind_types: Vec<DB::TypeMetadata>,
},
}
impl<DB> StatementCacheKey<DB>
where
DB: Backend,
DB::QueryBuilder: Default,
DB::TypeMetadata: Clone,
{
#[allow(unreachable_pub)]
pub fn for_source<T>(
source: &T,
bind_types: &[DB::TypeMetadata],
backend: &DB,
) -> QueryResult<Self>
where
T: QueryFragment<DB> + QueryId,
{
match T::query_id() {
Some(id) => Ok(StatementCacheKey::Type(id)),
None => {
let sql = Self::construct_sql(source, backend)?;
Ok(StatementCacheKey::Sql {
sql,
bind_types: bind_types.into(),
})
}
}
}
#[allow(unreachable_pub)]
pub fn sql<T: QueryFragment<DB>>(&self, source: &T, backend: &DB) -> QueryResult<Cow<'_, str>> {
match *self {
StatementCacheKey::Type(_) => Self::construct_sql(source, backend).map(Cow::Owned),
StatementCacheKey::Sql { ref sql, .. } => Ok(Cow::Borrowed(sql)),
}
}
fn construct_sql<T: QueryFragment<DB>>(source: &T, backend: &DB) -> QueryResult<String> {
let mut query_builder = DB::QueryBuilder::default();
source.to_sql(&mut query_builder, backend)?;
Ok(query_builder.finish())
}
}