krabka_connect_postgres/
offset.rs

1use std::{fmt, str::FromStr};
2
3use krabka_connect::{OffsetMap, OffsetValue, SourceOffset};
4
5use crate::PostgresConnectError;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
8pub struct PgLsn(pub u64);
9
10impl PgLsn {
11    #[must_use]
12    pub fn to_source_offset(self, database: &str, slot: &str) -> SourceOffset {
13        let partition = OffsetMap::from([
14            (
15                "database".to_owned(),
16                OffsetValue::String(database.to_owned()),
17            ),
18            ("slot".to_owned(), OffsetValue::String(slot.to_owned())),
19        ]);
20        let position = OffsetMap::from([("lsn".to_owned(), OffsetValue::String(self.to_string()))]);
21
22        SourceOffset::new(partition.into(), position.into())
23    }
24
25    /// # Errors
26    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
27    pub fn from_source_offset(
28        offset: &SourceOffset,
29        expected_slot: &str,
30    ) -> Result<PgLsn, PostgresConnectError> {
31        match offset.partition.get("slot") {
32            Some(OffsetValue::String(slot)) if slot == expected_slot => {}
33            Some(OffsetValue::String(slot)) => {
34                return Err(PostgresConnectError::Offset(format!(
35                    "source offset slot {slot:?} does not match expected slot {expected_slot:?}"
36                )));
37            }
38            _ => {
39                return Err(PostgresConnectError::Offset(
40                    "source offset missing string slot".to_owned(),
41                ));
42            }
43        }
44
45        match offset.position.get("lsn") {
46            Some(OffsetValue::String(lsn)) => lsn.parse(),
47            _ => Err(PostgresConnectError::Offset(
48                "source offset missing string lsn".to_owned(),
49            )),
50        }
51    }
52}
53
54impl fmt::Display for PgLsn {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        let hi = self.0 >> 32;
57        let lo = self.0 & 0xffff_ffff;
58
59        write!(f, "{hi:X}/{lo:X}")
60    }
61}
62
63impl FromStr for PgLsn {
64    type Err = PostgresConnectError;
65
66    fn from_str(s: &str) -> Result<Self, Self::Err> {
67        let (hi, lo) = s
68            .split_once('/')
69            .ok_or_else(|| PostgresConnectError::Offset(format!("invalid Postgres LSN {s:?}")))?;
70
71        let hi = parse_lsn_half(hi, s)?;
72        let lo = parse_lsn_half(lo, s)?;
73
74        Ok(PgLsn((hi << 32) | lo))
75    }
76}
77
78fn parse_lsn_half(half: &str, lsn: &str) -> Result<u64, PostgresConnectError> {
79    let value = u64::from_str_radix(half, 16)
80        .map_err(|_| PostgresConnectError::Offset(format!("invalid Postgres LSN {lsn:?}")))?;
81
82    if value > 0xffff_ffff {
83        return Err(PostgresConnectError::Offset(format!(
84            "invalid Postgres LSN {lsn:?}: component exceeds 32 bits"
85        )));
86    }
87
88    Ok(value)
89}
90
91#[cfg(test)]
92mod tests {
93    use krabka_connect::{OffsetValue, SourceOffset};
94
95    use super::PgLsn;
96    use crate::PostgresConnectError;
97
98    #[test]
99    fn lsn_round_trips_postgres_text_form() {
100        let lsn: PgLsn = "16/B374D848".parse().expect("valid LSN");
101
102        assert2::assert!(lsn == PgLsn(0x16_b374_d848));
103        assert2::assert!(lsn.to_string() == "16/B374D848".to_owned());
104        assert2::assert!("FFFFFFFF/FFFFFFFF".parse::<PgLsn>().expect("max LSN") == PgLsn(u64::MAX));
105    }
106
107    #[test]
108    fn source_offset_round_trips_and_checks_slot() {
109        let lsn: PgLsn = "0/2A".parse().expect("valid LSN");
110        let offset = lsn.to_source_offset("app", "slot_a");
111
112        assert2::assert!(
113            PgLsn::from_source_offset(&offset, "slot_a").expect("matching slot") == lsn
114        );
115        match PgLsn::from_source_offset(&offset, "slot_b").expect_err("slot mismatch") {
116            PostgresConnectError::Offset(message) => {
117                for needle in ["does not match expected slot", "slot_a", "slot_b"] {
118                    assert2::assert!(message.contains(needle));
119                }
120            }
121            error => panic!("expected offset error, got {error:?}"),
122        }
123    }
124
125    #[test]
126    fn source_offset_rejects_missing_or_non_string_slot() {
127        let missing = SourceOffset::default();
128        let mut non_string = PgLsn(42).to_source_offset("app", "slot_a");
129        non_string
130            .partition
131            .0
132            .insert("slot".to_owned(), OffsetValue::Long(7));
133
134        for (_name, offset) in [("missing", missing), ("non_string", non_string)] {
135            match PgLsn::from_source_offset(&offset, "slot_a").expect_err("slot should fail") {
136                PostgresConnectError::Offset(message) => {
137                    assert2::assert!(message == "source offset missing string slot");
138                }
139                error => panic!("expected offset error, got {error:?}"),
140            }
141        }
142    }
143
144    #[test]
145    fn malformed_lsn_is_offset_error() {
146        for (_name, input) in [
147            ("missing_separator", "not-a-lsn"),
148            ("invalid_hex", "1/XYZ"),
149            ("component_overflow", "100000000/0"),
150        ] {
151            assert2::assert!(matches!(
152                input.parse::<PgLsn>(),
153                Err(PostgresConnectError::Offset(_))
154            ));
155        }
156    }
157}
krabka_connect_postgres/offset.rs

krabka_connect_postgres/
offset.rs
