479 lines
14 KiB
Rust
479 lines
14 KiB
Rust
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());
|
|
}
|
|
}
|