Multi-tenant (#161)

* Multi-tenant

* save

* clippy

* save

* savee

* routing tests

* only load schema if multi-tenant

* shortcut if not multi-tenant

Author: levkk

pgdog.toml

@@ -106,3 +106,6 @@ fingerprint = "f4814b6fadabc4c1" #[17618446160277259457]
 
 [[manual_query]]
 fingerprint = "04dc05f480b702d3"
+
+[multi_tenant]
+column = "tenant_id"

pgdog/src/backend/databases.rs

@@ -363,8 +363,14 @@ pub(crate) fn new_pool(
             }
         };
 
-        let cluster_config =
-            ClusterConfig::new(general, user, &shard_configs, sharded_tables, mirror_of);
+        let cluster_config = ClusterConfig::new(
+            general,
+            user,
+            &shard_configs,
+            sharded_tables,
+            mirror_of,
+            config.multi_tenant(),
+        );
 
         Some((
             User {

pgdog/src/backend/pool/cluster.rs

@@ -1,20 +1,23 @@
 //! A collection of replicas and a primary.
 
+use parking_lot::RwLock;
+use std::sync::Arc;
+use tokio::spawn;
+use tracing::{error, info};
+
 use crate::{
     backend::{
         databases::databases,
         replication::{ReplicationConfig, ShardedColumn},
-        ShardedTables,
+        Schema, ShardedTables,
     },
-    config::{General, PoolerMode, ShardedTable, User},
+    config::{General, MultiTenant, PoolerMode, ShardedTable, User},
     net::messages::BackendKeyData,
 };
 
 use super::{Address, Config, Error, Guard, Request, Shard};
 use crate::config::LoadBalancingStrategy;
 
-use std::ffi::CString;
-
 #[derive(Clone, Debug)]
 /// Database configuration.
 pub struct PoolConfig {
@@ -36,6 +39,8 @@ pub struct Cluster {
     sharded_tables: ShardedTables,
     replication_sharding: Option<String>,
     mirror_of: Option<String>,
+    schema: Arc<RwLock<Schema>>,
+    multi_tenant: Option<MultiTenant>,
 }
 
 /// Sharding configuration from the cluster.
@@ -63,6 +68,7 @@ pub struct ClusterConfig<'a> {
     pub sharded_tables: ShardedTables,
     pub replication_sharding: Option<String>,
     pub mirror_of: Option<&'a str>,
+    pub multi_tenant: &'a Option<MultiTenant>,
 }
 
 impl<'a> ClusterConfig<'a> {
@@ -72,6 +78,7 @@ impl<'a> ClusterConfig<'a> {
         shards: &'a [ClusterShardConfig],
         sharded_tables: ShardedTables,
         mirror_of: Option<&'a str>,
+        multi_tenant: &'a Option<MultiTenant>,
     ) -> Self {
         Self {
             name: &user.database,
@@ -83,6 +90,7 @@ impl<'a> ClusterConfig<'a> {
             shards,
             sharded_tables,
             mirror_of,
+            multi_tenant,
         }
     }
 }
@@ -100,6 +108,7 @@ impl Cluster {
             sharded_tables,
             replication_sharding,
             mirror_of,
+            multi_tenant,
         } = config;
 
         Self {
@@ -114,6 +123,8 @@ impl Cluster {
             sharded_tables,
             replication_sharding,
             mirror_of: mirror_of.map(|s| s.to_owned()),
+            schema: Arc::new(RwLock::new(Schema::default())),
+            multi_tenant: multi_tenant.clone(),
         }
     }
 
@@ -160,6 +171,8 @@ impl Cluster {
             sharded_tables: self.sharded_tables.clone(),
             replication_sharding: self.replication_sharding.clone(),
             mirror_of: self.mirror_of.clone(),
+            schema: self.schema.clone(),
+            multi_tenant: self.multi_tenant.clone(),
         }
     }
 
@@ -182,52 +195,6 @@ impl Cluster {
         self.mirror_of.as_deref()
     }
 
-    /// Plugin input.
-    ///
-    /// # Safety
-    ///
-    /// This allocates, so make sure to call `Config::drop` when you're done.
-    ///
-    pub unsafe fn plugin_config(&self) -> Result<pgdog_plugin::bindings::Config, Error> {
-        use pgdog_plugin::bindings::{Config, DatabaseConfig, Role_PRIMARY, Role_REPLICA};
-        let mut databases: Vec<DatabaseConfig> = vec![];
-        let name = CString::new(self.name.as_str()).map_err(|_| Error::NullBytes)?;
-
-        for (index, shard) in self.shards.iter().enumerate() {
-            if let Some(ref primary) = shard.primary {
-                // Ignore hosts with null bytes.
-                let host = if let Ok(host) = CString::new(primary.addr().host.as_str()) {
-                    host
-                } else {
-                    continue;
-                };
-                databases.push(DatabaseConfig::new(
-                    host,
-                    primary.addr().port,
-                    Role_PRIMARY,
-                    index,
-                ));
-            }
-
-            for replica in shard.replicas.pools() {
-                // Ignore hosts with null bytes.
-                let host = if let Ok(host) = CString::new(replica.addr().host.as_str()) {
-                    host
-                } else {
-                    continue;
-                };
-                databases.push(DatabaseConfig::new(
-                    host,
-                    replica.addr().port,
-                    Role_REPLICA,
-                    index,
-                ));
-            }
-        }
-
-        Ok(Config::new(name, &databases, self.shards.len()))
-    }
-
     /// Get the password the user should use to connect to the database.
     pub fn password(&self) -> &str {
         &self.password
@@ -280,6 +247,11 @@ impl Cluster {
         true
     }
 
+    /// Multi-tenant config.
+    pub fn multi_tenant(&self) -> &Option<MultiTenant> {
+        &self.multi_tenant
+    }
+
     /// Get replication configuration for this cluster.
     pub fn replication_sharding_config(&self) -> Option<ReplicationConfig> {
         self.replication_sharding
@@ -295,11 +267,42 @@ impl Cluster {
         }
     }
 
+    /// Update schema from primary.
+    async fn update_schema(&self) -> Result<(), crate::backend::Error> {
+        let mut server = self.primary(0, &Request::default()).await?;
+        let schema = Schema::load(&mut server).await?;
+        info!(
+            "loaded {} tables from schema [{}]",
+            schema.tables().len(),
+            server.addr()
+        );
+        *self.schema.write() = schema;
+        Ok(())
+    }
+
+    fn load_schema(&self) -> bool {
+        self.multi_tenant.is_some()
+    }
+
+    /// Get currently loaded schema.
+    pub fn schema(&self) -> Schema {
+        self.schema.read().clone()
+    }
+
     /// Launch the connection pools.
     pub(crate) fn launch(&self) {
         for shard in self.shards() {
             shard.launch();
         }
+
+        if self.load_schema() {
+            let me = self.clone();
+            spawn(async move {
+                if let Err(err) = me.update_schema().await {
+                    error!("error loading schema: {}", err);
+                }
+            });
+        }
     }
 
     /// Shutdown the connection pools.

pgdog/src/backend/pool/connection/mirror.rs

@@ -6,7 +6,8 @@ use tracing::{debug, error};
 use crate::backend::Cluster;
 use crate::config::config;
 use crate::frontend::client::timeouts::Timeouts;
-use crate::frontend::{PreparedStatements, Router};
+use crate::frontend::{PreparedStatements, Router, RouterContext};
+use crate::net::Parameters;
 use crate::state::State;
 use crate::{
     backend::pool::{Error as PoolError, Request},
@@ -37,6 +38,7 @@ pub(crate) struct Mirror {
     router: Router,
     cluster: Cluster,
     prepared_statements: PreparedStatements,
+    params: Parameters,
     state: State,
 }
 
@@ -50,6 +52,7 @@ impl Mirror {
             prepared_statements: PreparedStatements::new(),
             cluster: cluster.clone(),
             state: State::Idle,
+            params: Parameters::default(),
         };
 
         let config = config();
@@ -110,18 +113,21 @@ impl Mirror {
     pub(crate) async fn handle(&mut self, request: &MirrorRequest) -> Result<(), Error> {
         if !self.connection.connected() {
             // TODO: handle parsing errors.
-            if let Err(err) = self.router.query(
+            if let Ok(context) = RouterContext::new(
                 &request.buffer,
                 &self.cluster,
                 &mut self.prepared_statements,
+                &self.params,
             ) {
-                error!("mirror query parse error: {}", err);
-                return Ok(()); // Drop request.
-            }
+                if let Err(err) = self.router.query(context) {
+                    error!("mirror query parse error: {}", err);
+                    return Ok(()); // Drop request.
+                }
 
-            self.connection
-                .connect(&request.request, &self.router.route())
-                .await?;
+                self.connection
+                    .connect(&request.request, &self.router.route())
+                    .await?;
+            }
         }
 
         // TODO: handle streaming.

pgdog/src/backend/schema/mod.rs

@@ -2,7 +2,8 @@
 pub mod columns;
 pub mod relation;
 
-use std::collections::HashMap;
+use std::sync::Arc;
+use std::{collections::HashMap, ops::Deref};
 use tracing::debug;
 
 pub use relation::Relation;
@@ -11,10 +12,16 @@ use super::{pool::Request, Cluster, Error, Server};
 
 static SETUP: &str = include_str!("setup.sql");
 
+#[derive(Debug, Default)]
+struct Inner {
+    search_path: Vec<String>,
+    relations: HashMap<(String, String), Relation>,
+}
+
 /// Load schema from database.
 #[derive(Debug, Clone, Default)]
 pub struct Schema {
-    relations: HashMap<(String, String), Relation>,
+    inner: Arc<Inner>,
 }
 
 impl Schema {
@@ -31,7 +38,23 @@ impl Schema {
             })
             .collect();
 
-        Ok(Self { relations })
+        let search_path = server
+            .fetch_all::<String>("SHOW search_path")
+            .await?
+            .pop()
+            .unwrap_or(String::from("$user, public"))
+            .split(",")
+            .map(|p| p.trim().replace("\"", ""))
+            .collect();
+
+        let inner = Inner {
+            search_path,
+            relations,
+        };
+
+        Ok(Self {
+            inner: Arc::new(inner),
+        })
     }
 
     /// Load schema from primary database.
@@ -68,7 +91,7 @@ impl Schema {
                     .iter()
                     .filter(|table| table.schema() != "pgdog")
                 {
-                    let column_match = schema_table.columns.iter().find(|column| {
+                    let column_match = schema_table.columns.values().find(|column| {
                         column.column_name == table.column && column.data_type == "bigint"
                     });
                     if let Some(column_match) = column_match {
@@ -110,24 +133,41 @@ impl Schema {
     /// Get table by name.
     pub fn table(&self, name: &str, schema: Option<&str>) -> Option<&Relation> {
         let schema = schema.unwrap_or("public");
-        self.relations.get(&(name.to_string(), schema.to_string()))
+        self.inner
+            .relations
+            .get(&(name.to_string(), schema.to_string()))
     }
 
     /// Get all indices.
     pub fn tables(&self) -> Vec<&Relation> {
-        self.relations
+        self.inner
+            .relations
             .values()
             .filter(|value| value.is_table())
             .collect()
     }
 
     /// Get all sequences.
     pub fn sequences(&self) -> Vec<&Relation> {
-        self.relations
+        self.inner
+            .relations
             .values()
             .filter(|value| value.is_sequence())
             .collect()
     }
+
+    /// Get search path components.
+    pub fn search_path(&self) -> &[String] {
+        &self.inner.search_path
+    }
+}
+
+impl Deref for Schema {
+    type Target = HashMap<(String, String), Relation>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner.relations
+    }
 }
 
 #[cfg(test)]