diesel/connection/
instrumentation.rs

1use downcast_rs::Downcast;
2use std::fmt::{Debug, Display};
3use std::num::NonZeroU32;
4use std::ops::{Deref, DerefMut};
5
6static GLOBAL_INSTRUMENTATION: std::sync::RwLock<fn() -> Option<Box<dyn Instrumentation>>> =
7    std::sync::RwLock::new(|| None);
8
9/// A helper trait for opaque query representations
10/// which allows to get a `Display` and `Debug`
11/// representation of the underlying type without
12/// exposing type specific details
13pub trait DebugQuery: Debug + Display {}
14
15impl<T, DB> DebugQuery for crate::query_builder::DebugQuery<'_, T, DB> where Self: Debug + Display {}
16
17/// A helper type that allows printing out str slices
18///
19/// This type is necessary because it's not possible
20/// to cast from a reference of a unsized type like `&str`
21/// to a reference of a trait object even if that
22/// type implements all necessary traits
23#[diesel_derives::__diesel_public_if(
24    feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
25)]
26pub(crate) struct StrQueryHelper<'query> {
27    s: &'query str,
28}
29
30impl<'query> StrQueryHelper<'query> {
31    /// Construct a new `StrQueryHelper`
32    #[diesel_derives::__diesel_public_if(
33        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
34    )]
35    #[cfg(any(
36        feature = "postgres",
37        feature = "sqlite",
38        feature = "mysql",
39        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
40    ))]
41    pub(crate) fn new(s: &'query str) -> Self {
42        Self { s }
43    }
44}
45
46impl Debug for StrQueryHelper<'_> {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        Debug::fmt(self.s, f)
49    }
50}
51
52impl Display for StrQueryHelper<'_> {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        Display::fmt(&self.s, f)
55    }
56}
57
58impl DebugQuery for StrQueryHelper<'_> {}
59
60/// This enum describes possible connection events
61/// that can be handled by an [`Instrumentation`] implementation
62///
63/// Some fields might contain sensitive information, like login
64/// details for the database.
65///
66/// Diesel does not guarantee that future versions will
67/// emit the same events in the same order or timing.
68/// In addition the output of the [`Debug`] and [`Display`]
69/// implementation of the enum itself and any of its fields
70/// is not guarantee to be stable.
71//
72// This type is carefully designed
73// to avoid any potential overhead by
74// taking references for all things
75// and by not performing any additional
76// work until required.
77// In addition it's carefully designed
78// not to be dependent on the actual backend
79// type, as that makes it easier to reuse
80// `Instrumentation` implementations in
81// a different context
82#[derive(Debug)]
83#[non_exhaustive]
84pub enum InstrumentationEvent<'a> {
85    /// An event emitted by before starting
86    /// establishing a new connection
87    #[non_exhaustive]
88    StartEstablishConnection {
89        /// The database url the connection
90        /// tries to connect to
91        ///
92        /// This might contain sensitive information
93        /// like the database password
94        url: &'a str,
95    },
96    /// An event emitted after establishing a
97    /// new connection
98    #[non_exhaustive]
99    FinishEstablishConnection {
100        /// The database url the connection
101        /// tries is connected to
102        ///
103        /// This might contain sensitive information
104        /// like the database password
105        url: &'a str,
106        /// An optional error if the connection failed
107        error: Option<&'a crate::result::ConnectionError>,
108    },
109    /// An event that is emitted before executing
110    /// a query
111    #[non_exhaustive]
112    StartQuery {
113        /// A opaque representation of the query
114        ///
115        /// This type implements [`Debug`] and [`Display`],
116        /// but should be considered otherwise as opaque.
117        ///
118        /// The exact output of the [`Debug`] and [`Display`]
119        /// implementation is not considered as part of the
120        /// stable API.
121        query: &'a dyn DebugQuery,
122    },
123    /// An event that is emitted when a query
124    /// is cached in the connection internal
125    /// prepared statement cache
126    #[non_exhaustive]
127    CacheQuery {
128        /// SQL string of the cached query
129        sql: &'a str,
130    },
131    /// An event that is emitted after executing
132    /// a query
133    #[non_exhaustive]
134    FinishQuery {
135        /// A opaque representation of the query
136        ///
137        /// This type implements [`Debug`] and [`Display`],
138        /// but should be considered otherwise as opaque.
139        ///
140        /// The exact output of the [`Debug`] and [`Display`]
141        /// implementation is not considered as part of the
142        /// stable API.
143        query: &'a dyn DebugQuery,
144        /// An optional error if the connection failed
145        error: Option<&'a crate::result::Error>,
146    },
147    /// An event that is emitted while
148    /// starting a new transaction
149    #[non_exhaustive]
150    BeginTransaction {
151        /// Transaction level of the newly started
152        /// transaction
153        depth: NonZeroU32,
154    },
155    /// An event that is emitted while
156    /// committing a transaction
157    #[non_exhaustive]
158    CommitTransaction {
159        /// Transaction level of the to be committed
160        /// transaction
161        depth: NonZeroU32,
162    },
163    /// An event that is emitted while
164    /// rolling back a transaction
165    #[non_exhaustive]
166    RollbackTransaction {
167        /// Transaction level of the to be rolled
168        /// back transaction
169        depth: NonZeroU32,
170    },
171}
172
173// these constructors exist to
174// keep `#[non_exhaustive]` on all the variants
175// and to gate the constructors on the unstable feature
176#[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
177impl<'a> InstrumentationEvent<'a> {
178    /// Create a new `InstrumentationEvent::StartEstablishConnection` event
179    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
180    pub fn start_establish_connection(url: &'a str) -> Self {
181        Self::StartEstablishConnection { url }
182    }
183
184    /// Create a new `InstrumentationEvent::FinishEstablishConnection` event
185    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
186    pub fn finish_establish_connection(
187        url: &'a str,
188        error: Option<&'a crate::result::ConnectionError>,
189    ) -> Self {
190        Self::FinishEstablishConnection { url, error }
191    }
192
193    /// Create a new `InstrumentationEvent::StartQuery` event
194    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
195    pub fn start_query(query: &'a dyn DebugQuery) -> Self {
196        Self::StartQuery { query }
197    }
198
199    /// Create a new `InstrumentationEvent::CacheQuery` event
200    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
201    pub fn cache_query(sql: &'a str) -> Self {
202        Self::CacheQuery { sql }
203    }
204
205    /// Create a new `InstrumentationEvent::FinishQuery` event
206    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
207    pub fn finish_query(
208        query: &'a dyn DebugQuery,
209        error: Option<&'a crate::result::Error>,
210    ) -> Self {
211        Self::FinishQuery { query, error }
212    }
213
214    /// Create a new `InstrumentationEvent::BeginTransaction` event
215    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
216    pub fn begin_transaction(depth: NonZeroU32) -> Self {
217        Self::BeginTransaction { depth }
218    }
219
220    /// Create a new `InstrumentationEvent::RollbackTransaction` event
221    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
222    pub fn rollback_transaction(depth: NonZeroU32) -> Self {
223        Self::RollbackTransaction { depth }
224    }
225
226    /// Create a new `InstrumentationEvent::CommitTransaction` event
227    #[cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes")]
228    pub fn commit_transaction(depth: NonZeroU32) -> Self {
229        Self::CommitTransaction { depth }
230    }
231}
232
233/// A type that provides an connection `Instrumentation`
234///
235/// This trait is the basic building block for logging or
236/// otherwise instrumenting diesel connection types. It
237/// acts as callback that receives information about certain
238/// important connection states
239///
240/// For simple usages this trait is implemented for closures
241/// accepting a [`InstrumentationEvent`] as argument.
242///
243/// More complex usages and integrations with frameworks like
244/// `tracing` and `log` are supposed to be part of their own
245/// crates.
246pub trait Instrumentation: Downcast + Send + 'static {
247    /// The function that is invoked for each event
248    fn on_connection_event(&mut self, event: InstrumentationEvent<'_>);
249}
250downcast_rs::impl_downcast!(Instrumentation);
251
252/// Get an instance of the default [`Instrumentation`]
253///
254/// This function is mostly useful for crates implementing
255/// their own connection types
256pub fn get_default_instrumentation() -> Option<Box<dyn Instrumentation>> {
257    match GLOBAL_INSTRUMENTATION.read() {
258        Ok(f) => (*f)(),
259        Err(_) => None,
260    }
261}
262
263/// Set a custom constructor for the default [`Instrumentation`]
264/// used by new connections
265///
266/// ```rust
267/// use diesel::connection::{set_default_instrumentation, Instrumentation, InstrumentationEvent};
268///
269/// // a simple logger that prints all events to stdout
270/// fn simple_logger() -> Option<Box<dyn Instrumentation>> {
271///     // we need the explicit argument type there due
272///     // to bugs in rustc
273///     Some(Box::new(|event: InstrumentationEvent<'_>| {
274///         println!("{event:?}")
275///     }))
276/// }
277///
278/// set_default_instrumentation(simple_logger);
279/// ```
280pub fn set_default_instrumentation(
281    default: fn() -> Option<Box<dyn Instrumentation>>,
282) -> crate::QueryResult<()> {
283    match GLOBAL_INSTRUMENTATION.write() {
284        Ok(mut l) => {
285            *l = default;
286            Ok(())
287        }
288        Err(e) => Err(crate::result::Error::DatabaseError(
289            crate::result::DatabaseErrorKind::Unknown,
290            Box::new(e.to_string()),
291        )),
292    }
293}
294
295impl<F> Instrumentation for F
296where
297    F: FnMut(InstrumentationEvent<'_>) + Send + 'static,
298{
299    fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
300        (self)(event)
301    }
302}
303
304impl Instrumentation for Box<dyn Instrumentation> {
305    fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
306        self.deref_mut().on_connection_event(event)
307    }
308}
309
310impl<T> Instrumentation for Option<T>
311where
312    T: Instrumentation,
313{
314    fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
315        if let Some(i) = self {
316            i.on_connection_event(event)
317        }
318    }
319}
320
321#[diesel_derives::__diesel_public_if(
322    feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
323)]
324/// An optional dyn instrumentation.
325///
326/// For ease of use, this type implements [`Deref`] and [`DerefMut`] to `&dyn Instrumentation`,
327/// falling back to a no-op implementation if no instrumentation is set.
328///
329/// The DynInstrumentation type is useful because without it we actually did tend to return
330/// (accidentally) &mut Option<Box> as &mut dyn Instrumentation from connection.instrumentation(),
331/// so downcasting would have to be done in these two steps by the user, which is counter-intuitive.
332pub(crate) struct DynInstrumentation {
333    /// zst
334    no_instrumentation: NoInstrumentation,
335    inner: Option<Box<dyn Instrumentation>>,
336}
337
338impl Deref for DynInstrumentation {
339    type Target = dyn Instrumentation;
340
341    fn deref(&self) -> &Self::Target {
342        self.inner.as_deref().unwrap_or(&self.no_instrumentation)
343    }
344}
345
346impl DerefMut for DynInstrumentation {
347    fn deref_mut(&mut self) -> &mut Self::Target {
348        self.inner
349            .as_deref_mut()
350            .unwrap_or(&mut self.no_instrumentation)
351    }
352}
353
354impl DynInstrumentation {
355    /// Create a instance of the default instrumentation provider
356    #[diesel_derives::__diesel_public_if(
357        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
358    )]
359    pub(crate) fn default_instrumentation() -> Self {
360        Self {
361            inner: get_default_instrumentation(),
362            no_instrumentation: NoInstrumentation,
363        }
364    }
365
366    /// Create a noop instrumentation provider instance
367    #[diesel_derives::__diesel_public_if(
368        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
369    )]
370    pub(crate) fn none() -> Self {
371        Self {
372            inner: None,
373            no_instrumentation: NoInstrumentation,
374        }
375    }
376
377    /// register an event with the given instrumentation implementation
378    #[diesel_derives::__diesel_public_if(
379        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"
380    )]
381    pub(crate) fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) {
382        // This implementation is not necessary to be able to call this method on this object
383        // because of the already existing Deref impl.
384        // However it allows avoiding the dynamic dispatch to the stub value
385        if let Some(inner) = self.inner.as_deref_mut() {
386            inner.on_connection_event(event)
387        }
388    }
389}
390
391impl<I: Instrumentation> From<I> for DynInstrumentation {
392    fn from(instrumentation: I) -> Self {
393        Self {
394            inner: Some(unpack_instrumentation(Box::new(instrumentation))),
395            no_instrumentation: NoInstrumentation,
396        }
397    }
398}
399
400struct NoInstrumentation;
401
402impl Instrumentation for NoInstrumentation {
403    fn on_connection_event(&mut self, _: InstrumentationEvent<'_>) {}
404}
405
406/// Unwrap unnecessary boxing levels
407fn unpack_instrumentation(
408    mut instrumentation: Box<dyn Instrumentation>,
409) -> Box<dyn Instrumentation> {
410    loop {
411        match instrumentation.downcast::<Box<dyn Instrumentation>>() {
412            Ok(extra_boxed_instrumentation) => instrumentation = *extra_boxed_instrumentation,
413            Err(not_extra_boxed_instrumentation) => {
414                break not_extra_boxed_instrumentation;
415            }
416        }
417    }
418}