add auth_query_database config

.gitignore

@@ -10,3 +10,4 @@ lcov.info
 dev/.bash_history
 dev/cache
 !dev/cache/.keepme
+.bundle

CONFIG.md

@@ -297,6 +297,38 @@ Query to be sent to servers to obtain the hash used for md5 authentication. The
 established using the database configured in the pool. This parameter is inherited by every pool
 and can be redefined in pool configuration.
 
+#### Configuration example
+
+> This example uses the `postgres` database.
+
+add on `pgcat.toml`:
+
+```toml
+auth_query = "SELECT * FROM pgcat.user_lookup('$1');"
+auth_query_user = "connection_pooler"
+auth_query_password = "user-look-up-pass"
+auth_query_database = "postgres"
+```
+
+setup in the `postgres` database:
+
+```sql
+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;
+
+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;
+```
+
 ### auth_query_user
 ```
 path: pools.<pool_name>.auth_query_user
@@ -319,6 +351,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:
+      - "15432: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", 15432, "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,28 @@ impl Config {
             self.general.server_lifetime
         );
         info!("Sever round robin: {}", self.general.server_round_robin);
+
+        if self.general.auth_query.is_some() {
+            info!(
+                "Auth query configured: {}",
+                self.general.auth_query.clone().unwrap()
+            );
+        }
+
+        if self.general.auth_query_user.is_some() {
+            info!(
+                "Auth query user: {}",
+                self.general.auth_query_user.clone().unwrap()
+            );
+        }
+
+        if self.general.auth_query_database.is_some() {
+            info!(
+                "Auth query database: {}",
+                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(),