improv_pg_model/
pgoutput.rs

1//! A pure decoder for the `pgoutput` logical-replication protocol, version 1
2//! (the binary messages `pg_logical_slot_peek_binary_changes` returns with
3//! `proto_version '1'`), and [`Stream`], which groups the decoded changes by
4//! source transaction.
5//!
6//! Integers are big-endian; strings are NUL-terminated and assumed UTF-8 (a
7//! `UTF8` server encoding); tuple values arrive in text form.
8
9use std::collections::HashMap;
10
11/// A malformed or unexpected `pgoutput` message.
12#[derive(Debug, Clone, PartialEq, thiserror::Error)]
13pub enum DecodeError {
14    #[error("pgoutput: message truncated")]
15    Truncated,
16    #[error("pgoutput: {0} unexpected trailing bytes")]
17    Trailing(usize),
18    #[error("pgoutput: unknown message type {0:?}")]
19    UnknownMessage(char),
20    #[error("pgoutput: unknown tuple column kind {0:?}")]
21    UnknownColumn(char),
22    #[error("pgoutput: unexpected tuple marker {0:?}")]
23    BadMarker(char),
24    #[error("pgoutput: string is not NUL-terminated UTF-8")]
25    BadString,
26    #[error("pgoutput: change for relation {0} before its Relation message")]
27    UnknownRelation(u32),
28    #[error("pgoutput: {0} outside a transaction")]
29    OutsideTxn(&'static str),
30}
31
32/// A decode result.
33pub type Result<T> = std::result::Result<T, DecodeError>;
34
35/// One column value of a tuple.
36#[derive(Debug, Clone, PartialEq)]
37pub enum Datum {
38    /// SQL NULL (`'n'`).
39    Null,
40    /// An unchanged TOASTed value (`'u'`): not sent, keep the old one.
41    Unchanged,
42    /// A value in text output form (`'t'`).
43    Text(String),
44}
45
46/// A row: one [`Datum`] per column of its [`Relation`].
47pub type Tuple = Vec<Datum>;
48
49/// A column of a [`Relation`] message.
50#[derive(Debug, Clone, PartialEq)]
51pub struct Column {
52    /// Part of the replica identity key.
53    pub key: bool,
54    /// Column name.
55    pub name: String,
56    /// `pg_type` oid.
57    pub type_oid: u32,
58    /// `atttypmod`.
59    pub type_mod: i32,
60}
61
62/// A `Relation` message: the shape of the tuples of relation `id`.
63#[derive(Debug, Clone, PartialEq)]
64pub struct Relation {
65    /// The relation's oid.
66    pub id: u32,
67    /// Schema name.
68    pub namespace: String,
69    /// Relation name.
70    pub name: String,
71    /// `relreplident`: `d`efault, `n`othing, `f`ull or `i`ndex.
72    pub replica_identity: u8,
73    /// Columns, in tuple order.
74    pub columns: Vec<Column>,
75}
76
77impl Relation {
78    /// The value of column `name` in `tuple` (a tuple of this relation).
79    pub fn get<'t>(&self, tuple: &'t Tuple, name: &str) -> Option<&'t Datum> {
80        let i = self.columns.iter().position(|c| c.name == name)?;
81        tuple.get(i)
82    }
83}
84
85/// A decoded `pgoutput` message.
86#[allow(missing_docs)]
87#[derive(Debug, Clone, PartialEq)]
88pub enum Message {
89    Begin {
90        final_lsn: u64,
91        commit_ts: i64,
92        xid: u32,
93    },
94    Commit {
95        flags: u8,
96        commit_lsn: u64,
97        /// The end of the commit record: where a consumer resumes.
98        end_lsn: u64,
99        commit_ts: i64,
100    },
101    Origin {
102        lsn: u64,
103        name: String,
104    },
105    Relation(Relation),
106    Type {
107        id: u32,
108        namespace: String,
109        name: String,
110    },
111    Insert {
112        rel: u32,
113        new: Tuple,
114    },
115    Update {
116        rel: u32,
117        /// The old key (`'K'`) or whole old row (`'O'`), when sent.
118        old: Option<Tuple>,
119        new: Tuple,
120    },
121    Delete {
122        rel: u32,
123        old: Tuple,
124    },
125    Truncate {
126        options: u8,
127        rels: Vec<u32>,
128    },
129}
130
131struct Reader<'a>(&'a [u8]);
132
133impl<'a> Reader<'a> {
134    fn take(&mut self, n: usize) -> Result<&'a [u8]> {
135        if self.0.len() < n {
136            return Err(DecodeError::Truncated);
137        }
138        let (head, tail) = self.0.split_at(n);
139        self.0 = tail;
140        Ok(head)
141    }
142
143    fn array<const N: usize>(&mut self) -> Result<[u8; N]> {
144        let mut a = [0; N];
145        a.copy_from_slice(self.take(N)?);
146        Ok(a)
147    }
148
149    fn u8(&mut self) -> Result<u8> {
150        Ok(self.take(1)?[0])
151    }
152
153    fn u16(&mut self) -> Result<u16> {
154        self.array().map(u16::from_be_bytes)
155    }
156
157    fn u32(&mut self) -> Result<u32> {
158        self.array().map(u32::from_be_bytes)
159    }
160
161    fn i32(&mut self) -> Result<i32> {
162        self.array().map(i32::from_be_bytes)
163    }
164
165    fn u64(&mut self) -> Result<u64> {
166        self.array().map(u64::from_be_bytes)
167    }
168
169    fn i64(&mut self) -> Result<i64> {
170        self.array().map(i64::from_be_bytes)
171    }
172
173    fn string(&mut self) -> Result<String> {
174        let n = self
175            .0
176            .iter()
177            .position(|&c| c == 0)
178            .ok_or(DecodeError::BadString)?;
179        let s = std::str::from_utf8(&self.0[..n]).map_err(|_| DecodeError::BadString)?;
180        self.0 = &self.0[n + 1..];
181        Ok(s.to_string())
182    }
183
184    fn marker(&mut self, ok: &[u8]) -> Result<u8> {
185        let m = self.u8()?;
186        if ok.contains(&m) {
187            Ok(m)
188        } else {
189            Err(DecodeError::BadMarker(m as char))
190        }
191    }
192
193    fn tuple(&mut self) -> Result<Tuple> {
194        let n = self.u16()?;
195        let mut t = Vec::with_capacity(usize::from(n));
196        for _ in 0..n {
197            t.push(match self.u8()? {
198                b'n' => Datum::Null,
199                b'u' => Datum::Unchanged,
200                b't' => {
201                    let len = self.u32()? as usize;
202                    let bytes = self.take(len)?;
203                    let s = std::str::from_utf8(bytes).map_err(|_| DecodeError::BadString)?;
204                    Datum::Text(s.to_string())
205                }
206                c => return Err(DecodeError::UnknownColumn(c as char)),
207            });
208        }
209        Ok(t)
210    }
211}
212
213/// Decode one `pgoutput` (protocol 1) message.
214pub fn decode(buf: &[u8]) -> Result<Message> {
215    let mut r = Reader(buf);
216    // Struct fields are evaluated in the order written: the wire order.
217    let m = match r.u8()? {
218        b'B' => Message::Begin {
219            final_lsn: r.u64()?,
220            commit_ts: r.i64()?,
221            xid: r.u32()?,
222        },
223        b'C' => Message::Commit {
224            flags: r.u8()?,
225            commit_lsn: r.u64()?,
226            end_lsn: r.u64()?,
227            commit_ts: r.i64()?,
228        },
229        b'O' => Message::Origin {
230            lsn: r.u64()?,
231            name: r.string()?,
232        },
233        b'R' => {
234            let id = r.u32()?;
235            let namespace = r.string()?;
236            let name = r.string()?;
237            let replica_identity = r.u8()?;
238            let n = r.u16()?;
239            let mut columns = Vec::with_capacity(usize::from(n));
240            for _ in 0..n {
241                columns.push(Column {
242                    key: r.u8()? & 1 == 1,
243                    name: r.string()?,
244                    type_oid: r.u32()?,
245                    type_mod: r.i32()?,
246                });
247            }
248            Message::Relation(Relation {
249                id,
250                namespace,
251                name,
252                replica_identity,
253                columns,
254            })
255        }
256        b'Y' => Message::Type {
257            id: r.u32()?,
258            namespace: r.string()?,
259            name: r.string()?,
260        },
261        b'I' => {
262            let rel = r.u32()?;
263            r.marker(b"N")?;
264            Message::Insert {
265                rel,
266                new: r.tuple()?,
267            }
268        }
269        b'U' => {
270            let rel = r.u32()?;
271            let old = match r.marker(b"KON")? {
272                b'N' => None,
273                _ => {
274                    let old = r.tuple()?;
275                    r.marker(b"N")?;
276                    Some(old)
277                }
278            };
279            Message::Update {
280                rel,
281                old,
282                new: r.tuple()?,
283            }
284        }
285        b'D' => {
286            let rel = r.u32()?;
287            r.marker(b"KO")?;
288            Message::Delete {
289                rel,
290                old: r.tuple()?,
291            }
292        }
293        b'T' => {
294            let n = r.u32()?;
295            let options = r.u8()?;
296            // Not preallocated: `n` is untrusted.
297            let rels = (0..n).map(|_| r.u32()).collect::<Result<_>>()?;
298            Message::Truncate { options, rels }
299        }
300        c => return Err(DecodeError::UnknownMessage(c as char)),
301    };
302    if r.0.is_empty() {
303        Ok(m)
304    } else {
305        Err(DecodeError::Trailing(r.0.len()))
306    }
307}
308
309/// A committed source transaction: its row changes (`Insert`, `Update`,
310/// `Delete` and `Truncate` messages, in order) and its commit's end LSN.
311#[derive(Debug, Clone, PartialEq, Default)]
312pub struct Txn {
313    /// Source transaction id.
314    pub xid: u32,
315    /// End of its commit record.
316    pub end_lsn: u64,
317    /// Row changes, in order.
318    pub changes: Vec<Message>,
319}
320
321/// Groups a sequence of messages (one decoding session) into committed
322/// transactions. A transaction still open at the end is not in `done`.
323#[derive(Debug, Default)]
324pub struct Stream {
325    /// Every relation described so far, by oid.
326    pub relations: HashMap<u32, Relation>,
327    open: Option<Txn>,
328    /// Committed transactions, in commit order.
329    pub done: Vec<Txn>,
330}
331
332impl Stream {
333    /// Decode and absorb one message.
334    pub fn feed(&mut self, buf: &[u8]) -> Result<()> {
335        match decode(buf)? {
336            Message::Begin { xid, .. } => {
337                if self.open.is_some() {
338                    return Err(DecodeError::OutsideTxn("Begin"));
339                }
340                self.open = Some(Txn {
341                    xid,
342                    ..Txn::default()
343                });
344            }
345            Message::Commit { end_lsn, .. } => {
346                let mut t = self.open.take().ok_or(DecodeError::OutsideTxn("Commit"))?;
347                t.end_lsn = end_lsn;
348                self.done.push(t);
349            }
350            Message::Relation(rel) => {
351                self.relations.insert(rel.id, rel);
352            }
353            Message::Origin { .. } | Message::Type { .. } => {}
354            change => {
355                let rels = match &change {
356                    Message::Insert { rel, .. }
357                    | Message::Update { rel, .. }
358                    | Message::Delete { rel, .. } => std::slice::from_ref(rel),
359                    Message::Truncate { rels, .. } => rels.as_slice(),
360                    _ => unreachable!("every other message is matched above"),
361                };
362                if let Some(&id) = rels.iter().find(|id| !self.relations.contains_key(id)) {
363                    return Err(DecodeError::UnknownRelation(id));
364                }
365                let t = self
366                    .open
367                    .as_mut()
368                    .ok_or(DecodeError::OutsideTxn("row change"))?;
369                t.changes.push(change);
370            }
371        }
372        Ok(())
373    }
374}
375
376/// An LSN in Postgres' `X/Y` text form.
377pub fn format_lsn(lsn: u64) -> String {
378    format!("{:X}/{:X}", lsn >> 32, lsn & 0xFFFF_FFFF)
379}
380
381/// Parse an LSN from Postgres' `X/Y` text form.
382pub fn parse_lsn(s: &str) -> Option<u64> {
383    let (hi, lo) = s.split_once('/')?;
384    let hi = u32::from_str_radix(hi, 16).ok()?;
385    let lo = u32::from_str_radix(lo, 16).ok()?;
386    Some((u64::from(hi) << 32) | u64::from(lo))
387}
388
389#[cfg(test)]
390#[cfg_attr(coverage_nightly, coverage(off))]
391mod tests {
392    use super::*;
393
394    fn s(b: &mut Vec<u8>, x: &str) {
395        b.extend_from_slice(x.as_bytes());
396        b.push(0);
397    }
398
399    fn tuple(b: &mut Vec<u8>, t: &[Datum]) {
400        b.extend_from_slice(&(t.len() as u16).to_be_bytes());
401        for d in t {
402            match d {
403                Datum::Null => b.push(b'n'),
404                Datum::Unchanged => b.push(b'u'),
405                Datum::Text(v) => {
406                    b.push(b't');
407                    b.extend_from_slice(&(v.len() as u32).to_be_bytes());
408                    b.extend_from_slice(v.as_bytes());
409                }
410            }
411        }
412    }
413
414    fn txt(v: &str) -> Datum {
415        Datum::Text(v.to_string())
416    }
417
418    fn begin(xid: u32) -> Vec<u8> {
419        let mut b = vec![b'B'];
420        b.extend_from_slice(&0x0000_0001_0000_0100u64.to_be_bytes());
421        b.extend_from_slice(&(-5i64).to_be_bytes());
422        b.extend_from_slice(&xid.to_be_bytes());
423        b
424    }
425
426    fn commit(end: u64) -> Vec<u8> {
427        let mut b = vec![b'C', 0];
428        b.extend_from_slice(&(end - 8).to_be_bytes());
429        b.extend_from_slice(&end.to_be_bytes());
430        b.extend_from_slice(&7i64.to_be_bytes());
431        b
432    }
433
434    /// `public.sales(quarter text KEY, product text, price float8)`.
435    fn relation(id: u32) -> Vec<u8> {
436        let mut b = vec![b'R'];
437        b.extend_from_slice(&id.to_be_bytes());
438        s(&mut b, "public");
439        s(&mut b, "sales");
440        b.push(b'f');
441        b.extend_from_slice(&3u16.to_be_bytes());
442        for (key, name, oid) in [
443            (1u8, "quarter", 25u32),
444            (0, "product", 25),
445            (0, "price", 701),
446        ] {
447            b.push(key);
448            s(&mut b, name);
449            b.extend_from_slice(&oid.to_be_bytes());
450            b.extend_from_slice(&(-1i32).to_be_bytes());
451        }
452        b
453    }
454
455    fn insert(rel: u32, t: &[Datum]) -> Vec<u8> {
456        let mut b = vec![b'I'];
457        b.extend_from_slice(&rel.to_be_bytes());
458        b.push(b'N');
459        tuple(&mut b, t);
460        b
461    }
462
463    fn update(rel: u32, old: Option<(u8, &[Datum])>, new: &[Datum]) -> Vec<u8> {
464        let mut b = vec![b'U'];
465        b.extend_from_slice(&rel.to_be_bytes());
466        if let Some((m, t)) = old {
467            b.push(m);
468            tuple(&mut b, t);
469        }
470        b.push(b'N');
471        tuple(&mut b, new);
472        b
473    }
474
475    #[test]
476    fn begin_and_commit() {
477        assert_eq!(
478            decode(&begin(42)),
479            Ok(Message::Begin {
480                final_lsn: 0x1_0000_0100,
481                commit_ts: -5,
482                xid: 42
483            })
484        );
485        assert_eq!(
486            decode(&commit(0x200)),
487            Ok(Message::Commit {
488                flags: 0,
489                commit_lsn: 0x1F8,
490                end_lsn: 0x200,
491                commit_ts: 7
492            })
493        );
494    }
495
496    #[test]
497    fn relation_columns_and_lookup() {
498        let Ok(Message::Relation(r)) = decode(&relation(16384)) else {
499            panic!("not a relation")
500        };
501        assert_eq!((r.id, &*r.namespace, &*r.name), (16384, "public", "sales"));
502        assert_eq!(r.replica_identity, b'f');
503        let cols: Vec<_> = r
504            .columns
505            .iter()
506            .map(|c| (c.key, &*c.name, c.type_oid))
507            .collect();
508        assert_eq!(
509            cols,
510            [
511                (true, "quarter", 25),
512                (false, "product", 25),
513                (false, "price", 701)
514            ]
515        );
516        assert_eq!(r.columns[2].type_mod, -1);
517        let t = vec![txt("Q1"), Datum::Null, txt("10")];
518        assert_eq!(r.get(&t, "price"), Some(&txt("10")));
519        assert_eq!(r.get(&t, "product"), Some(&Datum::Null));
520        assert_eq!(r.get(&t, "nope"), None);
521    }
522
523    #[test]
524    fn insert_text_null_and_unchanged_toast() {
525        let t = [txt("Q1"), Datum::Null, Datum::Unchanged, txt("")];
526        assert_eq!(
527            decode(&insert(7, &t)),
528            Ok(Message::Insert {
529                rel: 7,
530                new: t.to_vec()
531            })
532        );
533        // Non-ASCII text survives.
534        let t = [txt("Größe €")];
535        assert_eq!(
536            decode(&insert(7, &t)),
537            Ok(Message::Insert {
538                rel: 7,
539                new: t.to_vec()
540            })
541        );
542    }
543
544    #[test]
545    fn update_without_with_old_row_and_with_key() {
546        let new = [txt("Q1"), txt("Widget"), txt("12")];
547        assert_eq!(
548            decode(&update(7, None, &new)),
549            Ok(Message::Update {
550                rel: 7,
551                old: None,
552                new: new.to_vec()
553            })
554        );
555        let old = [txt("Q1"), txt("Widget"), txt("10")];
556        for marker in *b"OK" {
557            assert_eq!(
558                decode(&update(7, Some((marker, &old)), &new)),
559                Ok(Message::Update {
560                    rel: 7,
561                    old: Some(old.to_vec()),
562                    new: new.to_vec()
563                })
564            );
565        }
566    }
567
568    #[test]
569    fn delete_truncate_origin_type() {
570        let old = [txt("Q4"), Datum::Null, Datum::Null];
571        let mut b = vec![b'D'];
572        b.extend_from_slice(&9u32.to_be_bytes());
573        b.push(b'K');
574        tuple(&mut b, &old);
575        assert_eq!(
576            decode(&b),
577            Ok(Message::Delete {
578                rel: 9,
579                old: old.to_vec()
580            })
581        );
582
583        let mut b = vec![b'T'];
584        b.extend_from_slice(&2u32.to_be_bytes());
585        b.push(1);
586        b.extend_from_slice(&5u32.to_be_bytes());
587        b.extend_from_slice(&6u32.to_be_bytes());
588        assert_eq!(
589            decode(&b),
590            Ok(Message::Truncate {
591                options: 1,
592                rels: vec![5, 6]
593            })
594        );
595
596        let mut b = vec![b'O'];
597        b.extend_from_slice(&3u64.to_be_bytes());
598        s(&mut b, "node1");
599        assert_eq!(
600            decode(&b),
601            Ok(Message::Origin {
602                lsn: 3,
603                name: "node1".into()
604            })
605        );
606
607        let mut b = vec![b'Y'];
608        b.extend_from_slice(&600u32.to_be_bytes());
609        s(&mut b, "public");
610        s(&mut b, "mood");
611        assert_eq!(
612            decode(&b),
613            Ok(Message::Type {
614                id: 600,
615                namespace: "public".into(),
616                name: "mood".into()
617            })
618        );
619    }
620
621    #[test]
622    fn malformed_messages_are_errors_not_panics() {
623        let msgs = [
624            begin(1),
625            commit(16),
626            relation(3),
627            insert(3, &[txt("a"), Datum::Null]),
628            update(3, Some((b'O', &[txt("a")])), &[txt("b")]),
629        ];
630        for m in &msgs {
631            for n in 0..m.len() {
632                assert!(decode(&m[..n]).is_err(), "prefix {n} of {m:?}");
633            }
634            let mut long = m.clone();
635            long.push(0);
636            assert_eq!(decode(&long), Err(DecodeError::Trailing(1)));
637        }
638        assert_eq!(decode(b"Z"), Err(DecodeError::UnknownMessage('Z')));
639        let mut b = insert(3, &[txt("a")]);
640        b[5] = b'X';
641        assert_eq!(decode(&b), Err(DecodeError::BadMarker('X')));
642        let mut b = insert(3, &[txt("a")]);
643        b[8] = b'b'; // binary values are protocol 2+ with `binary`
644        assert_eq!(decode(&b), Err(DecodeError::UnknownColumn('b')));
645        let mut b = insert(3, &[txt("a")]);
646        *b.last_mut().unwrap() = 0xFF;
647        assert_eq!(decode(&b), Err(DecodeError::BadString));
648        let mut b = vec![b'D'];
649        b.extend_from_slice(&3u32.to_be_bytes());
650        b.push(b'N');
651        assert_eq!(decode(&b), Err(DecodeError::BadMarker('N')));
652    }
653
654    #[test]
655    fn stream_groups_committed_transactions() {
656        let row = |q: &str| [txt(q), txt("Widget"), txt("10")];
657        let mut st = Stream::default();
658        for m in [
659            begin(10),
660            relation(3),
661            insert(3, &row("Q1")),
662            update(3, None, &row("Q2")),
663            commit(0x100),
664            begin(11),
665            commit(0x200),
666            begin(12),
667            insert(3, &row("Q3")),
668        ] {
669            st.feed(&m).unwrap();
670        }
671        assert_eq!(st.done.len(), 2, "the open transaction 12 is not done");
672        assert_eq!((st.done[0].xid, st.done[0].end_lsn), (10, 0x100));
673        assert_eq!(st.done[0].changes.len(), 2);
674        assert!(matches!(
675            st.done[0].changes[1],
676            Message::Update { rel: 3, .. }
677        ));
678        assert_eq!((st.done[1].xid, st.done[1].end_lsn), (11, 0x200));
679        assert!(st.done[1].changes.is_empty());
680        assert_eq!(st.relations[&3].name, "sales");
681        assert_eq!(st.feed(&begin(13)), Err(DecodeError::OutsideTxn("Begin")));
682    }
683
684    #[test]
685    fn stream_rejects_out_of_order_messages() {
686        let mut st = Stream::default();
687        assert_eq!(st.feed(&commit(8)), Err(DecodeError::OutsideTxn("Commit")));
688        assert_eq!(
689            st.feed(&insert(3, &[])),
690            Err(DecodeError::UnknownRelation(3))
691        );
692        st.feed(&relation(3)).unwrap();
693        assert_eq!(
694            st.feed(&insert(3, &[])),
695            Err(DecodeError::OutsideTxn("row change"))
696        );
697        let mut b = vec![b'T'];
698        b.extend_from_slice(&1u32.to_be_bytes());
699        b.push(0);
700        b.extend_from_slice(&4u32.to_be_bytes());
701        st.feed(&begin(1)).unwrap();
702        assert_eq!(st.feed(&b), Err(DecodeError::UnknownRelation(4)));
703    }
704
705    #[test]
706    fn lsn_text_round_trip() {
707        for lsn in [0, 1, 0x1_6B37_4D48, u64::MAX] {
708            assert_eq!(parse_lsn(&format_lsn(lsn)), Some(lsn));
709        }
710        assert_eq!(format_lsn(0x1_6B37_4D48), "1/6B374D48");
711        assert_eq!(parse_lsn("0/16B374D8"), Some(0x16B3_74D8));
712        for bad in ["", "1", "1/", "/1", "g/1", "1/2/3", "100000000/0"] {
713            assert_eq!(parse_lsn(bad), None, "{bad}");
714        }
715    }
716}
improv_pg_model/pgoutput.rs

improv_pg_model/
pgoutput.rs
