Fix default_role being ignored when query_parser_enabled was false

* Update bb8 to 0.8.6

To get https://github.com/djc/bb8/pull/186 and https://github.com/djc/bb8/pull/189
which fix potential deadlocks (https://github.com/djc/bb8/issues/154).

Also, this (https://github.com/djc/bb8/pull/225) was needed to prevent a connection
leak which was conveniently spotted in our integration tests.

* Ignore ./.bundle (created by dev console)

---------

Co-authored-by: Jose Fernandez (magec) <joseferper@gmail.com>

.gitignore

@@ -12,3 +12,4 @@ dev/cache
 !dev/cache/.keepme
 .venv
 **/__pycache__
+.bundle

Cargo.lock

@@ -192,12 +192,11 @@ checksum = "604178f6c5c21f02dc555784810edfb88d34ac2c73b2eae109655649ee73ce3d"
 
 [[package]]
 name = "bb8"
-version = "0.8.1"
+version = "0.8.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "98b4b0f25f18bcdc3ac72bdb486ed0acf7e185221fd4dc985bc15db5800b0ba2"
+checksum = "d89aabfae550a5c44b43ab941844ffcd2e993cb6900b342debf59e9ea74acdb8"
 dependencies = [
  "async-trait",
- "futures-channel",
  "futures-util",
  "parking_lot",
  "tokio",

Cargo.toml

@@ -8,7 +8,7 @@ edition = "2021"
 tokio = { version = "1", features = ["full"] }
 bytes = "1"
 md-5 = "0.10"
-bb8 = "0.8.1"
+bb8 = "=0.8.6"
 async-trait = "0.1"
 rand = "0.8"
 chrono = "0.4"

src/client.rs

@@ -881,6 +881,7 @@ 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

src/query_router.rs

@@ -386,6 +386,18 @@ 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,
+        }
+    }
+
     /// 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 {
@@ -428,8 +440,9 @@ impl QueryRouter {
                     };
 
                     let has_locks = !query.locks.is_empty();
+                    let has_mutation = Self::is_mutation_query(query);
 
-                    if has_locks {
+                    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.
@@ -1048,6 +1061,11 @@ 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
@@ -1113,6 +1131,26 @@ 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();

tests/ruby/load_balancing_spec.rb

@@ -56,6 +56,41 @@ 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
@@ -161,4 +196,3 @@ describe "Least Outstanding Queries Load Balancing" do
     end
   end
 end
-