allow for multiple sharding keys

Author: JanKaul

src/config.rs

@@ -457,7 +457,7 @@ pub struct Pool {
     pub sharding_function: ShardingFunction,
 
     #[serde(default = "Pool::default_automatic_sharding_key")]
-    pub automatic_sharding_key: Option<String>,
+    pub automatic_sharding_keys: Option<Vec<String>>,
 
     pub sharding_key_regex: Option<String>,
     pub shard_id_regex: Option<String>,
@@ -495,7 +495,7 @@ impl Pool {
         LoadBalancingMode::Random
     }
 
-    pub fn default_automatic_sharding_key() -> Option<String> {
+    pub fn default_automatic_sharding_key() -> Option<Vec<String>> {
         None
     }
 
@@ -539,23 +539,23 @@ impl Pool {
             }
         }
 
-        self.automatic_sharding_key = match &self.automatic_sharding_key {
-            Some(key) => {
-                // No quotes in the key so we don't have to compare quoted
-                // to unquoted idents.
-                let key = key.replace("\"", "");
-
-                if key.split(".").count() != 2 {
-                    error!(
-                        "automatic_sharding_key '{}' must be fully qualified, e.g. t.{}`",
-                        key, key
-                    );
-                    return Err(Error::BadConfig);
+        match &mut self.automatic_sharding_keys {
+            Some(keys) => {
+                for key in keys {
+                    // No quotes in the key so we don't have to compare quoted
+                    // to unquoted idents.
+                    let key = key.replace("\"", "");
+
+                    if key.split(".").count() != 2 {
+                        error!(
+                            "automatic_sharding_key '{}' must be fully qualified, e.g. t.{}`",
+                            key, key
+                        );
+                        return Err(Error::BadConfig);
+                    }
                 }
-
-                Some(key)
             }
-            None => None,
+            None => (),
         };
 
         for (_, user) in &self.users {
@@ -577,7 +577,7 @@ impl Default for Pool {
             query_parser_enabled: false,
             primary_reads_enabled: false,
             sharding_function: ShardingFunction::PgBigintHash,
-            automatic_sharding_key: None,
+            automatic_sharding_keys: None,
             connect_timeout: None,
             idle_timeout: None,
             sharding_key_regex: None,

src/pool.rs

@@ -105,7 +105,7 @@ pub struct PoolSettings {
     pub sharding_function: ShardingFunction,
 
     // Sharding key
-    pub automatic_sharding_key: Option<String>,
+    pub automatic_sharding_keys: Option<Vec<String>>,
 
     // Health check timeout
     pub healthcheck_timeout: u64,
@@ -142,7 +142,7 @@ impl Default for PoolSettings {
             query_parser_enabled: false,
             primary_reads_enabled: true,
             sharding_function: ShardingFunction::PgBigintHash,
-            automatic_sharding_key: None,
+            automatic_sharding_keys: None,
             healthcheck_delay: General::default_healthcheck_delay(),
             healthcheck_timeout: General::default_healthcheck_timeout(),
             ban_time: General::default_ban_time(),
@@ -421,7 +421,7 @@ impl ConnectionPool {
                         query_parser_enabled: pool_config.query_parser_enabled,
                         primary_reads_enabled: pool_config.primary_reads_enabled,
                         sharding_function: pool_config.sharding_function,
-                        automatic_sharding_key: pool_config.automatic_sharding_key.clone(),
+                        automatic_sharding_keys: pool_config.automatic_sharding_keys.clone(),
                         healthcheck_delay: config.general.healthcheck_delay,
                         healthcheck_timeout: config.general.healthcheck_timeout,
                         ban_time: config.general.ban_time,

src/query_router.rs

@@ -384,7 +384,7 @@ impl QueryRouter {
 
                 // Likely a read-only query
                 Query(query) => {
-                    match &self.pool_settings.automatic_sharding_key {
+                    match &self.pool_settings.automatic_sharding_keys {
                         Some(_) => {
                             // TODO: if we have multiple queries in the same message,
                             // we can either split them and execute them individually
@@ -565,17 +565,23 @@ impl QueryRouter {
         let mut result = Vec::new();
         let mut found = false;
 
-        let sharding_key = self
+        let sharding_keys = self
             .pool_settings
-            .automatic_sharding_key
+            .automatic_sharding_keys
             .as_ref()
             .unwrap()
-            .split(".")
-            .map(|ident| Ident::new(ident))
-            .collect::<Vec<Ident>>();
-
-        // Sharding key must be always fully qualified
-        assert_eq!(sharding_key.len(), 2);
+            .iter()
+            .map(|x| {
+                x.split(".")
+                    .map(|ident| Ident::new(ident))
+                    .collect::<Vec<Ident>>()
+            })
+            .collect::<Vec<Vec<Ident>>>();
+
+        for sharding_key in sharding_keys.iter() {
+            // Sharding key must be always fully qualified
+            assert_eq!(sharding_key.len(), 2);
+        }
 
         // This parses `sharding_key = 5`. But it's technically
         // legal to write `5 = sharding_key`. I don't judge the people
@@ -587,7 +593,10 @@ impl QueryRouter {
                 Expr::Identifier(ident) => {
                     // Only if we're dealing with only one table
                     // and there is no ambiguity
-                    if &ident.value == &sharding_key[1].value {
+                    if let Some(sharding_key) = sharding_keys
+                        .iter()
+                        .find(|key| &ident.value == &key[1].value)
+                    {
                         // Sharding key is unique enough, don't worry about
                         // table names.
                         if &sharding_key[0].value == "*" {
@@ -618,8 +627,13 @@ impl QueryRouter {
                     // The key is fully qualified in the query,
                     // it will exist or Postgres will throw an error.
                     if idents.len() == 2 {
-                        found = &sharding_key[0].value == &idents[0].value
-                            && &sharding_key[1].value == &idents[1].value;
+                        found = sharding_keys
+                            .iter()
+                            .find(|key| {
+                                &key[0].value == &idents[0].value
+                                    && &key[1].value == &idents[1].value
+                            })
+                            .is_some();
                     }
                     // TODO: key can have schema as well, e.g. public.data.id (len == 3)
                 }
@@ -1103,7 +1117,7 @@ mod test {
             query_parser_enabled: true,
             primary_reads_enabled: false,
             sharding_function: ShardingFunction::PgBigintHash,
-            automatic_sharding_key: Some(String::from("test.id")),
+            automatic_sharding_keys: Some(vec![String::from("test.id")]),
             healthcheck_delay: PoolSettings::default().healthcheck_delay,
             healthcheck_timeout: PoolSettings::default().healthcheck_timeout,
             ban_time: PoolSettings::default().ban_time,
@@ -1167,7 +1181,7 @@ mod test {
             query_parser_enabled: true,
             primary_reads_enabled: false,
             sharding_function: ShardingFunction::PgBigintHash,
-            automatic_sharding_key: None,
+            automatic_sharding_keys: None,
             healthcheck_delay: PoolSettings::default().healthcheck_delay,
             healthcheck_timeout: PoolSettings::default().healthcheck_timeout,
             ban_time: PoolSettings::default().ban_time,
@@ -1205,17 +1219,26 @@ mod test {
         QueryRouter::setup();
 
         let mut qr = QueryRouter::new();
-        qr.pool_settings.automatic_sharding_key = Some("data.id".to_string());
+        qr.pool_settings.automatic_sharding_keys =
+            Some(vec!["data.id".to_string(), "derived.data_id".to_string()]);
         qr.pool_settings.shards = 3;
 
         assert!(qr.infer(&simple_query("SELECT * FROM data WHERE id = 5")));
         assert_eq!(qr.shard(), 2);
 
+        assert!(qr.infer(&simple_query("SELECT * FROM derived WHERE data_id = 5")));
+        assert_eq!(qr.shard(), 2);
+
         assert!(qr.infer(&simple_query(
             "SELECT one, two, three FROM public.data WHERE id = 6"
         )));
         assert_eq!(qr.shard(), 0);
 
+        assert!(qr.infer(&simple_query(
+            "SELECT one, two, three FROM public.derived WHERE data_id = 6"
+        )));
+        assert_eq!(qr.shard(), 0);
+
         assert!(qr.infer(&simple_query(
             "SELECT * FROM data
             INNER JOIN t2 ON data.id = 5
@@ -1242,7 +1265,8 @@ mod test {
         assert_eq!(qr.shard(), 2);
 
         // Super unique sharding key
-        qr.pool_settings.automatic_sharding_key = Some("*.unique_enough_column_name".to_string());
+        qr.pool_settings.automatic_sharding_keys =
+            Some(vec!["*.unique_enough_column_name".to_string()]);
         assert!(qr.infer(&simple_query(
             "SELECT * FROM table_x WHERE unique_enough_column_name = 6"
         )));
@@ -1269,7 +1293,7 @@ mod test {
         bind.put(payload);
 
         let mut qr = QueryRouter::new();
-        qr.pool_settings.automatic_sharding_key = Some("data.id".to_string());
+        qr.pool_settings.automatic_sharding_keys = Some(vec!["data.id".to_string()]);
         qr.pool_settings.shards = 3;
 
         assert!(qr.infer(&simple_query(stmt)));