chore(deps): bump rustls from 0.21.5 to 0.21.11

Bumps [rustls](https://github.com/rustls/rustls) from 0.21.5 to 0.21.11.
- [Release notes](https://github.com/rustls/rustls/releases)
- [Changelog](https://github.com/rustls/rustls/blob/main/CHANGELOG.md)
- [Commits](https://github.com/rustls/rustls/compare/v/0.21.5...v/0.21.11)

---
updated-dependencies:
- dependency-name: rustls
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>

.circleci/pgcat.toml

@@ -59,7 +59,6 @@ admin_password = "admin_pass"
 # session: one server connection per connected client
 # transaction: one server connection per client transaction
 pool_mode = "transaction"
-prepared_statements_cache_size = 500
 
 # If the client doesn't specify, route traffic to
 # this role by default.
@@ -142,7 +141,6 @@ query_parser_enabled = true
 query_parser_read_write_splitting = true
 primary_reads_enabled = true
 sharding_function = "pg_bigint_hash"
-prepared_statements_cache_size = 500
 
 [pools.simple_db.users.0]
 username = "simple_user"

.circleci/run_tests.sh

@@ -26,7 +26,6 @@ PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard1 -i
 PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard2 -i
 
 # Start Toxiproxy
-kill -9 $(pgrep toxiproxy) || true
 LOG_LEVEL=error toxiproxy-server &
 sleep 1
 
@@ -107,7 +106,7 @@ cd ../..
 # These tests will start and stop the pgcat server so it will need to be restarted after the tests
 #
 pip3 install -r tests/python/requirements.txt
-pytest || exit 1
+python3 tests/python/tests.py || exit 1
 
 
 #
@@ -178,6 +177,3 @@ killall pgcat -s SIGINT
 
 # Allow for graceful shutdown
 sleep 1
-
-kill -9 $(pgrep toxiproxy)
-sleep 1

.github/workflows/build-and-push.yaml

@@ -23,17 +23,14 @@ jobs:
 
     steps:
       - name: Checkout Repository
-        uses: actions/checkout@v4
-
-      - name: Set up QEMU
-        uses: docker/setup-qemu-action@v3
+        uses: actions/checkout@v3
 
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
+        uses: docker/setup-buildx-action@v2
 
       - name: Determine tags
         id: metadata
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@v4
         with:
           images: ${{ env.registry }}/${{ env.image-name }}
           tags: |
@@ -45,18 +42,15 @@ jobs:
             type=raw,value=latest,enable={{ is_default_branch }}
 
       - name: Log in to the Container registry
-        uses: docker/login-action@v3
+        uses: docker/login-action@v2.1.0
         with:
           registry: ${{ env.registry }}
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Build and push ${{ env.image-name }}
-        uses: docker/build-push-action@v6
+        uses: docker/build-push-action@v3
         with:
-          context: .
-          platforms: linux/amd64,linux/arm64
-          provenance: false
           push: true
           tags: ${{ steps.metadata.outputs.tags }}
           labels: ${{ steps.metadata.outputs.labels }}

.github/workflows/chart-lint-test.yaml

@@ -22,7 +22,7 @@ jobs:
       # Python is required because `ct lint` runs Yamale (https://github.com/23andMe/Yamale) and
       # yamllint (https://github.com/adrienverge/yamllint) which require Python
       - name: Set up Python
-        uses: actions/setup-python@v5.1.0
+        uses: actions/setup-python@v4.1.0
         with:
           python-version: 3.7
 
@@ -43,7 +43,7 @@ jobs:
         run: ct lint --config ct.yaml
 
       - name: Create kind cluster
-        uses: helm/kind-action@v1.10.0
+        uses: helm/kind-action@v1.7.0
         if: steps.list-changed.outputs.changed == 'true'
 
       - name: Run chart-testing (install)

.github/workflows/chart-release.yaml

@@ -32,7 +32,7 @@ jobs:
           version: v3.13.0
 
       - name: Run chart-releaser
-        uses: helm/chart-releaser-action@a917fd15b20e8b64b94d9158ad54cd6345335584 # v1.6.0
+        uses: helm/chart-releaser-action@be16258da8010256c6e82849661221415f031968 # v1.5.0
         with:
           charts_dir: charts
           config: cr.yaml

.github/workflows/publish-deb-package.yml

@@ -1,9 +1,6 @@
 name: pgcat package (deb)
 
 on:
-  push:
-    tags:
-      - v*
   workflow_dispatch:
     inputs:
       packageVersion:
@@ -19,14 +16,6 @@ jobs:
     runs-on: ${{ matrix.os }}
     steps:
     - uses: actions/checkout@v3
-    - name: Set package version
-      if: github.event_name == 'push'  # For push event
-      run: |
-        TAG=${{ github.ref_name }}
-        echo "packageVersion=${TAG#v}" >> "$GITHUB_ENV"
-    - name: Set package version (manual dispatch)
-      if: github.event_name == 'workflow_dispatch'  # For manual dispatch
-      run: echo "packageVersion=${{ github.event.inputs.packageVersion }}" >> "$GITHUB_ENV"
     - uses: actions-rs/toolchain@v1
       with:
         toolchain: stable
@@ -50,10 +39,10 @@ jobs:
           export ARCH=arm64
         fi
 
-        bash utilities/deb.sh ${{ env.packageVersion }}
+        bash utilities/deb.sh ${{ inputs.packageVersion }}
 
         deb-s3 upload \
           --lock \
           --bucket apt.postgresml.org \
-          pgcat-${{ env.packageVersion }}-ubuntu22.04-${ARCH}.deb \
+          pgcat-${{ inputs.packageVersion }}-ubuntu22.04-${ARCH}.deb \
           --codename $(lsb_release -cs)

.gitignore

@@ -10,6 +10,4 @@ lcov.info
 dev/.bash_history
 dev/cache
 !dev/cache/.keepme
-.venv
-**/__pycache__
-.bundle
+.venv

CONFIG.md

@@ -36,11 +36,10 @@ Port at which prometheus exporter listens on.
 ### connect_timeout
 ```
 path: general.connect_timeout
-default: 1000 # milliseconds
+default: 5000 # milliseconds
 ```
 
-How long the client waits to obtain a server connection before aborting (ms).
-This is similar to PgBouncer's `query_wait_timeout`.
+How long to wait before aborting a server connection (ms).
 
 ### idle_timeout
 ```
@@ -298,19 +297,6 @@ Load balancing mode
 `random` selects the server at random
 `loc` selects the server with the least outstanding busy connections
 
-### checkout_failure_limit
-```
-path: pools.<pool_name>.checkout_failure_limit
-default: 0 (disabled)
-```
-
-`Maximum number of checkout failures a client is allowed before it
-gets disconnected. This is needed to prevent persistent client/server
-imbalance in high availability setups where multiple PgCat instances are placed
-behind a single load balancer. If for any reason a client lands on a PgCat instance that has 
-a large number of connected clients, it might get stuck in perpetual checkout failure loop especially
-in session mode
-`
 ### default_role
 ```
 path: pools.<pool_name>.default_role
@@ -322,45 +308,6 @@ If the client doesn't specify, PgCat routes traffic to this role by default.
 `replica` round-robin between replicas only without touching the primary,
 `primary` all queries go to the primary unless otherwise specified.
 
-### db_activity_based_routing
-```
-path: pools.<pool_name>.db_activity_based_routing
-default: false
-```
-
-If enabled, PgCat will route queries to the primary if the queried table was recently written to.
-Only relevant when `query_parser_enabled` *and* `query_parser_read_write_splitting` is enabled.
-
-##### Considerations:
-- *This feature is experimental and may not work as expected.*
-- This feature only works when the same PgCat instance is used for both reads and writes to the database.
-- This feature is not relevant when the primary is not part of the pool of databases used for load balancing of read queries.
-- If more than one PgCat instance is used for HA purposes, this feature will not work as expected. A way to still make it work is by using sticky sessions.
-
-### db_activity_based_ms_init_delay
-```
-path: pools.<pool_name>.db_activity_based_ms_init_delay
-default: 100
-```
-
-The delay in milliseconds before the first activity-based routing check is performed.
-
-### db_activity_ttl
-```
-path: pools.<pool_name>.db_activity_ttl
-default: 900
-```
-
-The time in seconds after which a DB is considered inactive when no queries/updates are performed to it.
-
-### table_mutation_cache_ms_ttl
-```
-path: pools.<pool_name>.table_mutation_cache_ms_ttl
-default: 50
-```
-
-The time in milliseconds after a write to a table that all queries to that table will be routed to the primary.
-
 ### prepared_statements_cache_size
 ```
 path: general.prepared_statements_cache_size
@@ -515,18 +462,10 @@ path: pools.<pool_name>.users.<user_index>.pool_size
 default: 9
 ```
 
-Maximum number of server connections that can be established for this 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.
 
-### min_pool_size
-```
-path: pools.<pool_name>.users.<user_index>.min_pool_size
-default: 0
-```
-
-Minimum number of idle server connections to retain for this pool.
-
 ### statement_timeout
 ```
 path: pools.<pool_name>.users.<user_index>.statement_timeout
@@ -536,16 +475,6 @@ default: 0
 Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
 0 means it is disabled.
 
-### connect_timeout
-```
-path: pools.<pool_name>.users.<user_index>.connect_timeout
-default: <UNSET> # milliseconds
-```
-
-How long the client waits to obtain a server connection before aborting (ms).
-This is similar to PgBouncer's `query_wait_timeout`.
-If unset, uses the `connect_timeout` defined globally.
-
 ## `pools.<pool_name>.shards.<shard_index>` Section
 
 ### servers
@@ -573,3 +502,4 @@ default: "shard0"
 ```
 
 Database name (e.g. "postgres")
+

CONTRIBUTING.md

@@ -6,32 +6,6 @@ Thank you for contributing! Just a few tips here:
 2. Run the test suite (e.g. `pgbench`) to make sure everything still works. The tests are in `.circleci/run_tests.sh`.
 3. Performance is important, make sure there are no regressions in your branch vs. `main`.
 
-## How to run the integration tests locally and iterate on them
-We have integration tests written in Ruby, Python, Go and Rust.
-Below are the steps to run them in a developer-friendly way that allows iterating and quick turnaround.
-Hear me out, this should be easy, it will involve opening a shell into a container with all the necessary dependancies available for you and you can modify the test code and immediately rerun your test in the interactive shell.
-
-
-Quite simply, make sure you have docker installed and then run
-`./start_test_env.sh`
-
-That is it!
-
-Within this test environment you can modify the file in your favorite IDE and rerun the tests without having to bootstrap the entire environment again.
-
-Once the environment is ready, you can run the tests by running
-Ruby:   `cd /app/tests/ruby && bundle exec ruby <test_name>.rb --format documentation`
-Python: `cd /app/ && pytest`
-Rust:   `cd /app/tests/rust && cargo run`
-Go:     `cd /app/tests/go && /usr/local/go/bin/go test`
-
-You can also rebuild PgCat directly within the environment and the tests will run against the newly built binary
-To rebuild PgCat, just run `cargo build` within the container under `/app`
-
-![Animated gif showing how to run tests](https://github.com/user-attachments/assets/2258fde3-2aed-4efb-bdc5-e4f12dcd4d33)
-
-
-
 Happy hacking!
 
 ## TODOs

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgcat"
-version = "1.3.0"
+version = "1.1.2-dev4"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -8,7 +8,7 @@ edition = "2021"
 tokio = { version = "1", features = ["full"] }
 bytes = "1"
 md-5 = "0.10"
-bb8 = "=0.8.6"
+bb8 = "0.8.1"
 async-trait = "0.1"
 rand = "0.8"
 chrono = "0.4"
@@ -19,7 +19,7 @@ serde_derive = "1"
 regex = "1"
 num_cpus = "1"
 once_cell = "1"
-sqlparser = { version = "0.52", features = ["visitor"] }
+sqlparser = { version = "0.41", features = ["visitor"] }
 log = "0.4"
 arc-swap = "1"
 parking_lot = "0.12.1"
@@ -29,9 +29,7 @@ base64 = "0.21"
 stringprep = "0.1"
 tokio-rustls = "0.24"
 rustls-pemfile = "1"
-http-body-util = "0.1.2"
-hyper = { version = "1.4.1", features = ["full"] }
-hyper-util = { version = "0.1.7", features = ["tokio"] }
+hyper = { version = "0.14", features = ["full"] }
 phf = { version = "0.11.1", features = ["macros"] }
 exitcode = "1.1.2"
 futures = "0.3"
@@ -55,10 +53,6 @@ tracing-subscriber = { version = "0.3.17", features = [
     "std",
 ] }
 lru = "0.12.0"
-mini-moka = "0.10.3"
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"
-
-[dev-dependencies]
-serial_test = "*"

Dockerfile

@@ -1,4 +1,4 @@
-FROM rust:1.81.0-slim-bookworm AS builder
+FROM rust:1-slim-bookworm AS builder
 
 RUN apt-get update && \
     apt-get install -y build-essential
@@ -19,4 +19,3 @@ COPY --from=builder /app/pgcat.toml /etc/pgcat/pgcat.toml
 WORKDIR /etc/pgcat
 ENV RUST_LOG=info
 CMD ["pgcat"]
-STOPSIGNAL SIGINT

Dockerfile.ci

@@ -1,4 +1,4 @@
-FROM cimg/rust:1.81.0
+FROM cimg/rust:1.67.1
 COPY --from=sclevine/yj /bin/yj /bin/yj
 RUN /bin/yj -h
 RUN sudo apt-get update && \

README.md

@@ -231,7 +231,7 @@ User.find_by_email("test@example.com")
 ```sql
 -- Grab a bunch of users from shard 1
 SET SHARD TO '1';
-SELECT * FROM users LIMIT 10;
+SELECT * FROM users LIMT 10;
 
 -- Find by id
 SET SHARDING KEY TO '1234';
@@ -268,8 +268,6 @@ psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES'
 
 Additionally, Prometheus statistics are available at `/metrics` via HTTP.
 
-We also have a [basic Grafana dashboard](https://github.com/postgresml/pgcat/blob/main/grafana_dashboard.json) based on Prometheus metrics that you can import into Grafana and build on it or use it for monitoring.
-
 ### Live configuration reloading
 
 The config can be reloaded by sending a `kill -s SIGHUP` to the process or by querying `RELOAD` to the admin database. All settings except the `host` and `port` can be reloaded without restarting the pooler, including sharding and replicas configurations.

charts/pgcat/Chart.yaml

@@ -2,7 +2,7 @@ apiVersion: v2
 name: pgcat
 description: A Helm chart for PgCat a PostgreSQL pooler and proxy (like PgBouncer) with support for sharding, load balancing, failover and mirroring.
 maintainers:
-  - name: PostgresML
-    email: team@postgresml.org
-appVersion: "1.3.0"
-version: 0.2.5
+  - name: Wildcard
+    email: support@w6d.io
+appVersion: "1.1.1"
+version: 0.1.0

charts/pgcat/templates/secret.yaml

@@ -15,7 +15,6 @@ stringData:
     connect_timeout = {{ .Values.configuration.general.connect_timeout }}
     idle_timeout = {{ .Values.configuration.general.idle_timeout | int }}
     server_lifetime = {{ .Values.configuration.general.server_lifetime | int }}
-    server_tls = {{ .Values.configuration.general.server_tls }}
     idle_client_in_transaction_timeout = {{ .Values.configuration.general.idle_client_in_transaction_timeout | int }}
     healthcheck_timeout = {{ .Values.configuration.general.healthcheck_timeout }}
     healthcheck_delay = {{ .Values.configuration.general.healthcheck_delay }}
@@ -51,10 +50,6 @@ stringData:
     query_parser_enabled = {{ default true $pool.query_parser_enabled }}
     query_parser_read_write_splitting = {{ default true $pool.query_parser_read_write_splitting }}
     primary_reads_enabled = {{ default true $pool.primary_reads_enabled }}
-    db_activity_based_routing = {{ default false $pool.db_activity_based_routing }}
-    db_activity_based_ms_init_delay = {{ default 100 $pool.db_activity_based_ms_init_delay }}
-    db_activity_ttl = {{ default 900 $pool.db_activity_ttl }}
-    table_mutation_cache_ttl = {{ default 50 $pool.table_mutation_cache_ttl }}
     sharding_function = {{ default "pg_bigint_hash" $pool.sharding_function | quote }}
 
     {{-   range $index, $user := $pool.users }}
@@ -63,21 +58,11 @@ stringData:
     ##
     [pools.{{ $pool.name | quote }}.users.{{ $index }}]
     username = {{ $user.username | quote }}
-    {{- if $user.password }}
     password = {{ $user.password | quote }}
-    {{- else if and $user.passwordSecret.name $user.passwordSecret.key }}
-    {{- $secret := (lookup "v1" "Secret" $.Release.Namespace $user.passwordSecret.name) }}
-    {{- if $secret }}
-    {{- $password := index $secret.data $user.passwordSecret.key | b64dec }}
-    password = {{ $password | quote }}
-    {{- end }}
-    {{- end }}
     pool_size = {{ $user.pool_size }}
-    statement_timeout = {{ default 0 $user.statement_timeout }}
-    min_pool_size = {{ default 3 $user.min_pool_size }}
-    {{- if $user.server_lifetime }}
-    server_lifetime = {{ $user.server_lifetime }}
-    {{- end }}
+    statement_timeout = {{ $user.statement_timeout }}
+    min_pool_size = 3
+    server_lifetime = 60000
     {{-     if and $user.server_username $user.server_password }}
     server_username = {{ $user.server_username | quote }}
     server_password = {{ $user.server_password | quote }}

charts/pgcat/values.yaml

@@ -170,16 +170,13 @@ configuration:
     connect_timeout: 5000
 
     # How long an idle connection with a server is left open (ms).
-    idle_timeout: 30000  # milliseconds
+    idle_timeout: 30000 # milliseconds
 
     # Max connection lifetime before it's closed, even if actively used.
-    server_lifetime: 86400000  # 24 hours
-
-    # Whether to use TLS for server connections or not.
-    server_tls: false
+    server_lifetime: 86400000 # 24 hours
 
     # How long a client is allowed to be idle while in a transaction (ms).
-    idle_client_in_transaction_timeout: 0  # milliseconds
+    idle_client_in_transaction_timeout: 0 # milliseconds
 
     # @param configuration.general.healthcheck_timeout How much time to give `SELECT 1` health check query to return with a result (ms).
     healthcheck_timeout: 1000
@@ -243,15 +240,7 @@ configuration:
   ## 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"
   ## @param [object]
-  pools:
-    [{
-      name: "simple", pool_mode: "transaction",
-      users: [{username: "user", password: "pass", pool_size: 5, statement_timeout: 0}],
-      shards: [{
-        servers: [{host: "postgres", port: 5432, role: "primary"}],
-        database: "postgres"
-      }]
-    }]
+  pools: []
     # - ## default values
     #   ##
     #   ##
@@ -298,22 +287,6 @@ configuration:
     #  ## @param configuration.poolsPostgres.query_parser_read_write_splitting
     #  query_parser_read_write_splitting: true
 
-    # ## Db activity based routing. If enabled, we'll route queries to the primary if the table was recently mutated.
-    # ## @param configuration.poolsPostgres.db_activity_based_routing
-    # db_activity_based_routing: false
-
-    # ## DB activity based init delay. How long to wait before starting to route queries to the primary after a table mutation.
-    # ## @param configuration.poolsPostgres.db_activity_based_ms_init_delay
-    # db_activity_based_ms_init_delay: 100
-
-    # ## DB activity TTL. How long before marking the DB as inactive after no mutations or queries.
-    # ## @param configuration.poolsPostgres.db_activity_ttl
-    # db_activity_ttl: 900
-
-    # ## Table mutation cache TTL. How long to keep track of table mutations.
-    # ## @param configuration.poolsPostgres.table_mutation_cache_ttl
-    # table_mutation_cache_ttl: 50
-
     #  ## 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.
@@ -334,9 +307,7 @@ configuration:
     #  ## Credentials for users that may connect to this cluster
     #  ## @param users [array]
     #  ## @param users[0].username Name of the env var (required)
-    #  ## @param users[0].password Value for the env var (required) leave empty to use existing secret see passwordSecret.name and passwordSecret.key
-    #  ## @param users[0].passwordSecret.name Name of the secret containing the password
-    #  ## @param users[0].passwordSecret.key Key in the secret containing the password
+    #  ## @param users[0].password Value for the env var (required)
     #  ## @param users[0].pool_size Maximum number of server connections that can be established for this user
     #  ## @param users[0].statement_timeout Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
     #  users: []

cr.yaml

@@ -1,2 +1 @@
 sign: false
-pages_branch: main

pgcat.service

@@ -11,7 +11,6 @@ RestartSec=1
 Environment=RUST_LOG=info
 LimitNOFILE=65536
 ExecStart=/usr/bin/pgcat /etc/pgcat.toml
-ExecReload=/bin/kill -SIGHUP $MAINPID
 
 [Install]
 WantedBy=multi-user.target

pgcat.toml

@@ -179,7 +179,7 @@ primary_reads_enabled = true
 # `random`: picks a shard at random
 # `random_healthy`: picks a shard at random favoring shards with the least number of recent errors
 # `shard_<number>`: e.g. shard_0, shard_4, etc. picks a specific shard, everytime
-# default_shard = "shard_0"
+# no_shard_specified_behavior = "shard_0"
 
 # 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?

postinst

@@ -7,7 +7,3 @@ systemctl enable pgcat
 if ! id pgcat 2> /dev/null; then
 	useradd -s /usr/bin/false pgcat
 fi
-
-if [ -f /etc/pgcat.toml ]; then
-	systemctl start pgcat
-fi

src/admin.rs

@@ -55,12 +55,7 @@ where
 
     let query_parts: Vec<&str> = query.trim_end_matches(';').split_whitespace().collect();
 
-    match query_parts
-        .first()
-        .unwrap_or(&"")
-        .to_ascii_uppercase()
-        .as_str()
-    {
+    match query_parts[0].to_ascii_uppercase().as_str() {
         "BAN" => {
             trace!("BAN");
             ban(stream, query_parts).await
@@ -89,12 +84,7 @@ where
             trace!("SHUTDOWN");
             shutdown(stream).await
         }
-        "SHOW" => match query_parts
-            .get(1)
-            .unwrap_or(&"")
-            .to_ascii_uppercase()
-            .as_str()
-        {
+        "SHOW" => match query_parts[1].to_ascii_uppercase().as_str() {
             "HELP" => {
                 trace!("SHOW HELP");
                 show_help(stream).await

src/auth_passthrough.rs

@@ -1,4 +1,3 @@
-use crate::config::AuthType;
 use crate::errors::Error;
 use crate::pool::ConnectionPool;
 use crate::server::Server;
@@ -72,7 +71,6 @@ impl AuthPassthrough {
     pub async fn fetch_hash(&self, address: &crate::config::Address) -> Result<String, Error> {
         let auth_user = crate::config::User {
             username: self.user.clone(),
-            auth_type: AuthType::MD5,
             password: Some(self.password.clone()),
             server_username: None,
             server_password: None,

src/client.rs

@@ -14,9 +14,7 @@ use tokio::sync::mpsc::Sender;
 
 use crate::admin::{generate_server_parameters_for_admin, handle_admin};
 use crate::auth_passthrough::refetch_auth_hash;
-use crate::config::{
-    get_config, get_idle_client_in_transaction_timeout, Address, AuthType, PoolMode,
-};
+use crate::config::{get_config, get_idle_client_in_transaction_timeout, Address, PoolMode};
 use crate::constants::*;
 use crate::messages::*;
 use crate::plugins::PluginOutput;
@@ -465,8 +463,8 @@ where
             .count()
             == 1;
 
+        // Kick any client that's not admin while we're in admin-only mode.
         if !admin && admin_only {
-            // Kick any client that's not admin while we're in admin-only mode.
             debug!(
                 "Rejecting non-admin connection to {} when in admin only mode",
                 pool_name
@@ -483,76 +481,72 @@ where
         let process_id: i32 = rand::random();
         let secret_key: i32 = rand::random();
 
+        // Perform MD5 authentication.
+        // TODO: Add SASL support.
+        let salt = md5_challenge(&mut write).await?;
+
+        let code = match read.read_u8().await {
+            Ok(p) => p,
+            Err(_) => {
+                return Err(Error::ClientSocketError(
+                    "password code".into(),
+                    client_identifier,
+                ))
+            }
+        };
+
+        // PasswordMessage
+        if code as char != 'p' {
+            return Err(Error::ProtocolSyncError(format!(
+                "Expected p, got {}",
+                code as char
+            )));
+        }
+
+        let len = match read.read_i32().await {
+            Ok(len) => len,
+            Err(_) => {
+                return Err(Error::ClientSocketError(
+                    "password message length".into(),
+                    client_identifier,
+                ))
+            }
+        };
+
+        let mut password_response = vec![0u8; (len - 4) as usize];
+
+        match read.read_exact(&mut password_response).await {
+            Ok(_) => (),
+            Err(_) => {
+                return Err(Error::ClientSocketError(
+                    "password message".into(),
+                    client_identifier,
+                ))
+            }
+        };
+
         let mut prepared_statements_enabled = false;
 
         // Authenticate admin user.
         let (transaction_mode, mut server_parameters) = if admin {
             let config = get_config();
-            // TODO: Add SASL support.
-            // Perform MD5 authentication.
-            match config.general.admin_auth_type {
-                AuthType::Trust => (),
-                AuthType::MD5 => {
-                    let salt = md5_challenge(&mut write).await?;
 
-                    let code = match read.read_u8().await {
-                        Ok(p) => p,
-                        Err(_) => {
-                            return Err(Error::ClientSocketError(
-                                "password code".into(),
-                                client_identifier,
-                            ))
-                        }
-                    };
+            // Compare server and client hashes.
+            let password_hash = md5_hash_password(
+                &config.general.admin_username,
+                &config.general.admin_password,
+                &salt,
+            );
 
-                    // PasswordMessage
-                    if code as char != 'p' {
-                        return Err(Error::ProtocolSyncError(format!(
-                            "Expected p, got {}",
-                            code as char
-                        )));
-                    }
+            if password_hash != password_response {
+                let error = Error::ClientGeneralError("Invalid password".into(), client_identifier);
 
-                    let len = match read.read_i32().await {
-                        Ok(len) => len,
-                        Err(_) => {
-                            return Err(Error::ClientSocketError(
-                                "password message length".into(),
-                                client_identifier,
-                            ))
-                        }
-                    };
+                warn!("{}", error);
+                wrong_password(&mut write, username).await?;
 
-                    let mut password_response = vec![0u8; (len - 4) as usize];
-
-                    match read.read_exact(&mut password_response).await {
-                        Ok(_) => (),
-                        Err(_) => {
-                            return Err(Error::ClientSocketError(
-                                "password message".into(),
-                                client_identifier,
-                            ))
-                        }
-                    };
-
-                    // Compare server and client hashes.
-                    let password_hash = md5_hash_password(
-                        &config.general.admin_username,
-                        &config.general.admin_password,
-                        &salt,
-                    );
-
-                    if password_hash != password_response {
-                        let error =
-                            Error::ClientGeneralError("Invalid password".into(), client_identifier);
-
-                        warn!("{}", error);
-                        wrong_password(&mut write, username).await?;
-
-                        return Err(error);
-                    }
-                }
+                return Err(error);
             }
+
             (false, generate_server_parameters_for_admin())
         }
         // Authenticate normal user.
@@ -579,143 +573,92 @@ where
             // Obtain the hash to compare, we give preference to that written in cleartext in config
             // if there is nothing set in cleartext and auth passthrough (auth_query) is configured, we use the hash obtained
             // when the pool was created. If there is no hash there, we try to fetch it one more time.
-            match pool.settings.user.auth_type {
-                AuthType::Trust => (),
-                AuthType::MD5 => {
-                    // Perform MD5 authentication.
-                    // TODO: Add SASL support.
-                    let salt = md5_challenge(&mut write).await?;
+            let password_hash = if let Some(password) = &pool.settings.user.password {
+                Some(md5_hash_password(username, password, &salt))
+            } else {
+                if !get_config().is_auth_query_configured() {
+                    wrong_password(&mut write, username).await?;
+                    return Err(Error::ClientAuthImpossible(username.into()));
+                }
 
-                    let code = match read.read_u8().await {
-                        Ok(p) => p,
-                        Err(_) => {
-                            return Err(Error::ClientSocketError(
-                                "password code".into(),
-                                client_identifier,
-                            ))
-                        }
-                    };
+                let mut hash = (*pool.auth_hash.read()).clone();
 
-                    // PasswordMessage
-                    if code as char != 'p' {
-                        return Err(Error::ProtocolSyncError(format!(
-                            "Expected p, got {}",
-                            code as char
-                        )));
-                    }
+                if hash.is_none() {
+                    warn!(
+                        "Query auth configured \
+                          but no hash password found \
+                          for pool {}. Will try to refetch it.",
+                        pool_name
+                    );
 
-                    let len = match read.read_i32().await {
-                        Ok(len) => len,
-                        Err(_) => {
-                            return Err(Error::ClientSocketError(
-                                "password message length".into(),
-                                client_identifier,
-                            ))
-                        }
-                    };
-
-                    let mut password_response = vec![0u8; (len - 4) as usize];
-
-                    match read.read_exact(&mut password_response).await {
-                        Ok(_) => (),
-                        Err(_) => {
-                            return Err(Error::ClientSocketError(
-                                "password message".into(),
-                                client_identifier,
-                            ))
-                        }
-                    };
-
-                    let password_hash = if let Some(password) = &pool.settings.user.password {
-                        Some(md5_hash_password(username, password, &salt))
-                    } else {
-                        if !get_config().is_auth_query_configured() {
-                            wrong_password(&mut write, username).await?;
-                            return Err(Error::ClientAuthImpossible(username.into()));
-                        }
-
-                        let mut hash = (*pool.auth_hash.read()).clone();
-
-                        if hash.is_none() {
-                            warn!(
-                                "Query auth configured \
-                                  but no hash password found \
-                                  for pool {}. Will try to refetch it.",
-                                pool_name
-                            );
-
-                            match refetch_auth_hash(&pool).await {
-                                Ok(fetched_hash) => {
-                                    warn!(
-                                        "Password for {}, obtained. Updating.",
-                                        client_identifier
-                                    );
-
-                                    {
-                                        let mut pool_auth_hash = pool.auth_hash.write();
-                                        *pool_auth_hash = Some(fetched_hash.clone());
-                                    }
-
-                                    hash = Some(fetched_hash);
-                                }
-
-                                Err(err) => {
-                                    wrong_password(&mut write, username).await?;
-
-                                    return Err(Error::ClientAuthPassthroughError(
-                                        err.to_string(),
-                                        client_identifier,
-                                    ));
-                                }
-                            }
-                        };
-
-                        Some(md5_hash_second_pass(&hash.unwrap(), &salt))
-                    };
-
-                    // Once we have the resulting hash, we compare with what the client gave us.
-                    // If they do not match and auth query is set up, we try to refetch the hash one more time
-                    // to see if the password has changed since the pool was created.
-                    //
-                    // @TODO: we could end up fetching again the same password twice (see above).
-                    if password_hash.unwrap() != password_response {
-                        warn!(
-                            "Invalid password {}, will try to refetch it.",
-                            client_identifier
-                        );
-
-                        let fetched_hash = match refetch_auth_hash(&pool).await {
-                            Ok(fetched_hash) => fetched_hash,
-                            Err(err) => {
-                                wrong_password(&mut write, username).await?;
-
-                                return Err(err);
-                            }
-                        };
-
-                        let new_password_hash = md5_hash_second_pass(&fetched_hash, &salt);
-
-                        // Ok password changed in server an auth is possible.
-                        if new_password_hash == password_response {
-                            warn!(
-                                "Password for {}, changed in server. Updating.",
-                                client_identifier
-                            );
+                    match refetch_auth_hash(&pool).await {
+                        Ok(fetched_hash) => {
+                            warn!("Password for {}, obtained. Updating.", client_identifier);
 
                             {
                                 let mut pool_auth_hash = pool.auth_hash.write();
-                                *pool_auth_hash = Some(fetched_hash);
+                                *pool_auth_hash = Some(fetched_hash.clone());
                             }
-                        } else {
+
+                            hash = Some(fetched_hash);
+                        }
+
+                        Err(err) => {
                             wrong_password(&mut write, username).await?;
-                            return Err(Error::ClientGeneralError(
-                                "Invalid password".into(),
+
+                            return Err(Error::ClientAuthPassthroughError(
+                                err.to_string(),
                                 client_identifier,
                             ));
                         }
                     }
+                };
+
+                Some(md5_hash_second_pass(&hash.unwrap(), &salt))
+            };
+
+            // Once we have the resulting hash, we compare with what the client gave us.
+            // If they do not match and auth query is set up, we try to refetch the hash one more time
+            // to see if the password has changed since the pool was created.
+            //
+            // @TODO: we could end up fetching again the same password twice (see above).
+            if password_hash.unwrap() != password_response {
+                warn!(
+                    "Invalid password {}, will try to refetch it.",
+                    client_identifier
+                );
+
+                let fetched_hash = match refetch_auth_hash(&pool).await {
+                    Ok(fetched_hash) => fetched_hash,
+                    Err(err) => {
+                        wrong_password(&mut write, username).await?;
+
+                        return Err(err);
+                    }
+                };
+
+                let new_password_hash = md5_hash_second_pass(&fetched_hash, &salt);
+
+                // Ok password changed in server an auth is possible.
+                if new_password_hash == password_response {
+                    warn!(
+                        "Password for {}, changed in server. Updating.",
+                        client_identifier
+                    );
+
+                    {
+                        let mut pool_auth_hash = pool.auth_hash.write();
+                        *pool_auth_hash = Some(fetched_hash);
+                    }
+                } else {
+                    wrong_password(&mut write, username).await?;
+                    return Err(Error::ClientGeneralError(
+                        "Invalid password".into(),
+                        client_identifier,
+                    ));
                 }
             }
+
             let transaction_mode = pool.settings.pool_mode == PoolMode::Transaction;
             prepared_statements_enabled =
                 transaction_mode && pool.prepared_statement_cache.is_some();
@@ -859,8 +802,6 @@ where
         // e.g. primary, replica, which shard.
         let mut query_router = QueryRouter::new();
 
-        let mut checkout_failure_count: u64 = 0;
-
         self.stats.register(self.stats.clone());
 
         // Result returned by one of the plugins.
@@ -883,7 +824,6 @@ where
         };
 
         query_router.update_pool_settings(&pool.settings);
-        query_router.set_default_role();
 
         // Our custom protocol loop.
         // We expect the client to either start a transaction with regular queries
@@ -1110,25 +1050,7 @@ where
                         query_router.role(),
                         err
                     );
-                    checkout_failure_count += 1;
-                    if let Some(limit) = pool.settings.checkout_failure_limit {
-                        if checkout_failure_count >= limit {
-                            error!(
-                                "Checkout failure limit reached ({} / {}) - disconnecting client",
-                                checkout_failure_count, limit
-                            );
-                            error_response_terminal(
-                                &mut self.write,
-                                &format!(
-                                    "checkout failure limit reached ({} / {})",
-                                    checkout_failure_count, limit
-                                ),
-                            )
-                            .await?;
-                            self.stats.disconnect();
-                            return Ok(());
-                        }
-                    }
+
                     continue;
                 }
             };

src/config.rs

@@ -38,12 +38,12 @@ pub enum Role {
     Mirror,
 }
 
-impl std::fmt::Display for Role {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Role::Primary => write!(f, "primary"),
-            Role::Replica => write!(f, "replica"),
-            Role::Mirror => write!(f, "mirror"),
+impl ToString for Role {
+    fn to_string(&self) -> String {
+        match *self {
+            Role::Primary => "primary".to_string(),
+            Role::Replica => "replica".to_string(),
+            Role::Mirror => "mirror".to_string(),
         }
     }
 }
@@ -208,9 +208,6 @@ impl Address {
 pub struct User {
     pub username: String,
     pub password: Option<String>,
-
-    #[serde(default = "User::default_auth_type")]
-    pub auth_type: AuthType,
     pub server_username: Option<String>,
     pub server_password: Option<String>,
     pub pool_size: u32,
@@ -228,7 +225,6 @@ impl Default for User {
         User {
             username: String::from("postgres"),
             password: None,
-            auth_type: AuthType::MD5,
             server_username: None,
             server_password: None,
             pool_size: 15,
@@ -243,10 +239,6 @@ impl Default for User {
 }
 
 impl User {
-    pub fn default_auth_type() -> AuthType {
-        AuthType::MD5
-    }
-
     fn validate(&self) -> Result<(), Error> {
         if let Some(min_pool_size) = self.min_pool_size {
             if min_pool_size > self.pool_size {
@@ -342,9 +334,6 @@ pub struct General {
     pub admin_username: String,
     pub admin_password: String,
 
-    #[serde(default = "General::default_admin_auth_type")]
-    pub admin_auth_type: AuthType,
-
     #[serde(default = "General::default_validate_config")]
     pub validate_config: bool,
 
@@ -359,10 +348,6 @@ impl General {
         "0.0.0.0".into()
     }
 
-    pub fn default_admin_auth_type() -> AuthType {
-        AuthType::MD5
-    }
-
     pub fn default_port() -> u16 {
         5432
     }
@@ -471,7 +456,6 @@ impl Default for General {
             verify_server_certificate: false,
             admin_username: String::from("admin"),
             admin_password: String::from("admin"),
-            admin_auth_type: AuthType::MD5,
             validate_config: true,
             auth_query: None,
             auth_query_user: None,
@@ -492,20 +476,11 @@ pub enum PoolMode {
     Session,
 }
 
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Copy, Hash)]
-pub enum AuthType {
-    #[serde(alias = "trust", alias = "Trust")]
-    Trust,
-
-    #[serde(alias = "md5", alias = "MD5")]
-    MD5,
-}
-
-impl std::fmt::Display for PoolMode {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            PoolMode::Transaction => write!(f, "transaction"),
-            PoolMode::Session => write!(f, "session"),
+impl ToString for PoolMode {
+    fn to_string(&self) -> String {
+        match *self {
+            PoolMode::Transaction => "transaction".to_string(),
+            PoolMode::Session => "session".to_string(),
         }
     }
 }
@@ -518,13 +493,12 @@ pub enum LoadBalancingMode {
     #[serde(alias = "loc", alias = "LOC", alias = "least_outstanding_connections")]
     LeastOutstandingConnections,
 }
-
-impl std::fmt::Display for LoadBalancingMode {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            LoadBalancingMode::Random => write!(f, "random"),
+impl ToString for LoadBalancingMode {
+    fn to_string(&self) -> String {
+        match *self {
+            LoadBalancingMode::Random => "random".to_string(),
             LoadBalancingMode::LeastOutstandingConnections => {
-                write!(f, "least_outstanding_connections")
+                "least_outstanding_connections".to_string()
             }
         }
     }
@@ -558,14 +532,6 @@ pub struct Pool {
     /// Close idle connections that have been opened for longer than this.
     pub idle_timeout: Option<u64>,
 
-    /// Maximum number of checkout failures a client is allowed before it
-    /// gets disconnected. This is needed to prevent persistent client/server
-    /// imbalance in high availability setups where multiple PgCat instances are placed
-    /// behind a single load balancer. If for any reason a client lands on a PgCat instance that has
-    /// a large number of connected clients, it might get stuck in perpetual checkout failure loop especially
-    /// in session mode
-    pub checkout_failure_limit: Option<u64>,
-
     /// Close server connections that have been opened for longer than this.
     /// Only applied to idle connections. If the connection is actively used for
     /// longer than this period, the pool will not interrupt it.
@@ -597,19 +563,6 @@ pub struct Pool {
     #[serde(default = "Pool::default_prepared_statements_cache_size")]
     pub prepared_statements_cache_size: usize,
 
-    // Support for query routing based on database activity
-    #[serde(default = "Pool::default_db_activity_based_routing")]
-    pub db_activity_based_routing: bool,
-
-    #[serde(default = "Pool::default_db_activity_init_delay")]
-    pub db_activity_init_delay: u64,
-
-    #[serde(default = "Pool::default_db_activity_ttl")]
-    pub db_activity_ttl: u64,
-
-    #[serde(default = "Pool::default_table_mutation_cache_ms_ttl")]
-    pub table_mutation_cache_ms_ttl: u64,
-
     pub plugins: Option<Plugins>,
     pub shards: BTreeMap<String, Shard>,
     pub users: BTreeMap<String, User>,
@@ -663,25 +616,6 @@ impl Pool {
         0
     }
 
-    pub fn default_db_activity_based_routing() -> bool {
-        false
-    }
-
-    pub fn default_db_activity_init_delay() -> u64 {
-        // 100 milliseconds
-        100
-    }
-
-    pub fn default_db_activity_ttl() -> u64 {
-        // 15 minutes
-        15 * 60
-    }
-
-    pub fn default_table_mutation_cache_ms_ttl() -> u64 {
-        // 50 milliseconds
-        50
-    }
-
     pub fn validate(&mut self) -> Result<(), Error> {
         match self.default_role.as_ref() {
             "any" => (),
@@ -764,23 +698,6 @@ impl Pool {
             user.validate()?;
         }
 
-        if self.db_activity_based_routing {
-            if self.db_activity_init_delay == 0 {
-                error!("db_activity_init_delay must be greater than 0");
-                return Err(Error::BadConfig);
-            }
-
-            if self.table_mutation_cache_ms_ttl == 0 {
-                error!("table_mutation_cache_ms_ttl must be greater than 0");
-                return Err(Error::BadConfig);
-            }
-
-            if self.db_activity_ttl == 0 {
-                error!("db_activity_ttl must be greater than 0");
-                return Err(Error::BadConfig);
-            }
-        }
-
         Ok(())
     }
 }
@@ -790,7 +707,6 @@ impl Default for Pool {
         Pool {
             pool_mode: Self::default_pool_mode(),
             load_balancing_mode: Self::default_load_balancing_mode(),
-            checkout_failure_limit: None,
             default_role: String::from("any"),
             query_parser_enabled: false,
             query_parser_max_length: None,
@@ -811,10 +727,6 @@ impl Default for Pool {
             cleanup_server_connections: true,
             log_client_parameter_status_changes: false,
             prepared_statements_cache_size: Self::default_prepared_statements_cache_size(),
-            db_activity_based_routing: Self::default_db_activity_based_routing(),
-            db_activity_init_delay: Self::default_db_activity_init_delay(),
-            db_activity_ttl: Self::default_db_activity_ttl(),
-            table_mutation_cache_ms_ttl: Self::default_table_mutation_cache_ms_ttl(),
             plugins: None,
             shards: BTreeMap::from([(String::from("1"), Shard::default())]),
             users: BTreeMap::default(),
@@ -1087,17 +999,15 @@ impl Config {
     pub fn fill_up_auth_query_config(&mut self) {
         for (_name, pool) in self.pools.iter_mut() {
             if pool.auth_query.is_none() {
-                pool.auth_query.clone_from(&self.general.auth_query);
+                pool.auth_query = self.general.auth_query.clone();
             }
 
             if pool.auth_query_user.is_none() {
-                pool.auth_query_user
-                    .clone_from(&self.general.auth_query_user);
+                pool.auth_query_user = self.general.auth_query_user.clone();
             }
 
             if pool.auth_query_password.is_none() {
-                pool.auth_query_password
-                    .clone_from(&self.general.auth_query_password);
+                pool.auth_query_password = self.general.auth_query_password.clone();
             }
         }
     }
@@ -1245,7 +1155,7 @@ impl Config {
             "Default max server lifetime: {}ms",
             self.general.server_lifetime
         );
-        info!("Server round robin: {}", self.general.server_round_robin);
+        info!("Sever round robin: {}", self.general.server_round_robin);
         match self.general.tls_certificate.clone() {
             Some(tls_certificate) => {
                 info!("TLS certificate: {}", tls_certificate);
@@ -1307,17 +1217,6 @@ impl Config {
                 None => self.general.idle_timeout,
             };
             info!("[pool: {}] Idle timeout: {}ms", pool_name, idle_timeout);
-            match pool_config.checkout_failure_limit {
-                Some(checkout_failure_limit) => {
-                    info!(
-                        "[pool: {}] Checkout failure limit: {}",
-                        pool_name, checkout_failure_limit
-                    );
-                }
-                None => {
-                    info!("[pool: {}] Checkout failure limit: not set", pool_name);
-                }
-            };
             info!(
                 "[pool: {}] Sharding function: {}",
                 pool_name,
@@ -1362,22 +1261,6 @@ impl Config {
                 "[pool: {}] Cleanup server connections: {}",
                 pool_name, pool_config.cleanup_server_connections
             );
-            info!(
-                "[pool: {}] DB activity based routing: {}",
-                pool_name, pool_config.db_activity_based_routing
-            );
-            info!(
-                "[pool: {}] DB activity init delay: {}",
-                pool_name, pool_config.db_activity_init_delay
-            );
-            info!(
-                "[pool: {}] DB activity TTL: {}",
-                pool_name, pool_config.db_activity_ttl
-            );
-            info!(
-                "[pool: {}] Table mutation cache TTL: {}",
-                pool_name, pool_config.table_mutation_cache_ms_ttl
-            );
             info!(
                 "[pool: {}] Log client parameter status changes: {}",
                 pool_name, pool_config.log_client_parameter_status_changes

src/logger.rs

@@ -3,7 +3,7 @@ use tracing_subscriber;
 use tracing_subscriber::EnvFilter;
 
 pub fn init(args: &Args) {
-    // Initialize a default filter, and then override the builtin default "warning" with our
+    // Iniitalize a default filter, and then override the builtin default "warning" with our
     // commandline, (default: "info")
     let filter = EnvFilter::from_default_env().add_directive(args.log_level.into());
 

src/messages.rs

@@ -733,10 +733,6 @@ pub fn configure_socket(stream: &TcpStream) {
         }
         Err(err) => error!("Could not configure socket: {}", err),
     }
-    match sock_ref.set_nodelay(true) {
-        Ok(_) => (),
-        Err(err) => error!("Could not configure TCP_NODELAY for socket: {}", err),
-    }
 }
 
 pub trait BytesMutReader {

src/pool.rs

@@ -152,14 +152,6 @@ pub struct PoolSettings {
     /// Random or LeastOutstandingConnections.
     pub load_balancing_mode: LoadBalancingMode,
 
-    /// Maximum number of checkout failures a client is allowed before it
-    /// gets disconnected. This is needed to prevent persistent client/server
-    /// imbalance in high availability setups where multiple PgCat instances are placed
-    /// behind a single load balancer. If for any reason a client lands on a PgCat instance that has
-    /// a large number of connected clients, it might get stuck in perpetual checkout failure loop especially
-    /// in session mode
-    pub checkout_failure_limit: Option<u64>,
-
     // Number of shards.
     pub shards: usize,
 
@@ -182,18 +174,6 @@ pub struct PoolSettings {
     // Read from the primary as well or not.
     pub primary_reads_enabled: bool,
 
-    // Automatic primary/replica selection based on recent activity.
-    pub db_activity_based_routing: bool,
-
-    // DB activity init delay
-    pub db_activity_init_delay: u64,
-
-    // DB activity TTL
-    pub db_activity_ttl: u64,
-
-    // Table mutation cache TTL
-    pub table_mutation_cache_ms_ttl: u64,
-
     // Sharding function.
     pub sharding_function: ShardingFunction,
 
@@ -235,7 +215,6 @@ impl Default for PoolSettings {
         PoolSettings {
             pool_mode: PoolMode::Transaction,
             load_balancing_mode: LoadBalancingMode::Random,
-            checkout_failure_limit: None,
             shards: 1,
             user: User::default(),
             db: String::default(),
@@ -244,10 +223,6 @@ impl Default for PoolSettings {
             query_parser_max_length: None,
             query_parser_read_write_splitting: false,
             primary_reads_enabled: true,
-            db_activity_based_routing: false,
-            db_activity_init_delay: 100,
-            db_activity_ttl: 15 * 60,
-            table_mutation_cache_ms_ttl: 50,
             sharding_function: ShardingFunction::PgBigintHash,
             automatic_sharding_key: None,
             healthcheck_delay: General::default_healthcheck_delay(),
@@ -546,7 +521,6 @@ impl ConnectionPool {
                             None => pool_config.pool_mode,
                         },
                         load_balancing_mode: pool_config.load_balancing_mode,
-                        checkout_failure_limit: pool_config.checkout_failure_limit,
                         // shards: pool_config.shards.clone(),
                         shards: shard_ids.len(),
                         user: user.clone(),
@@ -563,10 +537,6 @@ impl ConnectionPool {
                             .query_parser_read_write_splitting,
                         primary_reads_enabled: pool_config.primary_reads_enabled,
                         sharding_function: pool_config.sharding_function,
-                        db_activity_based_routing: pool_config.db_activity_based_routing,
-                        db_activity_init_delay: pool_config.db_activity_init_delay,
-                        db_activity_ttl: pool_config.db_activity_ttl,
-                        table_mutation_cache_ms_ttl: pool_config.table_mutation_cache_ms_ttl,
                         automatic_sharding_key: pool_config.automatic_sharding_key.clone(),
                         healthcheck_delay: config.general.healthcheck_delay,
                         healthcheck_timeout: config.general.healthcheck_timeout,
@@ -843,7 +813,7 @@ impl ConnectionPool {
             }
         }
 
-        client_stats.checkout_error();
+        client_stats.checkout_success();
 
         Err(Error::AllServersDown)
     }

src/prometheus.rs

@@ -1,41 +1,23 @@
-use http_body_util::Full;
-use hyper::body;
-use hyper::body::Bytes;
-
-use hyper::server::conn::http1;
-use hyper::service::service_fn;
-use hyper::{Method, Request, Response, StatusCode};
-use hyper_util::rt::TokioIo;
+use hyper::service::{make_service_fn, service_fn};
+use hyper::{Body, Method, Request, Response, Server, StatusCode};
 use log::{debug, error, info};
 use phf::phf_map;
 use std::collections::HashMap;
 use std::fmt;
 use std::net::SocketAddr;
 use std::sync::atomic::Ordering;
-use tokio::net::TcpListener;
+use std::sync::Arc;
 
 use crate::config::Address;
 use crate::pool::{get_all_pools, PoolIdentifier};
-use crate::stats::get_server_stats;
 use crate::stats::pool::PoolStats;
+use crate::stats::{get_server_stats, ServerStats};
 
 struct MetricHelpType {
     help: &'static str,
     ty: &'static str,
 }
 
-struct ServerPrometheusStats {
-    bytes_received: u64,
-    bytes_sent: u64,
-    transaction_count: u64,
-    query_count: u64,
-    error_count: u64,
-    active_count: u64,
-    idle_count: u64,
-    login_count: u64,
-    tested_count: u64,
-}
-
 // reference for metric types: https://prometheus.io/docs/concepts/metric_types/
 // counters only increase
 // gauges can arbitrarily increase or decrease
@@ -138,46 +120,22 @@ static METRIC_HELP_AND_TYPES_LOOKUP: phf::Map<&'static str, MetricHelpType> = ph
     },
     "servers_bytes_received" => MetricHelpType {
         help: "Volume in bytes of network traffic received by server",
-        ty: "counter",
+        ty: "gauge",
     },
     "servers_bytes_sent" => MetricHelpType {
         help: "Volume in bytes of network traffic sent by server",
-        ty: "counter",
+        ty: "gauge",
     },
     "servers_transaction_count" => MetricHelpType {
         help: "Number of transactions executed by server",
-        ty: "counter",
+        ty: "gauge",
     },
     "servers_query_count" => MetricHelpType {
         help: "Number of queries executed by server",
-        ty: "counter",
+        ty: "gauge",
     },
     "servers_error_count" => MetricHelpType {
         help: "Number of errors",
-        ty: "counter",
-    },
-    "servers_idle_count" => MetricHelpType {
-        help: "Number of server connection in idle state",
-        ty: "gauge",
-    },
-    "servers_active_count" => MetricHelpType {
-        help: "Number of server connection in active state",
-        ty: "gauge",
-    },
-    "servers_tested_count" => MetricHelpType {
-        help: "Number of server connection in tested state",
-        ty: "gauge",
-    },
-    "servers_login_count" => MetricHelpType {
-        help: "Number of server connection in login state",
-        ty: "gauge",
-    },
-    "servers_is_banned" => MetricHelpType {
-        help: "0 if server is not banned, 1 if server is banned",
-        ty: "gauge",
-    },
-    "servers_is_paused" => MetricHelpType {
-        help: "0 if server is not paused, 1 if server is paused",
         ty: "gauge",
     },
     "databases_pool_size" => MetricHelpType {
@@ -200,17 +158,18 @@ struct PrometheusMetric<Value: fmt::Display> {
 
 impl<Value: fmt::Display> fmt::Display for PrometheusMetric<Value> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let mut sorted_labels: Vec<_> = self.labels.iter().collect();
-        sorted_labels.sort_by_key(|&(key, _)| key);
-        let formatted_labels = sorted_labels
+        let formatted_labels = self
+            .labels
             .iter()
             .map(|(key, value)| format!("{}=\"{}\"", key, value))
             .collect::<Vec<_>>()
             .join(",");
         write!(
             f,
-            "{name}{{{formatted_labels}}} {value}",
+            "# HELP {name} {help}\n# TYPE {name} {ty}\n{name}{{{formatted_labels}}} {value}\n",
             name = format_args!("pgcat_{}", self.name),
+            help = self.help,
+            ty = self.ty,
             formatted_labels = formatted_labels,
             value = self.value
         )
@@ -244,9 +203,7 @@ impl<Value: fmt::Display> PrometheusMetric<Value> {
         labels.insert("shard", address.shard.to_string());
         labels.insert("role", address.role.to_string());
         labels.insert("pool", address.pool_name.clone());
-        labels.insert("index", address.address_index.to_string());
         labels.insert("database", address.database.to_string());
-        labels.insert("username", address.username.clone());
 
         Self::from_name(&format!("databases_{}", name), value, labels)
     }
@@ -261,9 +218,7 @@ impl<Value: fmt::Display> PrometheusMetric<Value> {
         labels.insert("shard", address.shard.to_string());
         labels.insert("role", address.role.to_string());
         labels.insert("pool", address.pool_name.clone());
-        labels.insert("index", address.address_index.to_string());
         labels.insert("database", address.database.to_string());
-        labels.insert("username", address.username.clone());
 
         Self::from_name(&format!("servers_{}", name), value, labels)
     }
@@ -274,9 +229,7 @@ impl<Value: fmt::Display> PrometheusMetric<Value> {
         labels.insert("shard", address.shard.to_string());
         labels.insert("pool", address.pool_name.clone());
         labels.insert("role", address.role.to_string());
-        labels.insert("index", address.address_index.to_string());
         labels.insert("database", address.database.to_string());
-        labels.insert("username", address.username.clone());
 
         Self::from_name(&format!("stats_{}", name), value, labels)
     }
@@ -288,20 +241,9 @@ impl<Value: fmt::Display> PrometheusMetric<Value> {
 
         Self::from_name(&format!("pools_{}", name), value, labels)
     }
-
-    fn get_header(&self) -> String {
-        format!(
-            "\n# HELP {name} {help}\n# TYPE {name} {ty}",
-            name = format_args!("pgcat_{}", self.name),
-            help = self.help,
-            ty = self.ty,
-        )
-    }
 }
 
-async fn prometheus_stats(
-    request: Request<body::Incoming>,
-) -> Result<Response<Full<Bytes>>, hyper::http::Error> {
+async fn prometheus_stats(request: Request<Body>) -> Result<Response<Body>, hyper::http::Error> {
     match (request.method(), request.uri().path()) {
         (&Method::GET, "/metrics") => {
             let mut lines = Vec::new();
@@ -309,7 +251,6 @@ async fn prometheus_stats(
             push_pool_stats(&mut lines);
             push_server_stats(&mut lines);
             push_database_stats(&mut lines);
-            lines.push("".to_string()); // Ensure to end the stats with a line terminator as required by the specification.
 
             Response::builder()
                 .header("content-type", "text/plain; version=0.0.4")
@@ -323,7 +264,6 @@ async fn prometheus_stats(
 
 // Adds metrics shown in a SHOW STATS admin command.
 fn push_address_stats(lines: &mut Vec<String>) {
-    let mut grouped_metrics: HashMap<String, Vec<PrometheusMetric<u64>>> = HashMap::new();
     for (_, pool) in get_all_pools() {
         for shard in 0..pool.shards() {
             for server in 0..pool.servers(shard) {
@@ -333,10 +273,7 @@ fn push_address_stats(lines: &mut Vec<String>) {
                     if let Some(prometheus_metric) =
                         PrometheusMetric::<u64>::from_address(address, &key, value)
                     {
-                        grouped_metrics
-                            .entry(key)
-                            .or_default()
-                            .push(prometheus_metric);
+                        lines.push(prometheus_metric.to_string());
                     } else {
                         debug!("Metric {} not implemented for {}", key, address.name());
                     }
@@ -344,53 +281,33 @@ fn push_address_stats(lines: &mut Vec<String>) {
             }
         }
     }
-    for (_key, metrics) in grouped_metrics {
-        if !metrics.is_empty() {
-            lines.push(metrics[0].get_header());
-            for metric in metrics {
-                lines.push(metric.to_string());
-            }
-        }
-    }
 }
 
 // Adds relevant metrics shown in a SHOW POOLS admin command.
 fn push_pool_stats(lines: &mut Vec<String>) {
-    let mut grouped_metrics: HashMap<String, Vec<PrometheusMetric<u64>>> = HashMap::new();
     let pool_stats = PoolStats::construct_pool_lookup();
     for (pool_id, stats) in pool_stats.iter() {
         for (name, value) in stats.clone() {
             if let Some(prometheus_metric) =
                 PrometheusMetric::<u64>::from_pool(pool_id.clone(), &name, value)
             {
-                grouped_metrics
-                    .entry(name)
-                    .or_default()
-                    .push(prometheus_metric);
+                lines.push(prometheus_metric.to_string());
             } else {
                 debug!("Metric {} not implemented for ({})", name, *pool_id);
             }
         }
     }
-    for (_key, metrics) in grouped_metrics {
-        if !metrics.is_empty() {
-            lines.push(metrics[0].get_header());
-            for metric in metrics {
-                lines.push(metric.to_string());
-            }
-        }
-    }
 }
 
 // Adds relevant metrics shown in a SHOW DATABASES admin command.
 fn push_database_stats(lines: &mut Vec<String>) {
-    let mut grouped_metrics: HashMap<String, Vec<PrometheusMetric<u32>>> = HashMap::new();
     for (_, pool) in get_all_pools() {
         let pool_config = pool.settings.clone();
         for shard in 0..pool.shards() {
             for server in 0..pool.servers(shard) {
                 let address = pool.address(shard, server);
                 let pool_state = pool.pool_state(shard, server);
+
                 let metrics = vec![
                     ("pool_size", pool_config.user.pool_size),
                     ("current_connections", pool_state.connections),
@@ -399,10 +316,7 @@ fn push_database_stats(lines: &mut Vec<String>) {
                     if let Some(prometheus_metric) =
                         PrometheusMetric::<u32>::from_database_info(address, key, value)
                     {
-                        grouped_metrics
-                            .entry(key.to_string())
-                            .or_default()
-                            .push(prometheus_metric);
+                        lines.push(prometheus_metric.to_string());
                     } else {
                         debug!("Metric {} not implemented for {}", key, address.name());
                     }
@@ -410,73 +324,45 @@ fn push_database_stats(lines: &mut Vec<String>) {
             }
         }
     }
-    for (_key, metrics) in grouped_metrics {
-        if !metrics.is_empty() {
-            lines.push(metrics[0].get_header());
-            for metric in metrics {
-                lines.push(metric.to_string());
-            }
-        }
-    }
 }
 
 // Adds relevant metrics shown in a SHOW SERVERS admin command.
 fn push_server_stats(lines: &mut Vec<String>) {
     let server_stats = get_server_stats();
-    let mut prom_stats = HashMap::<String, ServerPrometheusStats>::new();
+    let mut server_stats_by_addresses = HashMap::<String, Arc<ServerStats>>::new();
     for (_, stats) in server_stats {
-        let entry = prom_stats
-            .entry(stats.address_name())
-            .or_insert(ServerPrometheusStats {
-                bytes_received: 0,
-                bytes_sent: 0,
-                transaction_count: 0,
-                query_count: 0,
-                error_count: 0,
-                active_count: 0,
-                idle_count: 0,
-                login_count: 0,
-                tested_count: 0,
-            });
-        entry.bytes_received += stats.bytes_received.load(Ordering::Relaxed);
-        entry.bytes_sent += stats.bytes_sent.load(Ordering::Relaxed);
-        entry.transaction_count += stats.transaction_count.load(Ordering::Relaxed);
-        entry.query_count += stats.query_count.load(Ordering::Relaxed);
-        entry.error_count += stats.error_count.load(Ordering::Relaxed);
-        match stats.state.load(Ordering::Relaxed) {
-            crate::stats::ServerState::Login => entry.login_count += 1,
-            crate::stats::ServerState::Active => entry.active_count += 1,
-            crate::stats::ServerState::Tested => entry.tested_count += 1,
-            crate::stats::ServerState::Idle => entry.idle_count += 1,
-        }
+        server_stats_by_addresses.insert(stats.address_name(), stats);
     }
-    let mut grouped_metrics: HashMap<String, Vec<PrometheusMetric<u64>>> = HashMap::new();
+
     for (_, pool) in get_all_pools() {
         for shard in 0..pool.shards() {
             for server in 0..pool.servers(shard) {
                 let address = pool.address(shard, server);
-                if let Some(server_info) = prom_stats.get(&address.name()) {
+                if let Some(server_info) = server_stats_by_addresses.get(&address.name()) {
                     let metrics = [
-                        ("bytes_received", server_info.bytes_received),
-                        ("bytes_sent", server_info.bytes_sent),
-                        ("transaction_count", server_info.transaction_count),
-                        ("query_count", server_info.query_count),
-                        ("error_count", server_info.error_count),
-                        ("idle_count", server_info.idle_count),
-                        ("active_count", server_info.active_count),
-                        ("login_count", server_info.login_count),
-                        ("tested_count", server_info.tested_count),
-                        ("is_banned", if pool.is_banned(address) { 1 } else { 0 }),
-                        ("is_paused", if pool.paused() { 1 } else { 0 }),
+                        (
+                            "bytes_received",
+                            server_info.bytes_received.load(Ordering::Relaxed),
+                        ),
+                        ("bytes_sent", server_info.bytes_sent.load(Ordering::Relaxed)),
+                        (
+                            "transaction_count",
+                            server_info.transaction_count.load(Ordering::Relaxed),
+                        ),
+                        (
+                            "query_count",
+                            server_info.query_count.load(Ordering::Relaxed),
+                        ),
+                        (
+                            "error_count",
+                            server_info.error_count.load(Ordering::Relaxed),
+                        ),
                     ];
                     for (key, value) in metrics {
                         if let Some(prometheus_metric) =
                             PrometheusMetric::<u64>::from_server_info(address, key, value)
                         {
-                            grouped_metrics
-                                .entry(key.to_string())
-                                .or_default()
-                                .push(prometheus_metric);
+                            lines.push(prometheus_metric.to_string());
                         } else {
                             debug!("Metric {} not implemented for {}", key, address.name());
                         }
@@ -485,46 +371,17 @@ fn push_server_stats(lines: &mut Vec<String>) {
             }
         }
     }
-    for (_key, metrics) in grouped_metrics {
-        if !metrics.is_empty() {
-            lines.push(metrics[0].get_header());
-            for metric in metrics {
-                lines.push(metric.to_string());
-            }
-        }
-    }
 }
 
 pub async fn start_metric_server(http_addr: SocketAddr) {
-    let listener = TcpListener::bind(http_addr);
-    let listener = match listener.await {
-        Ok(listener) => listener,
-        Err(e) => {
-            error!("Failed to bind prometheus server to HTTP address: {}.", e);
-            return;
-        }
-    };
+    let http_service_factory =
+        make_service_fn(|_conn| async { Ok::<_, hyper::Error>(service_fn(prometheus_stats)) });
+    let server = Server::bind(&http_addr).serve(http_service_factory);
     info!(
         "Exposing prometheus metrics on http://{}/metrics.",
         http_addr
     );
-    loop {
-        let stream = match listener.accept().await {
-            Ok((stream, _)) => stream,
-            Err(e) => {
-                error!("Error accepting connection: {}", e);
-                continue;
-            }
-        };
-        let io = TokioIo::new(stream);
-
-        tokio::task::spawn(async move {
-            if let Err(err) = http1::Builder::new()
-                .serve_connection(io, service_fn(prometheus_stats))
-                .await
-            {
-                eprintln!("Error serving HTTP connection for metrics: {:?}", err);
-            }
-        });
+    if let Err(e) = server.await {
+        error!("Failed to run HTTP server: {}.", e);
     }
 }

src/query_router.rs

@@ -2,7 +2,6 @@
 /// or implied query characteristics.
 use bytes::{Buf, BytesMut};
 use log::{debug, error};
-use mini_moka::sync::Cache;
 use once_cell::sync::OnceCell;
 use regex::{Regex, RegexSet};
 use sqlparser::ast::Statement::{Delete, Insert, Query, StartTransaction, Update};
@@ -12,7 +11,6 @@ use sqlparser::ast::{
 };
 use sqlparser::dialect::PostgreSqlDialect;
 use sqlparser::parser::Parser;
-use std::sync::OnceLock;
 
 use crate::config::Role;
 use crate::errors::Error;
@@ -23,7 +21,6 @@ use crate::sharding::Sharder;
 
 use std::collections::BTreeSet;
 use std::io::Cursor;
-use std::time::Duration;
 use std::{cmp, mem};
 
 /// Regexes used to parse custom commands.
@@ -69,18 +66,6 @@ static CUSTOM_SQL_REGEX_SET: OnceCell<RegexSet> = OnceCell::new();
 // Get the value inside the custom command.
 static CUSTOM_SQL_REGEX_LIST: OnceCell<Vec<Regex>> = OnceCell::new();
 
-#[derive(Debug, Clone, PartialEq)]
-enum DatabaseActivityState {
-    Active,
-    Initializing,
-}
-
-// A moka cache for the databases
-// the key is the database name and the value is the database activity state
-static DATABASE_ACTIVITY_CACHE: OnceLock<Cache<String, DatabaseActivityState>> = OnceLock::new();
-// A moka cache for the tables, the key is the db_table.
-static TABLE_MUTATIONS_CACHE: OnceLock<Cache<String, bool>> = OnceLock::new();
-
 /// The query router.
 pub struct QueryRouter {
     /// Which shard we should be talking to right now.
@@ -102,12 +87,6 @@ pub struct QueryRouter {
     placeholders: Vec<i16>,
 }
 
-struct ExtractedExprsAndTables<'a> {
-    exprs: Vec<Expr>,
-    table_names: Vec<Vec<Ident>>,
-    assignments_opt: Option<&'a Vec<Assignment>>,
-}
-
 impl QueryRouter {
     /// One-time initialization of regexes
     /// that parse our custom SQL protocol.
@@ -407,53 +386,6 @@ impl QueryRouter {
         }
     }
 
-    /// Determines if a query is a mutation or not.
-    fn is_mutation_query(q: &sqlparser::ast::Query) -> bool {
-        use sqlparser::ast::*;
-
-        match q.body.as_ref() {
-            SetExpr::Insert(_) => true,
-            SetExpr::Update(_) => true,
-            SetExpr::Query(q) => Self::is_mutation_query(q),
-            _ => false,
-        }
-    }
-
-    fn database_activity_cache(&self) -> Cache<String, DatabaseActivityState> {
-        DATABASE_ACTIVITY_CACHE
-            .get_or_init(|| {
-                Cache::builder()
-                    .time_to_idle(Duration::from_secs(self.pool_settings.db_activity_ttl))
-                    .build()
-            })
-            .clone()
-    }
-
-    /// Check database activity state and reset it if necessary
-    fn database_activity_state(&self, db: &String) -> DatabaseActivityState {
-        let cache = self.database_activity_cache();
-
-        // Exists in cache
-        if cache.contains_key(db) {
-            return cache.get(db).unwrap();
-        }
-
-        // Not in cache
-        debug!("Adding database to cache: {}", db);
-
-        cache.insert(db.to_string(), DatabaseActivityState::Initializing);
-
-        // Set a timer to update the cache
-        let db = db.clone();
-        let db_activity_init_delay = self.pool_settings.db_activity_init_delay;
-        tokio::spawn(async move {
-            tokio::time::sleep(Duration::from_millis(db_activity_init_delay)).await;
-            cache.insert(db, DatabaseActivityState::Active);
-        });
-
-        DatabaseActivityState::Initializing
-    }
-
     /// Try to infer which server to connect to based on the contents of the query.
     pub fn infer(&mut self, ast: &Vec<sqlparser::ast::Statement>) -> Result<(), Error> {
         if !self.pool_settings.query_parser_read_write_splitting {
@@ -468,23 +400,9 @@ impl QueryRouter {
             return Err(Error::QueryRouterParserError("empty query".into()));
         }
 
-        let mut primary_set_based_on_activity = false;
         let mut visited_write_statement = false;
         let mut prev_inferred_shard = None;
 
-        if self.pool_settings.db_activity_based_routing {
-            let db = self.pool_settings.db.clone();
-            let state = self.database_activity_state(&db);
-            debug!("Database activity state: {:?}", state);
-
-            if let DatabaseActivityState::Initializing = state {
-                debug!("Database is initializing, going to primary");
-
-                self.active_role = Some(Role::Primary);
-                primary_set_based_on_activity = true;
-            }
-        }
-
         for q in ast {
             match q {
                 // All transactions go to the primary, probably a write.
@@ -495,22 +413,6 @@ impl QueryRouter {
 
                 // Likely a read-only query
                 Query(query) => {
-                    if primary_set_based_on_activity {
-                        // If we already set the role based on activity, we don't need to do it again
-                        continue;
-                    }
-
-                    if self.pool_settings.db_activity_based_routing {
-                        // Check if the tables in the query have been written to recently
-                        if self.query_handles_tables_in_mutation_cache(query) {
-                            debug!("Query handles tables in mutation cache, going to primary");
-
-                            self.active_role = Some(Role::Primary);
-                            primary_set_based_on_activity = true;
-                            continue;
-                        }
-                    }
-
                     match &self.pool_settings.automatic_sharding_key {
                         Some(_) => {
                             // TODO: if we have multiple queries in the same message,
@@ -525,13 +427,8 @@ impl QueryRouter {
                         None => (),
                     };
 
-                    let has_locks = !query.locks.is_empty();
-                    let has_mutation = Self::is_mutation_query(query);
-
-                    if has_locks || has_mutation {
-                        self.active_role = Some(Role::Primary);
-                    } else if !visited_write_statement {
-                        // If we already visited a write statement, we should be going to the primary.
+                    // If we already visited a write statement, we should be going to the primary.
+                    if !visited_write_statement {
                         self.active_role = match self.primary_reads_enabled() {
                             false => Some(Role::Replica), // If primary should not be receiving reads, use a replica.
                             true => None,                 // Any server role is fine in this case.
@@ -541,13 +438,6 @@ impl QueryRouter {
 
                 // Likely a write
                 _ => {
-                    debug!("Write statement found, going to primary");
-
-                    if self.pool_settings.db_activity_based_routing {
-                        // add all of the query tables to the mutation cache
-                        self.update_mutation_cache_on_write(q);
-                    }
-
                     match &self.pool_settings.automatic_sharding_key {
                         Some(_) => {
                             // TODO: similar to the above, if we have multiple queries in the
@@ -590,69 +480,62 @@ impl QueryRouter {
         Ok(())
     }
 
-    fn table_mutations_cache(&self) -> Cache<String, bool> {
-        TABLE_MUTATIONS_CACHE
-            .get_or_init(|| {
-                Cache::builder()
-                    .time_to_live(Duration::from_millis(
-                        self.pool_settings.table_mutation_cache_ms_ttl,
-                    ))
-                    .build()
-            })
-            .clone()
-    }
-
-    fn query_handles_tables_in_mutation_cache(&self, query: &sqlparser::ast::Query) -> bool {
-        let table_mutations_cache = self.table_mutations_cache();
-        debug!("Checking if query handles tables in mutation cache");
-        debug!("Table mutations cache: {:?}", table_mutations_cache);
-
-        for tables in self.table_names(query) {
-            for table in tables {
-                if table_mutations_cache.contains_key(&self.table_mutation_cache_key(table)) {
-                    return true;
-                }
-            }
-        }
-
-        false
-    }
-    fn extract_exprs_and_table_names<'a>(
-        &'a self,
-        q: &'a Statement,
-    ) -> Option<ExtractedExprsAndTables<'a>> {
+    fn infer_shard_on_write(&mut self, q: &Statement) -> Result<Option<usize>, Error> {
         let mut exprs = Vec::new();
+
+        // Collect all table names from the query.
         let mut table_names = Vec::new();
-        let mut assignments_opt = None;
 
         match q {
-            Insert(i) => {
+            Insert {
+                or,
+                into: _,
+                table_name,
+                columns,
+                overwrite: _,
+                source,
+                partitioned,
+                after_columns,
+                table: _,
+                on: _,
+                returning: _,
+                ignore: _,
+            } => {
                 // Not supported in postgres.
-                assert!(i.or.is_none());
-                assert!(i.partitioned.is_none());
-                assert!(i.after_columns.is_empty());
+                assert!(or.is_none());
+                assert!(partitioned.is_none());
+                assert!(after_columns.is_empty());
 
-                Self::process_table(&i.table_name, &mut table_names);
-                if let Some(source) = &i.source {
-                    Self::process_query(source, &mut exprs, &mut table_names, &Some(&i.columns));
+                Self::process_table(table_name, &mut table_names);
+                if let Some(source) = source {
+                    Self::process_query(source, &mut exprs, &mut table_names, &Some(columns));
                 }
             }
-            Delete(d) => {
-                if let Some(expr) = &d.selection {
+            Delete {
+                tables,
+                from,
+                using,
+                selection,
+                returning: _,
+                order_by: _,
+                limit: _,
+            } => {
+                if let Some(expr) = selection {
                     exprs.push(expr.clone());
                 }
 
-                // Multi-tables delete are not supported in postgres.
-                assert!(d.tables.is_empty());
+                // Multi tables delete are not supported in postgres.
+                assert!(tables.is_empty());
 
-                if let Some(using_tbl_with_join) = &d.using {
+                Self::process_tables_with_join(from, &mut exprs, &mut table_names);
+                if let Some(using_tbl_with_join) = using {
                     Self::process_tables_with_join(
                         using_tbl_with_join,
                         &mut exprs,
                         &mut table_names,
                     );
                 }
-                Self::process_selection(&d.selection, &mut exprs);
+                Self::process_selection(selection, &mut exprs);
             }
             Update {
                 table,
@@ -666,55 +549,14 @@ impl QueryRouter {
                     Self::process_table_with_join(from_tbl, &mut exprs, &mut table_names);
                 }
                 Self::process_selection(selection, &mut exprs);
-
-                assignments_opt = Some(assignments);
-            }
-            _ => return None,
-        };
-
-        Some(ExtractedExprsAndTables {
-            exprs,
-            table_names,
-            assignments_opt,
-        })
-    }
-
-    fn infer_shard_on_write(&mut self, q: &Statement) -> Result<Option<usize>, Error> {
-        if let Some(extracted) = self.extract_exprs_and_table_names(q) {
-            let exprs = extracted.exprs;
-            let table_names = extracted.table_names;
-            let assignments_opt = extracted.assignments_opt;
-
-            if let Some(assignments) = assignments_opt {
                 self.assignment_parser(assignments)?;
             }
-
-            Ok(self.infer_shard_from_exprs(exprs, table_names))
-        } else {
-            Ok(None)
-        }
-    }
-
-    fn update_mutation_cache_on_write(&self, q: &Statement) {
-        if let Some(extracted) = self.extract_exprs_and_table_names(q) {
-            debug!("Updating mutation cache on write");
-
-            let table_names = extracted.table_names;
-            debug!("Table names in mutation query: {:?}", table_names);
-            let table_mutations_cache = self.table_mutations_cache();
-            for tables in table_names {
-                for table in tables {
-                    table_mutations_cache.insert(self.table_mutation_cache_key(table), true);
-                }
+            _ => {
+                return Ok(None);
             }
-        }
-    }
+        };
 
-    // combines the database name and table name into a single string
-    // to be used as the key in the table mutation cache
-    // e.g. "mydb.mytable"
-    fn table_mutation_cache_key(&self, table: Ident) -> String {
-        format!("{}.{}", self.pool_settings.db, table.value)
+        Ok(self.infer_shard_from_exprs(exprs, table_names))
     }
 
     fn process_query(
@@ -963,13 +805,7 @@ impl QueryRouter {
 
         for a in assignments {
             if sharding_key[0].value == "*"
-                && sharding_key[1].value
-                    == a.target
-                        .to_string()
-                        .split('.')
-                        .last()
-                        .unwrap()
-                        .to_lowercase()
+                && sharding_key[1].value == a.id.last().unwrap().value.to_lowercase()
             {
                 return Err(Error::QueryRouterParserError(
                     "Sharding key cannot be updated.".into(),
@@ -1102,18 +938,6 @@ impl QueryRouter {
         self.infer_shard_from_exprs(exprs, table_names)
     }
 
-    /// get table names from query
-    fn table_names(&self, query: &sqlparser::ast::Query) -> Vec<Vec<Ident>> {
-        let mut exprs = Vec::new();
-
-        let mut table_names = Vec::new();
-        Self::process_query(query, &mut exprs, &mut table_names, &None);
-
-        debug!("Table names in query: {:?}", table_names);
-
-        table_names
-    }
-
     fn infer_shard_from_exprs(
         &mut self,
         exprs: Vec<Expr>,
@@ -1220,11 +1044,6 @@ impl QueryRouter {
         self.active_shard
     }
 
-    /// Set active_role as the default_role specified in the pool.
-    pub fn set_default_role(&mut self) {
-        self.active_role = self.pool_settings.default_role;
-    }
-
     /// Get the current desired server role we should be talking to.
     pub fn role(&self) -> Option<Role> {
         self.active_role
@@ -1281,7 +1100,6 @@ mod test {
     use crate::messages::simple_query;
     use crate::sharding::ShardingFunction;
     use bytes::BufMut;
-    use serial_test::serial;
 
     #[test]
     fn test_defaults() {
@@ -1291,26 +1109,6 @@ mod test {
         assert_eq!(qr.role(), None);
     }
 
-    #[test]
-    fn test_split_cte_queries() {
-        QueryRouter::setup();
-        let mut qr = QueryRouter::new();
-        qr.pool_settings.query_parser_read_write_splitting = true;
-        qr.pool_settings.query_parser_enabled = true;
-
-        let query = simple_query(
-            "WITH t AS (
-                SELECT id FROM users WHERE name ILIKE '%ja%'
-            )
-            UPDATE user_languages
-               SET settings = '{}'
-              FROM t WHERE t.id = user_id;",
-        );
-        let ast = qr.parse(&query).unwrap();
-        assert!(qr.infer(&ast).is_ok());
-        assert_eq!(qr.role(), Some(Role::Primary));
-    }
-
     #[test]
     fn test_infer_replica() {
         QueryRouter::setup();
@@ -1360,29 +1158,6 @@ mod test {
         }
     }
 
-    #[test]
-    fn test_select_for_update() {
-        QueryRouter::setup();
-        let mut qr = QueryRouter::new();
-        qr.pool_settings.query_parser_read_write_splitting = true;
-
-        let queries_in_primary_role = vec![
-            simple_query("BEGIN"), // Transaction start
-            simple_query("SELECT * FROM items WHERE id = 5 FOR UPDATE"),
-            simple_query("UPDATE items SET name = 'pumpkin' WHERE id = 5"),
-        ];
-
-        for query in queries_in_primary_role {
-            assert!(qr.infer(&qr.parse(&query).unwrap()).is_ok());
-            assert_eq!(qr.role(), Some(Role::Primary));
-        }
-
-        // query without lock do not change role
-        let query = simple_query("SELECT * FROM items WHERE id = 5");
-        assert!(qr.infer(&qr.parse(&query).unwrap()).is_ok());
-        assert_eq!(qr.role(), None);
-    }
-
     #[test]
     fn test_infer_primary_reads_enabled() {
         QueryRouter::setup();
@@ -1597,19 +1372,6 @@ mod test {
         assert!(!qr.query_parser_enabled());
     }
 
-    #[test]
-    fn test_query_parser() {
-        QueryRouter::setup();
-        let mut qr = QueryRouter::new();
-        qr.pool_settings.query_parser_read_write_splitting = true;
-
-        let query = simple_query("SELECT req_tab_0.*  FROM validation req_tab_0  WHERE  array['http://www.w3.org/ns/shacl#ValidationResult'] && req_tab_0.type::text[] AND ( (  (req_tab_0.focusnode = 'DataSource_Credilogic_DataSourceAddress_144959227') )  )");
-        assert!(qr.infer(&qr.parse(&query).unwrap()).is_ok());
-
-        let query = simple_query("WITH EmployeeSalaries AS (SELECT Department, Salary FROM Employees) SELECT Department, AVG(Salary) AS AverageSalary FROM EmployeeSalaries GROUP BY Department;");
-        assert!(qr.infer(&qr.parse(&query).unwrap()).is_ok());
-    }
-
     #[test]
     fn test_update_from_pool_settings() {
         QueryRouter::setup();
@@ -1617,7 +1379,6 @@ mod test {
         let pool_settings = PoolSettings {
             pool_mode: PoolMode::Transaction,
             load_balancing_mode: crate::config::LoadBalancingMode::Random,
-            checkout_failure_limit: None,
             shards: 2,
             user: crate::config::User::default(),
             default_role: Some(Role::Replica),
@@ -1638,10 +1399,6 @@ mod test {
             auth_query_password: None,
             auth_query_user: None,
             db: "test".to_string(),
-            db_activity_based_routing: PoolSettings::default().db_activity_based_routing,
-            db_activity_init_delay: PoolSettings::default().db_activity_init_delay,
-            db_activity_ttl: PoolSettings::default().db_activity_ttl,
-            table_mutation_cache_ms_ttl: PoolSettings::default().table_mutation_cache_ms_ttl,
             plugins: None,
         };
         let mut qr = QueryRouter::new();
@@ -1700,7 +1457,6 @@ mod test {
         let pool_settings = PoolSettings {
             pool_mode: PoolMode::Transaction,
             load_balancing_mode: crate::config::LoadBalancingMode::Random,
-            checkout_failure_limit: Some(10),
             shards: 5,
             user: crate::config::User::default(),
             default_role: Some(Role::Replica),
@@ -1721,10 +1477,6 @@ mod test {
             auth_query_password: None,
             auth_query_user: None,
             db: "test".to_string(),
-            db_activity_based_routing: PoolSettings::default().db_activity_based_routing,
-            db_activity_init_delay: PoolSettings::default().db_activity_init_delay,
-            db_activity_ttl: PoolSettings::default().db_activity_ttl,
-            table_mutation_cache_ms_ttl: PoolSettings::default().table_mutation_cache_ms_ttl,
             plugins: None,
         };
 
@@ -2140,150 +1892,4 @@ mod test {
 
         assert_eq!(res, Ok(PluginOutput::Allow));
     }
-
-    #[tokio::test]
-    #[serial]
-    async fn test_db_activity_based_routing_initializing_state() {
-        QueryRouter::setup();
-        let mut qr = QueryRouter::new();
-        qr.pool_settings.db_activity_based_routing = true;
-        qr.pool_settings.query_parser_read_write_splitting = true;
-        qr.pool_settings.query_parser_enabled = true;
-        qr.pool_settings.db = "test_table_mutation_cache".to_string();
-
-        qr.database_activity_cache()
-            .invalidate(&qr.pool_settings.db.clone());
-
-        let query = simple_query("SELECT * FROM some_table");
-        let ast = qr.parse(&query).unwrap();
-
-        // Initially, the database activity should be in the "Initializing" state
-        let state = qr.database_activity_state(&qr.pool_settings.db.clone());
-        assert_eq!(state, DatabaseActivityState::Initializing);
-
-        // Check that the router chooses the primary role due to "Initializing" state
-        assert!(qr.infer(&ast).is_ok());
-        assert_eq!(qr.role(), Some(Role::Primary));
-    }
-
-    #[tokio::test]
-    #[serial]
-    async fn test_db_activity_based_routing_active_state() {
-        QueryRouter::setup();
-        let mut qr = QueryRouter::new();
-        qr.pool_settings.db_activity_based_routing = true;
-        qr.pool_settings.query_parser_read_write_splitting = true;
-        qr.pool_settings.query_parser_enabled = true;
-        qr.pool_settings.db = "test_table_mutation_cache".to_string();
-
-        let db_name = qr.pool_settings.db.clone();
-        let cache = qr.database_activity_cache();
-        cache.insert(db_name.clone(), DatabaseActivityState::Active);
-
-        let query = simple_query("SELECT * FROM some_table");
-        let ast = qr.parse(&query).unwrap();
-
-        // Check that the router can choose a replica role when in "Active" state
-        assert!(qr.infer(&ast).is_ok());
-        assert_eq!(qr.role(), None); // Default should allow replica due to active state
-    }
-
-    #[tokio::test]
-    #[serial]
-    async fn test_table_mutation_cache_on_write() {
-        QueryRouter::setup();
-        let mut qr = QueryRouter::new();
-        qr.pool_settings.db_activity_based_routing = true;
-        qr.pool_settings.table_mutation_cache_ms_ttl = 20_000; // 20 seconds in milliseconds
-        qr.pool_settings.query_parser_enabled = true;
-        qr.pool_settings.query_parser_read_write_splitting = true;
-        qr.pool_settings.db = "test_table_mutation_cache".to_string();
-
-        qr.database_activity_cache()
-            .invalidate(&qr.pool_settings.db.clone());
-
-        let query = simple_query("UPDATE some_table SET col1 = 'value' WHERE col2 = 1");
-        let ast = qr.parse(&query).unwrap();
-
-        // Simulate the mutation query which should populate the mutation cache
-        assert!(qr.infer(&ast).is_ok());
-        assert_eq!(qr.role(), Some(Role::Primary));
-
-        let table_cache_key = qr.table_mutation_cache_key(Ident::new("some_table"));
-        let cache = qr.table_mutations_cache();
-
-        // Ensure the table mutation cache contains the table with recent write
-        assert!(cache.contains_key(&table_cache_key));
-    }
-
-    #[tokio::test]
-    #[serial]
-    async fn test_db_activity_based_routing_multi_query() {
-        use super::*;
-        use crate::messages::simple_query;
-        use tokio::time::Duration;
-
-        QueryRouter::setup();
-        let mut qr = QueryRouter::new();
-
-        // Configure the pool settings for db_activity_based_routing
-        qr.pool_settings.query_parser_read_write_splitting = true;
-        qr.pool_settings.query_parser_enabled = true;
-        qr.pool_settings.db_activity_based_routing = true;
-        qr.pool_settings.db = "test_db_activity_routing".to_string();
-
-        qr.database_activity_cache()
-            .invalidate(&qr.pool_settings.db.clone());
-
-        // First query when database is initializing
-        let query = simple_query("SELECT * FROM test_table");
-        let ast = qr.parse(&query).unwrap();
-        assert!(qr.infer(&ast).is_ok());
-        // Should route to primary because database is initializing
-        assert_eq!(qr.role(), Some(Role::Primary));
-
-        // Wait for the initialization delay to pass
-        tokio::time::sleep(Duration::from_millis(
-            qr.pool_settings.db_activity_init_delay * 2,
-        ))
-        .await;
-
-        // Next query after database is active
-        let query = simple_query("SELECT * FROM test_table");
-        let ast = qr.parse(&query).unwrap();
-        qr.active_role = None; // Reset the active_role
-        assert!(qr.infer(&ast).is_ok());
-        // Should route to replica because database is active and no recent mutations
-        assert_eq!(qr.role(), None);
-
-        // Simulate a write query to update the mutation cache
-        let query = simple_query("INSERT INTO test_table (id, name) VALUES (1, 'test')");
-        let ast = qr.parse(&query).unwrap();
-        qr.active_role = None; // Reset the active_role
-        assert!(qr.infer(&ast).is_ok());
-        // Should route to primary because it's a write operation
-        assert_eq!(qr.role(), Some(Role::Primary));
-
-        // Immediately run a read query on the same table
-        let query = simple_query("SELECT * FROM test_table WHERE id = 1");
-        let ast = qr.parse(&query).unwrap();
-        qr.active_role = None; // Reset the active_role
-        assert!(qr.infer(&ast).is_ok());
-        // Should route to primary because the table was recently mutated
-        assert_eq!(qr.role(), Some(Role::Primary));
-
-        // Wait for the mutation cache TTL to expire
-        tokio::time::sleep(Duration::from_millis(
-            qr.pool_settings.table_mutation_cache_ms_ttl * 2,
-        ))
-        .await;
-
-        // Run the read query again after cache expiration
-        let query = simple_query("SELECT * FROM test_table WHERE id = 1");
-        let ast = qr.parse(&query).unwrap();
-        qr.active_role = None; // Reset the active_role
-        assert!(qr.infer(&ast).is_ok());
-        // Should route to replica because mutation cache has expired
-        assert_eq!(qr.role(), None);
-    }
 }

src/sharding.rs

@@ -14,11 +14,11 @@ pub enum ShardingFunction {
     Sha1,
 }
 
-impl std::fmt::Display for ShardingFunction {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            ShardingFunction::PgBigintHash => write!(f, "pg_bigint_hash"),
-            ShardingFunction::Sha1 => write!(f, "sha1"),
+impl ToString for ShardingFunction {
+    fn to_string(&self) -> String {
+        match *self {
+            ShardingFunction::PgBigintHash => "pg_bigint_hash".to_string(),
+            ShardingFunction::Sha1 => "sha1".to_string(),
         }
     }
 }

start_test_env.sh

@@ -1,34 +0,0 @@
-GREEN="\033[0;32m"
-RED="\033[0;31m"
-BLUE="\033[0;34m"
-RESET="\033[0m"
-
-
-cd tests/docker/
-docker compose kill main || true
-docker compose build main
-docker compose down
-docker compose up -d
-# wait for the container to start
-while ! docker compose exec main ls; do
-    echo "Waiting for test environment to start"
-    sleep 1
-done
-echo "==================================="
-docker compose exec -e LOG_LEVEL=error -d main toxiproxy-server
-docker compose exec --workdir /app main cargo build
-docker compose exec -d --workdir /app main ./target/debug/pgcat ./.circleci/pgcat.toml
-docker compose exec --workdir /app/tests/ruby main bundle install
-docker compose exec --workdir /app/tests/python main pip3 install -r requirements.txt
-echo "Interactive test environment ready"
-echo "To run integration tests, you can use the following commands:"
-echo -e "   ${BLUE}Ruby:   ${RED}cd /app/tests/ruby && bundle exec ruby tests.rb --format documentation${RESET}"
-echo -e "   ${BLUE}Python: ${RED}cd /app/ && pytest ${RESET}"
-echo -e "   ${BLUE}Rust:   ${RED}cd /app/tests/rust && cargo run ${RESET}"
-echo -e "   ${BLUE}Go:     ${RED}cd /app/tests/go && /usr/local/go/bin/go test${RESET}"
-echo "the source code for tests are directly linked to the source code in the container so you can modify the code and run the tests again"
-echo "You can rebuild PgCat from within the container by running"
-echo -e "  ${GREEN}cargo build${RESET}"
-echo "and then run the tests again"
-echo "==================================="
-docker compose exec --workdir /app/tests main bash

tests/docker/Dockerfile

@@ -1,4 +1,4 @@
-FROM rust:1.81.0-slim-bookworm 
+FROM rust:bullseye
 
 COPY --from=sclevine/yj /bin/yj /bin/yj
 RUN /bin/yj -h

tests/docker/docker-compose.yml

@@ -1,3 +1,4 @@
+version: "3"
 services:
   pg1:
     image: postgres:14
@@ -47,8 +48,6 @@ services:
   main:
     build: .
     command: ["bash", "/app/tests/docker/run.sh"]
-    environment:
-      - INTERACTIVE_TEST_ENVIRONMENT=true
     volumes:
       - ../../:/app/
       - /app/target/

tests/docker/run.sh

@@ -5,38 +5,6 @@ rm /app/*.profraw || true
 rm /app/pgcat.profdata || true
 rm -rf /app/cov || true
 
-# Prepares the interactive test environment
-# 
-if [ -n "$INTERACTIVE_TEST_ENVIRONMENT" ]; then
-    ports=(5432 7432 8432 9432 10432)
-    for port in "${ports[@]}"; do
-        is_it_up=0
-        attempts=0
-        while [ $is_it_up -eq 0 ]; do
-            PGPASSWORD=postgres psql -h 127.0.0.1 -p $port -U postgres -c '\q' > /dev/null 2>&1
-            if [ $? -eq 0 ]; then
-                echo "PostgreSQL on port $port is up."
-                is_it_up=1
-            else
-                attempts=$((attempts+1))
-                if [ $attempts -gt 10 ]; then
-                    echo "PostgreSQL on port $port is down, giving up."
-                    exit 1
-                fi
-                echo "PostgreSQL on port $port is down, waiting for it to start."
-                sleep 1
-            fi
-        done
-    done
-    PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432  -U postgres -f /app/tests/sharding/query_routing_setup.sql
-    PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 7432  -U postgres -f /app/tests/sharding/query_routing_setup.sql
-    PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 8432  -U postgres -f /app/tests/sharding/query_routing_setup.sql
-    PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 9432  -U postgres -f /app/tests/sharding/query_routing_setup.sql
-    PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 10432 -U postgres -f /app/tests/sharding/query_routing_setup.sql
-    sleep 100000000000000000
-    exit 0
-fi
-
 export LLVM_PROFILE_FILE="/app/pgcat-%m-%p.profraw"
 export RUSTC_BOOTSTRAP=1
 export CARGO_INCREMENTAL=0

tests/python/requirements.txt

@@ -1,3 +1,2 @@
-pytest
 psycopg2==2.9.3
-psutil==5.9.1
+psutil==5.9.1

tests/python/test_auth.py

@@ -1,71 +0,0 @@
-import utils
-import signal
-
-class TestTrustAuth:
-    @classmethod
-    def setup_method(cls):
-        config= """
-        [general]
-        host = "0.0.0.0"
-        port = 6432
-        admin_username = "admin_user"
-        admin_password = ""
-        admin_auth_type = "trust"
-
-        [pools.sharded_db.users.0]
-        username = "sharding_user"
-        password = "sharding_user"
-        auth_type = "trust"
-        pool_size = 10
-        min_pool_size = 1
-        pool_mode = "transaction"
-
-        [pools.sharded_db.shards.0]
-        servers = [
-          [ "127.0.0.1", 5432, "primary" ],
-        ]
-        database = "shard0"
-        """
-        utils.pgcat_generic_start(config)
-
-    @classmethod
-    def teardown_method(self):
-        utils.pg_cat_send_signal(signal.SIGTERM)
-
-    def test_admin_trust_auth(self):
-        conn, cur = utils.connect_db_trust(admin=True)
-        cur.execute("SHOW POOLS")
-        res = cur.fetchall()
-        print(res)
-        utils.cleanup_conn(conn, cur)
-
-    def test_normal_trust_auth(self):
-        conn, cur = utils.connect_db_trust(autocommit=False)
-        cur.execute("SELECT 1")
-        res = cur.fetchall()
-        print(res)
-        utils.cleanup_conn(conn, cur)
-
-class TestMD5Auth:
-    @classmethod
-    def setup_method(cls):
-        utils.pgcat_start()
-
-    @classmethod
-    def teardown_method(self):
-        utils.pg_cat_send_signal(signal.SIGTERM)
-
-    def test_normal_db_access(self):
-        conn, cur = utils.connect_db(autocommit=False)
-        cur.execute("SELECT 1")
-        res = cur.fetchall()
-        print(res)
-        utils.cleanup_conn(conn, cur)
-
-    def test_admin_db_access(self):
-        conn, cur = utils.connect_db(admin=True)
-
-        cur.execute("SHOW POOLS")
-        res = cur.fetchall()
-        print(res)
-        utils.cleanup_conn(conn, cur)

tests/python/tests.py

@@ -1,12 +1,84 @@
-
+from typing import Tuple
+import psycopg2
+import psutil
+import os
 import signal
 import time
 
-import psycopg2
-import utils
-
 SHUTDOWN_TIMEOUT = 5
 
+PGCAT_HOST = "127.0.0.1"
+PGCAT_PORT = "6432"
+
+
+def pgcat_start():
+    pg_cat_send_signal(signal.SIGTERM)
+    os.system("./target/debug/pgcat .circleci/pgcat.toml &")
+    time.sleep(2)
+
+
+def pg_cat_send_signal(signal: signal.Signals):
+    try:
+        for proc in psutil.process_iter(["pid", "name"]):
+            if "pgcat" == proc.name():
+                os.kill(proc.pid, signal)
+    except Exception as e:
+        # The process can be gone when we send this signal
+        print(e)
+
+    if signal == signal.SIGTERM:
+        # Returns 0 if pgcat process exists
+        time.sleep(2)
+        if not os.system('pgrep pgcat'):
+            raise Exception("pgcat not closed after SIGTERM")
+
+
+def connect_db(
+    autocommit: bool = True,
+    admin: bool = False,
+) -> Tuple[psycopg2.extensions.connection, psycopg2.extensions.cursor]:
+
+    if admin:
+        user = "admin_user"
+        password = "admin_pass"
+        db = "pgcat"
+    else:
+        user = "sharding_user"
+        password = "sharding_user"
+        db = "sharded_db"
+
+    conn = psycopg2.connect(
+        f"postgres://{user}:{password}@{PGCAT_HOST}:{PGCAT_PORT}/{db}?application_name=testing_pgcat",
+        connect_timeout=2,
+    )
+    conn.autocommit = autocommit
+    cur = conn.cursor()
+
+    return (conn, cur)
+
+
+def cleanup_conn(conn: psycopg2.extensions.connection, cur: psycopg2.extensions.cursor):
+    cur.close()
+    conn.close()
+
+
+def test_normal_db_access():
+    pgcat_start()
+    conn, cur = connect_db(autocommit=False)
+    cur.execute("SELECT 1")
+    res = cur.fetchall()
+    print(res)
+    cleanup_conn(conn, cur)
+
+
+def test_admin_db_access():
+    conn, cur = connect_db(admin=True)
+
+    cur.execute("SHOW POOLS")
+    res = cur.fetchall()
+    print(res)
+    cleanup_conn(conn, cur)
+
 
 def test_shutdown_logic():
 
@@ -14,17 +86,17 @@ def test_shutdown_logic():
     # NO ACTIVE QUERIES SIGINT HANDLING
 
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and send query (not in transaction)
-    conn, cur = utils.connect_db()
+    conn, cur = connect_db()
 
     cur.execute("BEGIN;")
     cur.execute("SELECT 1;")
     cur.execute("COMMIT;")
 
     # Send sigint to pgcat
-    utils.pg_cat_send_signal(signal.SIGINT)
+    pg_cat_send_signal(signal.SIGINT)
     time.sleep(1)
 
     # Check that any new queries fail after sigint since server should close with no active transactions
@@ -36,18 +108,18 @@ def test_shutdown_logic():
         # Fail if query execution succeeded
         raise Exception("Server not closed after sigint")
 
-    utils.cleanup_conn(conn, cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(conn, cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
     # NO ACTIVE QUERIES ADMIN SHUTDOWN COMMAND
 
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and begin transaction
-    conn, cur = utils.connect_db()
-    admin_conn, admin_cur = utils.connect_db(admin=True)
+    conn, cur = connect_db()
+    admin_conn, admin_cur = connect_db(admin=True)
 
     cur.execute("BEGIN;")
     cur.execute("SELECT 1;")
@@ -66,24 +138,24 @@ def test_shutdown_logic():
         # Fail if query execution succeeded
         raise Exception("Server not closed after sigint")
 
-    utils.cleanup_conn(conn, cur)
-    utils.cleanup_conn(admin_conn, admin_cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(conn, cur)
+    cleanup_conn(admin_conn, admin_cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
     # HANDLE TRANSACTION WITH SIGINT
 
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and begin transaction
-    conn, cur = utils.connect_db()
+    conn, cur = connect_db()
 
     cur.execute("BEGIN;")
     cur.execute("SELECT 1;")
 
     # Send sigint to pgcat while still in transaction
-    utils.pg_cat_send_signal(signal.SIGINT)
+    pg_cat_send_signal(signal.SIGINT)
     time.sleep(1)
 
     # Check that any new queries succeed after sigint since server should still allow transaction to complete
@@ -93,18 +165,18 @@ def test_shutdown_logic():
         # Fail if query fails since server closed
         raise Exception("Server closed while in transaction", e.pgerror)
 
-    utils.cleanup_conn(conn, cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(conn, cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
     # HANDLE TRANSACTION WITH ADMIN SHUTDOWN COMMAND
 
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and begin transaction
-    conn, cur = utils.connect_db()
-    admin_conn, admin_cur = utils.connect_db(admin=True)
+    conn, cur = connect_db()
+    admin_conn, admin_cur = connect_db(admin=True)
 
     cur.execute("BEGIN;")
     cur.execute("SELECT 1;")
@@ -122,30 +194,30 @@ def test_shutdown_logic():
         # Fail if query fails since server closed
         raise Exception("Server closed while in transaction", e.pgerror)
 
-    utils.cleanup_conn(conn, cur)
-    utils.cleanup_conn(admin_conn, admin_cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(conn, cur)
+    cleanup_conn(admin_conn, admin_cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
     # NO NEW NON-ADMIN CONNECTIONS DURING SHUTDOWN
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and begin transaction
-    transaction_conn, transaction_cur = utils.connect_db()
+    transaction_conn, transaction_cur = connect_db()
 
     transaction_cur.execute("BEGIN;")
     transaction_cur.execute("SELECT 1;")
 
     # Send sigint to pgcat while still in transaction
-    utils.pg_cat_send_signal(signal.SIGINT)
+    pg_cat_send_signal(signal.SIGINT)
     time.sleep(1)
 
     start = time.perf_counter()
     try:
-        conn, cur = utils.connect_db()
+        conn, cur = connect_db()
         cur.execute("SELECT 1;")
-        utils.cleanup_conn(conn, cur)
+        cleanup_conn(conn, cur)
     except psycopg2.OperationalError as e:
         time_taken = time.perf_counter() - start
         if time_taken > 0.1:
@@ -155,49 +227,49 @@ def test_shutdown_logic():
     else:
         raise Exception("Able connect to database during shutdown")
 
-    utils.cleanup_conn(transaction_conn, transaction_cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(transaction_conn, transaction_cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
     # ALLOW NEW ADMIN CONNECTIONS DURING SHUTDOWN
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and begin transaction
-    transaction_conn, transaction_cur = utils.connect_db()
+    transaction_conn, transaction_cur = connect_db()
 
     transaction_cur.execute("BEGIN;")
     transaction_cur.execute("SELECT 1;")
 
     # Send sigint to pgcat while still in transaction
-    utils.pg_cat_send_signal(signal.SIGINT)
+    pg_cat_send_signal(signal.SIGINT)
     time.sleep(1)
 
     try:
-        conn, cur = utils.connect_db(admin=True)
+        conn, cur = connect_db(admin=True)
         cur.execute("SHOW DATABASES;")
-        utils.cleanup_conn(conn, cur)
+        cleanup_conn(conn, cur)
     except psycopg2.OperationalError as e:
         raise Exception(e)
 
-    utils.cleanup_conn(transaction_conn, transaction_cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(transaction_conn, transaction_cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
     # ADMIN CONNECTIONS CONTINUING TO WORK AFTER SHUTDOWN
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and begin transaction
-    transaction_conn, transaction_cur = utils.connect_db()
+    transaction_conn, transaction_cur = connect_db()
     transaction_cur.execute("BEGIN;")
     transaction_cur.execute("SELECT 1;")
 
-    admin_conn, admin_cur = utils.connect_db(admin=True)
+    admin_conn, admin_cur = connect_db(admin=True)
     admin_cur.execute("SHOW DATABASES;")
 
     # Send sigint to pgcat while still in transaction
-    utils.pg_cat_send_signal(signal.SIGINT)
+    pg_cat_send_signal(signal.SIGINT)
     time.sleep(1)
 
     try:
@@ -205,24 +277,24 @@ def test_shutdown_logic():
     except psycopg2.OperationalError as e:
         raise Exception("Could not execute admin command:", e)
 
-    utils.cleanup_conn(transaction_conn, transaction_cur)
-    utils.cleanup_conn(admin_conn, admin_cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(transaction_conn, transaction_cur)
+    cleanup_conn(admin_conn, admin_cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
     # HANDLE SHUTDOWN TIMEOUT WITH SIGINT
 
     # Start pgcat
-    utils.pgcat_start()
+    pgcat_start()
 
     # Create client connection and begin transaction, which should prevent server shutdown unless shutdown timeout is reached
-    conn, cur = utils.connect_db()
+    conn, cur = connect_db()
 
     cur.execute("BEGIN;")
     cur.execute("SELECT 1;")
 
     # Send sigint to pgcat while still in transaction
-    utils.pg_cat_send_signal(signal.SIGINT)
+    pg_cat_send_signal(signal.SIGINT)
 
     # pgcat shutdown timeout is set to SHUTDOWN_TIMEOUT seconds, so we sleep for SHUTDOWN_TIMEOUT + 1 seconds
     time.sleep(SHUTDOWN_TIMEOUT + 1)
@@ -236,7 +308,12 @@ def test_shutdown_logic():
         # Fail if query execution succeeded
         raise Exception("Server not closed after sigint and expected timeout")
 
-    utils.cleanup_conn(conn, cur)
-    utils.pg_cat_send_signal(signal.SIGTERM)
+    cleanup_conn(conn, cur)
+    pg_cat_send_signal(signal.SIGTERM)
 
     # - - - - - - - - - - - - - - - - - -
+
+
+test_normal_db_access()
+test_admin_db_access()
+test_shutdown_logic()

tests/python/utils.py

@@ -1,110 +0,0 @@
-import os
-import signal
-import time
-from typing import Tuple
-import tempfile
-
-import psutil
-import psycopg2
-
-PGCAT_HOST = "127.0.0.1"
-PGCAT_PORT = "6432"
-
-
-def _pgcat_start(config_path: str):
-    pg_cat_send_signal(signal.SIGTERM)
-    os.system(f"./target/debug/pgcat {config_path} &")
-    time.sleep(2)
-
-
-def pgcat_start():
-    _pgcat_start(config_path='.circleci/pgcat.toml')
-
-
-def pgcat_generic_start(config: str):
-    tmp = tempfile.NamedTemporaryFile()
-    with open(tmp.name, 'w') as f:
-        f.write(config)
-    _pgcat_start(config_path=tmp.name)
-
-
-def glauth_send_signal(signal: signal.Signals):
-    try:
-        for proc in psutil.process_iter(["pid", "name"]):
-            if proc.name() == "glauth":
-                os.kill(proc.pid, signal)
-    except Exception as e:
-        # The process can be gone when we send this signal
-        print(e)
-
-    if signal == signal.SIGTERM:
-        # Returns 0 if pgcat process exists
-        time.sleep(2)
-        if not os.system('pgrep glauth'):
-            raise Exception("glauth not closed after SIGTERM")
-
-
-def pg_cat_send_signal(signal: signal.Signals):
-    try:
-        for proc in psutil.process_iter(["pid", "name"]):
-            if "pgcat" == proc.name():
-                os.kill(proc.pid, signal)
-    except Exception as e:
-        # The process can be gone when we send this signal
-        print(e)
-
-    if signal == signal.SIGTERM:
-        # Returns 0 if pgcat process exists
-        time.sleep(2)
-        if not os.system('pgrep pgcat'):
-            raise Exception("pgcat not closed after SIGTERM")
-
-
-def connect_db(
-    autocommit: bool = True,
-    admin: bool = False,
-) -> Tuple[psycopg2.extensions.connection, psycopg2.extensions.cursor]:
-
-    if admin:
-        user = "admin_user"
-        password = "admin_pass"
-        db = "pgcat"
-    else:
-        user = "sharding_user"
-        password = "sharding_user"
-        db = "sharded_db"
-
-    conn = psycopg2.connect(
-        f"postgres://{user}:{password}@{PGCAT_HOST}:{PGCAT_PORT}/{db}?application_name=testing_pgcat",
-        connect_timeout=2,
-    )
-    conn.autocommit = autocommit
-    cur = conn.cursor()
-
-    return (conn, cur)
-
-def connect_db_trust(
-    autocommit: bool = True,
-    admin: bool = False,
-) -> Tuple[psycopg2.extensions.connection, psycopg2.extensions.cursor]:
-
-    if admin:
-        user = "admin_user"
-        db = "pgcat"
-    else:
-        user = "sharding_user"
-        db = "sharded_db"
-
-    conn = psycopg2.connect(
-        f"postgres://{user}@{PGCAT_HOST}:{PGCAT_PORT}/{db}?application_name=testing_pgcat",
-        connect_timeout=2,
-    )
-    conn.autocommit = autocommit
-    cur = conn.cursor()
-
-    return (conn, cur)
-
-
-def cleanup_conn(conn: psycopg2.extensions.connection, cur: psycopg2.extensions.cursor):
-    cur.close()
-    conn.close()

tests/ruby/Gemfile.lock

@@ -1,33 +1,22 @@
 GEM
   remote: https://rubygems.org/
   specs:
-    activemodel (7.1.4)
-      activesupport (= 7.1.4)
-    activerecord (7.1.4)
-      activemodel (= 7.1.4)
-      activesupport (= 7.1.4)
-      timeout (>= 0.4.0)
-    activesupport (7.1.4)
-      base64
-      bigdecimal
+    activemodel (7.0.4.1)
+      activesupport (= 7.0.4.1)
+    activerecord (7.0.4.1)
+      activemodel (= 7.0.4.1)
+      activesupport (= 7.0.4.1)
+    activesupport (7.0.4.1)
       concurrent-ruby (~> 1.0, >= 1.0.2)
-      connection_pool (>= 2.2.5)
-      drb
       i18n (>= 1.6, < 2)
       minitest (>= 5.1)
-      mutex_m
       tzinfo (~> 2.0)
     ast (2.4.2)
-    base64 (0.2.0)
-    bigdecimal (3.1.8)
-    concurrent-ruby (1.3.4)
-    connection_pool (2.4.1)
+    concurrent-ruby (1.1.10)
     diff-lcs (1.5.0)
-    drb (2.2.1)
-    i18n (1.14.5)
+    i18n (1.12.0)
       concurrent-ruby (~> 1.0)
-    minitest (5.25.1)
-    mutex_m (0.2.0)
+    minitest (5.17.0)
     parallel (1.22.1)
     parser (3.1.2.0)
       ast (~> 2.4.1)
@@ -35,8 +24,7 @@ GEM
     pg (1.3.2)
     rainbow (3.1.1)
     regexp_parser (2.3.1)
-    rexml (3.3.6)
-      strscan
+    rexml (3.2.5)
     rspec (3.11.0)
       rspec-core (~> 3.11.0)
       rspec-expectations (~> 3.11.0)
@@ -62,12 +50,10 @@ GEM
     rubocop-ast (1.17.0)
       parser (>= 3.1.1.0)
     ruby-progressbar (1.11.0)
-    strscan (3.1.0)
-    timeout (0.4.1)
     toml (0.3.0)
       parslet (>= 1.8.0, < 3.0.0)
     toxiproxy (2.0.1)
-    tzinfo (2.0.6)
+    tzinfo (2.0.5)
       concurrent-ruby (~> 1.0)
     unicode-display_width (2.1.0)
 

tests/ruby/admin_spec.rb

@@ -91,27 +91,6 @@ describe "Admin" do
     end
   end
 
-  [
-    "SHOW ME THE MONEY", 
-    "SHOW ME THE WAY", 
-    "SHOW UP", 
-    "SHOWTIME",
-    "HAMMER TIME",
-    "SHOWN TO BE TRUE",
-    "SHOW ",
-    "SHOW     ",
-    "SHOW 1",
-    ";;;;;"
-  ].each do |cmd|
-    describe "Bad command #{cmd}" do
-      it "does not panic and responds with PG::SystemError" do
-        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
-        expect { admin_conn.async_exec(cmd) }.to raise_error(PG::SystemError).with_message(/Unsupported/)
-        admin_conn.close
-      end
-    end
-  end
-
   describe "PAUSE" do
     it "pauses all pools" do
       admin_conn = PG::connect(processes.pgcat.admin_connection_string)

tests/ruby/load_balancing_spec.rb

@@ -56,41 +56,6 @@ describe "Random Load Balancing" do
       end
     end
   end
-
-  context "when all replicas are down " do
-    let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5, "transaction", "random", "debug", {"default_role" => "replica"}) }
-
-    it "unbans them automatically to prevent false positives in health checks that could make all replicas unavailable" do
-      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
-      failed_count = 0
-      number_of_replicas = processes[:replicas].length
-
-      # Take down all replicas
-      processes[:replicas].each(&:take_down)
-
-      (number_of_replicas + 1).times do |n|
-        conn.async_exec("SELECT 1 + 2")
-      rescue
-        conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
-        failed_count += 1
-      end
-
-      expect(failed_count).to eq(number_of_replicas + 1)
-      failed_count = 0
-
-      # Ban_time is configured to 60 so this reset will only work
-      # if the replicas are unbanned automatically
-      processes[:replicas].each(&:reset)
-
-      number_of_replicas.times do
-        conn.async_exec("SELECT 1 + 2")
-      rescue
-        conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
-        failed_count += 1
-      end
-      expect(failed_count).to eq(0)
-    end
-  end
 end
 
 describe "Least Outstanding Queries Load Balancing" do
@@ -196,3 +161,4 @@ describe "Least Outstanding Queries Load Balancing" do
     end
   end
 end
+

tests/ruby/misc_spec.rb

@@ -188,102 +188,6 @@ describe "Miscellaneous" do
     end
   end
 
-  describe "Checkout failure limit" do
-    context "when no checkout failure limit is set" do
-      before do 
-        new_configs = processes.pgcat.current_config
-        new_configs["general"]["connect_timeout"] = 200
-        new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
-        processes.pgcat.update_config(new_configs)
-        processes.pgcat.reload_config
-        sleep 0.5
-      end
-  
-      it "does not disconnect client" do
-        Array.new(5) do
-          Thread.new do
-            conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
-            for i in 0..4
-              begin
-                conn.async_exec("SELECT pg_sleep(0.5);")
-                expect(conn.status).to eq(PG::CONNECTION_OK)
-              rescue PG::SystemError
-                expect(conn.status).to eq(PG::CONNECTION_OK)
-              end
-            end
-            conn.close
-          end
-        end.each(&:join)
-      end
-    end
-
-    context "when checkout failure limit is set high" do
-      before do 
-        new_configs = processes.pgcat.current_config
-        new_configs["general"]["connect_timeout"] = 200
-        new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
-        new_configs["pools"]["sharded_db"]["checkout_failure_limit"] = 10000
-        processes.pgcat.update_config(new_configs)
-        processes.pgcat.reload_config
-        sleep 0.5
-      end
-  
-      it "does not disconnect client" do
-        Array.new(5) do
-          Thread.new do
-            conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
-            for i in 0..4
-              begin
-                conn.async_exec("SELECT pg_sleep(0.5);")
-                expect(conn.status).to eq(PG::CONNECTION_OK)
-              rescue PG::SystemError
-                expect(conn.status).to eq(PG::CONNECTION_OK)
-              end
-            end
-            conn.close
-          end
-        end.each(&:join)
-      end
-    end
-
-    context "when checkout failure limit is set low" do
-      before do 
-        new_configs = processes.pgcat.current_config
-        new_configs["general"]["connect_timeout"] = 200
-        new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
-        new_configs["pools"]["sharded_db"]["checkout_failure_limit"] = 2
-        processes.pgcat.update_config(new_configs)
-        processes.pgcat.reload_config
-        sleep 0.5
-      end
-  
-      it "disconnects client after reaching limit" do
-        Array.new(5) do
-          Thread.new do
-            conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
-            checkout_failure_count = 0
-            for i in 0..4
-              begin
-                conn.async_exec("SELECT pg_sleep(1);")
-                expect(conn.status).to eq(PG::CONNECTION_OK)
-              rescue PG::SystemError
-                checkout_failure_count += 1
-                expect(conn.status).to eq(PG::CONNECTION_OK)
-              rescue PG::ConnectionBad
-                expect(checkout_failure_count).to eq(2)
-                expect(conn.status).to eq(PG::CONNECTION_BAD)
-                break
-              end
-            end
-            conn.close
-          end
-        end.each(&:join)
-        puts processes.pgcat.logs
-
-      end
-    end
-  end
-  
   describe "Server version reporting" do
     it "reports correct version for normal and admin databases" do
       server_conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))

tests/rust/src/main.rs

@@ -15,11 +15,13 @@ async fn test_prepared_statements() {
     for _ in 0..5 {
         let pool = pool.clone();
         let handle = tokio::task::spawn(async move {
-            for i in 0..1000 {
-                match sqlx::query(&format!("SELECT {:?}", i % 5)).fetch_all(&pool).await {
+            for _ in 0..1000 {
+                match sqlx::query("SELECT one").fetch_all(&pool).await {
                     Ok(_) => (),
                     Err(err) => {
-                        panic!("prepared statement error: {}", err);
+                        if err.to_string().contains("prepared statement") {
+                            panic!("prepared statement error: {}", err);
+                        }
                     }
                 }
             }