fix: resolve cross-shard type OIDs on first use instead of failing the query

integration/rust/tests/integration/cross_shard_oid_drift.rs

@@ -9,6 +9,13 @@ struct Composite {
     b: String,
 }
 
+#[derive(sqlx::Type, Debug, Clone, PartialEq)]
+#[sqlx(type_name = "test_oid_drift_later_composite")]
+struct LaterComposite {
+    a: String,
+    b: String,
+}
+
 #[tokio::test]
 async fn test_oid_drift() {
     let conn = connections_sqlx().await.pop().unwrap();
@@ -74,3 +81,85 @@ async fn test_oid_drift() {
 
     admin.execute("RELOAD").await.unwrap();
 }
+
+/// A type created after PgDog loaded its OID mappings (here: with DDL reloads
+/// disabled) is resolved on first use, so the very first query using it
+/// already gets the canonical OID.
+#[tokio::test]
+async fn test_oid_drift_type_created_later() {
+    let conn = connections_sqlx().await.pop().unwrap();
+    let admin = admin_sqlx().await;
+
+    admin
+        .execute("SET canonicalize_type_information TO true")
+        .await
+        .unwrap();
+    admin
+        .execute("SET reload_schema_on_ddl TO false")
+        .await
+        .unwrap();
+    // Make sure the mappings are loaded before the type exists.
+    admin.execute("RELOAD").await.unwrap();
+    conn.execute("SELECT 1").await.unwrap();
+
+    conn.execute("DROP TABLE IF EXISTS test_oid_drift_later")
+        .await
+        .unwrap();
+    conn.execute("DROP TYPE IF EXISTS test_oid_drift_later_composite 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_later_composite AS (a text, b text)")
+        .await
+        .unwrap();
+    conn.execute(
+        "CREATE TABLE test_oid_drift_later (customer_id BIGINT, composite test_oid_drift_later_composite)",
+    )
+    .await
+    .unwrap();
+
+    let expected_oid: Oid = sqlx::query_scalar(
+        "SELECT oid FROM pg_type WHERE typname = 'test_oid_drift_later_composite'",
+    )
+    .fetch_one(&conn)
+    .await
+    .unwrap();
+
+    // The client learned the type's OID from shard 0 and sends it in Parse;
+    // shards where the OID differs must accept it on the first try.
+    let composite = LaterComposite {
+        a: String::from("a"),
+        b: String::from("b"),
+    };
+    for i in 1..=20_i64 {
+        sqlx::query("INSERT INTO test_oid_drift_later VALUES ($1, $2)")
+            .bind(i)
+            .bind(&composite)
+            .execute(&conn)
+            .await
+            .unwrap();
+    }
+
+    // Reads from every shard, including the ones where the OID differs,
+    // must not fail even once.
+    for customer_id in 1..=20_i64 {
+        let rows = sqlx::query("SELECT composite FROM test_oid_drift_later WHERE customer_id = $1")
+            .bind(customer_id)
+            .fetch_all(&conn)
+            .await
+            .unwrap();
+        assert_eq!(rows.len(), 1, "customer {customer_id}");
+        assert_eq!(
+            rows[0].column(0).type_info().oid(),
+            Some(expected_oid),
+            "customer {customer_id}"
+        );
+        let decoded: LaterComposite = rows[0].get(0);
+        assert_eq!(decoded, composite);
+    }
+
+    admin.execute("RELOAD").await.unwrap();
+}

pgdog/src/backend/error.rs

@@ -24,6 +24,9 @@ pub(crate) enum Error {
     #[error("unexpected transaction status: {0}")]
     UnexpectedTransactionStatus(char),
 
+    #[error("type information is not available")]
+    TypeInfoUnavailable,
+
     #[error("{0}")]
     ConnectionError(Box<ErrorResponse>),
 
@@ -116,9 +119,6 @@ pub(crate) 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/pool/cluster/schema_loader.rs

@@ -35,11 +35,7 @@ impl SchemaLoader for FromServer {
             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
-                        })
+                        .run_until_cancelled(async { canonical_oids.load(&shard).await })
                         .await;
 
                     match result {

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

@@ -468,6 +468,23 @@ impl Binding {
         }
     }
 
+    /// Any held server returned a type the canonical OID
+    /// mappings don't know about. Resets the flags.
+    pub(crate) fn take_oids_stale(&mut self) -> bool {
+        match self {
+            Binding::Direct(server, ..) => server.take_oids_stale(),
+            Binding::MultiShard(servers, _) => {
+                // Reset every server's flag, not just the first stale one.
+                let mut stale = false;
+                for server in servers.iter_mut() {
+                    stale |= server.take_oids_stale();
+                }
+                stale
+            }
+            _ => false,
+        }
+    }
+
     pub(super) fn dirty(&mut self) {
         match self {
             Binding::Direct(server, ..) => server.mark_dirty(true),

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, OidMappings, Oids, Shard};
 pub(crate) use state::State;
 pub(crate) use stats::Stats;
 

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

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