Merge pull request #712 from EspressoSystems/jb/configurable-sqlx-log…

…ging

Add an option for setting the threshold at which statements are considered slow

Cargo.toml

@@ -36,6 +36,7 @@ no-storage = []
 # Enable the availability data source backed by a Postgres database.
 sql-data-source = [
 	"include_dir",
+	"log",
 	"refinery",
 	"refinery-core",
 	"sqlx",
@@ -104,6 +105,7 @@ atomic_store = { git = "https://github.com/EspressoSystems/atomicstore.git", tag
 
 # Dependencies enabled by feature "sql-data-source".
 include_dir = { version = "0.7", optional = true }
+log = { version = "0.4", optional = true }
 refinery = { version = "0.8", features = ["tokio-postgres"], optional = true }
 refinery-core = { version = "0.8", optional = true }
 sqlx = { version = "0.8", features = [

src/data_source/sql.rs

@@ -26,10 +26,9 @@ pub use anyhow::Error;
 use hotshot_types::traits::node_implementation::NodeType;
 pub use refinery::Migration;
 
-pub use sql::{Postgres, Transaction};
+pub use sql::{Config, Transaction};
 
 pub type Builder<Types, Provider> = fetching::Builder<Types, SqlStorage, Provider>;
-pub type Config = sql::Config<Postgres>;
 
 impl Config {
     /// Connect to the database with this config.

src/data_source/storage/sql.rs

@@ -27,15 +27,16 @@ use chrono::Utc;
 use futures::future::FutureExt;
 use hotshot_types::traits::metrics::Metrics;
 use itertools::Itertools;
+use log::LevelFilter;
 use sqlx::{
     pool::{Pool, PoolOptions},
     postgres::{PgConnectOptions, PgSslMode},
-    ConnectOptions, Connection, Row,
+    ConnectOptions, Row,
 };
 use std::{cmp::min, fmt::Debug, str::FromStr, time::Duration};
 
 pub extern crate sqlx;
-pub use sqlx::{Database, Postgres, Sqlite};
+pub use sqlx::{Database, Sqlite};
 
 mod db;
 mod migrate;
@@ -185,11 +186,8 @@ fn add_custom_migrations(
 
 /// Postgres client config.
 #[derive(Clone, Debug)]
-pub struct Config<DB>
-where
-    DB: Database,
-{
-    db_opt: <DB::Connection as Connection>::Options,
+pub struct Config {
+    db_opt: PgConnectOptions,
     pool_opt: PoolOptions<Db>,
     schema: String,
     reset: bool,
@@ -199,7 +197,7 @@ where
     archive: bool,
 }
 
-impl Default for Config<Postgres> {
+impl Default for Config {
     fn default() -> Self {
         PgConnectOptions::default()
             .host("localhost")
@@ -208,7 +206,7 @@ impl Default for Config<Postgres> {
     }
 }
 
-impl From<PgConnectOptions> for Config<Postgres> {
+impl From<PgConnectOptions> for Config {
     fn from(db_opt: PgConnectOptions) -> Self {
         Self {
             db_opt,
@@ -223,15 +221,15 @@ impl From<PgConnectOptions> for Config<Postgres> {
     }
 }
 
-impl FromStr for Config<Postgres> {
+impl FromStr for Config {
     type Err = <PgConnectOptions as FromStr>::Err;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         Ok(PgConnectOptions::from_str(s)?.into())
     }
 }
 
-impl Config<Postgres> {
+impl Config {
     /// Set the hostname of the database server.
     ///
     /// The default is `localhost`.
@@ -275,9 +273,7 @@ impl Config<Postgres> {
         self.db_opt = self.db_opt.ssl_mode(PgSslMode::Require);
         self
     }
-}
 
-impl<DB: Database> Config<DB> {
     /// Set the name of the schema to use for queries.
     ///
     /// The default schema is named `hotshot` and is created via the default migrations.
@@ -375,6 +371,16 @@ impl<DB: Database> Config<DB> {
         self.pool_opt = self.pool_opt.max_connections(max);
         self
     }
+
+    /// Log at WARN level any time a SQL statement takes longer than `threshold`.
+    ///
+    /// The default threshold is 1s.
+    pub fn slow_statement_threshold(mut self, threshold: Duration) -> Self {
+        self.db_opt = self
+            .db_opt
+            .log_slow_statements(LevelFilter::Warn, threshold);
+        self
+    }
 }
 
 /// Storage for the APIs provided in this crate, backed by a remote PostgreSQL database.
@@ -395,7 +401,7 @@ pub struct Pruner {
 
 impl SqlStorage {
     /// Connect to a remote database.
-    pub async fn connect<DB: Database>(mut config: Config<DB>) -> Result<Self, Error> {
+    pub async fn connect(mut config: Config) -> Result<Self, Error> {
         let schema = config.schema.clone();
         let pool = config
             .pool_opt
@@ -409,7 +415,7 @@ impl SqlStorage {
                 }
                 .boxed()
             })
-            .connect(config.db_opt.to_url_lossy().as_ref())
+            .connect_with(config.db_opt)
             .await?;
 
         // Create or connect to the schema for this query service.
@@ -678,7 +684,7 @@ pub mod testing {
     use portpicker::pick_unused_port;
     use refinery::Migration;
 
-    use super::{Config, Postgres};
+    use super::Config;
     use crate::testing::sleep;
 
     #[derive(Debug)]
@@ -745,7 +751,7 @@ pub mod testing {
             self.port
         }
 
-        pub fn config(&self) -> Config<Postgres> {
+        pub fn config(&self) -> Config {
             Config::default()
                 .user("postgres")
                 .password("password")