feat: query-based sharding key lookup for sharded tables

pgdog-config/src/core.rs

@@ -120,6 +120,49 @@ impl ConfigAndUsers {
         self.config.check();
         self.users.check(&self.config);
         self.validate_server_auth()?;
+        self.validate_sharded_tables()?;
+        Ok(())
+    }
+
+    fn validate_sharded_tables(&self) -> Result<(), Error> {
+        for table in &self.config.sharded_tables {
+            let name = table.name.as_deref().unwrap_or("*");
+            match (&table.lookup, &table.query) {
+                (Some(_), None) | (None, Some(_)) => {
+                    return Err(Error::ParseError(format!(
+                        "sharded table \"{}\", column \"{}\": \"lookup\" and \"query\" must be set together",
+                        name, table.column,
+                    )));
+                }
+                (Some(lookup), Some(query)) => {
+                    if query.contains("$1") {
+                        return Err(Error::ParseError(format!(
+                            "sharded table \"{}\", column \"{}\": lookup \"query\" takes no parameters; it loads the whole lookup table (two columns: sharding key value, translated value)",
+                            name, table.column,
+                        )));
+                    }
+
+                    // Lookup queries run on a single shard picked round-robin,
+                    // so the lookup table must have the same data on all shards.
+                    let lookup_name = lookup.rsplit('.').next().unwrap_or(lookup);
+                    let omnisharded = self
+                        .config
+                        .omnisharded_tables
+                        .iter()
+                        .filter(|tables| tables.database == table.database)
+                        .flat_map(|tables| tables.tables.iter())
+                        .any(|t| t == lookup || t == lookup_name);
+                    if !omnisharded {
+                        return Err(Error::ParseError(format!(
+                            "sharded table \"{}\", column \"{}\": lookup table \"{}\" must be listed in [[omnisharded_tables]] for database \"{}\"",
+                            name, table.column, lookup, table.database,
+                        )));
+                    }
+                }
+                (None, None) => (),
+            }
+        }
+
         Ok(())
     }
 

pgdog-config/src/sharding.rs

@@ -86,6 +86,29 @@ pub struct ShardedTableConfig {
     ///
     /// <https://docs.pgdog.dev/configuration/pgdog.toml/sharded_tables/#shard-by-list-and-range>
     pub mapping: Option<Vec<ShardedMappingConfig>>,
+
+    /// Name of the table used to translate sharding key values before hashing.
+    /// The table is loaded into PgDog's memory using `query`; writes to it
+    /// observed by PgDog invalidate and reload the in-memory copy.
+    ///
+    /// **Note:** The table must be listed in `[[omnisharded_tables]]` for the
+    /// same database: it's loaded from a single shard picked round-robin, so
+    /// every shard must have the same data. Must be set together with `query`.
+    #[serde(default)]
+    pub lookup: Option<String>,
+
+    /// Query loading the lookup table into memory, returning two columns:
+    /// the sharding key value, and the value that's hashed to pick a shard,
+    /// e.g. `SELECT id, COALESCE(parent_organization_id, id) FROM organizations`.
+    /// Takes no parameters.
+    ///
+    /// **Note:** The query must return a row for every value that will ever
+    /// route through the sharded column — including values that translate to
+    /// themselves. Values without a row fail with an error: absence can't be
+    /// told apart from a row added after the table was loaded, so it's never
+    /// used for routing. Must be set together with `lookup`.
+    #[serde(default)]
+    pub query: Option<String>,
 }
 
 impl ShardedTableConfig {
@@ -599,9 +622,143 @@ pub struct QueryParser {
 
 #[cfg(test)]
 mod tests {
-    use crate::Config;
+    use crate::{Config, ConfigAndUsers};
+
+    use super::{DataType, QueryParserEngine, QueryParserLevel};
+
+    #[test]
+    fn sharded_table_reads_lookup_from_config() {
+        let source = r#"
+[[sharded_tables]]
+database = "houston"
+column = "organization_id"
+data_type = "varchar"
+lookup = "org_family_roots"
+query = "SELECT organization_id, root_organization_id FROM org_family_roots"
+
+[[omnisharded_tables]]
+database = "houston"
+tables = ["org_family_roots"]
+"#;
+
+        let config: Config = toml::from_str(source).unwrap();
+
+        assert_eq!(config.sharded_tables.len(), 1);
+        let table = &config.sharded_tables[0];
+        assert_eq!(table.database, "houston");
+        assert_eq!(table.column, "organization_id");
+        assert_eq!(table.data_type, DataType::Varchar);
+        assert_eq!(table.lookup.as_deref(), Some("org_family_roots"));
+        assert_eq!(
+            table.query.as_deref(),
+            Some("SELECT organization_id, root_organization_id FROM org_family_roots")
+        );
+
+        let mut config = ConfigAndUsers {
+            config,
+            ..Default::default()
+        };
+        config.check().unwrap();
+    }
+
+    #[test]
+    fn sharded_table_lookup_requires_query() {
+        let source = r#"
+[[sharded_tables]]
+database = "houston"
+column = "organization_id"
+lookup = "org_family_roots"
+"#;
+
+        let mut config = ConfigAndUsers {
+            config: toml::from_str(source).unwrap(),
+            ..Default::default()
+        };
+        let err = config.check().unwrap_err();
+        assert!(err.to_string().contains("must be set together"));
+    }
 
-    use super::{QueryParserEngine, QueryParserLevel};
+    #[test]
+    fn sharded_table_query_requires_lookup() {
+        let source = r#"
+[[sharded_tables]]
+database = "houston"
+column = "organization_id"
+query = "SELECT organization_id, root_organization_id FROM org_family_roots"
+"#;
+
+        let mut config = ConfigAndUsers {
+            config: toml::from_str(source).unwrap(),
+            ..Default::default()
+        };
+        let err = config.check().unwrap_err();
+        assert!(err.to_string().contains("must be set together"));
+    }
+
+    #[test]
+    fn sharded_table_lookup_query_rejects_parameters() {
+        let source = r#"
+[[sharded_tables]]
+database = "houston"
+column = "organization_id"
+lookup = "org_family_roots"
+query = "SELECT root_organization_id FROM org_family_roots WHERE organization_id = $1"
+
+[[omnisharded_tables]]
+database = "houston"
+tables = ["org_family_roots"]
+"#;
+
+        let mut config = ConfigAndUsers {
+            config: toml::from_str(source).unwrap(),
+            ..Default::default()
+        };
+        let err = config.check().unwrap_err();
+        assert!(err.to_string().contains("takes no parameters"));
+    }
+
+    #[test]
+    fn sharded_table_lookup_requires_omnisharded() {
+        let source = r#"
+[[sharded_tables]]
+database = "houston"
+column = "organization_id"
+lookup = "org_family_roots"
+query = "SELECT organization_id, root_organization_id FROM org_family_roots"
+
+[[omnisharded_tables]]
+database = "other_database"
+tables = ["org_family_roots"]
+"#;
+
+        let mut config = ConfigAndUsers {
+            config: toml::from_str(source).unwrap(),
+            ..Default::default()
+        };
+        let err = config.check().unwrap_err();
+        assert!(err.to_string().contains("omnisharded_tables"));
+    }
+
+    #[test]
+    fn sharded_table_lookup_omnisharded_matches_qualified_name() {
+        let source = r#"
+[[sharded_tables]]
+database = "houston"
+column = "organization_id"
+lookup = "app.org_family_roots"
+query = "SELECT organization_id, root_organization_id FROM app.org_family_roots"
+
+[[omnisharded_tables]]
+database = "houston"
+tables = ["org_family_roots"]
+"#;
+
+        let mut config = ConfigAndUsers {
+            config: toml::from_str(source).unwrap(),
+            ..Default::default()
+        };
+        config.check().unwrap();
+    }
 
     #[test]
     fn query_parser_reads_default_values_from_config() {

pgdog/src/backend/databases.rs

@@ -24,7 +24,7 @@ use crate::config::PoolerMode;
 use crate::frontend::PreparedStatements;
 use crate::frontend::client::query_engine::two_pc::Manager;
 use crate::frontend::router::parser::Cache;
-use crate::frontend::router::sharding::{Mapping, ShardedTable};
+use crate::frontend::router::sharding::{LookupCache, Mapping, ShardedTable};
 use crate::{
     backend::pool::PoolConfig,
     config::{
@@ -74,6 +74,20 @@ pub fn replace_databases(new_databases: Databases, reload: bool) -> Result<(), E
     }
     // 3. Launch new databases first.
     new_databases.launch();
+
+    // Drop loaded sharding key lookup maps if the lookup configuration
+    // changed; they reload cold under the new config on next use.
+    let mut lookups = Vec::new();
+    for cluster in new_databases.all().values() {
+        let schema = cluster.sharding_schema();
+        for table in schema.tables.tables() {
+            if let (Some(lookup), Some(query)) = (&table.lookup, &table.query) {
+                lookups.push((lookup.clone(), query.clone()));
+            }
+        }
+    }
+    LookupCache::get().update_config(lookups);
+
     DATABASES.store(new_databases);
     // 4. Shutdown all databases.
     old_databases.shutdown();
@@ -516,6 +530,8 @@ fn resolve_sharded_table(
         centroid_probes: config.centroid_probes,
         hasher: config.hasher.clone(),
         mapping: mapping.flatten(),
+        lookup: config.lookup.clone(),
+        query: config.query.clone(),
     }
 }
 

pgdog/src/backend/replication/sharded_tables.rs

@@ -21,6 +21,8 @@ struct Inner {
     common_mapping: Option<CommonMapping>,
     omnisharded_sticky: bool,
     system_catalogs: SystemCatalogsBehavior,
+    /// At least one table translates sharding keys through a lookup table.
+    has_lookups: bool,
 }
 
 #[derive(Debug)]
@@ -85,6 +87,8 @@ impl ShardedTables {
             _ => None,
         };
 
+        let has_lookups = tables.iter().any(|table| table.lookup.is_some());
+
         Self {
             inner: Arc::new(Inner {
                 tables,
@@ -95,6 +99,7 @@ impl ShardedTables {
                 common_mapping,
                 omnisharded_sticky,
                 system_catalogs,
+                has_lookups,
             }),
         }
     }
@@ -103,6 +108,11 @@ impl ShardedTables {
         &self.inner.tables
     }
 
+    /// At least one table translates sharding keys through a lookup table.
+    pub fn has_lookups(&self) -> bool {
+        self.inner.has_lookups
+    }
+
     pub fn omnishards(&self) -> &HashMap<String, bool> {
         &self.inner.omnisharded
     }

pgdog/src/frontend/client/query_engine/end_transaction.rs

@@ -9,6 +9,11 @@ impl QueryEngine {
         rollback: bool,
         extended: bool,
     ) -> Result<(), Error> {
+        // The transaction is over on all shards (2PC transactions land
+        // here after phase 2). Invalidate lookup tables it wrote to
+        // before the client learns the transaction is done.
+        self.settle_lookup_invalidations();
+
         let bytes_sent = if extended {
             self.extended_transaction_reply(context, false, rollback)
                 .await?

pgdog/src/frontend/client/query_engine/mod.rs

@@ -4,7 +4,7 @@ use crate::{
     frontend::{
         BufferedQuery, Client, ClientComms, Command, Error, Router, RouterContext, Stats,
         client::query_engine::{hooks::QueryEngineHooks, route_query::ClusterCheck},
-        router::{Route, parser::Shard},
+        router::{Route, parser::Shard, sharding::LookupCache},
     },
     net::{ErrorResponse, Message, Parameters},
     state::State,
@@ -66,6 +66,10 @@ pub struct QueryEngine {
     // They will remain pinned to their connection until they unpin manually
     // or disconnect.
     manual_lock: bool,
+    // Lookup tables written to by the current statement or transaction.
+    // Their cached sharding key translations are flushed when the write
+    // completes, before the client sees the result.
+    lookup_invalidations: Vec<String>,
 }
 
 impl QueryEngine {
@@ -89,9 +93,30 @@ impl QueryEngine {
             router: Router::default(),
             advisory_locks: AdvisoryLocks::default(),
             manual_lock: false,
+            lookup_invalidations: Vec::new(),
         })
     }
 
+    /// Flush cached sharding key translations for lookup tables written
+    /// to by the statement or transaction that just completed. Called
+    /// before the client is told the write is done, so every query routed
+    /// after that reads the new mapping.
+    fn settle_lookup_invalidations(&mut self) {
+        for lookup in self.lookup_invalidations.drain(..) {
+            LookupCache::get().flush_lookup_table(&lookup);
+        }
+    }
+}
+
+impl Drop for QueryEngine {
+    fn drop(&mut self) {
+        // The client disconnected with a lookup table write whose outcome
+        // we never observed; invalidate conservatively.
+        self.settle_lookup_invalidations();
+    }
+}
+
+impl QueryEngine {
     pub fn from_client(client: &Client) -> Result<Self, Error> {
         Self::new(&client.params, &client.comms, client.admin)
     }

pgdog/src/frontend/client/query_engine/query.rs

@@ -233,6 +233,12 @@ impl QueryEngine {
 
             if !context.in_transaction() {
                 self.stats.transaction(two_pc_auto);
+
+                // The statement or transaction is complete on all shards.
+                // Invalidate lookup tables it wrote to before the client
+                // learns the write is done, so every query routed after
+                // that reads the new mapping.
+                self.settle_lookup_invalidations();
             }
         }