add auth_query_database config

This commit adds the new config auth_query_database to allow use a
different database to get the hashes from the database instance.

I updated the example folder with functional code to validate the new
feature and updated the config docs.

Signed-off-by: Sebastian Webber <[email protected]>

Author: sebastianwebber

CONFIG.md

@@ -319,6 +319,16 @@ Password to be used for connecting to servers to obtain the hash used for md5 au
 specified in `auth_query_user`. The connection will be established using the database configured in the pool.
 This parameter is inherited by every pool and can be redefined in pool configuration.
 
+### auth_query_database
+```
+path: pools.<pool_name>.auth_query_database
+default: <UNSET>
+example: "postgres"
+```
+
+Database to be used for connecting to servers to obtain the hash used for md5 authentication by sending the query
+specified in `auth_query_query`. This parameter is inherited by every pool and can be redefined in pool configuration.
+
 ### automatic_sharding_key
 ```
 path: pools.<pool_name>.automatic_sharding_key

docker-compose.yml

@@ -5,6 +5,10 @@ services:
     environment:
       POSTGRES_PASSWORD: postgres
       POSTGRES_HOST_AUTH_METHOD: md5
+    ports:
+      - "54321:5432"
+    volumes:
+      - "${PWD}/examples/docker:/docker-entrypoint-initdb.d"
   pgcat:
     build: .
     command:

examples/docker/01-setup.sql

@@ -0,0 +1,2 @@
+ALTER SYSTEM SET password_encryption to 'md5';
+SELECT pg_reload_conf();

examples/docker/02-add-users-and-db.sql

@@ -0,0 +1,3 @@
+CREATE ROLE myappadmin PASSWORD 'myappadmin' LOGIN;
+CREATE DATABASE myapp owner myappadmin;
+GRANT ALL PRIVILEGES ON DATABASE myapp TO myappadmin;

examples/docker/03-query-auth.sql

@@ -0,0 +1,17 @@
+CREATE ROLE connection_pooler PASSWORD 'user-look-up-pass' LOGIN;
+CREATE SCHEMA IF NOT EXISTS pgcat;
+
+CREATE OR REPLACE FUNCTION pgcat.user_lookup(i_username text)
+    RETURNS table ("user" text, hash text) AS $$
+SELECT usename as user, passwd as hash FROM pg_catalog.pg_shadow
+WHERE usename = i_username;
+$$ LANGUAGE sql SECURITY DEFINER;
+
+-- usage:
+-- SELECT * FROM pgcat.user_lookup('$1');
+
+GRANT CONNECT ON DATABASE postgres TO  connection_pooler;
+GRANT USAGE ON SCHEMA pgcat TO connection_pooler;
+
+REVOKE ALL ON FUNCTION pgcat.user_lookup(text) FROM public, connection_pooler;
+GRANT EXECUTE ON FUNCTION pgcat.user_lookup(text) TO connection_pooler;

examples/docker/pgcat.toml

@@ -1,123 +1,42 @@
-#
-# PgCat config example.
-#
-
-#
-# General pooler settings
+## pgcat.toml
 [general]
-# What IP to run on, 0.0.0.0 means accessible from everywhere.
 host = "0.0.0.0"
-
-# Port to run on, same as PgBouncer used in this example.
 port = 6432
-
-# Whether to enable prometheus exporter or not.
 enable_prometheus_exporter = true
-
-# Port at which prometheus exporter listens on.
 prometheus_exporter_port = 9930
 
-# How long to wait before aborting a server connection (ms).
-connect_timeout = 5000
-
-# How much time to give `SELECT 1` health check query to return with a result (ms).
-healthcheck_timeout = 1000
-
-# How long to keep connection available for immediate re-use, without running a healthcheck query on it
-healthcheck_delay = 30000
-
-# How much time to give clients during shutdown before forcibly killing client connections (ms).
-shutdown_timeout = 60000
-
-# For how long to ban a server if it fails a health check (seconds).
-ban_time = 60 # seconds
-
-# If we should log client connections
-log_client_connections = false
+log_client_connections = true
+log_client_disconnections = true
 
-# If we should log client disconnections
-log_client_disconnections = false
+server_tls = false
+verify_server_certificate = false
+#tls_certificate = "/etc/postgresql/certificate/server.crt"
+#tls_private_key = "/etc/postgresql/certificate/server.key"
 
-# TLS
-# tls_certificate = "server.cert"
-# tls_private_key = "server.key"
-
-# Credentials to access the virtual administrative database (pgbouncer or pgcat)
-# Connecting to that database allows running commands like `SHOW POOLS`, `SHOW DATABASES`, etc..
 admin_username = "postgres"
 admin_password = "postgres"
 
-# pool
-# configs are structured as pool.<pool_name>
-# the pool_name is what clients use as database name when connecting
-# For the example below a client can connect using "postgres://sharding_user:sharding_user@pgcat_host:pgcat_port/sharded"
-[pools.postgres]
-# Pool mode (see PgBouncer docs for more).
-# session: one server connection per connected client
-# transaction: one server connection per client transaction
-pool_mode = "transaction"
+auth_query = "SELECT * FROM pgcat.user_lookup('$1');"
+auth_query_user = "connection_pooler"
+auth_query_password = "user-look-up-pass"
+auth_query_database = "postgres"
 
-# If the client doesn't specify, route traffic to
-# this role by default.
-#
-# any: round-robin between primary and replicas,
-# replica: round-robin between replicas only without touching the primary,
-# primary: all queries go to the primary unless otherwise specified.
+## session mode
+[pools.myapp]
+pool_mode = "session"
 default_role = "any"
 
-# Query parser. If enabled, we'll attempt to parse
-# every incoming query to determine if it's a read or a write.
-# If it's a read query, we'll direct it to a replica. Otherwise, if it's a write,
-# we'll direct it to the primary.
-query_parser_enabled = true
-
-# If the query parser is enabled and this setting is enabled, the primary will be part of the pool of databases used for
-# load balancing of read queries. Otherwise, the primary will only be used for write
-# queries. The primary can always be explicitly selected with our custom protocol.
-primary_reads_enabled = true
+# query_parser_enabled = true
 
-# So what if you wanted to implement a different hashing function,
-# or you've already built one and you want this pooler to use it?
-#
-# Current options:
-#
-# pg_bigint_hash: PARTITION BY HASH (Postgres hashing function)
-# sha1: A hashing function based on SHA1
-#
 sharding_function = "pg_bigint_hash"
 
-# Credentials for users that may connect to this cluster
-[pools.postgres.users.0]
-username = "postgres"
-password = "postgres"
-# Maximum number of server connections that can be established for this user
-# The maximum number of connection from a single Pgcat process to any database in the cluster
-# 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.
-statement_timeout = 0
-
-# Shard 0
-[pools.postgres.shards.0]
-# [ host, port, role ]
-servers = [
-    [ "postgres", 5432, "primary" ],
-    [ "postgres", 5432, "replica" ]
-]
-# Database name (e.g. "postgres")
-database = "postgres"
-
-[pools.postgres.shards.1]
-servers = [
-    [ "postgres", 5432, "primary" ],
-    [ "postgres", 5432, "replica" ],
-]
-database = "postgres"
+[pools.myapp.users.0]
+username = "myappadmin"
+pool_size = 23
 
-[pools.postgres.shards.2]
+[pools.myapp.shards.0]
 servers = [
-    [ "postgres", 5432, "primary" ],
-    [ "postgres", 5432, "replica" ],
+#  [ "localhost", 54321, "primary" ], # to use without docker
+    [ "postgres", 5432, "primary" ],  # to use with docker
 ]
-database = "postgres"
+database = "myapp"

pgcat.toml

@@ -195,6 +195,10 @@ sharding_function = "pg_bigint_hash"
 # This parameter is inherited by every pool and can be redefined in pool configuration.
 # auth_query_password = "sharding_user"
 
+#Database to be used for connecting to servers to obtain the hash used for md5 authentication by sending the query
+#specified in `auth_query_query`. This parameter is inherited by every pool and can be redefined in pool configuration.
+# auth_query_database = "shared_db"
+
 # Automatically parse this from queries and route queries to the right shard!
 # automatic_sharding_key = "data.id"
 

src/auth_passthrough.rs

@@ -8,15 +8,17 @@ pub struct AuthPassthrough {
     password: String,
     query: String,
     user: String,
+    database: Option<String>,
 }
 
 impl AuthPassthrough {
     /// Initializes an AuthPassthrough.
-    pub fn new(query: &str, user: &str, password: &str) -> Self {
+    pub fn new(query: &str, user: &str, password: &str, database: Option<String>) -> Self {
         AuthPassthrough {
             password: password.to_string(),
             query: query.to_string(),
             user: user.to_string(),
+            database: database,
         }
     }
 
@@ -28,6 +30,7 @@ impl AuthPassthrough {
                 pool_config.auth_query.as_ref().unwrap(),
                 pool_config.auth_query_user.as_ref().unwrap(),
                 pool_config.auth_query_password.as_ref().unwrap(),
+                pool_config.auth_query_database.clone(),
             ));
         }
 
@@ -64,7 +67,7 @@ impl AuthPassthrough {
     /// ```
     /// use pgcat::auth_passthrough::AuthPassthrough;
     /// use pgcat::config::Address;
-    /// let auth_passthrough = AuthPassthrough::new("SELECT * FROM public.user_lookup('$1');", "postgres", "postgres");
+    /// let auth_passthrough = AuthPassthrough::new("SELECT * FROM public.user_lookup('$1');", "postgres", "postgres", None);
     /// auth_passthrough.fetch_hash(&Address::default());
     /// ```
     ///
@@ -87,7 +90,11 @@ impl AuthPassthrough {
 
         let auth_query = self.query.replace("$1", user);
 
-        match Server::exec_simple_query(address, &auth_user, &auth_query).await {
+        let mut query_address = address.clone();
+        if self.database.is_some() {
+            query_address.database = self.database.clone().unwrap();
+        }
+        match Server::exec_simple_query(&query_address, &auth_user, &auth_query).await {
             Ok(password_data) => {
                 if password_data.len() == 2 && password_data.first().unwrap() == user {
                     if let Some(stripped_hash) = password_data

src/config.rs

@@ -322,6 +322,7 @@ pub struct General {
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
+    pub auth_query_database: Option<String>,
 
     #[serde(default)]
     pub prepared_statements: bool,
@@ -448,6 +449,7 @@ impl Default for General {
             auth_query: None,
             auth_query_user: None,
             auth_query_password: None,
+            auth_query_database: None,
             server_lifetime: Self::default_server_lifetime(),
             server_round_robin: Self::default_server_round_robin(),
             validate_config: true,
@@ -538,6 +540,7 @@ pub struct Pool {
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
+    pub auth_query_database: Option<String>,
 
     #[serde(default = "Pool::default_cleanup_server_connections")]
     pub cleanup_server_connections: bool,
@@ -674,6 +677,7 @@ impl Default for Pool {
             auth_query: None,
             auth_query_user: None,
             auth_query_password: None,
+            auth_query_database: None,
             server_lifetime: None,
             plugins: None,
             cleanup_server_connections: true,
@@ -876,6 +880,10 @@ impl Config {
             if pool.auth_query_password.is_none() {
                 pool.auth_query_password = self.general.auth_query_password.clone();
             }
+
+            if pool.auth_query_database.is_none() {
+                pool.auth_query_database = self.general.auth_query_database.clone();
+            }
         }
     }
 }
@@ -1011,6 +1019,21 @@ impl Config {
             self.general.server_lifetime
         );
         info!("Sever round robin: {}", self.general.server_round_robin);
+
+        if self.is_auth_query_configured() {
+            info!(
+                "Auth query configured: {}",
+                self.general.auth_query.clone().unwrap()
+            );
+
+            if self.general.auth_query_database.is_some() {
+                info!(
+                    "Auth query will use the db: {}",
+                    self.general.auth_query_database.clone().unwrap()
+                );
+            }
+        }
+
         match self.general.tls_certificate.clone() {
             Some(tls_certificate) => {
                 info!("TLS certificate: {}", tls_certificate);

src/pool.rs

@@ -142,6 +142,7 @@ pub struct PoolSettings {
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
+    pub auth_query_database: Option<String>,
 
     /// Plugins
     pub plugins: Option<Plugins>,
@@ -169,6 +170,7 @@ impl Default for PoolSettings {
             auth_query: None,
             auth_query_user: None,
             auth_query_password: None,
+            auth_query_database: None,
             plugins: None,
         }
     }
@@ -474,6 +476,7 @@ impl ConnectionPool {
                         auth_query: pool_config.auth_query.clone(),
                         auth_query_user: pool_config.auth_query_user.clone(),
                         auth_query_password: pool_config.auth_query_password.clone(),
+                        auth_query_database: pool_config.auth_query_database.clone(),
                         plugins: match pool_config.plugins {
                             Some(ref plugins) => Some(plugins.clone()),
                             None => config.plugins.clone(),

src/query_router.rs

@@ -1176,6 +1176,7 @@ mod test {
             auth_query: None,
             auth_query_password: None,
             auth_query_user: None,
+            auth_query_database: None,
             db: "test".to_string(),
             plugins: None,
         };
@@ -1251,6 +1252,7 @@ mod test {
             auth_query: None,
             auth_query_password: None,
             auth_query_user: None,
+            auth_query_database: None,
             db: "test".to_string(),
             plugins: None,
         };

src/server.rs

@@ -1194,7 +1194,10 @@ impl Server {
     ) -> Result<Vec<String>, Error> {
         let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new()));
 
-        debug!("Connecting to server to obtain auth hashes.");
+        debug!(
+            "Connecting to server to obtain auth hashes from db '{}'.",
+            &address.database
+        );
         let mut server = Server::startup(
             address,
             user,