feat: maintenance mode per DB (#1014)

fix #975 

Allow the maintenance mode to be enabled/disabled on a per-database
basis.

Author: levkk

integration/rust/tests/integration/maintenance_mode.rs

@@ -120,7 +120,7 @@ async fn test_maintenance_mode_parsing() {
     let result = admin.simple_query("MAINTENANCE INVALID").await;
     assert!(result.is_err());
 
-    let result = admin.simple_query("MAINTENANCE ON EXTRA").await;
+    let result = admin.simple_query("MAINTENANCE ON DB_NAME EXTRA").await;
     assert!(result.is_err());
 }
 
@@ -156,6 +156,67 @@ async fn test_maintenance_mode_concurrent_operations() {
     tokio::try_join!(admin_task, client_task).unwrap();
 }
 
+#[tokio::test]
+#[serial]
+async fn test_maintenance_mode_single_database() {
+    let admin = admin_tokio().await;
+    let failover = connection_failover().await;
+    let mut pools = connections_sqlx().await;
+    let pgdog = pools.remove(0); // "pgdog" database
+
+    // Clean slate.
+    admin.simple_query("MAINTENANCE OFF").await.unwrap();
+    admin
+        .simple_query("MAINTENANCE OFF failover")
+        .await
+        .unwrap();
+
+    // Both databases work normally.
+    pgdog.execute("SELECT 1").await.unwrap();
+    failover.execute("SELECT 1").await.unwrap();
+
+    // Put only the `failover` database into maintenance mode.
+    admin.simple_query("MAINTENANCE ON failover").await.unwrap();
+
+    // Queries to `pgdog` must keep working, unaffected by `failover`.
+    tokio::time::timeout(Duration::from_secs(1), pgdog.execute("SELECT 2"))
+        .await
+        .expect("pgdog query should not block while only failover is in maintenance")
+        .unwrap();
+
+    // Meanwhile, a query to `failover` should block until maintenance is lifted.
+    let failover_blocked = failover.clone();
+    let blocked = tokio::spawn(async move {
+        failover_blocked.execute("SELECT 2").await.unwrap();
+    });
+
+    sleep(Duration::from_millis(100)).await;
+    assert!(
+        !blocked.is_finished(),
+        "failover query should be blocked during maintenance"
+    );
+
+    // `pgdog` still works while the `failover` query is blocked.
+    tokio::time::timeout(Duration::from_secs(1), pgdog.execute("SELECT 3"))
+        .await
+        .expect("pgdog query should not block")
+        .unwrap();
+
+    // Lift maintenance on `failover`; the blocked query should now complete.
+    admin
+        .simple_query("MAINTENANCE OFF failover")
+        .await
+        .unwrap();
+
+    tokio::time::timeout(Duration::from_secs(5), blocked)
+        .await
+        .expect("failover query should complete after maintenance is lifted")
+        .unwrap();
+
+    pgdog.close().await;
+    failover.close().await;
+}
+
 #[tokio::test]
 #[serial]
 async fn test_maintenance_mode_transaction_behavior() {

pgdog/src/admin/maintenance_mode.rs

@@ -1,38 +1,55 @@
 //! Turn maintenance mode on/off.
+//!
+//! Maintenance mode is special: it's completely independent from the config
+//! and will hold true during config changes, e.g. when some databases disappear
+//! and new ones are added.
+//!
+//! This is useful when changing the sharding config online, for example.
+//!
 
 use crate::backend::maintenance_mode;
 
 use super::prelude::*;
 
-/// Turn maintenance mode on/off.
+/// Turn maintenance mode on/off, optionally for a single database.
 #[derive(Default)]
 pub struct MaintenanceMode {
     enable: bool,
+    database: Option<String>,
 }
 
 #[async_trait]
 impl Command for MaintenanceMode {
     fn parse(sql: &str) -> Result<Self, Error> {
         let parts = sql.split(" ").collect::<Vec<_>>();
 
-        match parts[..] {
-            ["maintenance", "on"] => Ok(Self { enable: true }),
-            ["maintenance", "off"] => Ok(Self { enable: false }),
-            _ => Err(Error::Syntax),
-        }
+        let (enable, database) = match parts[..] {
+            ["maintenance", "on"] => (true, None),
+            ["maintenance", "off"] => (false, None),
+            ["maintenance", "on", database] => (true, Some(database.to_string())),
+            ["maintenance", "off", database] => (false, Some(database.to_string())),
+            _ => return Err(Error::Syntax),
+        };
+
+        Ok(Self { enable, database })
     }
 
     async fn execute(&self) -> Result<Vec<Message>, Error> {
+        let database = self.database.as_deref();
         if self.enable {
-            maintenance_mode::start();
+            maintenance_mode::start(database);
         } else {
-            maintenance_mode::stop();
+            maintenance_mode::stop(database);
         }
 
         Ok(vec![])
     }
 
     fn name(&self) -> String {
-        format!("MAINTENANCE {}", if self.enable { "ON" } else { "OFF" })
+        let state = if self.enable { "ON" } else { "OFF" };
+        match &self.database {
+            Some(database) => format!("MAINTENANCE {} {}", state, database),
+            None => format!("MAINTENANCE {}", state),
+        }
     }
 }