Allow shard setting with comments (#293)

What
Allows shard selection by the client to come in via comments like /* shard_id: 1 */ select * from foo;

Why
We're using a setup in Ruby that makes it tough or impossible to inject commands on the connection to set the shard before it gets to the "real" SQL being run. Instead we have an updated PG adapter that allows injection of comments before each executed SQL statement. We need this support in pgcat in order to keep some complex shard picking logic in Ruby code while using pgcat for connection management.

Local Testing
Run postgres and pgcat with the default options. Run psql < tests/sharding/query_routing_setup.sql to setup the database for the tests and run ./tests/pgbench/external_shard_test.sh as often as needed to exercise the shard setting comment test.

src/config.rs

@@ -2,6 +2,7 @@
 use arc_swap::ArcSwap;
 use log::{error, info};
 use once_cell::sync::Lazy;
+use regex::Regex;
 use serde_derive::{Deserialize, Serialize};
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::hash::Hash;
@@ -342,8 +343,15 @@ pub struct Pool {
     #[serde(default = "Pool::default_automatic_sharding_key")]
     pub automatic_sharding_key: Option<String>,
 
+    pub sharding_key_regex: Option<String>,
+    pub shard_id_regex: Option<String>,
+    pub regex_search_limit: Option<usize>,
+
     pub shards: BTreeMap<String, Shard>,
     pub users: BTreeMap<String, User>,
+    // Note, don't put simple fields below these configs. There's a compatability issue with TOML that makes it
+    // incompatible to have simple fields in TOML after complex objects. See
+    // https://users.rust-lang.org/t/why-toml-to-string-get-error-valueaftertable/85903
 }
 
 impl Pool {
@@ -387,6 +395,18 @@ impl Pool {
             shard.validate()?;
         }
 
+        for (option, name) in [
+            (&self.shard_id_regex, "shard_id_regex"),
+            (&self.sharding_key_regex, "sharding_key_regex"),
+        ] {
+            if let Some(regex) = option {
+                if let Err(parse_err) = Regex::new(regex.as_str()) {
+                    error!("{} is not a valid Regex: {}", name, parse_err);
+                    return Err(Error::BadConfig);
+                }
+            }
+        }
+
         Ok(())
     }
 }
@@ -405,6 +425,9 @@ impl Default for Pool {
             automatic_sharding_key: None,
             connect_timeout: None,
             idle_timeout: None,
+            sharding_key_regex: None,
+            shard_id_regex: None,
+            regex_search_limit: Some(1000),
         }
     }
 }