Allow pause/resuming all pools (#566)

support pausing all pools

src/admin.rs

@@ -74,11 +74,11 @@ where
         }
         "PAUSE" => {
             trace!("PAUSE");
-            pause(stream, query_parts[1]).await
+            pause(stream, query_parts).await
         }
         "RESUME" => {
             trace!("RESUME");
-            resume(stream, query_parts[1]).await
+            resume(stream, query_parts).await
         }
         "SHUTDOWN" => {
             trace!("SHUTDOWN");
@@ -797,19 +797,33 @@ where
 }
 
 /// Pause a pool. It won't pass any more queries to the backends.
-async fn pause<T>(stream: &mut T, query: &str) -> Result<(), Error>
+async fn pause<T>(stream: &mut T, tokens: Vec<&str>) -> Result<(), Error>
 where
     T: tokio::io::AsyncWrite + std::marker::Unpin,
 {
-    let parts: Vec<&str> = query.split(",").map(|part| part.trim()).collect();
+    let parts: Vec<&str> = match tokens.len() == 2 {
+        true => tokens[1].split(",").map(|part| part.trim()).collect(),
+        false => Vec::new(),
+    };
 
-    if parts.len() != 2 {
+    match parts.len() {
-        error_response(
+        0 => {
-            stream,
+            for (_, pool) in get_all_pools() {
-            "PAUSE requires a database and a user, e.g. PAUSE my_db,my_user",
+                pool.pause();
-        )
+            }
-        .await
+
-    } else {
+            let mut res = BytesMut::new();
+
+            res.put(command_complete("PAUSE"));
+
+            // ReadyForQuery
+            res.put_u8(b'Z');
+            res.put_i32(5);
+            res.put_u8(b'I');
+
+            write_all_half(stream, &res).await
+        }
+        2 => {
             let database = parts[0];
             let user = parts[1];
 
@@ -841,22 +855,38 @@ where
                 }
             }
         }
+        _ => error_response(stream, "usage: PAUSE [db, user]").await,
+    }
 }
 
 /// Resume a pool. Queries are allowed again.
-async fn resume<T>(stream: &mut T, query: &str) -> Result<(), Error>
+async fn resume<T>(stream: &mut T, tokens: Vec<&str>) -> Result<(), Error>
 where
     T: tokio::io::AsyncWrite + std::marker::Unpin,
 {
-    let parts: Vec<&str> = query.split(",").map(|part| part.trim()).collect();
+    let parts: Vec<&str> = match tokens.len() == 2 {
+        true => tokens[1].split(",").map(|part| part.trim()).collect(),
+        false => Vec::new(),
+    };
 
-    if parts.len() != 2 {
+    match parts.len() {
-        error_response(
+        0 => {
-            stream,
+            for (_, pool) in get_all_pools() {
-            "RESUME requires a database and a user, e.g. RESUME my_db,my_user",
+                pool.resume();
-        )
+            }
-        .await
+
-    } else {
+            let mut res = BytesMut::new();
+
+            res.put(command_complete("RESUME"));
+
+            // ReadyForQuery
+            res.put_u8(b'Z');
+            res.put_i32(5);
+            res.put_u8(b'I');
+
+            write_all_half(stream, &res).await
+        }
+        2 => {
             let database = parts[0];
             let user = parts[1];
 
@@ -888,6 +918,8 @@ where
                 }
             }
         }
+        _ => error_response(stream, "usage: RESUME [db, user]").await,
+    }
 }
 
 /// Send response packets for shutdown.

tests/ruby/admin_spec.rb

@@ -90,4 +90,28 @@ describe "Admin" do
       expect(results["pool_mode"]).to eq("transaction")
     end
   end
+
+  describe "PAUSE" do
+    it "pauses all pools" do
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+      results = admin_conn.async_exec("SHOW DATABASES").to_a
+      expect(results.map{ |r| r["paused"] }.uniq).to eq(["0"])
+
+      admin_conn.async_exec("PAUSE")
+
+      results = admin_conn.async_exec("SHOW DATABASES").to_a
+      expect(results.map{ |r| r["paused"] }.uniq).to eq(["1"])
+
+      admin_conn.async_exec("RESUME")
+
+      results = admin_conn.async_exec("SHOW DATABASES").to_a
+      expect(results.map{ |r| r["paused"] }.uniq).to eq(["0"])
+    end
+
+    it "handles errors" do
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+      expect { admin_conn.async_exec("PAUSE foo").to_a }.to raise_error(PG::SystemError)
+      expect { admin_conn.async_exec("PAUSE foo,bar").to_a }.to raise_error(PG::SystemError)
+    end
+  end
 end