Skip to content

Commit

Permalink
Merge pull request diesel-rs#3864 from weiznich/feature/connection_in…
Browse files Browse the repository at this point in the history
…strumentation_support

Introduce support for connection instrumentation
  • Loading branch information
weiznich authored Dec 22, 2023
2 parents d627da4 + 7df81d9 commit ca3ba76
Show file tree
Hide file tree
Showing 17 changed files with 1,063 additions and 95 deletions.
2 changes: 1 addition & 1 deletion diesel/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "diesel"
version = "2.1.1"
version = "2.1.4"
license = "MIT OR Apache-2.0"
description = "A safe, extensible ORM and Query Builder for PostgreSQL, SQLite, and MySQL"
readme = "README.md"
Expand Down
315 changes: 315 additions & 0 deletions diesel/src/connection/instrumentation.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,315 @@
use std::fmt::Debug;
use std::fmt::Display;
use std::num::NonZeroU32;
use std::ops::DerefMut;

static GLOBAL_INSTRUMENTATION: std::sync::RwLock<fn() -> Option<Box<dyn Instrumentation>>> =
std::sync::RwLock::new(|| None);

/// A helper trait for opaque query representations
/// which allows to get a `Display` and `Debug`
/// representation of the underlying type without
/// exposing type specific details
pub trait DebugQuery: Debug + Display {}

impl<T, DB> DebugQuery for crate::query_builder::DebugQuery<'_, T, DB> where Self: Debug + Display {}

/// A helper type that allows printing out str slices
///
/// This type is necessary because it's not possible
/// to cast from a reference of a unsized type like `&str`
/// to a reference of a trait object even if that
/// type implements all necessary traits
#[diesel_derives::__diesel_public_if(
feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
)]
pub(crate) struct StrQueryHelper<'query> {
s: &'query str,
}

impl<'query> StrQueryHelper<'query> {
/// Construct a new `StrQueryHelper`
#[diesel_derives::__diesel_public_if(
feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
)]
#[cfg(any(
feature = "postgres",
feature = "sqlite",
feature = "mysql",
feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
))]
pub(crate) fn new(s: &'query str) -> Self {
Self { s }
}
}

impl Debug for StrQueryHelper<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
Debug::fmt(self.s, f)
}
}

impl Display for StrQueryHelper<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
Display::fmt(&self.s, f)
}
}

impl DebugQuery for StrQueryHelper<'_> {}

/// This enum describes possible connection events
/// that can be handled by an [`Instrumentation`] implementation
///
/// Some fields might contain sensitive information, like login
/// details for the database.
///
/// Diesel does not guarantee that future versions will
/// emit the same events in the same order or timing.
/// In addition the output of the [`Debug`] and [`Display`]
/// implementation of the enum itself and any of its fields
/// is not guarantee to be stable.
//
// This types is carefully designed
// to avoid any potential overhead by
// taking references for all things
// and by not performing any additional
// work until required.
// In addition it's carefully designed
// not to be dependent on the actual backend
// type, as that makes it easier to to reuse
// `Instrumentation` implementations in
// different a different context
#[derive(Debug)]
#[non_exhaustive]
pub enum InstrumentationEvent<'a> {
/// An event emitted by before starting
/// establishing a new connection
#[non_exhaustive]
StartEstablishConnection {
/// The database url the connection
/// tries to connect to
///
/// This might contain sensitive information
/// like the database password
url: &'a str,
},
/// An event emitted after establishing a
/// new connection
#[non_exhaustive]
FinishEstablishConnection {
/// The database url the connection
/// tries is connected to
///
/// This might contain sensitive information
/// like the database password
url: &'a str,
/// An optional error if the connection failed
error: Option<&'a crate::result::ConnectionError>,
},
/// An event that is emitted before executing
/// a query
#[non_exhaustive]
StartQuery {
/// A opaque representation of the query
///
/// This type implements [`Debug`] and [`Display`],
/// but should be considered otherwise as opaque.
///
/// The exact output of the [`Debug`] and [`Display`]
/// implementation is not considered as part of the
/// stable API.
query: &'a dyn DebugQuery,
},
/// An event that is emitted when a query
/// is cached in the connection internal
/// prepared statement cache
#[non_exhaustive]
CacheQuery {
/// SQL string of the cached query
sql: &'a str,
},
/// An event that is emitted after executing
/// a query
#[non_exhaustive]
FinishQuery {
/// A opaque representation of the query
///
/// This type implements [`Debug`] and [`Display`],
/// but should be considered otherwise as opaque.
///
/// The exact output of the [`Debug`] and [`Display`]
/// implementation is not considered as part of the
/// stable API.
query: &'a dyn DebugQuery,
/// An optional error if the connection failed
error: Option<&'a crate::result::Error>,
},
/// An event that is emitted while
/// starting a new transaction
#[non_exhaustive]
BeginTransaction {
/// Transaction level of the newly started
/// transaction
depth: NonZeroU32,
},
/// An event that is emitted while
/// committing a transaction
#[non_exhaustive]
CommitTransaction {
/// Transaction level of the to be committed
/// transaction
depth: NonZeroU32,
},
/// An event that is emitted while
/// rolling back a transaction
#[non_exhaustive]
RollbackTransaction {
/// Transaction level of the to be rolled
/// back transaction
depth: NonZeroU32,
},
}

// these constructors exist to
// keep `#[non_exhaustive]` on all the variants
// and to gate the constructors on the unstable feature
impl<'a> InstrumentationEvent<'a> {
/// Create a new `InstrumentationEvent::StartEstablishConnection` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn start_establish_connection(url: &'a str) -> Self {
Self::StartEstablishConnection { url }
}

/// Create a new `InstrumentationEvent::FinishEstablishConnection` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn finish_establish_connection(
url: &'a str,
error: Option<&'a crate::result::ConnectionError>,
) -> Self {
Self::FinishEstablishConnection { url, error }
}

/// Create a new `InstrumentationEvent::StartQuery` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn start_query(query: &'a dyn DebugQuery) -> Self {
Self::StartQuery { query }
}

/// Create a new `InstrumentationEvent::CacheQuery` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn cache_query(sql: &'a str) -> Self {
Self::CacheQuery { sql }
}

/// Create a new `InstrumentationEvent::FinishQuery` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn finish_query(
query: &'a dyn DebugQuery,
error: Option<&'a crate::result::Error>,
) -> Self {
Self::FinishQuery { query, error }
}

/// Create a new `InstrumentationEvent::BeginTransaction` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn begin_transaction(depth: NonZeroU32) -> Self {
Self::BeginTransaction { depth }
}

/// Create a new `InstrumentationEvent::RollbackTransaction` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn rollback_transaction(depth: NonZeroU32) -> Self {
Self::RollbackTransaction { depth }
}

/// Create a new `InstrumentationEvent::CommitTransaction` event
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub fn commit_transaction(depth: NonZeroU32) -> Self {
Self::CommitTransaction { depth }
}
}

/// A type that provides an connection `Instrumentation`
///
/// This trait is the basic building block for logging or
/// otherwise instrumenting diesel connection types. It
/// acts as callback that receives information about certain
/// important connection states
///
/// For simple usages this trait is implemented for closures
/// accepting a [`InstrumentationEvent`] as argument.
///
/// More complex usages and integrations with frameworks like
/// `tracing` and `log` are supposed to be part of their own
/// crates.
pub trait Instrumentation: Send + 'static {
/// The function that is invoced for each event
fn on_connection_event(&mut self, event: InstrumentationEvent<'_>);
}

/// Get an instance of the default [`Instrumentation`]
///
/// This function is mostly useful for crates implementing
/// their own connection types
pub fn get_default_instrumentation() -> Option<Box<dyn Instrumentation>> {
match GLOBAL_INSTRUMENTATION.read() {
Ok(f) => (*f)(),
Err(_) => None,
}
}

/// Set a custom constructor for the default [`Instrumentation`]
/// used by new connections
///
/// ```rust
/// use diesel::connection::{set_default_instrumentation, Instrumentation, InstrumentationEvent};
///
/// // a simple logger that prints all events to stdout
/// fn simple_logger() -> Option<Box<dyn Instrumentation>> {
/// // we need the explicit argument type there due
/// // to bugs in rustc
/// Some(Box::new(|event: InstrumentationEvent<'_>| println!("{event:?}")))
/// }
///
/// set_default_instrumentation(simple_logger);
/// ```
pub fn set_default_instrumentation(
default: fn() -> Option<Box<dyn Instrumentation>>,
) -> crate::QueryResult<()> {
match GLOBAL_INSTRUMENTATION.write() {
Ok(mut l) => {
*l = default;
Ok(())
}
Err(e) => Err(crate::result::Error::DatabaseError(
crate::result::DatabaseErrorKind::Unknown,
Box::new(e.to_string()),
)),
}
}

impl<F> Instrumentation for F
where
F: FnMut(InstrumentationEvent<'_>) + Send + 'static,
{
fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
(self)(event)
}
}

impl Instrumentation for Box<dyn Instrumentation> {
fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
self.deref_mut().on_connection_event(event)
}
}

impl<T> Instrumentation for Option<T>
where
T: Instrumentation,
{
fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
if let Some(i) = self {
i.on_connection_event(event)
}
}
}
18 changes: 18 additions & 0 deletions diesel/src/connection/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
//! Types related to database connections
pub(crate) mod instrumentation;
#[cfg(all(
not(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"),
any(feature = "sqlite", feature = "postgres", feature = "mysql")
Expand All @@ -15,6 +16,12 @@ use crate::query_builder::{Query, QueryFragment, QueryId};
use crate::result::*;
use std::fmt::Debug;

#[doc(inline)]
pub use self::instrumentation::{
get_default_instrumentation, set_default_instrumentation, DebugQuery, Instrumentation,
InstrumentationEvent,
};
#[doc(inline)]
pub use self::transaction_manager::{
AnsiTransactionManager, InTransactionStatus, TransactionDepthChange, TransactionManager,
TransactionManagerStatus, ValidTransactionManagerStatus,
Expand All @@ -28,6 +35,9 @@ pub(crate) use self::private::ConnectionSealed;
#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub use self::private::MultiConnectionHelper;

#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
pub use self::instrumentation::StrQueryHelper;

#[cfg(all(
not(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"),
any(feature = "sqlite", feature = "postgres", feature = "mysql")
Expand Down Expand Up @@ -381,6 +391,14 @@ where
fn transaction_state(
&mut self,
) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData;

#[diesel_derives::__diesel_public_if(
feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
)]
fn instrumentation(&mut self) -> &mut dyn Instrumentation;

/// Set a specific [`Instrumentation`] implementation for this connection
fn set_instrumentation(&mut self, instrumentation: impl Instrumentation);
}

/// The specific part of a [`Connection`] which actually loads data from the database
Expand Down
Loading

0 comments on commit ca3ba76

Please sign in to comment.