1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

extern crate diesel;
#[macro_use]
extern crate log;

use std::ops::Deref;
use std::time::{Duration, Instant};

use diesel::prelude::*;
use diesel::debug_query;
use diesel::backend::{Backend, UsesAnsiSavepointSyntax};
use diesel::connection::{AnsiTransactionManager, SimpleConnection};
use diesel::deserialize::QueryableByName;
use diesel::query_builder::{AsQuery, QueryFragment, QueryId};
use diesel::sql_types::HasSqlType;

/// Wraps a diesel `Connection` to time and log each query using
/// the configured logger for the `log` crate.
///
/// Currently, this produces a `debug` log on every query,
/// an `info` on queries that take longer than 1 second,
/// and a `warn`ing on queries that take longer than 5 seconds.
/// These thresholds will be configurable in a future version.
pub struct LoggingConnection<C: Connection>(C);

impl<C: Connection> LoggingConnection<C> {
    pub fn new(conn: C) -> Self {
        LoggingConnection(conn)
    }
}

impl<C: Connection> Deref for LoggingConnection<C> {
    type Target = C;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<C> SimpleConnection for LoggingConnection<C>
where
    C: Connection + Send + 'static,
{
    fn batch_execute(&self, query: &str) -> QueryResult<()> {
        self.0.batch_execute(query)
    }
}

impl<C: Connection> Connection for LoggingConnection<C>
where
    C: Connection<TransactionManager = AnsiTransactionManager> + Send + 'static,
    C::Backend: UsesAnsiSavepointSyntax,
    <C::Backend as Backend>::QueryBuilder: Default,
{
    type Backend = C::Backend;
    type TransactionManager = C::TransactionManager;

    fn establish(database_url: &str) -> ConnectionResult<Self> {
        Ok(LoggingConnection(C::establish(database_url)?))
    }

    fn execute(&self, query: &str) -> QueryResult<usize> {
        let start_time = Instant::now();
        let result = self.0.execute(query);
        let duration = start_time.elapsed();
        log_query(query, duration);
        result
    }

    fn query_by_index<T, U>(&self, source: T) -> QueryResult<Vec<U>>
    where
        T: AsQuery,
        T::Query: QueryFragment<Self::Backend> + QueryId,
        Self::Backend: HasSqlType<T::SqlType>,
        U: Queryable<T::SqlType, Self::Backend>,
    {
        let query = source.as_query();
        let debug_query = debug_query::<Self::Backend, _>(&query).to_string();
        let start_time = Instant::now();
        let result = self.0.query_by_index(query);
        let duration = start_time.elapsed();
        log_query(&debug_query, duration);
        result
    }

    fn query_by_name<T, U>(&self, source: &T) -> QueryResult<Vec<U>>
    where
        T: QueryFragment<Self::Backend> + QueryId,
        U: QueryableByName<Self::Backend>,
    {
        let debug_query = debug_query::<Self::Backend, _>(&source).to_string();
        let start_time = Instant::now();
        let result = self.0.query_by_name(source);
        let duration = start_time.elapsed();
        log_query(&debug_query, duration);
        result
    }

    fn execute_returning_count<T>(&self, source: &T) -> QueryResult<usize>
    where
        T: QueryFragment<Self::Backend> + QueryId,
    {
        let debug_query = debug_query::<Self::Backend, _>(&source).to_string();
        let start_time = Instant::now();
        let result = self.0.execute_returning_count(source);
        let duration = start_time.elapsed();
        log_query(&debug_query, duration);
        result
    }

    fn transaction_manager(&self) -> &Self::TransactionManager {
        self.0.transaction_manager()
    }
}

fn log_query(query: &str, duration: Duration) {
    if duration.as_secs() >= 5 {
        warn!(
            "Slow query ran in {:.2} seconds: {}",
            duration_to_secs(duration),
            query
        );
    } else if duration.as_secs() >= 1 {
        info!(
            "Slow query ran in {:.2} seconds: {}",
            duration_to_secs(duration),
            query
        );
    } else {
        debug!("Query ran in {:.1} ms: {}", duration_to_ms(duration), query);
    }
}

const NANOS_PER_MILLI: u32 = 1_000_000;
const MILLIS_PER_SEC: u32 = 1_000;

fn duration_to_secs(duration: Duration) -> f32 {
    duration_to_ms(duration) / MILLIS_PER_SEC as f32
}

fn duration_to_ms(duration: Duration) -> f32 {
    (duration.as_secs() as u32 * 1000) as f32
        + (duration.subsec_nanos() as f32 / NANOS_PER_MILLI as f32)
}