fix: canonicalize type OIDs embedded in binary array and composite values

integration/rust/tests/integration/cross_shard_oid_drift.rs

@@ -74,3 +74,78 @@ async fn test_oid_drift() {
 
     admin.execute("RELOAD").await.unwrap();
 }
+
+#[derive(sqlx::Type, Debug, Clone, Copy, PartialEq)]
+#[sqlx(type_name = "test_oid_drift_mood", rename_all = "lowercase")]
+enum Mood {
+    Sad,
+    Ok,
+    Happy,
+}
+
+/// Binary arrays embed the element type's OID, so arrays of custom types
+/// have to be rewritten in both directions, not just the RowDescription.
+#[tokio::test]
+async fn test_oid_drift_arrays() {
+    let conn = connections_sqlx().await.pop().unwrap();
+    let admin = admin_sqlx().await;
+
+    conn.execute("DROP TABLE IF EXISTS test_oid_drift_arrays")
+        .await
+        .unwrap();
+    conn.execute("DROP TYPE IF EXISTS test_oid_drift_mood CASCADE")
+        .await
+        .unwrap();
+    // Intentionally cause the OID of the type to differ between shards
+    conn.execute("/* pgdog_shard: 1 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;")
+        .await
+        .unwrap();
+    conn.execute("CREATE TYPE test_oid_drift_mood AS ENUM ('sad', 'ok', 'happy')")
+        .await
+        .unwrap();
+    conn.execute(
+        "CREATE TABLE test_oid_drift_arrays (customer_id BIGINT, moods test_oid_drift_mood[])",
+    )
+    .await
+    .unwrap();
+    admin
+        .execute("SET canonicalize_type_information TO true")
+        .await
+        .unwrap();
+    admin.execute("RELOAD").await.unwrap();
+
+    let canonical_oid: Oid =
+        sqlx::query_scalar("SELECT oid FROM pg_type WHERE typname = 'test_oid_drift_mood'")
+            .fetch_one(&conn)
+            .await
+            .unwrap();
+
+    let moods = vec![Mood::Sad, Mood::Ok, Mood::Happy];
+    for i in 1..=20_i64 {
+        // Binary array parameter, element OID as learned from shard 0.
+        sqlx::query("INSERT INTO test_oid_drift_arrays VALUES ($1, $2)")
+            .bind(i)
+            .bind(&moods)
+            .execute(&conn)
+            .await
+            .unwrap();
+    }
+
+    for customer_id in 1..=20_i64 {
+        let row = sqlx::query("SELECT moods FROM test_oid_drift_arrays WHERE customer_id = $1")
+            .bind(customer_id)
+            .fetch_one(&conn)
+            .await
+            .unwrap();
+
+        // The element OID inside the binary array payload is shard 0's.
+        let raw = row.try_get_raw(0).unwrap().as_bytes().unwrap().to_vec();
+        let element_oid = u32::from_be_bytes([raw[8], raw[9], raw[10], raw[11]]);
+        assert_eq!(element_oid, canonical_oid.0, "customer {customer_id}");
+
+        let decoded: Vec<Mood> = row.get(0);
+        assert_eq!(decoded, moods);
+    }
+
+    admin.execute("RELOAD").await.unwrap();
+}

pgdog/src/backend/pool/mod.rs

@@ -42,7 +42,7 @@ pub(crate) use password::Password;
 pub(crate) use pool_impl::Pool;
 pub(crate) use request::Request;
 pub(crate) use role::PoolRole;
-pub(crate) use shard::{CanonicalOids, Oids, Shard};
+pub(crate) use shard::{CanonicalOids, Oids, PayloadRewriter, Shard};
 pub(crate) use state::State;
 pub(crate) use stats::Stats;
 

pgdog/src/backend/pool/shard/mod.rs

@@ -29,7 +29,9 @@ pub(crate) mod role_detector;
 
 use failover_signal::{FailoverSignal, FailoverSignalWatcher};
 use monitor::*;
-pub(crate) use oids::{CanonicalOids, Oids};
+#[cfg(test)]
+pub(crate) use oids::TypeKind;
+pub(crate) use oids::{CanonicalOids, Oids, PayloadRewriter};
 use role_detector::*;
 
 #[cfg_attr(test, derive(Default))]

pgdog/src/backend/pool/shard/oids.rs

@@ -1,3 +1,10 @@
+//! Canonical type OID mappings.
+//!
+//! Types created with `CREATE TYPE` (or by extensions) get a different OID
+//! on each shard. Clients cache type information by OID, so PgDog presents
+//! shard 0's OIDs to clients and translates them on the way to and from
+//! the other shards.
+
 use super::{Request, Shard};
 use crate::{
     backend::{Error, Server},
@@ -8,6 +15,60 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use tracing::info;
 
+mod payload;
+pub(crate) use payload::PayloadRewriter;
+
+/// What a type's binary representation looks like, as far as
+/// embedded type OIDs are concerned.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum TypeKind {
+    /// Array: binary values carry the element type OID.
+    Array { element: u32 },
+    /// Composite: binary values carry the OID of every field.
+    Composite,
+    /// Domain: encoded like its base type.
+    Domain { base: u32 },
+    /// Everything else: no embedded OIDs.
+    Other,
+}
+
+impl TypeKind {
+    /// From `pg_type` columns.
+    fn from_catalog(typtype: &str, typcategory: &str, typelem: u32, typbasetype: u32) -> Self {
+        match (typtype, typcategory) {
+            (_, "A") if typelem != 0 => Self::Array { element: typelem },
+            ("c", _) => Self::Composite,
+            ("d", _) if typbasetype != 0 => Self::Domain { base: typbasetype },
+            _ => Self::Other,
+        }
+    }
+}
+
+/// A type on the shard: `schema.name`, its OID and kind.
+type TypeRow = (String, u32, TypeKind);
+
+/// The canonical set of types, from shard 0.
+#[derive(Debug, Default)]
+pub(crate) struct CanonicalTypes {
+    by_name: HashMap<String, u32>,
+    kinds: Arc<HashMap<u32, TypeKind>>,
+}
+
+impl FromIterator<TypeRow> for CanonicalTypes {
+    fn from_iter<I: IntoIterator<Item = TypeRow>>(iter: I) -> Self {
+        let mut by_name = HashMap::new();
+        let mut kinds = HashMap::new();
+        for (name, oid, kind) in iter {
+            by_name.insert(name, oid);
+            kinds.insert(oid, kind);
+        }
+        Self {
+            by_name,
+            kinds: Arc::new(kinds),
+        }
+    }
+}
+
 #[derive(Debug)]
 /// The mapping from a shards type OID to a canonical one
 pub(crate) struct Oids {
@@ -34,8 +95,11 @@ impl Oids {
                 let canonical = self.canonical_oids.oids.wait().await;
                 let mut canonical_to_shard = HashMap::new();
                 let mut shard_to_canonical = HashMap::new();
-                for (type_name, oid) in oids {
+                let mut shard_kinds = HashMap::new();
+                for (type_name, oid, kind) in oids {
+                    shard_kinds.insert(oid, kind);
                     let canonical = canonical
+                        .by_name
                         .get(&type_name)
                         .copied()
                         .ok_or(Error::MissingCanonicalOid(type_name))?;
@@ -57,6 +121,8 @@ impl Oids {
                 Ok(OidMappings {
                     canonical_to_shard,
                     shard_to_canonical,
+                    shard_kinds,
+                    canonical_kinds: Arc::clone(&canonical.kinds),
                 })
             })
             .await
@@ -78,12 +144,24 @@ impl Oids {
 
     #[cfg(test)]
     pub(crate) fn from_canonical(canonical_to_shard: HashMap<u32, u32>) -> Arc<Self> {
+        Self::from_canonical_with_kinds(canonical_to_shard, HashMap::new(), HashMap::new())
+    }
+
+    /// Mappings plus the type kinds on both sides, simulating what's loaded from `pg_type`.
+    #[cfg(test)]
+    pub(crate) fn from_canonical_with_kinds(
+        canonical_to_shard: HashMap<u32, u32>,
+        shard_kinds: HashMap<u32, TypeKind>,
+        canonical_kinds: HashMap<u32, TypeKind>,
+    ) -> Arc<Self> {
         let shard_to_canonical = canonical_to_shard.iter().map(|(&k, &v)| (v, k)).collect();
         Arc::new(Self {
             canonical_oids: Default::default(),
             mappings: SetOnceCell::from(OidMappings {
                 canonical_to_shard,
                 shard_to_canonical,
+                shard_kinds,
+                canonical_kinds: Arc::new(canonical_kinds),
             }),
         })
     }
@@ -102,11 +180,40 @@ impl Default for Oids {
 pub(crate) struct OidMappings {
     pub(crate) canonical_to_shard: HashMap<u32, u32>,
     pub(crate) shard_to_canonical: HashMap<u32, u32>,
+    /// Kinds of the shard's types, by shard OID.
+    shard_kinds: HashMap<u32, TypeKind>,
+    /// Kinds of the canonical types, by canonical OID.
+    canonical_kinds: Arc<HashMap<u32, TypeKind>>,
+}
+
+impl OidMappings {
+    /// Whether canonicalization has anything to do on this shard at all.
+    pub(crate) fn is_identity(&self) -> bool {
+        self.shard_to_canonical.is_empty()
+    }
+
+    /// The shard's OID for a canonical type OID.
+    pub(crate) fn shard_oid(&self, canonical: u32) -> u32 {
+        self.canonical_to_shard
+            .get(&canonical)
+            .copied()
+            .unwrap_or(canonical)
+    }
+
+    /// Rewriter for values coming from the shard (DataRow).
+    pub(crate) fn to_canonical(&self) -> PayloadRewriter<'_> {
+        PayloadRewriter::new(&self.shard_kinds, &self.shard_to_canonical)
+    }
+
+    /// Rewriter for values going to the shard (Bind parameters).
+    pub(crate) fn to_shard(&self) -> PayloadRewriter<'_> {
+        PayloadRewriter::new(&self.canonical_kinds, &self.canonical_to_shard)
+    }
 }
 
 #[derive(Debug, Default)]
 pub(crate) struct CanonicalOids {
-    oids: SetOnceCell<HashMap<String, u32>>,
+    oids: SetOnceCell<CanonicalTypes>,
 }
 
 impl CanonicalOids {
@@ -118,20 +225,51 @@ impl CanonicalOids {
     }
 }
 
-async fn load_oids(
-    server: &mut Server,
-) -> Result<impl Iterator<Item = (String, u32)> + use<>, Error> {
+async fn load_oids(server: &mut Server) -> Result<impl Iterator<Item = TypeRow> + use<>, Error> {
     // OIDs < 10,000 are reserved for PG's internal use and are assumed to be stable
     Ok(server
         .fetch_all::<DataRow>(
-            "SELECT nspname || '.' || typname, pg_type.oid FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid WHERE pg_type.oid >= 10000",
+            "SELECT nspname || '.' || typname, pg_type.oid, typtype::text, typcategory::text, typelem, typbasetype \
+             FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid \
+             WHERE pg_type.oid >= 10000",
         )
         .await?
         .into_iter()
         .map(|row| {
+            let name = row.get_text(0).expect("selected 6 columns");
+            let oid = row.get_int(1, true).expect("selected 6 columns") as u32;
+            let typtype = row.get_text(2).expect("selected 6 columns");
+            let typcategory = row.get_text(3).expect("selected 6 columns");
+            let typelem = row.get_int(4, true).expect("selected 6 columns") as u32;
+            let typbasetype = row.get_int(5, true).expect("selected 6 columns") as u32;
             (
-                row.get_text(0).expect("selected 2 columns"),
-                row.get_int(1, true).expect("selected 2 columns") as u32,
+                name,
+                oid,
+                TypeKind::from_catalog(&typtype, &typcategory, typelem, typbasetype),
             )
         }))
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_type_kind_from_catalog() {
+        assert_eq!(
+            TypeKind::from_catalog("b", "A", 16400, 0),
+            TypeKind::Array { element: 16400 }
+        );
+        assert_eq!(TypeKind::from_catalog("c", "C", 0, 0), TypeKind::Composite);
+        assert_eq!(
+            TypeKind::from_catalog("d", "N", 0, 23),
+            TypeKind::Domain { base: 23 }
+        );
+        assert_eq!(TypeKind::from_catalog("e", "E", 0, 0), TypeKind::Other);
+        // A domain over an array is category A but not itself an array.
+        assert_eq!(
+            TypeKind::from_catalog("d", "A", 16400, 16399),
+            TypeKind::Array { element: 16400 }
+        );
+    }
+}