Functional sharding

Author: levkk

pgdog/src/backend/pool/cluster.rs

@@ -210,6 +210,7 @@ mod test {
                     name: Some("sharded".into()),
                     column: "id".into(),
                     primary: true,
+                    shard: None,
                 }]),
                 shards: vec![Shard::default(), Shard::default()],
                 ..Default::default()

pgdog/src/config/mod.rs

@@ -528,12 +528,17 @@ pub struct ShardedTable {
     /// This table is the primary sharding anchor (e.g. "users").
     #[serde(default)]
     pub primary: bool,
+    /// Hardcoded shard number.
+    pub shard: Option<usize>,
 }
 
 /// Queries with manual routing rules.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)]
 pub struct ManualQuery {
+    /// Query fingerprint.
     pub fingerprint: String,
+    /// Which shard this query should go to.
+    pub shard: Option<usize>,
 }
 
 #[cfg(test)]

pgdog/src/frontend/router/parser/query.rs

@@ -179,9 +179,12 @@ impl QueryParser {
                 let fingerprint = fingerprint(query).map_err(Error::PgQuery)?;
                 let manual_route = databases().manual_query(&fingerprint.hex).cloned();
 
-                // TODO: check routing logic required by config.
-                if manual_route.is_some() {
-                    route.set_shard(round_robin::next() % cluster.shards().len());
+                if let Some(manual_route) = manual_route {
+                    if let Some(shard) = manual_route.shard {
+                        route.set_shard(shard);
+                    } else {
+                        route.set_shard(round_robin::next() % cluster.shards().len());
+                    }
                 }
             }
         }