Fix port; add user pool mode (#395)

* Fix port; add user pool mode * will probably break our session/transaction mode tests
postgresml · Apr 5, 2023 · a62f6b0 · a62f6b0
1 parent 89e15f0
commit a62f6b0
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 6 deletions.
diff --git a/pgcat.toml b/pgcat.toml
@@ -134,6 +134,7 @@ password = "sharding_user"
 # is the sum of pool_size across all users.
 pool_size = 9
 
+
 # Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
 # 0 means it is disabled.
 statement_timeout = 0

diff --git a/src/auth_passthrough.rs b/src/auth_passthrough.rs
@@ -73,6 +73,7 @@ impl AuthPassthrough {
             password: Some(self.password.clone()),
             pool_size: 1,
             statement_timeout: 0,
+            pool_mode: None,
         };
 
         let user = &address.username;

diff --git a/src/config.rs b/src/config.rs
@@ -179,6 +179,7 @@ pub struct User {
     pub username: String,
     pub password: Option<String>,
     pub pool_size: u32,
+    pub pool_mode: Option<PoolMode>,
     #[serde(default)] // 0
     pub statement_timeout: u64,
 }
@@ -190,6 +191,7 @@ impl Default for User {
             password: None,
             pool_size: 15,
             statement_timeout: 0,
+            pool_mode: None,
         }
     }
 }
@@ -201,7 +203,7 @@ pub struct General {
     pub host: String,
 
     #[serde(default = "General::default_port")]
-    pub port: i16,
+    pub port: u16,
 
     pub enable_prometheus_exporter: Option<bool>,
     pub prometheus_exporter_port: i16,
@@ -261,7 +263,7 @@ impl General {
         "0.0.0.0".into()
     }
 
-    pub fn default_port() -> i16 {
+    pub fn default_port() -> u16 {
         5432
     }
 
@@ -356,6 +358,7 @@ pub enum PoolMode {
     #[serde(alias = "session", alias = "Session")]
     Session,
 }
+
 impl ToString for PoolMode {
     fn to_string(&self) -> String {
         match *self {
@@ -816,8 +819,9 @@ impl Config {
                     .to_string()
             );
             info!(
-                "[pool: {}] Pool mode: {:?}",
-                pool_name, pool_config.pool_mode
+                "[pool: {}] Default pool mode: {}",
+                pool_name,
+                pool_config.pool_mode.to_string()
             );
             info!(
                 "[pool: {}] Load Balancing mode: {:?}",
@@ -868,7 +872,16 @@ impl Config {
                 info!(
                     "[pool: {}][user: {}] Statement timeout: {}",
                     pool_name, user.1.username, user.1.statement_timeout
-                )
+                );
+                info!(
+                    "[pool: {}][user: {}] Pool mode: {}",
+                    pool_name,
+                    user.1.username,
+                    match user.1.pool_mode {
+                        Some(pool_mode) => pool_mode.to_string(),
+                        None => pool_config.pool_mode.to_string(),
+                    }
+                );
             }
         }
     }

diff --git a/src/pool.rs b/src/pool.rs
@@ -382,7 +382,10 @@ impl ConnectionPool {
                     server_info: Arc::new(RwLock::new(BytesMut::new())),
                     auth_hash: pool_auth_hash,
                     settings: PoolSettings {
-                        pool_mode: pool_config.pool_mode,
+                        pool_mode: match user.pool_mode {
+                            Some(pool_mode) => pool_mode,
+                            None => pool_config.pool_mode,
+                        },
                         load_balancing_mode: pool_config.load_balancing_mode,
                         // shards: pool_config.shards.clone(),
                         shards: shard_ids.len(),