Merge branch 'main' of github.com:postgresml/pgcat into mostafa_fix_prepared_stmts

Make iterating on integration tests easier (#789 )
Writing and iterating on integration tests are cumbersome, having to wait 10 minutes for the test-suite to run just to see if your test works or not is unacceptable. In this PR, I added a detailed workflow for writing tests that should shorten the feedback cycle of modifying tests to be as low as a few seconds. It will involve opening a shell into a long-lived container that has all the setup and dependencies necessary and then running your desired tests directly there. I added a convenience script that bootstraps the environment and then opens an interactive shell into the container and you can then run tests immediately in an environment that is more or less identical to what we have running in CircleCI
2026-03-23 17:36:28 +00:00 · 2024-09-03 18:11:32 -05:00 · 2024-09-03 11:15:53 -05:00 · 2024-09-01 16:14:44 -05:00 · 2024-09-01 15:41:33 -05:00 · 2024-09-01 15:31:27 -05:00
27 changed files with 3302 additions and 585 deletions
--- a/.circleci/pgcat.toml
+++ b/.circleci/pgcat.toml
@@ -59,6 +59,7 @@ 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.
@@ -141,6 +142,7 @@ 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"
--- a/.github/workflows/build-and-push.yaml
+++ b/.github/workflows/build-and-push.yaml
@@ -23,14 +23,17 @@ jobs:

    steps:
      - name: Checkout Repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3

      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v2
+        uses: docker/setup-buildx-action@v3

      - name: Determine tags
        id: metadata
-        uses: docker/metadata-action@v4
+        uses: docker/metadata-action@v5
        with:
          images: ${{ env.registry }}/${{ env.image-name }}
          tags: |
@@ -42,15 +45,18 @@ jobs:
            type=raw,value=latest,enable={{ is_default_branch }}

      - name: Log in to the Container registry
-        uses: docker/login-action@v2.1.0
+        uses: docker/login-action@v3
        with:
          registry: ${{ env.registry }}
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}

      - name: Build and push ${{ env.image-name }}
-        uses: docker/build-push-action@v3
+        uses: docker/build-push-action@v6
        with:
+          context: .
+          platforms: linux/amd64,linux/arm64
+          provenance: false
          push: true
          tags: ${{ steps.metadata.outputs.tags }}
          labels: ${{ steps.metadata.outputs.labels }}
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -6,6 +6,32 @@ 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 && python3 tests/python/tests.py`
+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
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -146,6 +146,12 @@ dependencies = [
 "syn 2.0.26",
 ]

+[[package]]
+name = "atomic-waker"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
+
 [[package]]
 name = "atomic_enum"
 version = "0.2.0"
@@ -542,29 +548,23 @@ checksum = "b6c80984affa11d98d1b88b66ac8853f143217b399d3c74116778ff8fdb4ed2e"

 [[package]]
 name = "h2"
-version = "0.3.20"
+version = "0.4.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "97ec8491ebaf99c8eaa73058b045fe58073cd6be7f596ac993ced0b0a0c01049"
+checksum = "524e8ac6999421f49a846c2d4411f337e53497d8ec55d67753beffa43c5d9205"
 dependencies = [
+ "atomic-waker",
 "bytes",
 "fnv",
 "futures-core",
 "futures-sink",
- "futures-util",
 "http",
- "indexmap 1.9.3",
+ "indexmap",
 "slab",
 "tokio",
 "tokio-util",
 "tracing",
 ]

-[[package]]
-name = "hashbrown"
-version = "0.12.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
-
 [[package]]
 name = "hashbrown"
 version = "0.14.0"
@@ -609,9 +609,9 @@ dependencies = [

 [[package]]
 name = "http"
-version = "0.2.9"
+version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bd6effc99afb63425aff9b05836f029929e345a6148a14b7ecd5ab67af944482"
+checksum = "21b9ddb458710bc376481b842f5da65cdf31522de232c1ca8146abce2a358258"
 dependencies = [
 "bytes",
 "fnv",
@@ -620,12 +620,24 @@ dependencies = [

 [[package]]
 name = "http-body"
-version = "0.4.5"
+version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d5f38f16d184e36f2408a55281cd658ecbd3ca05cce6d6510a176eca393e26d1"
+checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184"
 dependencies = [
 "bytes",
 "http",
+]
+
+[[package]]
+name = "http-body-util"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "793429d76616a256bcb62c2a2ec2bed781c8307e797e2598c50010f2bee2544f"
+dependencies = [
+ "bytes",
+ "futures-util",
+ "http",
+ "http-body",
 "pin-project-lite",
 ]

@@ -643,13 +655,12 @@ checksum = "c4a1e36c821dbe04574f602848a19f742f4fb3c98d40449f11bcad18d6b17421"

 [[package]]
 name = "hyper"
-version = "0.14.27"
+version = "1.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ffb1cfd654a8219eaef89881fdb3bb3b1cdc5fa75ded05d6933b2b382e395468"
+checksum = "50dfd22e0e76d0f662d429a5f80fcaf3855009297eab6a0a9f8543834744ba05"
 dependencies = [
 "bytes",
 "futures-channel",
- "futures-core",
 "futures-util",
 "h2",
 "http",
@@ -658,13 +669,26 @@ dependencies = [
 "httpdate",
 "itoa",
 "pin-project-lite",
- "socket2 0.4.9",
+ "smallvec",
 "tokio",
- "tower-service",
- "tracing",
 "want",
 ]

+[[package]]
+name = "hyper-util"
+version = "0.1.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cde7055719c54e36e95e8719f95883f22072a48ede39db7fc17a4e1d5281e9b9"
+dependencies = [
+ "bytes",
+ "futures-util",
+ "http",
+ "http-body",
+ "hyper",
+ "pin-project-lite",
+ "tokio",
+]
+
 [[package]]
 name = "iana-time-zone"
 version = "0.1.57"
@@ -709,16 +733,6 @@ dependencies = [
 "unicode-normalization",
 ]

-[[package]]
-name = "indexmap"
-version = "1.9.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
-dependencies = [
- "autocfg",
- "hashbrown 0.12.3",
-]
-
 [[package]]
 name = "indexmap"
 version = "2.0.0"
@@ -726,7 +740,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d5477fe2230a79769d8dc68e0eabf5437907c0457a5614a9e8dddb67f65eb65d"
 dependencies = [
 "equivalent",
- "hashbrown 0.14.0",
+ "hashbrown",
 ]

 [[package]]
@@ -848,7 +862,7 @@ version = "0.12.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1efa59af2ddfad1854ae27d75009d538d0998b4b2fd47083e743ac1a10e46c60"
 dependencies = [
- "hashbrown 0.14.0",
+ "hashbrown",
 ]

 [[package]]
@@ -1020,7 +1034,7 @@ checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94"

 [[package]]
 name = "pgcat"
-version = "1.1.2-dev4"
+version = "1.2.0"
 dependencies = [
 "arc-swap",
 "async-trait",
@@ -1034,7 +1048,9 @@ dependencies = [
 "fallible-iterator",
 "futures",
 "hmac",
+ "http-body-util",
 "hyper",
+ "hyper-util",
 "itertools",
 "jemallocator",
 "log",
@@ -1478,9 +1494,9 @@ dependencies = [

 [[package]]
 name = "smallvec"
-version = "1.11.0"
+version = "1.13.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "62bb4feee49fdd9f707ef802e22365a35de4b7b299de4763d44bfea899442ff9"
+checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67"

 [[package]]
 name = "socket2"
@@ -1741,19 +1757,13 @@ version = "0.19.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f8123f27e969974a3dfba720fdb560be359f57b44302d280ba72e76a74480e8a"
 dependencies = [
- "indexmap 2.0.0",
+ "indexmap",
 "serde",
 "serde_spanned",
 "toml_datetime",
 "winnow",
 ]

-[[package]]
-name = "tower-service"
-version = "0.3.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b6bc1c9ce2b5135ac7f93c72918fc37feb872bdc6a5533a8b85eb4b86bfdae52"
-
 [[package]]
 name = "tracing"
 version = "0.1.37"
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pgcat"
-version = "1.1.2-dev4"
+version = "1.2.0"
 edition = "2021"

 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -29,7 +29,9 @@ base64 = "0.21"
 stringprep = "0.1"
 tokio-rustls = "0.24"
 rustls-pemfile = "1"
-hyper = { version = "0.14", features = ["full"] }
+http-body-util = "0.1.2"
+hyper = { version = "1.4.1", features = ["full"] }
+hyper-util = { version = "0.1.7", features = ["tokio"] }
 phf = { version = "0.11.1", features = ["macros"] }
 exitcode = "1.1.2"
 futures = "0.3"
--- a/Dockerfile.ci
+++ b/Dockerfile.ci
@@ -1,4 +1,4 @@
-FROM cimg/rust:1.67.1
+FROM cimg/rust:1.79.0
 COPY --from=sclevine/yj /bin/yj /bin/yj
 RUN /bin/yj -h
 RUN sudo apt-get update && \
--- a/README.md
+++ b/README.md
@@ -268,6 +268,8 @@ 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.
--- a/charts/pgcat/Chart.yaml
+++ b/charts/pgcat/Chart.yaml
@@ -4,5 +4,5 @@ description: A Helm chart for PgCat a PostgreSQL pooler and proxy (like PgBounce
 maintainers:
  - name: Wildcard
    email: support@w6d.io
-appVersion: "1.1.1"
-version: 0.1.0
+appVersion: "1.2.0"
+version: 0.2.0
--- a/charts/pgcat/values.yaml
+++ b/charts/pgcat/values.yaml
@@ -170,13 +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
+    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
@@ -240,7 +240,15 @@ 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: []
+  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"
+      }]
+    }]
    # - ## default values
    #   ##
    #   ##
--- a/grafana_dashboard.json
+++ b/grafana_dashboard.json
--- a/src/admin.rs
+++ b/src/admin.rs
@@ -55,7 +55,12 @@ where

    let query_parts: Vec<&str> = query.trim_end_matches(';').split_whitespace().collect();

-    match query_parts[0].to_ascii_uppercase().as_str() {
+    match query_parts
+        .first()
+        .unwrap_or(&"")
+        .to_ascii_uppercase()
+        .as_str()
+    {
        "BAN" => {
            trace!("BAN");
            ban(stream, query_parts).await
@@ -84,7 +89,12 @@ where
            trace!("SHUTDOWN");
            shutdown(stream).await
        }
-        "SHOW" => match query_parts[1].to_ascii_uppercase().as_str() {
+        "SHOW" => match query_parts
+            .get(1)
+            .unwrap_or(&"")
+            .to_ascii_uppercase()
+            .as_str()
+        {
            "HELP" => {
                trace!("SHOW HELP");
                show_help(stream).await
--- a/src/client.rs
+++ b/src/client.rs
@@ -1204,12 +1204,9 @@ where
                        if !server.in_transaction() {
                            // Report transaction executed statistics.
                            self.stats.transaction();
-                            server.stats().transaction(
-                                Instant::now()
-                                    .duration_since(server.transaction_start().into())
-                                    .as_millis() as u64,
-                                self.server_parameters.get_application_name(),
-                            );
+                            server
+                                .stats()
+                                .transaction(self.server_parameters.get_application_name());

                            // Release server back to the pool if we are in transaction mode.
                            // If we are in session mode, we keep the server until the client disconnects.
@@ -1463,12 +1460,9 @@ where

                        if !server.in_transaction() {
                            self.stats.transaction();
-                            server.stats().transaction(
-                                Instant::now()
-                                    .duration_since(server.transaction_start().into())
-                                    .as_millis() as u64,
-                                self.server_parameters.get_application_name(),
-                            );
+                            server
+                                .stats()
+                                .transaction(self.server_parameters.get_application_name());

                            // Release server back to the pool if we are in transaction mode.
                            // If we are in session mode, we keep the server until the client disconnects.
@@ -1517,12 +1511,9 @@ where

                        if !server.in_transaction() {
                            self.stats.transaction();
-                            server.stats().transaction(
-                                Instant::now()
-                                    .duration_since(server.transaction_start().into())
-                                    .as_millis() as u64,
-                                self.server_parameters.get_application_name(),
-                            );
+                            server
+                                .stats()
+                                .transaction(self.server_parameters.get_application_name());

                            // Release server back to the pool if we are in transaction mode.
                            // If we are in session mode, we keep the server until the client disconnects.
@@ -1738,14 +1729,13 @@ where
    /// and also the pool's statement cache. Add it to extended protocol data.
    fn buffer_parse(&mut self, message: BytesMut, pool: &ConnectionPool) -> Result<(), Error> {
        // Avoid parsing if prepared statements not enabled
-        if !self.prepared_statements_enabled {
+        let client_given_name = Parse::get_name(&message)?;
+        if !self.prepared_statements_enabled || client_given_name.is_empty() {
            debug!("Anonymous parse message");
            self.extended_protocol_data_buffer
                .push_back(ExtendedProtocolData::create_new_parse(message, None));
            return Ok(());
        }
-
-        let client_given_name = Parse::get_name(&message)?;
        let parse: Parse = (&message).try_into()?;

        // Compute the hash of the parse statement
@@ -1783,15 +1773,14 @@ where
    /// saved in the client cache.
    async fn buffer_bind(&mut self, message: BytesMut) -> Result<(), Error> {
        // Avoid parsing if prepared statements not enabled
-        if !self.prepared_statements_enabled {
+        let client_given_name = Bind::get_name(&message)?;
+        if !self.prepared_statements_enabled || client_given_name.is_empty() {
            debug!("Anonymous bind message");
            self.extended_protocol_data_buffer
                .push_back(ExtendedProtocolData::create_new_bind(message, None));
            return Ok(());
        }

-        let client_given_name = Bind::get_name(&message)?;
-
        match self.prepared_statements.get(&client_given_name) {
            Some((rewritten_parse, _)) => {
                let message = Bind::rename(message, &rewritten_parse.name)?;
@@ -1843,7 +1832,8 @@ where
        }

        let describe: Describe = (&message).try_into()?;
-        if describe.target == 'P' {
+        let client_given_name = describe.statement_name.clone();
+        if describe.target == 'P' || client_given_name.is_empty() {
            debug!("Portal describe message");
            self.extended_protocol_data_buffer
                .push_back(ExtendedProtocolData::create_new_describe(message, None));
@@ -1851,8 +1841,6 @@ where
            return Ok(());
        }

-        let client_given_name = describe.statement_name.clone();
-
        match self.prepared_statements.get(&client_given_name) {
            Some((rewritten_parse, _)) => {
                let describe = describe.rename(&rewritten_parse.name);
--- a/src/messages.rs
+++ b/src/messages.rs
@@ -821,10 +821,10 @@ impl ExtendedProtocolData {
 pub struct Parse {
    code: char,
    #[allow(dead_code)]
-    len: i32,
+    len: u32,
    pub name: String,
    query: String,
-    num_params: i16,
+    num_params: u16,
    param_types: Vec<i32>,
 }

@@ -834,12 +834,11 @@ impl TryFrom<&BytesMut> for Parse {
    fn try_from(buf: &BytesMut) -> Result<Parse, Error> {
        let mut cursor = Cursor::new(buf);
        let code = cursor.get_u8() as char;
-        let len = cursor.get_i32();
+        let len = cursor.get_u32();
        let name = cursor.read_string()?;
        let query = cursor.read_string()?;
-        let num_params = cursor.get_i16();
+        let num_params = cursor.get_u16();
        let mut param_types = Vec::new();
-
        for _ in 0..num_params {
            param_types.push(cursor.get_i32());
        }
@@ -875,10 +874,10 @@ impl TryFrom<Parse> for BytesMut {
            + 4 * parse.num_params as usize;

        bytes.put_u8(parse.code as u8);
-        bytes.put_i32(len as i32);
+        bytes.put_u32(len as u32);
        bytes.put_slice(name);
        bytes.put_slice(query);
-        bytes.put_i16(parse.num_params);
+        bytes.put_u16(parse.num_params);
        for param in parse.param_types {
            bytes.put_i32(param);
        }
@@ -945,14 +944,14 @@ impl Parse {
 pub struct Bind {
    code: char,
    #[allow(dead_code)]
-    len: i64,
+    len: u64,
    portal: String,
    pub prepared_statement: String,
-    num_param_format_codes: i16,
+    num_param_format_codes: u16,
    param_format_codes: Vec<i16>,
-    num_param_values: i16,
+    num_param_values: u16,
    param_values: Vec<(i32, BytesMut)>,
-    num_result_column_format_codes: i16,
+    num_result_column_format_codes: u16,
    result_columns_format_codes: Vec<i16>,
 }

@@ -962,17 +961,17 @@ impl TryFrom<&BytesMut> for Bind {
    fn try_from(buf: &BytesMut) -> Result<Bind, Error> {
        let mut cursor = Cursor::new(buf);
        let code = cursor.get_u8() as char;
-        let len = cursor.get_i32();
+        let len = cursor.get_u32();
        let portal = cursor.read_string()?;
        let prepared_statement = cursor.read_string()?;
-        let num_param_format_codes = cursor.get_i16();
+        let num_param_format_codes = cursor.get_u16();
        let mut param_format_codes = Vec::new();

        for _ in 0..num_param_format_codes {
            param_format_codes.push(cursor.get_i16());
        }

-        let num_param_values = cursor.get_i16();
+        let num_param_values = cursor.get_u16();
        let mut param_values = Vec::new();

        for _ in 0..num_param_values {
@@ -994,7 +993,7 @@ impl TryFrom<&BytesMut> for Bind {
            }
        }

-        let num_result_column_format_codes = cursor.get_i16();
+        let num_result_column_format_codes = cursor.get_u16();
        let mut result_columns_format_codes = Vec::new();

        for _ in 0..num_result_column_format_codes {
@@ -1003,7 +1002,7 @@ impl TryFrom<&BytesMut> for Bind {

        Ok(Bind {
            code,
-            len: len as i64,
+            len: len as u64,
            portal,
            prepared_statement,
            num_param_format_codes,
@@ -1042,19 +1041,19 @@ impl TryFrom<Bind> for BytesMut {
        len += 2 * bind.num_result_column_format_codes as usize;

        bytes.put_u8(bind.code as u8);
-        bytes.put_i32(len as i32);
+        bytes.put_u32(len as u32);
        bytes.put_slice(portal);
        bytes.put_slice(prepared_statement);
-        bytes.put_i16(bind.num_param_format_codes);
+        bytes.put_u16(bind.num_param_format_codes);
        for param_format_code in bind.param_format_codes {
            bytes.put_i16(param_format_code);
        }
-        bytes.put_i16(bind.num_param_values);
+        bytes.put_u16(bind.num_param_values);
        for (param_len, param) in bind.param_values {
            bytes.put_i32(param_len);
            bytes.put_slice(&param);
        }
-        bytes.put_i16(bind.num_result_column_format_codes);
+        bytes.put_u16(bind.num_result_column_format_codes);
        for result_column_format_code in bind.result_columns_format_codes {
            bytes.put_i16(result_column_format_code);
        }
@@ -1068,7 +1067,7 @@ impl Bind {
    pub fn get_name(buf: &BytesMut) -> Result<String, Error> {
        let mut cursor = Cursor::new(buf);
        // Skip the code and length
-        cursor.advance(mem::size_of::<u8>() + mem::size_of::<i32>());
+        cursor.advance(mem::size_of::<u8>() + mem::size_of::<u32>());
        cursor.read_string()?;
        cursor.read_string()
    }
@@ -1078,17 +1077,17 @@ impl Bind {
        let mut cursor = Cursor::new(&buf);
        // Read basic data from the cursor
        let code = cursor.get_u8();
-        let current_len = cursor.get_i32();
+        let current_len = cursor.get_u32();
        let portal = cursor.read_string()?;
        let prepared_statement = cursor.read_string()?;

        // Calculate new length
-        let new_len = current_len + new_name.len() as i32 - prepared_statement.len() as i32;
+        let new_len = current_len + new_name.len() as u32 - prepared_statement.len() as u32;

        // Begin building the response buffer
        let mut response_buf = BytesMut::with_capacity(new_len as usize + 1);
        response_buf.put_u8(code);
-        response_buf.put_i32(new_len);
+        response_buf.put_u32(new_len);

        // Put the portal and new name into the buffer
        // Note: panic if the provided string contains null byte
@@ -1112,7 +1111,7 @@ pub struct Describe {
    code: char,

    #[allow(dead_code)]
-    len: i32,
+    len: u32,
    pub target: char,
    pub statement_name: String,
 }
@@ -1123,7 +1122,7 @@ impl TryFrom<&BytesMut> for Describe {
    fn try_from(bytes: &BytesMut) -> Result<Describe, Error> {
        let mut cursor = Cursor::new(bytes);
        let code = cursor.get_u8() as char;
-        let len = cursor.get_i32();
+        let len = cursor.get_u32();
        let target = cursor.get_u8() as char;
        let statement_name = cursor.read_string()?;

@@ -1146,7 +1145,7 @@ impl TryFrom<Describe> for BytesMut {
        let len = 4 + 1 + statement_name.len();

        bytes.put_u8(describe.code as u8);
-        bytes.put_i32(len as i32);
+        bytes.put_u32(len as u32);
        bytes.put_u8(describe.target as u8);
        bytes.put_slice(statement_name);

--- a/src/prometheus.rs
+++ b/src/prometheus.rs
@@ -1,23 +1,41 @@
-use hyper::service::{make_service_fn, service_fn};
-use hyper::{Body, Method, Request, Response, Server, StatusCode};
+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 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 std::sync::Arc;
+use tokio::net::TcpListener;

 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
@@ -120,22 +138,46 @@ 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: "gauge",
+        ty: "counter",
    },
    "servers_bytes_sent" => MetricHelpType {
        help: "Volume in bytes of network traffic sent by server",
-        ty: "gauge",
+        ty: "counter",
    },
    "servers_transaction_count" => MetricHelpType {
        help: "Number of transactions executed by server",
-        ty: "gauge",
+        ty: "counter",
    },
    "servers_query_count" => MetricHelpType {
        help: "Number of queries executed by server",
-        ty: "gauge",
+        ty: "counter",
    },
    "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 {
@@ -203,7 +245,9 @@ 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("user", address.username.clone());

        Self::from_name(&format!("databases_{}", name), value, labels)
    }
@@ -218,8 +262,9 @@ 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("user", address.username.clone());
        Self::from_name(&format!("servers_{}", name), value, labels)
    }

@@ -229,7 +274,9 @@ 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("user", address.username.clone());

        Self::from_name(&format!("stats_{}", name), value, labels)
    }
@@ -243,7 +290,9 @@ impl<Value: fmt::Display> PrometheusMetric<Value> {
    }
 }

-async fn prometheus_stats(request: Request<Body>) -> Result<Response<Body>, hyper::http::Error> {
+async fn prometheus_stats(
+    request: Request<body::Incoming>,
+) -> Result<Response<Full<Bytes>>, hyper::http::Error> {
    match (request.method(), request.uri().path()) {
        (&Method::GET, "/metrics") => {
            let mut lines = Vec::new();
@@ -329,34 +378,51 @@ fn push_database_stats(lines: &mut Vec<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 server_stats_by_addresses = HashMap::<String, Arc<ServerStats>>::new();
+    let mut prom_stats = HashMap::<String, ServerPrometheusStats>::new();
    for (_, stats) in server_stats {
-        server_stats_by_addresses.insert(stats.address_name(), 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,
+        }
    }

    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) = server_stats_by_addresses.get(&address.name()) {
+                if let Some(server_info) = prom_stats.get(&address.name()) {
                    let metrics = [
-                        (
-                            "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),
-                        ),
+                        ("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 }),
                    ];
                    for (key, value) in metrics {
                        if let Some(prometheus_metric) =
@@ -374,14 +440,35 @@ fn push_server_stats(lines: &mut Vec<String>) {
 }

 pub async fn start_metric_server(http_addr: SocketAddr) {
-    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);
+    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;
+        }
+    };
    info!(
        "Exposing prometheus metrics on http://{}/metrics.",
        http_addr
    );
-    if let Err(e) = server.await {
-        error!("Failed to run HTTP server: {}.", e);
+    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);
+            }
+        });
    }
 }
--- a/src/query_router.rs
+++ b/src/query_router.rs
@@ -427,8 +427,12 @@ impl QueryRouter {
                        None => (),
                    };

-                    // If we already visited a write statement, we should be going to the primary.
-                    if !visited_write_statement {
+                    let has_locks = !query.locks.is_empty();
+
+                    if has_locks {
+                        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.
                        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.
@@ -1158,6 +1162,29 @@ 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();
@@ -1372,6 +1399,19 @@ 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();
--- a/src/server.rs
+++ b/src/server.rs
@@ -15,7 +15,6 @@ use std::sync::Arc;
 use std::time::SystemTime;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, BufStream};
 use tokio::net::TcpStream;
-use tokio::time::Instant;
 use tokio_rustls::rustls::{OwnedTrustAnchor, RootCertStore};
 use tokio_rustls::{client::TlsStream, TlsConnector};

@@ -286,9 +285,6 @@ pub struct Server {
    /// Is the server inside a transaction or idle.
    in_transaction: bool,

-    /// The time the most recent transaction started.
-    transaction_start: Instant,
-
    /// Is there more data for the client to read.
    data_available: bool,

@@ -702,7 +698,6 @@ impl Server {
                            ))
                        }
                    };
-
                    trace!("Error: {}", error_code);

                    match error_code {
@@ -808,7 +803,6 @@ impl Server {
                        process_id,
                        secret_key,
                        in_transaction: false,
-                        transaction_start: Instant::now(),
                        in_copy_mode: false,
                        data_available: false,
                        bad: false,
@@ -941,7 +935,6 @@ impl Server {
                        // In transaction.
                        'T' => {
                            self.in_transaction = true;
-                            self.transaction_start = Instant::now();
                        }

                        // Idle, transaction over.
@@ -1019,6 +1012,12 @@ impl Server {
                            // which can leak between clients. This is a best effort to block bad clients
                            // from poisoning a transaction-mode pool by setting inappropriate session variables
                            match command.as_str() {
+                                "DISCARD ALL" => {
+                                    self.clear_prepared_statement_cache();
+                                }
+                                "DEALLOCATE ALL" => {
+                                    self.clear_prepared_statement_cache();
+                                }
                                "SET" => {
                                    // We don't detect set statements in transactions
                                    // No great way to differentiate between set and set local
@@ -1138,6 +1137,12 @@ impl Server {
        has_it
    }

+    fn clear_prepared_statement_cache(&mut self) {
+        if let Some(cache) = &mut self.prepared_statement_cache {
+            cache.clear();
+        }
+    }
+
    fn add_prepared_statement_to_cache(&mut self, name: &str) -> Option<String> {
        let cache = match &mut self.prepared_statement_cache {
            Some(cache) => cache,
@@ -1226,12 +1231,6 @@ impl Server {
        self.in_transaction
    }

-    /// The start time of the most recent transaction.
-    /// Will be stale if not in a transaction.
-    pub fn transaction_start(&self) -> Instant {
-        self.transaction_start
-    }
-
    /// Currently copying data from client to server or vice-versa.
    pub fn in_copy_mode(&self) -> bool {
        self.in_copy_mode
--- a/src/stats/server.rs
+++ b/src/stats/server.rs
@@ -187,12 +187,11 @@ impl ServerStats {
    /// we report each individual queries outside a transaction as a transaction
    /// We only count the initial BEGIN as a transaction, all queries within do not
    /// count as transactions
-    pub fn transaction(&self, milliseconds: u64, application_name: &str) {
+    pub fn transaction(&self, application_name: &str) {
        self.set_application(application_name.to_string());

        self.transaction_count.fetch_add(1, Ordering::Relaxed);
        self.address.stats.xact_count_add();
-        self.address.stats.xact_time_add(milliseconds);
    }

    /// Report data sent to a server
--- a/start_test_env.sh
+++ b/start_test_env.sh
@@ -0,0 +1,34 @@
+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 && python3 tests/python/tests.py${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
--- a/tests/docker/docker-compose.yml
+++ b/tests/docker/docker-compose.yml
@@ -1,4 +1,3 @@
-version: "3"
 services:
  pg1:
    image: postgres:14
@@ -48,6 +47,8 @@ services:
  main:
    build: .
    command: ["bash", "/app/tests/docker/run.sh"]
+    environment:
+      - INTERACTIVE_TEST_ENVIRONMENT=true
    volumes:
      - ../../:/app/
      - /app/target/
--- a/tests/docker/run.sh
+++ b/tests/docker/run.sh
@@ -5,6 +5,38 @@ 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
--- a/tests/ruby/admin_spec.rb
+++ b/tests/ruby/admin_spec.rb
@@ -91,6 +91,27 @@ 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)
--- a/tests/ruby/helpers/frontend_messages.rb
+++ b/tests/ruby/helpers/frontend_messages.rb
@@ -0,0 +1,145 @@
+
+class PostgresMessage
+  # Base class for common functionality
+
+  def encode_string(str)
+    "#{str}\0"  # Encode a string with a null terminator
+  end
+
+  def encode_int16(value)
+    [value].pack('n')  # Encode an Int16
+  end
+
+  def encode_int32(value)
+    [value].pack('N')  # Encode an Int32
+  end
+
+  def message_prefix(type, length)
+    "#{type}#{encode_int32(length)}"  # Message type and length prefix
+  end
+end
+
+class SimpleQueryMessage < PostgresMessage
+  attr_accessor :query
+
+  def initialize(query = "")
+    @query = query
+  end
+
+  def to_bytes
+    query_bytes = encode_string(@query)
+    length = 4 + query_bytes.size  # Length includes 4 bytes for length itself
+    message_prefix('Q', length) + query_bytes
+  end
+end
+
+class ParseMessage < PostgresMessage
+  attr_accessor :statement_name, :query, :parameter_types
+
+  def initialize(statement_name = "", query = "", parameter_types = [])
+    @statement_name = statement_name
+    @query = query
+    @parameter_types = parameter_types
+  end
+
+  def to_bytes
+    statement_name_bytes = encode_string(@statement_name)
+    query_bytes = encode_string(@query)
+    parameter_types_bytes = @parameter_types.pack('N*')
+    
+    length = 4 + statement_name_bytes.size + query_bytes.size + 2 + parameter_types_bytes.size
+    message_prefix('P', length) + statement_name_bytes + query_bytes + encode_int16(@parameter_types.size) + parameter_types_bytes
+  end
+end
+
+class BindMessage < PostgresMessage
+  attr_accessor :portal_name, :statement_name, :parameter_format_codes, :parameters, :result_column_format_codes
+
+  def initialize(portal_name = "", statement_name = "", parameter_format_codes = [], parameters = [], result_column_format_codes = [])
+    @portal_name = portal_name
+    @statement_name = statement_name
+    @parameter_format_codes = parameter_format_codes
+    @parameters = parameters
+    @result_column_format_codes = result_column_format_codes
+  end
+
+  def to_bytes
+    portal_name_bytes = encode_string(@portal_name)
+    statement_name_bytes = encode_string(@statement_name)
+    parameter_format_codes_bytes = @parameter_format_codes.pack('n*')
+    
+    parameters_bytes = @parameters.map do |param|
+      if param.nil?
+        encode_int32(-1)
+      else
+        encode_int32(param.bytesize) + param
+      end
+    end.join
+
+    result_column_format_codes_bytes = @result_column_format_codes.pack('n*')
+
+    length = 4 + portal_name_bytes.size + statement_name_bytes.size + 2 + parameter_format_codes_bytes.size + 2 + parameters_bytes.size + 2 + result_column_format_codes_bytes.size
+    message_prefix('B', length) + portal_name_bytes + statement_name_bytes + encode_int16(@parameter_format_codes.size) + parameter_format_codes_bytes + encode_int16(@parameters.size) + parameters_bytes + encode_int16(@result_column_format_codes.size) + result_column_format_codes_bytes
+  end
+end
+
+class DescribeMessage < PostgresMessage
+  attr_accessor :type, :name
+
+  def initialize(type = 'S', name = "")
+    @type = type
+    @name = name
+  end
+
+  def to_bytes
+    name_bytes = encode_string(@name)
+    length = 4 + 1 + name_bytes.size
+    message_prefix('D', length) + @type + name_bytes
+  end
+end
+
+
+class ExecuteMessage < PostgresMessage
+  attr_accessor :portal_name, :max_rows
+
+  def initialize(portal_name = "", max_rows = 0)
+    @portal_name = portal_name
+    @max_rows = max_rows
+  end
+
+  def to_bytes
+    portal_name_bytes = encode_string(@portal_name)
+    length = 4 + portal_name_bytes.size + 4
+    message_prefix('E', length) + portal_name_bytes + encode_int32(@max_rows)
+  end
+end
+
+class FlushMessage < PostgresMessage
+  def to_bytes
+    length = 4
+    message_prefix('H', length)
+  end
+end
+
+class SyncMessage < PostgresMessage
+  def to_bytes
+    length = 4
+    message_prefix('S', length)
+  end
+end
+
+class CloseMessage < PostgresMessage
+  attr_accessor :type, :name
+
+  def initialize(type = 'S', name = "")
+    @type = type
+    @name = name
+  end
+
+  def to_bytes
+    name_bytes = encode_string(@name)
+    length = 4 + 1 + name_bytes.size
+    message_prefix('C', length) + @type + name_bytes
+  end
+end
+
--- a/tests/ruby/helpers/pg_socket.rb
+++ b/tests/ruby/helpers/pg_socket.rb
@@ -1,5 +1,6 @@
 require 'socket'
 require 'digest/md5'
+require_relative 'frontend_messages'

 BACKEND_MESSAGE_CODES = {
  'Z' => "ReadyForQuery",
@@ -18,9 +19,13 @@ class PostgresSocket
    @host = host
    @socket = TCPSocket.new @host, @port
    @parameters = {}
-    @verbose = true
+    @verbose = false
  end

+  def send_message(message)
+    @socket.write(message.to_bytes)
+  end
+  
  def send_md5_password_message(username, password, salt)
    m = Digest::MD5.hexdigest(password + username)
    m = Digest::MD5.hexdigest(m + salt.map(&:chr).join(""))
@@ -113,107 +118,6 @@ class PostgresSocket
    log "[F] Sent CancelRequest message"
  end

-  def send_query_message(query)
-    query_size = query.length
-    message_size = 1 + 4 + query_size
-    message = []
-    message << "Q".ord
-    message << [message_size].pack('l>').unpack('CCCC') # 4
-    message << query.split('').map(&:ord) # 2, 11
-    message << 0 # 1, 12
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent Q message (#{query})"
-  end
-
-  def send_parse_message(query)
-    query_size = query.length
-    message_size = 2 + 2 + 4 + query_size
-    message = []
-    message << "P".ord
-    message << [message_size].pack('l>').unpack('CCCC') # 4
-    message << 0 # unnamed statement
-    message << query.split('').map(&:ord) # 2, 11
-    message << 0 # 1, 12
-    message << [0, 0]
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent P message (#{query})"
-  end
-
-  def send_bind_message
-    message = []
-    message << "B".ord
-    message << [12].pack('l>').unpack('CCCC') # 4
-    message << 0 # unnamed statement
-    message << 0 # unnamed statement
-    message << [0, 0] # 2
-    message << [0, 0] # 2
-    message << [0, 0] # 2
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent B message"
-  end
-
-  def send_describe_message(mode)
-    message = []
-    message << "D".ord
-    message << [6].pack('l>').unpack('CCCC') # 4
-    message << mode.ord
-    message << 0 # unnamed statement
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent D message"
-  end
-
-  def send_execute_message(limit=0)
-    message = []
-    message << "E".ord
-    message << [9].pack('l>').unpack('CCCC') # 4
-    message << 0 # unnamed statement
-    message << [limit].pack('l>').unpack('CCCC') # 4
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent E message"
-  end
-
-  def send_sync_message
-    message = []
-    message << "S".ord
-    message << [4].pack('l>').unpack('CCCC') # 4
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent S message"
-  end
-
-  def send_copydone_message
-    message = []
-    message << "c".ord
-    message << [4].pack('l>').unpack('CCCC') # 4
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent c message"
-  end
-
-  def send_copyfail_message
-    message = []
-    message << "f".ord
-    message << [5].pack('l>').unpack('CCCC') # 4
-    message << 0
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent f message"
-  end
-
-  def send_flush_message
-    message = []
-    message << "H".ord
-    message << [4].pack('l>').unpack('CCCC') # 4
-    message.flatten!
-    @socket.write(message.flatten.pack('C*'))
-    log "[F] Sent H message"
-  end
-
  def read_from_server()
    output_messages = []
    retry_count = 0
--- a/tests/ruby/protocol_spec.rb
+++ b/tests/ruby/protocol_spec.rb
@@ -16,10 +16,14 @@ describe "Portocol handling" do
  end

  def run_comparison(sequence, socket_a, socket_b)
-    sequence.each do |msg, *args|
-      socket_a.send(msg, *args)
-      socket_b.send(msg, *args)
-
+    sequence.each do |msg|
+      if msg.is_a?(Symbol)
+        socket_a.send(msg)
+        socket_b.send(msg)
+      else
+        socket_a.send_message(msg)
+        socket_b.send_message(msg)
+      end
      compare_messages(
        socket_a.read_from_server,
        socket_b.read_from_server
@@ -83,9 +87,9 @@ describe "Portocol handling" do

  context "Cancel Query" do
    let(:sequence) {
-      [
-        [:send_query_message, "SELECT pg_sleep(5)"],
-        [:cancel_query]
+      [        
+        SimpleQueryMessage.new("SELECT pg_sleep(5)"),
+        :cancel_query
      ]
    }

@@ -95,12 +99,12 @@ describe "Portocol handling" do
  xcontext "Simple query after parse" do
    let(:sequence) {
      [
-        [:send_parse_message, "SELECT 5"],
-        [:send_query_message, "SELECT 1"],
-        [:send_bind_message],
-        [:send_describe_message, "P"],
-        [:send_execute_message],
-        [:send_sync_message],
+        ParseMessage.new("", "SELECT 5", []),
+        SimpleQueryMessage.new("SELECT 1"),
+        BindMessage.new("", "", [], [], [0]),
+        DescribeMessage.new("P", ""),
+        ExecuteMessage.new("", 1),
+        SyncMessage.new
      ]
    }

@@ -111,8 +115,8 @@ describe "Portocol handling" do
  xcontext "Flush message" do
    let(:sequence) {
      [
-        [:send_parse_message, "SELECT 1"],
-        [:send_flush_message]
+        ParseMessage.new("", "SELECT 1", []),
+        FlushMessage.new
      ]
    }

@@ -122,9 +126,7 @@ describe "Portocol handling" do

  xcontext "Bind without parse" do
    let(:sequence) {
-      [
-        [:send_bind_message]
-      ]
+      [BindMessage.new("", "", [], [], [0])]
    }
    # This is known to fail.
    # Server responds immediately, Proxy buffers the message
@@ -133,23 +135,155 @@ describe "Portocol handling" do

  context "Simple message" do
    let(:sequence) {
-      [[:send_query_message, "SELECT 1"]]
+      [SimpleQueryMessage.new("SELECT 1")]
    }

    it_behaves_like "at parity with database"
  end
+  
+  10.times do |i|
+    context "Extended protocol" do
+      let(:sequence) {
+        [
+          ParseMessage.new("", "SELECT 1", []),
+          BindMessage.new("", "", [], [], [0]),
+          DescribeMessage.new("S", ""),
+          ExecuteMessage.new("", 1),
+          SyncMessage.new
+        ]
+      }

-  context "Extended protocol" do
-    let(:sequence) {
-      [
-        [:send_parse_message, "SELECT 1"],
-        [:send_bind_message],
-        [:send_describe_message, "P"],
-        [:send_execute_message],
-        [:send_sync_message],
-      ]
-    }
-
-    it_behaves_like "at parity with database"
+      it_behaves_like "at parity with database"
+    end
  end
-end
+
+  describe "Protocol-level prepared statements" do
+    let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 1, "transaction") } 
+    before do 
+      q_sock = PostgresSocket.new('localhost', processes.pgcat.port)
+      q_sock.send_startup_message("sharding_user", "sharded_db", "sharding_user")
+      table_query = "CREATE TABLE IF NOT EXISTS employees (employee_id SERIAL PRIMARY KEY, salary NUMERIC(10, 2) CHECK (salary > 0));"
+      q_sock.send_message(SimpleQueryMessage.new(table_query))
+      q_sock.close
+      
+      current_configs = processes.pgcat.current_config
+      current_configs["pools"]["sharded_db"]["prepared_statements_cache_size"] = 500
+      processes.pgcat.update_config(current_configs)
+      processes.pgcat.reload_config
+    end
+    after do 
+      q_sock = PostgresSocket.new('localhost', processes.pgcat.port)
+      q_sock.send_startup_message("sharding_user", "sharded_db", "sharding_user")
+      table_query = "DROP TABLE IF EXISTS employees;"
+      q_sock.send_message(SimpleQueryMessage.new(table_query))
+      q_sock.close
+    end
+
+    context "When unnamed prepared statements are used" do
+      it "does not cache them" do
+        socket = PostgresSocket.new('localhost', processes.pgcat.port)
+        socket.send_startup_message("sharding_user", "sharded_db", "sharding_user")
+      
+        socket.send_message(SimpleQueryMessage.new("DISCARD ALL"))
+        socket.read_from_server
+        
+        10.times do |i|
+          socket.send_message(ParseMessage.new("", "SELECT #{i}", []))
+          socket.send_message(BindMessage.new("", "", [], [], [0]))
+          socket.send_message(DescribeMessage.new("S", ""))
+          socket.send_message(ExecuteMessage.new("", 1))
+          socket.send_message(SyncMessage.new)
+          socket.read_from_server
+        end
+      
+        socket.send_message(SimpleQueryMessage.new("SELECT name, statement, prepare_time, parameter_types FROM pg_prepared_statements"))
+        result = socket.read_from_server
+        number_of_saved_statements = result.count { |m| m[:code] == 'D' }
+        expect(number_of_saved_statements).to eq(0)
+      end
+    end
+
+    context "When named prepared statements are used" do
+      it "caches them" do
+        socket = PostgresSocket.new('localhost', processes.pgcat.port)
+        socket.send_startup_message("sharding_user", "sharded_db", "sharding_user")
+
+        socket.send_message(SimpleQueryMessage.new("DISCARD ALL"))
+        socket.read_from_server
+        
+        3.times do
+          socket.send_message(ParseMessage.new("my_query", "SELECT * FROM employees WHERE employee_id in ($1,$2,$3)", [0,0,0]))
+          socket.send_message(BindMessage.new("", "my_query", [0,0,0], [0,0,0].map(&:to_s), [0,0,0,0,0,0]))
+          socket.send_message(SyncMessage.new)
+          socket.read_from_server
+        end
+      
+        3.times do
+          socket.send_message(ParseMessage.new("my_other_query", "SELECT * FROM employees WHERE salary in ($1,$2,$3)", [0,0,0]))
+          socket.send_message(BindMessage.new("", "my_other_query", [0,0,0], [0,0,0].map(&:to_s), [0,0,0,0,0,0]))
+          socket.send_message(SyncMessage.new)
+          socket.read_from_server
+        end
+
+        socket.send_message(SimpleQueryMessage.new("SELECT name, statement, prepare_time, parameter_types FROM pg_prepared_statements"))
+        result = socket.read_from_server
+        number_of_saved_statements = result.count { |m| m[:code] == 'D' }
+        expect(number_of_saved_statements).to eq(2)
+      end
+    end
+
+    context "When DISCARD ALL/DEALLOCATE ALL are called" do 
+      it "resets server and client caches" do 
+        socket = PostgresSocket.new('localhost', processes.pgcat.port)
+        socket.send_startup_message("sharding_user", "sharded_db", "sharding_user")
+  
+        20.times do |i|
+          socket.send_message(ParseMessage.new("my_query_#{i}", "SELECT * FROM employees WHERE employee_id in ($1,$2,$3)", [0,0,0]))
+        end
+      
+        20.times do |i|
+          socket.send_message(BindMessage.new("", "my_query_#{i}", [0,0,0], [0,0,0].map(&:to_s), [0,0]))
+        end 
+      
+        socket.send_message(SyncMessage.new)
+        socket.read_from_server
+        
+        socket.send_message(SimpleQueryMessage.new("DISCARD ALL"))
+        socket.read_from_server
+        responses = []
+        4.times do |i|
+          socket.send_message(ParseMessage.new("my_query_#{i}", "SELECT * FROM employees WHERE employee_id in ($1,$2,$3)", [0,0,0]))
+          socket.send_message(BindMessage.new("", "my_query_#{i}", [0,0,0], [0,0,0].map(&:to_s), [0,0]))
+          socket.send_message(SyncMessage.new)
+
+          responses += socket.read_from_server
+        end
+
+        errors = responses.select { |message| message[:code] == 'E' }
+        error_message = errors.map { |message| message[:bytes].map(&:chr).join("") }.join("\n")
+        raise StandardError, "Encountered the following errors: #{error_message}" if errors.length > 0
+      end
+    end
+    
+    context "Maximum number of bound paramters" do
+      it "does not crash" do
+        test_socket = PostgresSocket.new('localhost', processes.pgcat.port)
+        test_socket.send_startup_message("sharding_user", "sharded_db", "sharding_user")
+      
+        types = Array.new(65_535) { |i| 0 }
+
+        params = Array.new(65_535) { |i| "$#{i+1}" }.join(",")
+        test_socket.send_message(ParseMessage.new("my_query", "SELECT * FROM employees WHERE employee_id in (#{params})", types))
+
+        test_socket.send_message(BindMessage.new("my_query", "my_query", types, types.map(&:to_s), types))
+
+        test_socket.send_message(SyncMessage.new)
+        
+        # If the proxy crashes, this will raise an error
+        expect { test_socket.read_from_server }.to_not raise_error
+
+        test_socket.close
+      end
+    end
+  end
+end
--- a/tests/ruby/stats_spec.rb
+++ b/tests/ruby/stats_spec.rb
@@ -16,7 +16,7 @@ describe "Stats" do
      it "updates *_query_time and *_wait_time" do
        connections = Array.new(3) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
        connections.each do |c|
-          Thread.new { c.async_exec("BEGIN; SELECT pg_sleep(0.25); COMMIT;") }
+          Thread.new { c.async_exec("SELECT pg_sleep(0.25)") }
        end
        sleep(1)
        connections.map(&:close)
@@ -25,32 +25,10 @@ describe "Stats" do
        sleep(15.5)
        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
        results = admin_conn.async_exec("SHOW STATS")[0]
+        admin_conn.close
        expect(results["total_query_time"].to_i).to be_within(200).of(750)
        expect(results["avg_query_time"].to_i).to be_within(50).of(250)

-        expect(results["total_xact_time"].to_i).to be_within(200).of(750)
-        expect(results["avg_xact_time"].to_i).to be_within(50).of(250)
-
-        expect(results["total_wait_time"].to_i).to_not eq(0)
-        expect(results["avg_wait_time"].to_i).to_not eq(0)
-
-        connections = Array.new(3) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
-        connections.each do |c|
-          Thread.new { c.async_exec("SELECT pg_sleep(0.25);") }
-        end
-        sleep(1)
-        connections.map(&:close)
-
-        results = admin_conn.async_exec("SHOW STATS")[0]
-        admin_conn.close
-        # This should increase with more queries
-        expect(results["total_query_time"].to_i).to be_within(400).of(1500)
-        expect(results["avg_query_time"].to_i).to be_within(50).of(250)
-
-        # This should not increase as we did not run any additional transactions
-        expect(results["total_xact_time"].to_i).to be_within(200).of(750)
-        expect(results["avg_xact_time"].to_i).to be_within(50).of(250)
-
        expect(results["total_wait_time"].to_i).to_not eq(0)
        expect(results["avg_wait_time"].to_i).to_not eq(0)
      end
--- a/tests/rust/Cargo.lock
+++ b/tests/rust/Cargo.lock
--- a/tests/rust/src/main.rs
+++ b/tests/rust/src/main.rs
@@ -15,13 +15,11 @@ async fn test_prepared_statements() {
    for _ in 0..5 {
        let pool = pool.clone();
        let handle = tokio::task::spawn(async move {
-            for _ in 0..1000 {
-                match sqlx::query("SELECT one").fetch_all(&pool).await {
+            for i in 0..1000 {
+                match sqlx::query(&format!("SELECT {:?}", i % 5)).fetch_all(&pool).await {
                    Ok(_) => (),
                    Err(err) => {
-                        if err.to_string().contains("prepared statement") {
-                            panic!("prepared statement error: {}", err);
-                        }
+                        panic!("prepared statement error: {}", err);
                    }
                }
            }
Author	SHA1	Message	Date
Mostafa	3fc9e5dec1	Merge branch 'main' of github.com:postgresml/pgcat into mostafa_fix_prepared_stmts	2024-09-03 18:11:32 -05:00
Mostafa Abdelraouf	69af6cc5e5	Make iterating on integration tests easier (#789 ) Writing and iterating on integration tests are cumbersome, having to wait 10 minutes for the test-suite to run just to see if your test works or not is unacceptable. In this PR, I added a detailed workflow for writing tests that should shorten the feedback cycle of modifying tests to be as low as a few seconds. It will involve opening a shell into a long-lived container that has all the setup and dependencies necessary and then running your desired tests directly there. I added a convenience script that bootstraps the environment and then opens an interactive shell into the container and you can then run tests immediately in an environment that is more or less identical to what we have running in CircleCI	2024-09-03 11:15:53 -05:00
Mostafa	f7c5c0faf9	fix bind	2024-09-01 16:14:44 -05:00
Mostafa	982d03c374	fix syntax	2024-09-01 15:41:33 -05:00
Mostafa	686b7ca7c5	Fixes	2024-09-01 15:31:27 -05:00
Mostafa	7c55bf78fe	Add failing tests	2024-09-01 14:39:05 -05:00
Mostafa Abdelraouf	ca34597002	Fix broken integration test #740 (#787 )	2024-08-31 17:15:13 -05:00
Mostafa Abdelraouf	2def40ea6a	Add test case for issue 776 (#786 ) I am adding a tiny test that uses the SQL statement that was reported to break an older version of SQL parser library #776	2024-08-31 10:52:33 -05:00
Mostafa Abdelraouf	c05129018d	Improve Prometheus stats + Add Grafana dashboard (#785 ) We were missing some labels on metrics generated by the Prometheus exporter so I fixed that. There are still some gaps that I want to address with respect to the metrics we track but this seems like a good start. I also created a Grafana Dashboard and exported it to JSON. It is designed with the same metric names the Prometheus exporter uses.	2024-08-31 08:18:57 -05:00
Mostafa Abdelraouf	4a7a6a8e7a	Cut 1.2.0 release (#783 )	2024-08-30 08:30:16 -05:00
Mostafa Abdelraouf	29a476e190	QueryRouter: route to primary when locks exists (select for update) (#782 ) Authored-by: Javier Goday <jgoday@gmail.com>	2024-08-30 04:26:36 -05:00
KwongTN	81933b918d	Add linux/arm64 docker image build support (#774 )	2024-08-29 13:50:38 -05:00
Saraj Munjal	7cbc9178d8	Bump the hyper crate to v1.4.1 and rework prometheus server handling (#778 ) Bump hyper to v1.4.1 and rework prometheus server handling	2024-08-29 09:47:58 -05:00
Mostafa Abdelraouf	2c8b2f0776	Fix CI image build step (#780 ) The docker CI build image is failing due to this error 249.5 Finished release [optimized] target(s) in 2m 49s 249.5 Installing /home/circleci/.cargo/bin/rustfilt 249.5 Installed package `rustfilt v0.2.1` (executable `rustfilt`) 249.5 error: failed to compile `cargo-binutils v0.3.6`, intermediate artifacts can be found at `/tmp/cargo-installrWENQG` 249.5 249.5 Caused by: 249.5 package `cargo-platform v0.1.8` cannot be built because it requires rustc 1.73 or newer, while the currently active rustc version is 1.67.1 249.5 Try re-running cargo install with `--locked` 249.5 Summary Successfully installed rustfilt! Failed to install cargo-binutils (see error(s) above). 249.5 error: some crates failed to install So I am bumping the version up	2024-08-29 08:37:13 -05:00
Mostafa Abdelraouf	8f9a2b8e6f	Fix a Panic in admin commands (#779 ) We have a panic when we send SHOW or ;;;;;;;;;;;;;;;;; to admin database. This PR fixes these panics and adds a couple of tests	2024-08-28 21:29:40 -05:00