From 5cce9929dd59b2adf147e8e523328bf53cd3cf90 Mon Sep 17 00:00:00 2001
From: Felipe Gonzalez <gonzalezz_felipe@hotmail.com>
Date: Thu, 18 Apr 2024 16:49:38 -0300
Subject: [PATCH] chore: Add statement timeout to created users (#25)

* chore: Add statement timeout to created users

* Fix default timeout and update version

* Rename query_timeout to statement_timeout

* Rename query_timeout to statement_timeout on README

* Change match pattern to unwrap_or
---
 operator/Cargo.lock      |  2 +-
 operator/Cargo.toml      |  2 +-
 operator/README.md       |  1 +
 operator/src/config.rs   | 49 ++++++++++++++++++++++++++++++++++++++++
 operator/src/postgres.rs | 13 ++++++++++-
 5 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/operator/Cargo.lock b/operator/Cargo.lock
index f3fb5bf..26d33b1 100644
--- a/operator/Cargo.lock
+++ b/operator/Cargo.lock
@@ -615,7 +615,7 @@ checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
 
 [[package]]
 name = "ext-cardano-dbsync"
-version = "0.1.0"
+version = "0.1.1"
 dependencies = [
  "actix-web",
  "bech32",
diff --git a/operator/Cargo.toml b/operator/Cargo.toml
index 18f736b..6b0cc32 100644
--- a/operator/Cargo.toml
+++ b/operator/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ext-cardano-dbsync"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 default-run = "controller"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/operator/README.md b/operator/README.md
index 11294fe..e6dd318 100644
--- a/operator/README.md
+++ b/operator/README.md
@@ -24,6 +24,7 @@ This project is a Kubernetes custom controller to create users on dbsync's Postg
 | DB_MAX_CONNECTIONS | 2                                                                                       |
 | DCU_PER_SECOND     | preview=5,preprod=5,mainnet=5                                                           |
 | METRICS_DELAY      | 30                                                                                      |
+| STATEMENT_TIMEOUT  | 12000                                                                                   |
 
 
 ## Commands
diff --git a/operator/src/config.rs b/operator/src/config.rs
index 4346f3a..8091968 100644
--- a/operator/src/config.rs
+++ b/operator/src/config.rs
@@ -17,6 +17,7 @@ pub struct Config {
     pub dcu_per_second: HashMap<String, f64>,
 
     pub metrics_delay: Duration,
+    pub statement_timeout: u64,
 }
 
 impl Config {
@@ -63,12 +64,60 @@ impl Config {
                 .expect("METRICS_DELAY must be a number"),
         );
 
+        let statement_timeout = env::var("STATEMENT_TIMEOUT")
+            .unwrap_or("120000".to_string())
+            .parse::<u64>()
+            .expect("STATEMENT_TIMEOUT must be a number");
+
         Self {
             db_urls,
             db_names,
             db_max_connections,
             dcu_per_second,
             metrics_delay,
+            statement_timeout,
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_from_env() {
+        env::set_var("DB_URLS", "url1,url2");
+        env::set_var(
+            "DB_NAMES",
+            "preview=dbsync-preview,preprod=dbsync-preprod,mainnet=dbsync-mainnet",
+        );
+        env::set_var("DCU_PER_SECOND", "preview=5,preprod=5,mainnet=5");
+        env::set_var("METRICS_DELAY", "100");
+        env::set_var("STATEMENT_TIMEOUT", "100");
+
+        let config = Config::from_env();
+        assert_eq!(config.db_urls, vec!["url1".to_owned(), "url2".to_owned()]);
+        assert_eq!(
+            config.db_names,
+            HashMap::from([
+                ("preview".to_owned(), "dbsync-preview".to_owned()),
+                ("preprod".to_owned(), "dbsync-preprod".to_owned()),
+                ("mainnet".to_owned(), "dbsync-mainnet".to_owned())
+            ])
+        );
+        assert_eq!(
+            config.dcu_per_second,
+            HashMap::from([
+                ("preview".to_owned(), 5.0),
+                ("preprod".to_owned(), 5.0),
+                ("mainnet".to_owned(), 5.0)
+            ])
+        );
+        assert_eq!(config.statement_timeout, 100);
+
+        // Check default query timeout
+        env::remove_var("STATEMENT_TIMEOUT");
+        let config = Config::from_env();
+        assert_eq!(config.statement_timeout, 120000);
+    }
+}
diff --git a/operator/src/postgres.rs b/operator/src/postgres.rs
index fd0d17c..1839052 100644
--- a/operator/src/postgres.rs
+++ b/operator/src/postgres.rs
@@ -3,7 +3,7 @@ use std::str::FromStr;
 use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
 use tokio_postgres::{NoTls, Row};
 
-use crate::Error;
+use crate::{get_config, Error};
 
 #[derive(Clone)]
 pub struct Postgres {
@@ -32,6 +32,10 @@ impl Postgres {
         let query_create_user = format!("create user \"{username}\" with password '{password}';");
         let query_grant = format!("grant select on all tables in schema public to \"{username}\";");
 
+        let timeout = get_config().statement_timeout;
+        let query_set_timeout =
+            format!("alter role \"{username}\" set statement_timeout = '{timeout}';");
+
         let mut client = self.pool.get().await?;
         let tx = client.transaction().await?;
 
@@ -49,6 +53,13 @@ impl Postgres {
             return Err(Error::PgError(err.to_string()));
         }
 
+        let set_timeout_stmt = tx.prepare(&query_set_timeout).await?;
+        let set_timeout_result = tx.execute(&set_timeout_stmt, &[]).await;
+        if let Err(err) = set_timeout_result {
+            tx.rollback().await?;
+            return Err(Error::PgError(err.to_string()));
+        }
+
         tx.commit().await?;
         Ok(())
     }