Extended query protocol sharding (#339)

* Prepared stmt sharding

s

tests

* len check

* remove python test

* latest rust

* move that to debug for sure

* Add the actual tests

* latest image

* Update tests/ruby/sharding_spec.rb

Author: levkk

.dockerignore

@@ -2,3 +2,5 @@ target/
 tests/
 tracing/
 .circleci/
+.git/
+dev/

Dockerfile.ci

@@ -1,7 +1,7 @@
-FROM cimg/rust:1.62.0
+FROM cimg/rust:1.67.1
 RUN sudo apt-get update && \
 	sudo apt-get install -y \
-		psmisc postgresql-contrib-12 postgresql-client-12 libpq-dev \
+		psmisc postgresql-contrib-14 postgresql-client-14 libpq-dev \
 		ruby ruby-dev python3 python3-pip \
 		lcov llvm-11 iproute2 && \
 	sudo apt-get upgrade curl && \

dev/script/console

@@ -3,4 +3,10 @@
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 export HOST_UID="$(id -u)"
 export HOST_GID="$(id -g)"
-docker-compose -f "${DIR}/../docker-compose.yaml" run --rm pgcat-shell
+
+if [[ "${1}" == "down" ]]; then
+	docker-compose -f "${DIR}/../docker-compose.yaml" down
+	exit 0
+else
+	docker-compose -f "${DIR}/../docker-compose.yaml" run --rm pgcat-shell
+fi

src/client.rs

@@ -675,14 +675,42 @@ where
                 // allocate a connection, we wouldn't be able to send back an error message
                 // to the client so we buffer them and defer the decision to error out or not
                 // to when we get the S message
-                'P' | 'B' | 'D' | 'E' => {
+                'D' | 'E' => {
                     self.buffer.put(&message[..]);
                     continue;
                 }
+
+                'Q' => {
+                    if query_router.query_parser_enabled() {
+                        query_router.infer(&message);
+                    }
+                }
+
+                'P' => {
+                    self.buffer.put(&message[..]);
+
+                    if query_router.query_parser_enabled() {
+                        query_router.infer(&message);
+                    }
+
+                    continue;
+                }
+
+                'B' => {
+                    self.buffer.put(&message[..]);
+
+                    if query_router.query_parser_enabled() {
+                        query_router.infer_shard_from_bind(&message);
+                    }
+
+                    continue;
+                }
+
                 'X' => {
                     debug!("Client disconnecting");
                     return Ok(());
                 }
+
                 _ => (),
             }
 
@@ -711,11 +739,7 @@ where
             // Handle all custom protocol commands, if any.
             match query_router.try_execute_command(&message) {
                 // Normal query, not a custom command.
-                None => {
-                    if query_router.query_parser_enabled() {
-                        query_router.infer(&message);
-                    }
-                }
+                None => (),
 
                 // SET SHARD TO
                 Some((Command::SetShard, _)) => {
@@ -727,7 +751,7 @@ where
                         error_response(
                             &mut self.write,
                             &format!(
-                                "shard {} is more than configured {}, staying on shard {}",
+                                "shard {} is more than configured {}, staying on shard {} (shard numbers start at 0)",
                                 query_router.shard(),
                                 pool.shards(),
                                 current_shard,