Handle custom types with differing OIDs across shards

.schema/pgdog.schema.json

@@ -31,6 +31,7 @@
         "ban_timeout": 300000,
         "broadcast_address": null,
         "broadcast_port": 6433,
+        "canonicalize_type_information": false,
         "checkout_timeout": 5000,
         "client_connection_recovery": "drop",
         "client_idle_in_transaction_timeout": 9223372036854775807,
@@ -634,6 +635,11 @@
           "maximum": 65535,
           "minimum": 0
         },
+        "canonicalize_type_information": {
+          "description": "Controls whether PgDog maps each shard's type information to a canonical source.\n\n_Default:_ `false`\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/general/#canonicalize_type_information>",
+          "type": "boolean",
+          "default": false
+        },
         "checkout_timeout": {
           "description": "Maximum amount of time a client is allowed to wait for a connection from the pool.\n\n_Default:_ `5000`\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/general/#checkout_timeout>",
           "type": "integer",

integration/rust/Cargo.toml

@@ -6,6 +6,9 @@ edition = "2024"
 [lib]
 test = true
 
+[features]
+new_parser = []
+
 [dependencies]
 tokio-postgres = {version = "0.7.13", features = ["with-uuid-1"]}
 sqlx = { version = "0.8.6", features = ["postgres", "runtime-tokio", "tls-native-tls", "bigdecimal", "chrono", "json", "rust_decimal"]}

integration/rust/tests/integration/cross_shard_oid_drift.rs

@@ -0,0 +1,77 @@
+#![cfg(feature = "new_parser")]
+use crate::setup::{admin_sqlx, connections_sqlx};
+use sqlx::postgres::types::Oid;
+use sqlx::{Column, Executor, Row};
+
+#[derive(sqlx::Type, Debug, Clone, PartialEq)]
+#[sqlx(type_name = "test_oid_drift_composite")]
+struct Composite {
+    a: String,
+    b: String,
+}
+
+#[tokio::test]
+async fn test_oid_drift() {
+    let conn = connections_sqlx().await.pop().unwrap();
+    let admin = admin_sqlx().await;
+
+    // Intentionally cause the OID of the type to differ between shards
+    conn.execute("/* pgdog_shard: 0 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;")
+        .await
+        .unwrap();
+    conn.execute("DROP TYPE IF EXISTS test_oid_drift_composite CASCADE")
+        .await
+        .unwrap();
+    conn.execute("CREATE TYPE test_oid_drift_composite AS (a text, b text)")
+        .await
+        .unwrap();
+    conn.execute("DROP TABLE IF EXISTS test_oid_drift")
+        .await
+        .unwrap();
+    conn.execute(
+        "CREATE TABLE test_oid_drift (customer_id BIGINT, composite test_oid_drift_composite)",
+    )
+    .await
+    .unwrap();
+    admin
+        .execute("SET canonicalize_type_information TO true")
+        .await
+        .unwrap();
+
+    let composite = Composite {
+        a: String::from("a"),
+        b: String::from("b"),
+    };
+    for i in 1..=20 {
+        sqlx::query("INSERT INTO test_oid_drift VALUES ($1, $2)")
+            .bind(i)
+            .bind(&composite)
+            .execute(&conn)
+            .await
+            .unwrap();
+    }
+
+    let rows: Vec<Composite> = sqlx::query_scalar("SELECT composite FROM test_oid_drift")
+        .fetch_all(&conn)
+        .await
+        .unwrap();
+    assert_eq!(rows, vec![composite.clone(); 20]);
+
+    let simple_rows = conn
+        .fetch_all("SELECT composite FROM test_oid_drift")
+        .await
+        .unwrap();
+
+    let expected_oid: Oid =
+        sqlx::query_scalar("SELECT oid FROM pg_type WHERE typname = 'test_oid_drift_composite'")
+            .fetch_one(&conn)
+            .await
+            .unwrap();
+    let given_oid = simple_rows.first().unwrap().column(0).type_info().oid();
+    assert_eq!(given_oid, Some(expected_oid));
+
+    let simple_data: Vec<Composite> = simple_rows.into_iter().map(|row| row.get(0)).collect();
+    assert_eq!(simple_data, vec![composite; 20]);
+
+    admin.execute("RELOAD").await.unwrap();
+}

integration/rust/tests/integration/mod.rs

@@ -10,6 +10,7 @@ pub mod client_ids;
 pub mod connection_recovery;
 pub mod copy;
 pub mod cross_shard_disabled;
+mod cross_shard_oid_drift;
 pub mod distinct;
 pub mod explain;
 pub mod fake_transactions;

pgdog-config/src/general.rs

@@ -767,6 +767,14 @@ pub struct General {
     #[serde(default = "General::load_schema")]
     pub load_schema: LoadSchema,
 
+    /// Controls whether PgDog maps each shard's type information to a canonical source.
+    ///
+    /// _Default:_ `false`
+    ///
+    /// <https://docs.pgdog.dev/configuration/pgdog.toml/general/#canonicalize_type_information>
+    #[serde(default = "General::canonicalize_type_information")]
+    pub canonicalize_type_information: bool,
+
     /// Replication lag threshold (in bytes) at which PgDog will pause traffic automatically during a traffic cutover.
     ///
     /// _Default:_ `1000000`
@@ -915,6 +923,7 @@ impl Default for General {
             resharding_replication_retry_min_delay: Self::resharding_replication_retry_min_delay(),
             reload_schema_on_ddl: Self::reload_schema_on_ddl(),
             load_schema: Self::load_schema(),
+            canonicalize_type_information: Self::canonicalize_type_information(),
             cutover_replication_lag_threshold: Self::cutover_replication_lag_threshold(),
             cutover_traffic_stop_threshold: Self::cutover_traffic_stop_threshold(),
             cutover_last_transaction_delay: Self::cutover_last_transaction_delay(),
@@ -1327,6 +1336,10 @@ impl General {
         Self::env_enum_or_default("PGDOG_LOAD_SCHEMA")
     }
 
+    fn canonicalize_type_information() -> bool {
+        Self::env_or_default("PGDOG_CANONICALIZE_TYPE_INFORMATION", false)
+    }
+
     pub fn mirror_queue() -> usize {
         Self::env_or_default("PGDOG_MIRROR_QUEUE", 128)
     }

pgdog/benches/comment_parser.rs

@@ -1,7 +1,6 @@
 use brunch::{Bench, benches};
 #[cfg(not(feature = "new_parser"))]
 use pg_query::scan_raw;
-use pgdog::backend::ShardingSchema;
 use pgdog::frontend::router::parser::comment::parse_edge_comment;
 
 const QUERY_WITH_LEADING: &str =
@@ -13,20 +12,20 @@ const QUERY_NO_COMMENT: &str = "SELECT * FROM users WHERE id = $1 AND name = $2"
 #[cfg(feature = "new_parser")]
 benches!(
     Bench::new("parse_edge_comment(leading)")
-        .run(|| parse_edge_comment(QUERY_WITH_LEADING, &ShardingSchema::default())),
+        .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())),
     Bench::new("parse_edge_comment(trailing)")
-        .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &ShardingSchema::default())),
+        .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &Default::default())),
     Bench::new("parse_edge_comment(no comment)")
-        .run(|| parse_edge_comment(QUERY_NO_COMMENT, &ShardingSchema::default())),
+        .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())),
 );
 #[cfg(not(feature = "new_parser"))]
 benches!(
     Bench::new("parse_edge_comment(leading)")
-        .run(|| parse_edge_comment(QUERY_WITH_LEADING, &ShardingSchema::default())),
+        .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())),
     Bench::new("parse_edge_comment(trailing)")
-        .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &ShardingSchema::default())),
+        .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &Default::default())),
     Bench::new("parse_edge_comment(no comment)")
-        .run(|| parse_edge_comment(QUERY_NO_COMMENT, &ShardingSchema::default())),
+        .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())),
     Bench::new("scan_raw(leading)").run(|| scan_raw(QUERY_WITH_LEADING)),
     Bench::new("scan_raw(trailing)").run(|| scan_raw(QUERY_WITH_TRAILING)),
     Bench::new("scan_raw(no comment)").run(|| scan_raw(QUERY_NO_COMMENT)),

pgdog/src/admin/probe.rs

@@ -32,6 +32,7 @@ impl Command for Probe {
                 &Address::try_from(self.url.clone()).map_err(|_| Error::InvalidAddress)?,
                 ServerOptions::default(),
                 ConnectReason::Probe,
+                Default::default(),
             ),
         )
         .await?

pgdog/src/admin/set.rs

@@ -211,6 +211,10 @@ impl Command for Set {
                 config.config.general.connect_timeout = self.value.parse()?;
             }
 
+            "canonicalize_type_information" => {
+                config.config.general.canonicalize_type_information = Self::from_json(&self.value)?;
+            }
+
             _ => return Ok(vec![]),
         }
 

pgdog/src/backend/error.rs

@@ -143,6 +143,9 @@ pub enum Error {
 
     #[error("cannot ignore response for message type: {0}")]
     UnsupportedHandleIgnore(char),
+
+    #[error("missing canonical oid for type {0}")]
+    MissingCanonicalOid(String),
 }
 
 impl From<crate::frontend::Error> for Error {

pgdog/src/backend/mod.rs

@@ -21,7 +21,9 @@ pub mod validation;
 pub use connect_reason::ConnectReason;
 pub use disconnect_reason::DisconnectReason;
 pub use error::Error;
-pub use pool::{Cluster, ClusterShardConfig, LoadBalancer, Pool, Shard, ShardingSchema};
+pub(crate) use pool::{
+    CanonicalOids, Cluster, ClusterShardConfig, Oids, Pool, Shard, ShardingSchema,
+};
 pub use prepared_statements::PreparedStatements;
 pub use protocol::*;
 pub use pub_sub::{PubSubClient, PubSubListener};

pgdog/src/backend/pool/cluster.rs

@@ -23,7 +23,9 @@ use crate::{
     net::{Query, messages::FrontendPid},
 };
 
-use super::{Address, Config, Error, Guard, MirrorStats, Request, Shard, ShardConfig};
+use super::{
+    Address, CanonicalOids, Config, Error, Guard, MirrorStats, Request, Shard, ShardConfig,
+};
 use crate::config::LoadBalancingStrategy;
 use launch::Readiness;
 
@@ -85,6 +87,7 @@ pub struct Cluster {
     tls_client_certificate_required: bool,
     #[debug(skip)]
     schema_loader: Box<dyn SchemaLoader>,
+    canonical_oids: Option<Arc<CanonicalOids>>,
 }
 
 /// Sharding configuration from the cluster.
@@ -174,6 +177,7 @@ pub struct ClusterConfig<'a> {
     identity: &'a Option<String>,
     tls_client_certificate_required: bool,
     schema_cache: SchemaCache,
+    canonicalize_oids: bool,
 }
 
 impl<'a> ClusterConfig<'a> {
@@ -245,6 +249,7 @@ impl<'a> ClusterConfig<'a> {
             identity: &user.identity,
             tls_client_certificate_required: user.tls_client_certificate_required.unwrap_or(true),
             schema_cache,
+            canonicalize_oids: general.canonicalize_type_information,
         }
     }
 }
@@ -293,12 +298,14 @@ impl Cluster {
             identity,
             tls_client_certificate_required,
             schema_cache,
+            canonicalize_oids,
         } = config;
 
         let identifier = Arc::new(DatabaseUser {
             user: user.to_owned(),
             database: name.to_owned(),
         });
+        let canonical_oids = canonicalize_oids.then(|| schema_cache.canonical_oids(name));
 
         Self {
             identifier: identifier.clone(),
@@ -357,6 +364,7 @@ impl Cluster {
             identity: identity.clone(),
             tls_client_certificate_required,
             schema_loader: Box::new(schema_loader::FromServer),
+            canonical_oids,
         }
     }
 
@@ -411,7 +419,7 @@ impl Cluster {
     }
 
     /// Get all shards.
-    pub fn shards(&self) -> &[Shard] {
+    pub(crate) fn shards(&self) -> &[Shard] {
         &self.shards
     }
 
@@ -678,6 +686,11 @@ impl Cluster {
 
         Ok(())
     }
+
+    #[cfg(feature = "new_parser")]
+    pub(crate) fn is_canonicalizing_oids(&self) -> bool {
+        self.canonical_oids.is_some()
+    }
 }
 
 #[cfg(test)]
@@ -813,6 +826,11 @@ mod test {
                 rewrite: config.config.rewrite.clone(),
                 two_phase_commit: config.config.general.two_phase_commit,
                 two_phase_commit_auto: config.config.general.two_phase_commit_auto.unwrap_or(false),
+                canonical_oids: config
+                    .config
+                    .general
+                    .canonicalize_type_information
+                    .then(Default::default),
                 ..Default::default()
             }
         }

pgdog/src/backend/pool/cluster/schema_loader.rs

@@ -3,6 +3,7 @@ use crate::backend::pool::ee::schema_changed_hook;
 use crate::tasks;
 use crate::util::safe_sleep;
 use dyn_clone::DynClone;
+use std::sync::Arc;
 use std::time::Duration;
 use tokio::select;
 use tracing::error;
@@ -25,6 +26,42 @@ impl SchemaLoader for FromServer {
             return;
         }
 
+        // For now we treat shard 0 as the canonical OID source
+        if let Some(shard) = cluster.shards().first()
+            && let Some(canonical_oids) = &cluster.canonical_oids
+        {
+            let canonical_oids = Arc::clone(canonical_oids);
+            let shard = shard.clone();
+            tasks::spawn("load canonical oids", async move {
+                loop {
+                    let result = tasks::shutdown_signal()
+                        .run_until_cancelled(async {
+                            canonical_oids
+                                .load(&mut *shard.primary_or_replica(&Default::default()).await?)
+                                .await
+                        })
+                        .await;
+
+                    match result {
+                        Some(Ok(_)) | None => break,
+                        Some(Err(err)) => {
+                            if shard.online() {
+                                error!("error loading canonical type information: {err}");
+                                safe_sleep(Duration::from_millis(100)).await;
+                            } else {
+                                // Cluster is shutting down
+                                break;
+                            }
+                        }
+                    }
+                }
+            });
+        } else {
+            for shard in cluster.shards() {
+                shard.skip_loading_oids()
+            }
+        }
+
         for shard in cluster.shards() {
             let identifier = cluster.identifier();
             let shard = shard.clone();

pgdog/src/backend/pool/inner.rs

@@ -324,6 +324,7 @@ impl Inner {
 
         for conn in idle.iter_mut() {
             conn.stats_mut().set_pool_id(destination.id());
+            conn.replace_oids(&destination.inner().oids)
         }
 
         (idle, taken)
@@ -352,6 +353,7 @@ impl Inner {
             if moved.id() != self.id {
                 server.stats_mut().set_pool_id(moved.id());
                 server.stats().update();
+                server.replace_oids(&moved.inner().oids);
                 moved.lock().maybe_check_in(server, now, stats, true)?;
                 return Ok(result);
             }