dump some extra files

2026-03-08 22:34:03 +01:00
parent ed85af4a4f
commit 2358ca29c7
4 changed files with 1461 additions and 0 deletions
@@ -0,0 +1,478 @@
+use bytes::{Buf, BytesMut};
+use std::collections::HashMap;
+
+/// PostgreSQL logical replication message parser
+/// Handles pgoutput protocol messages from replication streams
+
+#[derive(Debug, Clone)]
+pub struct ColumnInfo {
+    pub flags: u8,              // 1 if part of key, 0 otherwise
+    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,
+    pub replica_identity: u8,
+    pub columns: Vec<ColumnInfo>,
+}
+
+#[derive(Debug, Clone)]
+pub enum ColumnValue {
+    Null,
+    UnchangedToasted,
+    Text(String),
+    Binary(Vec<u8>),
+}
+
+#[derive(Debug)]
+pub struct TupleData {
+    pub values: Vec<ColumnValue>,
+}
+
+#[derive(Debug)]
+pub enum LogicalMessage {
+    Begin {
+        final_lsn: u64,
+        timestamp: i64,
+        xid: u32,
+    },
+    Relation(RelationInfo),
+    Insert {
+        xid: u32,
+        rel_oid: u32,
+        new_tuple: TupleData,
+    },
+    Update {
+        xid: u32,
+        rel_oid: u32,
+        key_tuple: Option<TupleData>,
+        old_tuple: Option<TupleData>,
+        new_tuple: TupleData,
+    },
+    Delete {
+        xid: u32,
+        rel_oid: u32,
+        key_tuple: Option<TupleData>,
+        old_tuple: Option<TupleData>,
+    },
+    Commit {
+        flags: u8,
+        commit_lsn: u64,
+        end_lsn: u64,
+        timestamp: i64,
+    },
+    Type {
+        xid: u32,
+        type_oid: u32,
+        namespace: String,
+        type_name: String,
+    },
+    Truncate {
+        xid: u32,
+        num_relations: u32,
+        options: u8,
+        rel_oids: Vec<u32>,
+    },
+    Message {
+        xid: u32,
+        flags: u8,
+        lsn: u64,
+        prefix: String,
+        content: Vec<u8>,
+    },
+}
+
+pub struct MessageParser {
+    /// Map of relation OID -> RelationInfo for decoding tuples
+    relations: HashMap<u32, RelationInfo>,
+    /// Map of type OID -> type name for potential future use
+    types: HashMap<u32, (String, String)>,
+}
+
+impl MessageParser {
+    pub fn new() -> Self {
+        Self {
+            relations: HashMap::new(),
+            types: HashMap::new(),
+        }
+    }
+
+    /// Parse a single logical replication message from bytes
+    /// Assumes input starts with the message type byte
+    pub fn parse_message(&mut self, input: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
+        if input.is_empty() {
+            return Ok(None);
+        }
+
+        let msg_type = input[0] as char;
+
+        match msg_type {
+            'B' => self.parse_begin(input),
+            'I' => self.parse_insert(input),
+            'U' => self.parse_update(input),
+            'D' => self.parse_delete(input),
+            'C' => self.parse_commit(input),
+            'R' => {
+                self.parse_relation(input)?;
+                Ok(None) // Relation messages are stored internally
+            }
+            'Y' => self.parse_type(input),
+            'T' => self.parse_truncate(input),
+            'M' => self.parse_message_msg(input),
+            't' => Err("Tuple data outside of Insert/Update/Delete (unexpected)".to_string()),
+            _ => Err(format!("Unknown message type: {} (byte: {})", msg_type, input[0])),
+        }
+    }
+
+    fn read_string(buf: &mut BytesMut) -> Result<String, String> {
+        // PostgreSQL strings are null-terminated
+        let mut vec = Vec::new();
+        loop {
+            if buf.is_empty() {
+                return Err("EOF while reading string".to_string());
+            }
+            let byte = buf.get_u8();
+            if byte == 0 {
+                break;
+            }
+            vec.push(byte);
+        }
+        String::from_utf8(vec).map_err(|e| format!("Invalid UTF-8 in string: {}", e))
+    }
+
+    fn read_i64(buf: &mut BytesMut) -> Result<i64, String> {
+        if buf.len() < 8 {
+            return Err("Not enough bytes for i64".to_string());
+        }
+        Ok(buf.get_i64())
+    }
+
+    fn read_u64(buf: &mut BytesMut) -> Result<u64, String> {
+        if buf.len() < 8 {
+            return Err("Not enough bytes for u64".to_string());
+        }
+        Ok(buf.get_u64())
+    }
+
+    fn read_i32(buf: &mut BytesMut) -> Result<i32, String> {
+        if buf.len() < 4 {
+            return Err("Not enough bytes for i32".to_string());
+        }
+        Ok(buf.get_i32())
+    }
+
+    fn read_u32(buf: &mut BytesMut) -> Result<u32, String> {
+        if buf.len() < 4 {
+            return Err("Not enough bytes for u32".to_string());
+        }
+        Ok(buf.get_u32())
+    }
+
+    fn read_i16(buf: &mut BytesMut) -> Result<i16, String> {
+        if buf.len() < 2 {
+            return Err("Not enough bytes for i16".to_string());
+        }
+        Ok(buf.get_i16())
+    }
+
+    fn read_u8(buf: &mut BytesMut) -> Result<u8, String> {
+        if buf.is_empty() {
+            return Err("Not enough bytes for u8".to_string());
+        }
+        Ok(buf.get_u8())
+    }
+
+    fn parse_begin(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, 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<Option<LogicalMessage>, 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'
+        // xid only present for streamed transactions (version 2+), skip for now
+        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,
+            });
+        }
+
+        let relation = RelationInfo {
+            oid: rel_oid,
+            namespace,
+            name,
+            replica_identity,
+            columns,
+        };
+
+        self.relations.insert(rel_oid, relation);
+        Ok(())
+    }
+
+    fn parse_insert(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
+        buf.get_u8(); // 'I'
+        let xid = Self::read_u32(buf)?;
+        let rel_oid = Self::read_u32(buf)?;
+
+        // Next byte must be 'N' for new tuple
+        let marker = Self::read_u8(buf)?;
+        if marker != b'N' {
+            return Err(format!("Expected 'N' in INSERT, got: {}", marker as char));
+        }
+
+        let new_tuple = self.parse_tuple_data(buf, rel_oid)?;
+
+        Ok(Some(LogicalMessage::Insert {
+            xid,
+            rel_oid,
+            new_tuple,
+        }))
+    }
+
+    fn parse_update(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
+        buf.get_u8(); // 'U'
+        let xid = Self::read_u32(buf)?;
+        let rel_oid = Self::read_u32(buf)?;
+
+        let mut key_tuple = None;
+        let mut old_tuple = None;
+
+        // Optional K or O markers
+        if !buf.is_empty() {
+            match buf[0] as char {
+                'K' => {
+                    buf.get_u8(); // consume 'K'
+                    key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
+                }
+                'O' => {
+                    buf.get_u8(); // consume 'O'
+                    old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
+                }
+                _ => {} // neither K nor O, proceed to N
+            }
+        }
+
+        let marker = Self::read_u8(buf)?;
+        if marker != b'N' {
+            return Err(format!("Expected 'N' in UPDATE, got: {}", marker as char));
+        }
+
+        let new_tuple = self.parse_tuple_data(buf, rel_oid)?;
+
+        Ok(Some(LogicalMessage::Update {
+            xid,
+            rel_oid,
+            key_tuple,
+            old_tuple,
+            new_tuple,
+        }))
+    }
+
+    fn parse_delete(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
+        buf.get_u8(); // 'D'
+        let xid = Self::read_u32(buf)?;
+        let rel_oid = Self::read_u32(buf)?;
+
+        let mut key_tuple = None;
+        let mut old_tuple = None;
+
+        // Either K or O (but not both)
+        if !buf.is_empty() {
+            match buf[0] as char {
+                'K' => {
+                    buf.get_u8(); // consume 'K'
+                    key_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
+                }
+                'O' => {
+                    buf.get_u8(); // consume 'O'
+                    old_tuple = Some(self.parse_tuple_data(buf, rel_oid)?);
+                }
+                _ => return Err(format!("Expected K or O in DELETE, got: {}", buf[0] as char)),
+            }
+        }
+
+        Ok(Some(LogicalMessage::Delete {
+            xid,
+            rel_oid,
+            key_tuple,
+            old_tuple,
+        }))
+    }
+
+    fn parse_type(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
+        buf.get_u8(); // 'Y'
+        let type_oid = Self::read_u32(buf)?;
+        let namespace = Self::read_string(buf)?;
+        let type_name = Self::read_string(buf)?;
+        let xid = 0; // Type doesn't have xid in base protocol
+
+        self.types
+            .insert(type_oid, (namespace.clone(), type_name.clone()));
+
+        Ok(Some(LogicalMessage::Type {
+            xid,
+            type_oid,
+            namespace,
+            type_name,
+        }))
+    }
+
+    fn parse_truncate(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
+        buf.get_u8(); // 'T'
+        let xid = Self::read_u32(buf)?;
+        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 {
+            xid,
+            num_relations,
+            options,
+            rel_oids,
+        }))
+    }
+
+    fn parse_message_msg(&mut self, buf: &mut BytesMut) -> Result<Option<LogicalMessage>, String> {
+        buf.get_u8(); // 'M'
+        let xid = Self::read_u32(buf)?;
+        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!(
+                "Not enough bytes for message content: {} < {}",
+                buf.len(),
+                content_len
+            ));
+        }
+
+        let content = buf.split_to(content_len).to_vec();
+
+        Ok(Some(LogicalMessage::Message {
+            xid,
+            flags,
+            lsn,
+            prefix,
+            content,
+        }))
+    }
+
+    fn parse_tuple_data(&self, buf: &mut BytesMut, rel_oid: u32) -> Result<TupleData, String> {
+        let num_columns = Self::read_i16(buf)? as usize;
+        let mut values = Vec::with_capacity(num_columns);
+
+        let relation = self
+            .relations
+            .get(&rel_oid)
+            .ok_or_else(|| format!("Unknown relation OID: {}", rel_oid))?;
+
+        if num_columns != relation.columns.len() {
+            return Err(format!(
+                "Column count mismatch: {} vs {}",
+                num_columns,
+                relation.columns.len()
+            ));
+        }
+
+        for _ in 0..num_columns {
+            let marker = Self::read_u8(buf)?;
+
+            let value = match marker as char {
+                'n' => ColumnValue::Null,
+                'u' => ColumnValue::UnchangedToasted,
+                't' => {
+                    let len = Self::read_i32(buf)? as usize;
+                    if buf.len() < len {
+                        return Err("Not enough bytes for text column".to_string());
+                    }
+                    let data = buf.split_to(len).to_vec();
+                    let text = String::from_utf8(data)
+                        .map_err(|e| format!("Invalid UTF-8 in column: {}", e))?;
+                    ColumnValue::Text(text)
+                }
+                'b' => {
+                    let len = Self::read_i32(buf)? as usize;
+                    if buf.len() < len {
+                        return Err("Not enough bytes for binary column".to_string());
+                    }
+                    let data = buf.split_to(len).to_vec();
+                    ColumnValue::Binary(data)
+                }
+                _ => return Err(format!("Unknown column marker: {}", marker as char)),
+            };
+
+            values.push(value);
+        }
+
+        Ok(TupleData { values })
+    }
+
+    /// Get a relation by OID (used for cache invalidation mapping)
+    pub fn get_relation(&self, oid: u32) -> Option<&RelationInfo> {
+        self.relations.get(&oid)
+    }
+
+    /// Get all cached relations
+    pub fn relations(&self) -> &HashMap<u32, RelationInfo> {
+        &self.relations
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_message_parser_creation() {
+        let parser = MessageParser::new();
+        assert!(parser.relations.is_empty());
+    }
+}