Add support for multi-database / multi-user pools (#96)

* Add support for multi-database / multi-user pools

* Nothing

* cargo fmt

* CI

* remove test users

* rename pool

* Update tests to use admin user/pass

* more fixes

* Revert bad change

* Use PGDATABASE env var

* send server info in case of admin

Author: drdrsh

.circleci/pgcat.toml

@@ -5,21 +5,12 @@
 #
 # General pooler settings
 [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
 
-# How many connections to allocate per server.
-pool_size = 15
-
-# Pool mode (see PgBouncer docs for more).
-# session: one server connection per connected client
-# transaction: one server connection per client transaction
-pool_mode = "transaction"
-
 # How long to wait before aborting a server connection (ms).
 connect_timeout = 100
 
@@ -29,56 +20,27 @@ healthcheck_timeout = 100
 # For how long to ban a server if it fails a health check (seconds).
 ban_time = 60 # Seconds
 
-#
+# Reload config automatically if it changes.
 autoreload = true
 
+# TLS
 tls_certificate = ".circleci/server.cert"
 tls_private_key = ".circleci/server.key"
 
-#
-# User to use for authentication against the server.
-[user]
-name = "sharding_user"
-password = "sharding_user"
-
-
-#
-# Shards in the cluster
-[shards]
-
-# Shard 0
-[shards.0]
-
-# [ host, port, role ]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5433, "replica" ],
-    # [ "127.0.1.1", 5432, "replica" ],
-]
-# Database name (e.g. "postgres")
-database = "shard0"
-
-[shards.1]
-# [ host, port, role ]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5433, "replica" ],
-    # [ "127.0.1.1", 5432, "replica" ],
-]
-database = "shard1"
-
-[shards.2]
-# [ host, port, role ]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5433, "replica" ],
-    # [ "127.0.1.1", 5432, "replica" ],
-]
-database = "shard2"
-
+# 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 = "admin_user"
+admin_password = "admin_pass"
 
-# Settings for our query routing layer.
-[query_router]
+# 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.sharded_db]
+# Pool mode (see PgBouncer docs for more).
+# session: one server connection per connected client
+# transaction: one server connection per client transaction
+pool_mode = "transaction"
 
 # If the client doesn't specify, route traffic to
 # this role by default.
@@ -88,7 +50,6 @@ database = "shard2"
 # primary: all queries go to the primary unless otherwise specified.
 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,
@@ -109,3 +70,36 @@ primary_reads_enabled = true
 # sha1: A hashing function based on SHA1
 #
 sharding_function = "pg_bigint_hash"
+
+# Credentials for users that may connect to this cluster
+[pools.sharded_db.users.0]
+username = "sharding_user"
+password = "sharding_user"
+# 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
+
+# Shard 0
+[pools.sharded_db.shards.0]
+# [ host, port, role ]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ]
+]
+# Database name (e.g. "postgres")
+database = "shard0"
+
+[pools.sharded_db.shards.1]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ],
+]
+database = "shard1"
+
+[pools.sharded_db.shards.2]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ],
+]
+database = "shard2"

.circleci/run_tests.sh

@@ -32,6 +32,7 @@ toxiproxy-cli create -l 127.0.0.1:5433 -u 127.0.0.1:5432 postgres_replica
 start_pgcat "info"
 
 export PGPASSWORD=sharding_user
+export PGDATABASE=sharded_db
 
 # pgbench test
 pgbench -U sharding_user -i -h 127.0.0.1 -p 6432
@@ -47,7 +48,7 @@ sleep 1
 killall psql -s SIGINT
 
 # Reload pool (closing unused server connections)
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD'
+PGPASSWORD=admin_pass psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD'
 
 (psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'SELECT pg_sleep(50)' || true) &
 sleep 1
@@ -72,15 +73,17 @@ cd tests/ruby && \
 cd ../..
 
 # Admin tests
-psql -U sharding_user -e -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW STATS' > /dev/null
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD' > /dev/null
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW CONFIG' > /dev/null
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES' > /dev/null
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW LISTS' > /dev/null
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW POOLS' > /dev/null
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW VERSION' > /dev/null
-psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c "SET client_encoding TO 'utf8'" > /dev/null # will ignore
-(! psql -U sharding_user -e -h 127.0.0.1 -p 6432 -d random_db -c 'SHOW STATS' > /dev/null)
+export PGPASSWORD=admin_pass
+psql -U admin_user -e -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW STATS' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW CONFIG' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW LISTS' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW POOLS' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW VERSION' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c "SET client_encoding TO 'utf8'" > /dev/null # will ignore
+(! psql -U admin_user -e -h 127.0.0.1 -p 6432 -d random_db -c 'SHOW STATS' > /dev/null)
+export PGPASSWORD=sharding_user
 
 # Start PgCat in debug to demonstrate failover better
 start_pgcat "trace"

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "pgcat"
-version = "0.4.0-beta1"
+version = "0.6.0-alpha1"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

Cargo.toml

@@ -5,29 +5,20 @@
 #
 # General pooler settings
 [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
 
-# How many connections to allocate per server.
-pool_size = 15
-
-# Pool mode (see PgBouncer docs for more).
-# session: one server connection per connected client
-# transaction: one server connection per client transaction
-pool_mode = "transaction"
-
 # 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
 
 # For how long to ban a server if it fails a health check (seconds).
-ban_time = 60 # Seconds
+ban_time = 60 # seconds
 
 # Reload config automatically if it changes.
 autoreload = false
@@ -36,50 +27,20 @@ autoreload = false
 # tls_certificate = "server.cert"
 # tls_private_key = "server.key"
 
-#
-# User to use for authentication against the server.
-[user]
-name = "sharding_user"
-password = "sharding_user"
-
-
-#
-# Shards in the cluster
-[shards]
-
-# Shard 0
-[shards.0]
-
-# [ host, port, role ]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5432, "replica" ],
-    # [ "127.0.1.1", 5432, "replica" ],
-]
-# Database name (e.g. "postgres")
-database = "shard0"
-
-[shards.1]
-# [ host, port, role ]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5432, "replica" ],
-    # [ "127.0.1.1", 5432, "replica" ],
-]
-database = "shard1"
+# 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 = "user"
+admin_password = "pass"
 
-[shards.2]
-# [ host, port, role ]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5432, "replica" ],
-    # [ "127.0.1.1", 5432, "replica" ],
-]
-database = "shard2"
-
-
-# Settings for our query routing layer.
-[query_router]
+# 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.sharded]
+# Pool mode (see PgBouncer docs for more).
+# session: one server connection per connected client
+# transaction: one server connection per client transaction
+pool_mode = "transaction"
 
 # If the client doesn't specify, route traffic to
 # this role by default.
@@ -89,7 +50,6 @@ database = "shard2"
 # primary: all queries go to the primary unless otherwise specified.
 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,
@@ -110,3 +70,61 @@ primary_reads_enabled = true
 # sha1: A hashing function based on SHA1
 #
 sharding_function = "pg_bigint_hash"
+
+# Credentials for users that may connect to this cluster
+[pools.sharded.users.0]
+username = "sharding_user"
+password = "sharding_user"
+# 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
+
+[pools.sharded.users.1]
+username = "other_user"
+password = "other_user"
+pool_size = 21
+
+# Shard 0
+[pools.sharded.shards.0]
+# [ host, port, role ]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ]
+]
+# Database name (e.g. "postgres")
+database = "shard0"
+
+[pools.sharded.shards.1]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ],
+]
+database = "shard1"
+
+[pools.sharded.shards.2]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ],
+]
+database = "shard2"
+
+
+[pools.simple_db]
+pool_mode = "session"
+default_role = "primary"
+query_parser_enabled = true
+primary_reads_enabled = true
+sharding_function = "pg_bigint_hash"
+
+[pools.simple_db.users.0]
+username = "simple_user"
+password = "simple_user"
+pool_size = 5
+
+[pools.simple_db.shards.0]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ]
+]
+database = "some_db"