save

Author: levkk

integration/postgres_fdw/pgdog.toml

@@ -13,6 +13,20 @@ shard = 1
 host = "127.0.0.1"
 database_name = "shard_1_fdw"
 
+[[databases]]
+name = "pgdog"
+shard = 0
+host = "127.0.0.1"
+database_name = "shard_0_fdw"
+role = "replica"
+
+[[databases]]
+name = "pgdog"
+shard = 1
+host = "127.0.0.1"
+database_name = "shard_1_fdw"
+role = "replica"
+
 [[sharded_tables]]
 column = "user_id"
 database = "pgdog"

integration/postgres_fdw/users.toml

@@ -2,3 +2,8 @@
 name = "pgdog"
 password = "pgdog"
 database = "pgdog"
+
+[[users]]
+name = "lev"
+password = "lev"
+database = "pgdog"

pgdog-config/src/fdw.rs

@@ -3,9 +3,6 @@ use serde::{Deserialize, Serialize};
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Copy)]
 #[serde(deny_unknown_fields)]
 pub struct Fdw {
-    #[serde(default)]
-    pub enabled: bool,
-
     #[serde(default = "default_port")]
     pub port: u16,
 
@@ -16,7 +13,6 @@ pub struct Fdw {
 impl Default for Fdw {
     fn default() -> Self {
         Self {
-            enabled: bool::default(),
             port: default_port(),
             launch_timeout: default_launch_timeout(),
         }

pgdog-config/src/general.rs

@@ -5,7 +5,9 @@ use std::path::PathBuf;
 use std::time::Duration;
 
 use crate::pooling::ConnectionRecovery;
-use crate::{CopyFormat, QueryParserEngine, QueryParserLevel, SystemCatalogsBehavior};
+use crate::{
+    CopyFormat, CrossShardBackend, QueryParserEngine, QueryParserLevel, SystemCatalogsBehavior,
+};
 
 use super::auth::{AuthType, PassthoughAuth};
 use super::database::{LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy};
@@ -203,6 +205,9 @@ pub struct General {
     /// Copy format used for resharding.
     #[serde(default)]
     pub resharding_copy_format: CopyFormat,
+    /// Cross-shard backend.
+    #[serde(default = "General::cross_shard_backend")]
+    pub cross_shard_backend: CrossShardBackend,
 }
 
 impl Default for General {
@@ -274,6 +279,7 @@ impl Default for General {
             system_catalogs: Self::default_system_catalogs(),
             omnisharded_sticky: bool::default(),
             resharding_copy_format: CopyFormat::default(),
+            cross_shard_backend: Self::cross_shard_backend(),
         }
     }
 }
@@ -398,6 +404,10 @@ impl General {
         )
     }
 
+    fn cross_shard_backend() -> CrossShardBackend {
+        Self::env_enum_or_default("PGDOG_CROSS_SHARD_BACKEND")
+    }
+
     pub fn query_timeout(&self) -> Duration {
         Duration::from_millis(self.query_timeout)
     }

pgdog-config/src/sharding.rs

@@ -363,3 +363,41 @@ impl Display for CopyFormat {
         }
     }
 }
+
+#[derive(Serialize, Deserialize, Debug, Copy, Clone, PartialEq, Eq, Hash, Default)]
+#[serde(rename_all = "snake_case", deny_unknown_fields)]
+pub enum CrossShardBackend {
+    #[default]
+    PgDog,
+    Fdw,
+    Hybrid,
+}
+
+impl CrossShardBackend {
+    pub fn need_fdw(&self) -> bool {
+        matches!(self, Self::Fdw | Self::Hybrid)
+    }
+}
+
+impl Display for CrossShardBackend {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::PgDog => write!(f, "pgdog"),
+            Self::Fdw => write!(f, "fdw"),
+            Self::Hybrid => write!(f, "hybrid"),
+        }
+    }
+}
+
+impl FromStr for CrossShardBackend {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "pgdog" => Ok(Self::PgDog),
+            "fdw" => Ok(Self::Fdw),
+            "hybrid" => Ok(Self::Hybrid),
+            _ => Err(()),
+        }
+    }
+}

pgdog/src/backend/databases.rs

@@ -91,7 +91,7 @@ pub fn init() -> Result<(), Error> {
     let config = config();
 
     // Start postgres_fdw compatibility engine.
-    if config.config.fdw.enabled {
+    if config.config.general.cross_shard_backend.need_fdw() {
         PostgresLauncher::get().launch();
     }
 

pgdog/src/backend/pool/cluster.rs

@@ -1,7 +1,10 @@
 //! A collection of replicas and a primary.
 
 use parking_lot::Mutex;
-use pgdog_config::{PreparedStatements, QueryParserEngine, QueryParserLevel, Rewrite, RewriteMode};
+use pgdog_config::{
+    CrossShardBackend, PreparedStatements, QueryParserEngine, QueryParserLevel, Rewrite,
+    RewriteMode,
+};
 use std::{
     sync::{
         atomic::{AtomicBool, AtomicUsize, Ordering},
@@ -76,6 +79,7 @@ pub struct Cluster {
     lb_strategy: LoadBalancingStrategy,
     rw_split: ReadWriteSplit,
     fdw_lb: Option<FdwLoadBalancer>,
+    cross_shard_backend: CrossShardBackend,
 }
 
 /// Sharding configuration from the cluster.
@@ -148,6 +152,7 @@ pub struct ClusterConfig<'a> {
     pub query_parser_engine: QueryParserEngine,
     pub connection_recovery: ConnectionRecovery,
     pub lsn_check_interval: Duration,
+    pub cross_shard_backend: CrossShardBackend,
 }
 
 impl<'a> ClusterConfig<'a> {
@@ -195,6 +200,7 @@ impl<'a> ClusterConfig<'a> {
             query_parser_engine: general.query_parser_engine,
             connection_recovery: general.connection_recovery,
             lsn_check_interval: Duration::from_millis(general.lsn_check_interval),
+            cross_shard_backend: general.cross_shard_backend,
         }
     }
 }
@@ -228,6 +234,7 @@ impl Cluster {
             connection_recovery,
             lsn_check_interval,
             query_parser_engine,
+            cross_shard_backend,
         } = config;
 
         let identifier = Arc::new(DatabaseUser {
@@ -276,9 +283,12 @@ impl Cluster {
             lb_strategy,
             rw_split,
             fdw_lb: None,
+            cross_shard_backend,
         };
 
-        cluster.fdw_lb = FdwLoadBalancer::new(&cluster).ok();
+        if cross_shard_backend.need_fdw() {
+            cluster.fdw_lb = FdwLoadBalancer::new(&cluster).ok();
+        }
 
         cluster
     }
@@ -460,6 +470,10 @@ impl Cluster {
         true
     }
 
+    pub fn cross_shard_backend(&self) -> CrossShardBackend {
+        self.cross_shard_backend
+    }
+
     /// This database/user pair is responsible for schema management.
     pub fn schema_admin(&self) -> bool {
         self.schema_admin

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

@@ -30,7 +30,9 @@ impl QueryEngine {
         }
 
         let connect_route = connect_route.unwrap_or(context.client_request.route());
-        let connect_route = if context.params.is_postgres_fdw() {
+        let connect_route = if (context.params.is_postgres_fdw() || connect_route.is_cross_shard())
+            && self.backend.cluster()?.cross_shard_backend().need_fdw()
+        {
             lazy_static! {
                 static ref FDW_ROUTE: Route = Route::fdw_fallback();
             }