From b64b320095fad30d51b8a031ec0b6d87b6a1a20c Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <matthias@neon.tech>
Date: Mon, 23 Jun 2025 18:30:30 +0200
Subject: [PATCH] Make PgVersionId a newtype

This makes several casts & checks a lot easier to check, as `type` is
nothing but an alias and doesn't actually prevent casts.
---
 libs/postgres_versioninfo/src/lib.rs   | 53 ++++++++++++++++++++++++--
 safekeeper/src/control_file_upgrade.rs |  4 +-
 safekeeper/src/safekeeper.rs           | 14 ++++---
 safekeeper/src/state.rs                |  3 +-
 4 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/libs/postgres_versioninfo/src/lib.rs b/libs/postgres_versioninfo/src/lib.rs
index 481e30d015..01d28298ca 100644
--- a/libs/postgres_versioninfo/src/lib.rs
+++ b/libs/postgres_versioninfo/src/lib.rs
@@ -1,3 +1,4 @@
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use serde_repr::{Deserialize_repr, Serialize_repr};
 use std::fmt::{Display, Formatter};
 use std::str::FromStr;
@@ -14,7 +15,51 @@ pub enum PgMajorVersion {
     // !!! When you add a new PgMajorVersion, don't forget to update PgMajorVersion::ALL
 }
 
-pub type PgVersionId = u32;
+#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
+pub struct PgVersionId(u32);
+
+impl PgVersionId {
+    pub const UNKNOWN: PgVersionId = PgVersionId(0);
+
+    pub fn from_full_pg_version(version: u32) -> PgVersionId {
+        match version {
+            0 => PgVersionId(version), // unknown version
+            140000..180000 => PgVersionId(version),
+            _ => panic!("Invalid full PostgreSQL version ID {version}"),
+        }
+    }
+}
+
+impl Display for PgVersionId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        u32::fmt(&self.0, f)
+    }
+}
+
+impl Serialize for PgVersionId {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        u32::serialize(&self.0, serializer)
+    }
+}
+
+impl<'de> Deserialize<'de> for PgVersionId {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        u32::deserialize(deserializer).map(PgVersionId)
+    }
+
+    fn deserialize_in_place<D>(deserializer: D, place: &mut Self) -> Result<(), D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        u32::deserialize_in_place(deserializer, &mut place.0)
+    }
+}
 
 impl PgMajorVersion {
     /// Get the numerical representation of the represented Major Version
@@ -89,13 +134,13 @@ impl TryFrom<PgVersionId> for PgMajorVersion {
     type Error = InvalidPgVersion;
 
     fn try_from(value: PgVersionId) -> Result<Self, Self::Error> {
-        Ok(match value / 10000 {
+        Ok(match value.0 / 10000 {
             14 => PgMajorVersion::PG14,
             15 => PgMajorVersion::PG15,
             16 => PgMajorVersion::PG16,
             17 => PgMajorVersion::PG17,
             _ => {
-                panic!("invalid pg version ID {value}");
+                panic!("Unknown PgVersionId: {value}");
             }
         })
     }
@@ -103,7 +148,7 @@ impl TryFrom<PgVersionId> for PgMajorVersion {
 
 impl From<PgMajorVersion> for PgVersionId {
     fn from(value: PgMajorVersion) -> Self {
-        ((value as u32) * 10000) as PgVersionId
+        PgVersionId((value as u32) * 10000)
     }
 }
 
diff --git a/safekeeper/src/control_file_upgrade.rs b/safekeeper/src/control_file_upgrade.rs
index 588810cf6e..555cbe457b 100644
--- a/safekeeper/src/control_file_upgrade.rs
+++ b/safekeeper/src/control_file_upgrade.rs
@@ -445,13 +445,13 @@ pub fn upgrade_control_file(buf: &[u8], version: u32) -> Result<TimelinePersiste
     } else if version == 6 {
         info!("reading safekeeper control file version {}", version);
         let mut oldstate = TimelinePersistentState::des(&buf[..buf.len()])?;
-        if oldstate.server.pg_version != 0 {
+        if oldstate.server.pg_version != PgVersionId::UNKNOWN {
             return Ok(oldstate);
         }
 
         // set pg_version to the default v14
         info!("setting pg_version to 140005");
-        oldstate.server.pg_version = 140005;
+        oldstate.server.pg_version = PgVersionId::from_full_pg_version(140005);
 
         return Ok(oldstate);
     } else if version == 7 {
diff --git a/safekeeper/src/safekeeper.rs b/safekeeper/src/safekeeper.rs
index 582ceebf25..02dda7ae8e 100644
--- a/safekeeper/src/safekeeper.rs
+++ b/safekeeper/src/safekeeper.rs
@@ -9,6 +9,7 @@ use anyhow::{Context, Result, bail};
 use byteorder::{LittleEndian, ReadBytesExt};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use postgres_ffi::{MAX_SEND_SIZE, TimeLineID};
+use postgres_versioninfo::{PgMajorVersion, PgVersionId};
 use pq_proto::SystemId;
 use safekeeper_api::membership::{
     INVALID_GENERATION, MemberSet, SafekeeperGeneration as Generation, SafekeeperId,
@@ -29,7 +30,7 @@ use crate::{control_file, wal_storage};
 
 pub const SK_PROTO_VERSION_2: u32 = 2;
 pub const SK_PROTO_VERSION_3: u32 = 3;
-pub const UNKNOWN_SERVER_VERSION: u32 = 0;
+pub const UNKNOWN_SERVER_VERSION: PgVersionId = PgVersionId::UNKNOWN;
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
 pub struct TermLsn {
@@ -220,7 +221,7 @@ pub struct ProposerGreeting {
     pub timeline_id: TimelineId,
     pub mconf: membership::Configuration,
     /// Postgres server version
-    pub pg_version: u32,
+    pub pg_version: PgVersionId,
     pub system_id: SystemId,
     pub wal_seg_size: u32,
 }
@@ -231,7 +232,7 @@ pub struct ProposerGreetingV2 {
     /// proposer-acceptor protocol version
     pub protocol_version: u32,
     /// Postgres server version
-    pub pg_version: u32,
+    pub pg_version: PgVersionId,
     pub proposer_id: PgUuid,
     pub system_id: SystemId,
     pub timeline_id: TimelineId,
@@ -513,7 +514,7 @@ impl ProposerAcceptorMessage {
                         tenant_id,
                         timeline_id,
                         mconf,
-                        pg_version,
+                        pg_version: PgVersionId::from_full_pg_version(pg_version),
                         system_id,
                         wal_seg_size,
                     };
@@ -963,7 +964,8 @@ where
          * because safekeepers parse WAL headers and the format
          * may change between versions.
          */
-        if msg.pg_version / 10000 != self.state.server.pg_version / 10000
+        if PgMajorVersion::try_from(msg.pg_version).unwrap()
+            != PgMajorVersion::try_from(self.state.server.pg_version).unwrap()
             && self.state.server.pg_version != UNKNOWN_SERVER_VERSION
         {
             bail!(
@@ -1750,7 +1752,7 @@ mod tests {
                 }]),
             },
             server: ServerInfo {
-                pg_version: 140000,
+                pg_version: PgVersionId::from_full_pg_version(140000),
                 system_id: 0x1234567887654321,
                 wal_seg_size: 0x12345678,
             },
diff --git a/safekeeper/src/state.rs b/safekeeper/src/state.rs
index 7533005c35..22740b37ed 100644
--- a/safekeeper/src/state.rs
+++ b/safekeeper/src/state.rs
@@ -7,6 +7,7 @@ use std::time::SystemTime;
 
 use anyhow::{Result, bail};
 use postgres_ffi::WAL_SEGMENT_SIZE;
+use postgres_versioninfo::PgVersionId;
 use safekeeper_api::membership::Configuration;
 use safekeeper_api::models::{TimelineMembershipSwitchResponse, TimelineTermBumpResponse};
 use safekeeper_api::{INITIAL_TERM, ServerInfo, Term};
@@ -149,7 +150,7 @@ impl TimelinePersistentState {
             &TenantTimelineId::empty(),
             Configuration::empty(),
             ServerInfo {
-                pg_version: 170000, /* Postgres server version (major * 10000) */
+                pg_version: PgVersionId::from_full_pg_version(170000), /* Postgres server version (major * 10000) */
                 system_id: 0,       /* Postgres system identifier */
                 wal_seg_size: WAL_SEGMENT_SIZE as u32,
             },