From 56b2c63a4c11c207a55d9eebbcc0e7b5ea4237a8 Mon Sep 17 00:00:00 2001
From: pawurb <contact@pawelurbanek.com>
Date: Tue, 24 Sep 2024 23:18:45 +0200
Subject: [PATCH] Custom error impl

---
 Cargo.toml  | 15 ++++++++-------
 src/lib.rs  | 34 ++++++++++++++++++++++++++--------
 src/main.rs | 34 ++++++++++++++++++++++++++++------
 3 files changed, 62 insertions(+), 21 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index b27391d..e8b73e9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,16 +9,17 @@ readme = "README.md"
 repository = "https://github.com/pawurb/rust-pg-extras"
 version = "0.3.2"
 
-exclude = [
-  "docker-compose.yml.sample",
-  "live_tests.sh",
-]
+exclude = ["docker-compose.yml.sample", "live_tests.sh"]
 
 [dependencies]
 prettytable-rs = "0.10"
-sqlx = {version = "0.6", features = ["runtime-tokio-rustls", "postgres", "macros", "bigdecimal"]}
-thiserror = "1.0"
-tokio = {version = "1", features = ["full"]}
+sqlx = { version = "0.6", features = [
+  "runtime-tokio-rustls",
+  "postgres",
+  "macros",
+  "bigdecimal",
+] }
+tokio = { version = "1", features = ["full"] }
 
 [[bin]]
 name = "pg_extras"
diff --git a/src/lib.rs b/src/lib.rs
index 3e0a51f..8756426 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,6 @@
 use std::collections::HashMap;
-use std::env;
+use std::time::Duration;
+use std::{env, fmt};
 pub mod queries;
 pub use queries::all_locks::AllLocks;
 pub use queries::bloat::Bloat;
@@ -36,7 +37,6 @@ pub use queries::total_table_size::TotalTableSize;
 pub use queries::unused_indexes::UnusedIndexes;
 pub use queries::vacuum_stats::VacuumStats;
 use sqlx::postgres::PgPoolOptions;
-use thiserror::Error;
 
 #[macro_use]
 extern crate prettytable;
@@ -204,17 +204,31 @@ pub async fn db_settings() -> Result<Vec<DbSettings>, PgExtrasError> {
     get_rows(None).await
 }
 
-#[derive(Error, Debug)]
+#[derive(Debug, Clone)]
 #[non_exhaustive]
 pub enum PgExtrasError {
-    #[error("Both $DATABASE_URL and $PG_EXTRAS_DATABASE_URL are not set.")]
     MissingConfigVars(),
-    #[error("Cannot connect to database: '{0}'")]
     DbConnectionError(String),
-    #[error("Unknown pg-extras error: '{0}'")]
     Unknown(String),
 }
 
+impl fmt::Display for PgExtrasError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::MissingConfigVars() => write!(
+                f,
+                "Both $DATABASE_URL and $PG_EXTRAS_DATABASE_URL are not set."
+            ),
+            Self::DbConnectionError(e) => {
+                write!(f, "Cannot connect to database: '{}'", e)
+            }
+            Self::Unknown(e) => write!(f, "Unknown pg-extras error: '{}'", e),
+        }
+    }
+}
+
+impl std::error::Error for PgExtrasError {}
+
 async fn get_rows<T: Query>(
     params: Option<HashMap<String, String>>,
 ) -> Result<Vec<T>, PgExtrasError> {
@@ -228,6 +242,7 @@ async fn get_rows<T: Query>(
 
     let pool = match PgPoolOptions::new()
         .max_connections(5)
+        .acquire_timeout(Duration::from_secs(10))
         .connect(db_url()?.as_str())
         .await
     {
@@ -265,11 +280,14 @@ fn limit_params(limit: Option<String>) -> HashMap<String, String> {
 
 #[cfg(test)]
 mod tests {
+    use std::time::Duration;
+
     use super::*;
 
     async fn setup() -> Result<(), Box<dyn std::error::Error>> {
         let pool = PgPoolOptions::new()
             .max_connections(5)
+            .acquire_timeout(Duration::from_secs(10))
             .connect(db_url()?.as_str())
             .await?;
 
@@ -322,10 +340,10 @@ mod tests {
         Ok(())
     }
 
-    fn is_normal<T: Sized + Send + Sync + Unpin>() {}
-
     #[test]
     fn normal_types() {
+        fn is_normal<T: Sized + Send + Sync + Unpin>() {}
+
         is_normal::<NullIndexes>();
         is_normal::<Bloat>();
         is_normal::<Blocking>();
diff --git a/src/main.rs b/src/main.rs
index 7825eaf..67ab0e3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,8 +7,7 @@ use pg_extras::{
     vacuum_stats, PgExtrasError,
 };
 
-use std::env;
-use thiserror::Error;
+use std::{env, fmt};
 
 #[tokio::main]
 async fn main() {
@@ -136,19 +135,42 @@ async fn execute(command: Option<&String>) -> Result<(), PgExtrasCmdError> {
     Ok(())
 }
 
-#[derive(Error, Debug)]
+#[derive(Debug, Clone)]
 #[non_exhaustive]
 pub enum PgExtrasCmdError {
-    #[error("Unknown command '{0}'. Check https://github.com/pawurb/rust-pg-extras for list of available commands.")]
     UnknownCommand(String),
-    #[error("Missing command. Check https://github.com/pawurb/rust-pg-extras for list of available commands.")]
     MissingCommand,
-    #[error("{0}")]
     Other(PgExtrasError),
 }
 
+impl fmt::Display for PgExtrasCmdError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let check_command_text =
+            "Check https://github.com/pawurb/rust-pg-extras for list of available commands.";
+        match self {
+            Self::UnknownCommand(command) => {
+                write!(f, "Unknown command '{}'. {}", command, check_command_text)
+            }
+            Self::MissingCommand => {
+                write!(f, "Missing command. {}", check_command_text)
+            }
+            Self::Other(error) => {
+                write!(f, "{}", error)
+            }
+        }
+    }
+}
+
 impl From<pg_extras::PgExtrasError> for PgExtrasCmdError {
     fn from(error: pg_extras::PgExtrasError) -> Self {
         Self::Other(error)
     }
 }
+
+impl std::error::Error for PgExtrasCmdError {}
+
+#[test]
+fn normal_types() {
+    fn is_normal<T: Sized + Send + Sync + Unpin>() {}
+    is_normal::<PgExtrasCmdError>();
+}