use bytes::{Buf, BytesMut}; use std::collections::HashMap; /// PostgreSQL logical replication binary protocol parser (pgoutput, version 1). /// /// Parses the raw WAL bytes delivered by `START_REPLICATION … (proto_version '1')`. /// In protocol version 1 the transaction ID (XID) is **only** present in `Begin` /// messages; DML messages (Insert / Update / Delete / Truncate) carry no XID. /// The [`WalEmitter`](super::emitter::WalEmitter) tracks the current XID via the /// `Begin`/`Commit` boundary. // --------------------------------------------------------------------------- // Schema metadata // --------------------------------------------------------------------------- #[derive(Debug, Clone)] pub struct ColumnInfo { /// Bit 0: column is part of the replica identity key. pub flags: u8, pub name: String, pub type_oid: u32, pub type_modifier: i32, } #[derive(Debug, Clone)] pub struct RelationInfo { pub oid: u32, pub namespace: String, pub name: String, /// `d` = default (PK), `n` = nothing, `f` = full row, `i` = index. pub replica_identity: u8, pub columns: Vec, } // --------------------------------------------------------------------------- // Tuple data // --------------------------------------------------------------------------- #[derive(Debug, Clone)] pub enum ColumnValue { Null, UnchangedToasted, Text(String), Binary(Vec), } #[derive(Debug, Clone)] pub struct TupleData { pub values: Vec, } // --------------------------------------------------------------------------- // Logical messages (pgoutput v1) // --------------------------------------------------------------------------- #[derive(Debug)] pub enum LogicalMessage { /// Transaction start. `xid` identifies the transaction; DML messages that /// follow do not carry their own XID. Begin { /// Final LSN of the transaction (equals the Commit LSN). final_lsn: u64, /// Commit timestamp (microseconds since PG epoch 2000-01-01). timestamp: i64, /// Transaction ID. xid: u32, }, /// Schema descriptor sent before the first DML that references a relation /// in a session (and after schema changes). Stored internally by the parser. Relation(RelationInfo), /// Row inserted into `rel_oid`. Insert { rel_oid: u32, new_tuple: TupleData, }, /// Row updated in `rel_oid`. /// /// `key_tuple` is present when REPLICA IDENTITY is the primary key /// (default). `old_tuple` is present for REPLICA IDENTITY FULL. Update { rel_oid: u32, key_tuple: Option, old_tuple: Option, new_tuple: TupleData, }, /// Row deleted from `rel_oid`. Delete { rel_oid: u32, key_tuple: Option, old_tuple: Option, }, /// Transaction committed successfully. Commit { flags: u8, commit_lsn: u64, end_lsn: u64, timestamp: i64, }, /// One or more tables truncated. Truncate { num_relations: u32, options: u8, rel_oids: Vec, }, /// Custom data-type descriptor (pgoutput emits these before first use). Type { type_oid: u32, namespace: String, type_name: String, }, /// Application-level logical message emitted via `pg_logical_emit_message`. Message { /// Bit 0: message is transactional. flags: u8, lsn: u64, prefix: String, content: Vec, }, } // --------------------------------------------------------------------------- // Parser // --------------------------------------------------------------------------- /// Stateful pgoutput v1 parser. /// /// Maintains a cache of `RelationInfo` objects keyed by OID so that /// Insert/Update/Delete messages can be decoded without extra round-trips. pub struct MessageParser { relations: HashMap, _types: HashMap, } impl MessageParser { pub fn new() -> Self { Self { relations: HashMap::new(), _types: HashMap::new(), } } // ----------------------------------------------------------------------- // Public entry point // ----------------------------------------------------------------------- /// Parse the next logical message from `buf`. /// /// `Relation` messages are absorbed into internal state and return `None` /// so callers only see semantically meaningful messages. pub fn parse_message(&mut self, buf: &mut BytesMut) -> Result, String> { if buf.is_empty() { return Ok(None); } match buf[0] as char { 'B' => self.parse_begin(buf), 'C' => self.parse_commit(buf), 'R' => { self.parse_relation(buf)?; Ok(None) } 'I' => self.parse_insert(buf), 'U' => self.parse_update(buf), 'D' => self.parse_delete(buf), 'T' => self.parse_truncate(buf), 'Y' => self.parse_type(buf), 'M' => self.parse_message_msg(buf), other => Err(format!( "unknown pgoutput message type: '{}' (0x{:02x})", other, buf[0] )), } } /// Look up a relation by OID (returns `None` if not yet seen). pub fn get_relation(&self, oid: u32) -> Option<&RelationInfo> { self.relations.get(&oid) } // ----------------------------------------------------------------------- // Low-level readers // ----------------------------------------------------------------------- fn read_u8(buf: &mut BytesMut) -> Result { if buf.is_empty() { return Err("unexpected EOF reading u8".into()); } Ok(buf.get_u8()) } fn read_i16(buf: &mut BytesMut) -> Result { if buf.len() < 2 { return Err("unexpected EOF reading i16".into()); } Ok(buf.get_i16()) } fn read_u32(buf: &mut BytesMut) -> Result { if buf.len() < 4 { return Err("unexpected EOF reading u32".into()); } Ok(buf.get_u32()) } fn read_i32(buf: &mut BytesMut) -> Result { if buf.len() < 4 { return Err("unexpected EOF reading i32".into()); } Ok(buf.get_i32()) } fn read_u64(buf: &mut BytesMut) -> Result { if buf.len() < 8 { return Err("unexpected EOF reading u64".into()); } Ok(buf.get_u64()) } fn read_i64(buf: &mut BytesMut) -> Result { if buf.len() < 8 { return Err("unexpected EOF reading i64".into()); } Ok(buf.get_i64()) } /// Read a null-terminated UTF-8 string. fn read_string(buf: &mut BytesMut) -> Result { let mut bytes = Vec::new(); loop { if buf.is_empty() { return Err("unexpected EOF reading null-terminated string".into()); } let b = buf.get_u8(); if b == 0 { break; } bytes.push(b); } String::from_utf8(bytes).map_err(|e| format!("invalid UTF-8 in string: {e}")) } // ----------------------------------------------------------------------- // Message parsers // ----------------------------------------------------------------------- fn parse_begin(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'B' let final_lsn = Self::read_u64(buf)?; let timestamp = Self::read_i64(buf)?; let xid = Self::read_u32(buf)?; Ok(Some(LogicalMessage::Begin { final_lsn, timestamp, xid })) } fn parse_commit(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'C' let flags = Self::read_u8(buf)?; let commit_lsn = Self::read_u64(buf)?; let end_lsn = Self::read_u64(buf)?; let timestamp = Self::read_i64(buf)?; Ok(Some(LogicalMessage::Commit { flags, commit_lsn, end_lsn, timestamp })) } fn parse_relation(&mut self, buf: &mut BytesMut) -> Result<(), String> { buf.get_u8(); // 'R' let rel_oid = Self::read_u32(buf)?; let namespace = Self::read_string(buf)?; let name = Self::read_string(buf)?; let replica_identity = Self::read_u8(buf)?; let num_columns = Self::read_i16(buf)? as usize; let mut columns = Vec::with_capacity(num_columns); for _ in 0..num_columns { let flags = Self::read_u8(buf)?; let col_name = Self::read_string(buf)?; let type_oid = Self::read_u32(buf)?; let type_modifier = Self::read_i32(buf)?; columns.push(ColumnInfo { flags, name: col_name, type_oid, type_modifier }); } self.relations.insert( rel_oid, RelationInfo { oid: rel_oid, namespace, name, replica_identity, columns }, ); Ok(()) } fn parse_insert(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'I' // pgoutput v1: Relation OID comes directly after type byte (no XID). let rel_oid = Self::read_u32(buf)?; let marker = Self::read_u8(buf)?; if marker != b'N' { return Err(format!("INSERT: expected 'N' marker, got '{}'", marker as char)); } let new_tuple = self.parse_tuple_data(buf, rel_oid)?; Ok(Some(LogicalMessage::Insert { rel_oid, new_tuple })) } fn parse_update(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'U' // pgoutput v1: no XID in Update. let rel_oid = Self::read_u32(buf)?; let mut key_tuple = None; let mut old_tuple = None; if !buf.is_empty() { match buf[0] as char { 'K' => { buf.get_u8(); key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); } 'O' => { buf.get_u8(); old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); } _ => {} // no optional tuple, proceed to 'N' } } let marker = Self::read_u8(buf)?; if marker != b'N' { return Err(format!("UPDATE: expected 'N' marker, got '{}'", marker as char)); } let new_tuple = self.parse_tuple_data(buf, rel_oid)?; Ok(Some(LogicalMessage::Update { rel_oid, key_tuple, old_tuple, new_tuple })) } fn parse_delete(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'D' // pgoutput v1: no XID in Delete. let rel_oid = Self::read_u32(buf)?; let mut key_tuple = None; let mut old_tuple = None; if !buf.is_empty() { match buf[0] as char { 'K' => { buf.get_u8(); key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); } 'O' => { buf.get_u8(); old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?); } other => { return Err(format!("DELETE: expected 'K' or 'O', got '{}'", other)); } } } Ok(Some(LogicalMessage::Delete { rel_oid, key_tuple, old_tuple })) } fn parse_truncate(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'T' // pgoutput v1: no XID in Truncate. let num_relations = Self::read_u32(buf)?; let options = Self::read_u8(buf)?; let mut rel_oids = Vec::with_capacity(num_relations as usize); for _ in 0..num_relations { rel_oids.push(Self::read_u32(buf)?); } Ok(Some(LogicalMessage::Truncate { num_relations, options, rel_oids })) } fn parse_type(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'Y' // pgoutput v1: no XID in Type. let type_oid = Self::read_u32(buf)?; let namespace = Self::read_string(buf)?; let type_name = Self::read_string(buf)?; self._types.insert(type_oid, (namespace.clone(), type_name.clone())); Ok(Some(LogicalMessage::Type { type_oid, namespace, type_name })) } fn parse_message_msg(&mut self, buf: &mut BytesMut) -> Result, String> { buf.get_u8(); // 'M' // pgoutput v1: no XID in Message. let flags = Self::read_u8(buf)?; let lsn = Self::read_u64(buf)?; let prefix = Self::read_string(buf)?; let content_len = Self::read_i32(buf)? as usize; if buf.len() < content_len { return Err(format!( "Message: need {} bytes for content, have {}", content_len, buf.len() )); } let content = buf.split_to(content_len).to_vec(); Ok(Some(LogicalMessage::Message { flags, lsn, prefix, content })) } // ----------------------------------------------------------------------- // Tuple decoder // ----------------------------------------------------------------------- fn parse_tuple_data(&self, buf: &mut BytesMut, rel_oid: u32) -> Result { let num_columns = Self::read_i16(buf)? as usize; let relation = self .relations .get(&rel_oid) .ok_or_else(|| format!("unknown relation OID {rel_oid} — missing Relation message?"))?; if num_columns != relation.columns.len() { return Err(format!( "tuple has {} columns but relation '{}' has {}", num_columns, relation.name, relation.columns.len() )); } let mut values = Vec::with_capacity(num_columns); for _ in 0..num_columns { let marker = Self::read_u8(buf)?; let value = match marker { b'n' => ColumnValue::Null, b'u' => ColumnValue::UnchangedToasted, b't' => { let len = Self::read_i32(buf)? as usize; if buf.len() < len { return Err(format!( "text column: need {} bytes, have {}", len, buf.len() )); } let data = buf.split_to(len).to_vec(); let text = String::from_utf8(data) .map_err(|e| format!("non-UTF-8 text column: {e}"))?; ColumnValue::Text(text) } b'b' => { let len = Self::read_i32(buf)? as usize; if buf.len() < len { return Err(format!( "binary column: need {} bytes, have {}", len, buf.len() )); } ColumnValue::Binary(buf.split_to(len).to_vec()) } other => { return Err(format!( "unknown column marker '{}'", other as char )); } }; values.push(value); } Ok(TupleData { values }) } } impl Default for MessageParser { fn default() -> Self { Self::new() } } // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- #[cfg(test)] mod tests { use super::*; #[test] fn parser_starts_empty() { let p = MessageParser::new(); assert!(p.relations.is_empty()); } #[test] fn empty_buf_returns_none() { let mut p = MessageParser::new(); let mut buf = BytesMut::new(); assert!(p.parse_message(&mut buf).unwrap().is_none()); } }