From 1ebfa387fa5d47e3200b990f63f55bbe36adeff1 Mon Sep 17 00:00:00 2001
From: Matthias Beyer <matthias.beyer@atos.net>
Date: Fri, 23 Apr 2021 12:51:57 +0200
Subject: Add support for database connection timeout

This patch adds support for a database connection timeout, giving the user the
option to configure one, but hardcoding a default timeout to 30 seconds.

Signed-off-by: Matthias Beyer <matthias.beyer@atos.net>
---
 config.toml                 |  4 ++++
 src/cli.rs                  |  7 +++++++
 src/config/not_validated.rs |  4 ++++
 src/db/connection.rs        | 14 ++++++++++++--
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/config.toml b/config.toml
index 79f341a..a5a0024 100644
--- a/config.toml
+++ b/config.toml
@@ -127,6 +127,10 @@ database_user     = "pgdev"
 database_password = "password"
 database_name     = "butido"
 
+# Set a database connection timeout
+# If not set, this defaults to 30
+#database_connection_timeout = 30
+
 
 # Phases which can be configured in the packages
 
diff --git a/src/cli.rs b/src/cli.rs
index 4866ccd..4c629f8 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -69,6 +69,13 @@ pub fn cli<'a>() -> App<'a> {
             .value_name("NAME")
             .about("Overwrite the database name set via configuration. Can also be overriden via environment, but this setting has presendence.")
         )
+        .arg(Arg::new("database_connection_timeout")
+            .required(false)
+            .multiple(false)
+            .long("db-timeout")
+            .value_name("TIMEOUT")
+            .about("Overwrite the database connection timeout set via configuration. Can also be overriden via environment, but this setting has presendence.")
+        )
 
         .subcommand(App::new("generate-completions")
             .version(crate_version!())
diff --git a/src/config/not_validated.rs b/src/config/not_validated.rs
index 451a3d2..14399c4 100644
--- a/src/config/not_validated.rs
+++ b/src/config/not_validated.rs
@@ -96,6 +96,10 @@ pub struct NotValidatedConfiguration {
     #[serde(rename = "database_name")]
     database_name: String,
 
+    #[getset(get = "pub")]
+    #[serde(rename = "database_connection_timeout")]
+    database_connection_timeout: Option<u16>,
+
     #[getset(get = "pub")]
     docker: DockerConfig,
 
diff --git a/src/db/connection.rs b/src/db/connection.rs
index 33de9e7..fb18e32 100644
--- a/src/db/connection.rs
+++ b/src/db/connection.rs
@@ -34,17 +34,21 @@ pub struct DbConnectionConfig {
 
     #[getset(get = "pub")]
     database_name: String,
+
+    #[getset(get = "pub")]
+    database_connection_timeout: String,
 }
 
 impl Into<String> for DbConnectionConfig {
     fn into(self) -> String {
         format!(
-            "postgres://{user}:{password}@{host}:{port}/{name}",
+            "postgres://{user}:{password}@{host}:{port}/{name}?connect_timeout={timeout}",
             host = self.database_host,
             port = self.database_port,
             user = self.database_user,
             password = self.database_password,
-            name = self.database_name
+            name = self.database_name,
+            timeout = self.database_connection_timeout,
         )
     }
 }
@@ -66,6 +70,11 @@ pub fn parse_db_connection_config(config: &Configuration, cli: &ArgMatches) -> D
         config.database_password().to_string()
     });
     let database_name = find_value(cli, "database_name", || config.database_name().to_string());
+    let database_connection_timeout = find_value(cli, "database_connection_timeout",
+        || {
+            // hardcoded default of 30 seconds database timeout
+            config.database_connection_timeout().unwrap_or(30).to_string()
+        });
 
     DbConnectionConfig {
         database_host,
@@ -73,6 +82,7 @@ pub fn parse_db_connection_config(config: &Configuration, cli: &ArgMatches) -> D
         database_user,
         database_password,
         database_name,
+        database_connection_timeout,
     }
 }
 
-- 
cgit v1.2.3