livevue-rs/src/pg_replication/wal_parser.rs

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<ColumnInfo>,
}

// ---------------------------------------------------------------------------
// Tuple data
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub enum ColumnValue {
    Null,
    UnchangedToasted,
    Text(String),
    Binary(Vec<u8>),
}

#[derive(Debug, Clone)]
pub struct TupleData {
    pub values: Vec<ColumnValue>,
}

// ---------------------------------------------------------------------------
// 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<TupleData>,
        old_tuple: Option<TupleData>,
        new_tuple: TupleData,
    },

    /// Row deleted from `rel_oid`.
    Delete {
        rel_oid: u32,
        key_tuple: Option<TupleData>,
        old_tuple: Option<TupleData>,
    },

    /// 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<u32>,
    },

    /// 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<u8>,
    },
}

// ---------------------------------------------------------------------------
// 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<u32, RelationInfo>,
    _types: HashMap<u32, (String, String)>,
}

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<Option<LogicalMessage>, 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<u8, String> {
        if buf.is_empty() {
            return Err("unexpected EOF reading u8".into());
        }
        Ok(buf.get_u8())
    }

    fn read_i16(buf: &mut BytesMut) -> Result<i16, String> {
        if buf.len() < 2 {
            return Err("unexpected EOF reading i16".into());
        }
        Ok(buf.get_i16())
    }

    fn read_u32(buf: &mut BytesMut) -> Result<u32, String> {
        if buf.len() < 4 {
            return Err("unexpected EOF reading u32".into());
        }
        Ok(buf.get_u32())
    }

    fn read_i32(buf: &mut BytesMut) -> Result<i32, String> {
        if buf.len() < 4 {
            return Err("unexpected EOF reading i32".into());
        }
        Ok(buf.get_i32())
    }

    fn read_u64(buf: &mut BytesMut) -> Result<u64, String> {
        if buf.len() < 8 {
            return Err("unexpected EOF reading u64".into());
        }
        Ok(buf.get_u64())
    }

    fn read_i64(buf: &mut BytesMut) -> Result<i64, String> {
        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<String, String> {
        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<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'
        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<Option<LogicalMessage>, 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<Option<LogicalMessage>, 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<Option<LogicalMessage>, 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<Option<LogicalMessage>, 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<Option<LogicalMessage>, 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<Option<LogicalMessage>, 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<TupleData, String> {
        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());
    }
}