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.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

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

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);
+    }
+}

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(())
     }