link gif

I am adding a tiny test that uses the SQL statement that was reported to break an older version of SQL parser library

#776

.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"

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

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.

src/prometheus.rs

@@ -12,19 +12,30 @@ 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
@@ -127,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 {
@@ -210,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)
     }
@@ -225,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)
     }
 
@@ -236,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)
     }
@@ -338,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),
-                            "bytes_received",
+                        ("bytes_sent", server_info.bytes_sent),
-                            server_info.bytes_received.load(Ordering::Relaxed),
+                        ("transaction_count", server_info.transaction_count),
-                        ),
+                        ("query_count", server_info.query_count),
-                        ("bytes_sent", server_info.bytes_sent.load(Ordering::Relaxed)),
+                        ("error_count", server_info.error_count),
-                        (
+                        ("idle_count", server_info.idle_count),
-                            "transaction_count",
+                        ("active_count", server_info.active_count),
-                            server_info.transaction_count.load(Ordering::Relaxed),
+                        ("login_count", server_info.login_count),
-                        ),
+                        ("tested_count", server_info.tested_count),
-                        (
+                        ("is_banned", if pool.is_banned(address) { 1 } else { 0 }),
-                            "query_count",
+                        ("is_paused", if pool.paused() { 1 } else { 0 }),
-                            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) =

src/query_router.rs

@@ -1399,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();

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

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/

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

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 {
+            for i in 0..1000 {
-                match sqlx::query("SELECT one").fetch_all(&pool).await {
+                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);
-                        }
                     }
                 }
             }