#![allow(unused)] use core::str; use std::{ cmp, io::{self, Cursor, Read}, str::Utf8Error, }; use byteorder::{BigEndian, ReadBytesExt}; use bytes::Bytes; use rust_postgres::types::{Oid, Type}; use thiserror::Error; use super::trigger::LogicalReplicationSettings; const PRIMARY_KEEPALIVE_BYTE: u8 = b'k'; const X_LOG_DATA_BYTE: u8 = b'w'; /** * This implementation is inspired by Postgres replication functionality * from https://github.com/supabase/pg_replicate * * Original implementation: * - https://github.com/supabase/pg_replicate/blob/main/pg_replicate/src/conversions/cdc_event.rs * */ #[derive(Debug)] pub struct PrimaryKeepAliveBody { pub wal_end: u64, pub timestamp: i64, pub reply: bool, } impl PrimaryKeepAliveBody { pub fn new(wal_end: u64, timestamp: i64, reply: bool) -> PrimaryKeepAliveBody { PrimaryKeepAliveBody { wal_end, timestamp, reply } } } const BEGIN_BYTE: u8 = b'B'; const COMMIT_BYTE: u8 = b'C'; const ORIGIN_BYTE: u8 = b'O'; const RELATION_BYTE: u8 = b'R'; const TYPE_BYTE: u8 = b'Y'; const INSERT_BYTE: u8 = b'I'; const UPDATE_BYTE: u8 = b'U'; const DELETE_BYTE: u8 = b'D'; const TUPLE_NEW_BYTE: u8 = b'N'; const TUPLE_KEY_BYTE: u8 = b'K'; const TUPLE_OLD_BYTE: u8 = b'O'; const TUPLE_DATA_NULL_BYTE: u8 = b'n'; const TUPLE_DATA_TOAST_BYTE: u8 = b'u'; const TUPLE_DATA_TEXT_BYTE: u8 = b't'; const TUPLE_DATA_BINARY_BYTE: u8 = b'b'; const REPLICA_IDENTITY_DEFAULT_BYTE: i8 = 0x64; const REPLICA_IDENTITY_NOTHING_BYTE: i8 = 0x6E; const REPLICA_IDENTITY_FULL_BYTE: i8 = 0x66; const REPLICA_IDENTITY_INDEX_BYTE: i8 = 0x69; #[derive(Debug)] pub enum ReplicaIdentity { Default, Nothing, Full, Index, } #[derive(Debug)] pub struct Column { pub flags: i8, pub name: String, pub type_o_id: Option, pub type_modifier: i32, } impl Column { pub fn new(flags: i8, name: String, type_o_id: Option, type_modifier: i32) -> Self { Self { flags, name, type_o_id, type_modifier } } } pub type Columns = Vec; #[derive(Debug)] pub struct RelationBody { pub transaction_id: Option, pub o_id: Oid, pub namespace: String, pub name: String, pub replica_identity: ReplicaIdentity, pub columns: Columns, } impl RelationBody { pub fn new( transaction_id: Option, o_id: Oid, namespace: String, name: String, replica_identity: ReplicaIdentity, columns: Columns, ) -> Self { Self { transaction_id, o_id, namespace, name, replica_identity, columns } } } #[derive(Debug)] pub struct InsertBody { pub transaction_id: Option, pub o_id: Oid, pub tuple: Vec, } impl InsertBody { pub fn new(transaction_id: Option, o_id: Oid, tuple: Vec) -> Self { Self { transaction_id, o_id, tuple } } } #[derive(Debug)] pub struct UpdateBody { transaction_id: Option, pub o_id: Oid, pub old_tuple: Option>, pub key_tuple: Option>, pub new_tuple: Vec, } impl UpdateBody { pub fn new( transaction_id: Option, o_id: Oid, old_tuple: Option>, key_tuple: Option>, new_tuple: Vec, ) -> Self { Self { transaction_id, o_id, old_tuple, key_tuple, new_tuple } } } #[derive(Debug)] pub struct DeleteBody { transaction_id: Option, pub o_id: Oid, pub old_tuple: Option>, pub key_tuple: Option>, } impl DeleteBody { pub fn new( transaction_id: Option, o_id: Oid, old_tuple: Option>, key_tuple: Option>, ) -> Self { Self { transaction_id, o_id, old_tuple, key_tuple } } } #[derive(Debug)] pub enum TupleData { Null, UnchangedToast, Text(Bytes), Binary(Bytes), } impl TupleData { fn parse(buf: &mut Buffer) -> Result, ConversionError> { let number_of_columns = buf.read_i16::()?; let mut tuples = Vec::with_capacity(number_of_columns as usize); for _ in 0..number_of_columns { let byte = buf.read_u8()?; let tuple_data = match byte { TUPLE_DATA_NULL_BYTE => TupleData::Null, TUPLE_DATA_TOAST_BYTE => TupleData::UnchangedToast, TUPLE_DATA_TEXT_BYTE => { let len = buf.read_i32::()?; let mut data = vec![0; len as usize]; buf.read_exact(&mut data)?; TupleData::Text(data.into()) } TUPLE_DATA_BINARY_BYTE => { let len = buf.read_i32::()?; let mut data = vec![0; len as usize]; buf.read_exact(&mut data)?; TupleData::Binary(data.into()) } byte => { return Err(ConversionError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!("unknown replication message byte `{}`", byte), ))); } }; tuples.push(tuple_data); } Ok(tuples) } } #[derive(Debug)] pub enum TransactionBody { Insert(InsertBody), Update(UpdateBody), Delete(DeleteBody), } #[non_exhaustive] #[derive(Debug)] pub enum LogicalReplicationMessage { Begin, Commit, Relation(RelationBody), Type, Insert(InsertBody), Update(UpdateBody), Delete(DeleteBody), } #[derive(Debug)] pub struct XLogDataBody { pub wal_start: u64, pub wal_end: u64, pub timestamp: i64, pub data: Bytes, } #[derive(Error, Debug)] pub enum ConversionError { #[error("Error: {0}")] Io(#[from] io::Error), #[error("Utf8Error conversion: {0}")] Utf8(#[from] Utf8Error), } struct Buffer { bytes: Bytes, idx: usize, } impl Buffer { pub fn new(bytes: Bytes, idx: usize) -> Buffer { Buffer { bytes, idx } } fn slice(&self) -> &[u8] { &self.bytes[self.idx..] } fn read_cstr(&mut self) -> Result { match self.slice().iter().position(|&x| x == 0) { Some(pos) => { let start = self.idx; let end = start + pos; let cstr = str::from_utf8(&self.bytes[start..end])?.to_owned(); self.idx = end + 1; Ok(cstr) } None => Err(ConversionError::Io(io::Error::new( io::ErrorKind::UnexpectedEof, "unexpected EOF", ))), } } } impl Read for Buffer { fn read(&mut self, buf: &mut [u8]) -> io::Result { let len = { let slice = self.slice(); let len = cmp::min(slice.len(), buf.len()); buf[..len].copy_from_slice(&slice[..len]); len }; self.idx += len; Ok(len) } } impl XLogDataBody { pub fn new(wal_start: u64, wal_end: u64, timestamp: i64, data: Bytes) -> XLogDataBody { XLogDataBody { wal_start, wal_end, timestamp, data } } pub fn parse( self, logical_replication_settings: &LogicalReplicationSettings, ) -> Result { let mut buf = Buffer::new(self.data.clone(), 0); let byte = buf.read_u8()?; let logical_replication_message = match byte { BEGIN_BYTE => { buf.read_i64::()?; buf.read_i64::()?; buf.read_i32::()?; LogicalReplicationMessage::Begin } COMMIT_BYTE => { buf.read_i8()?; buf.read_u64::()?; buf.read_u64::()?; buf.read_i64::()?; LogicalReplicationMessage::Commit } RELATION_BYTE => { let transaction_id = match logical_replication_settings.streaming { true => Some(buf.read_i32::()?), false => None, }; let o_id = buf.read_u32::()?; let namespace = buf.read_cstr()?; let name = buf.read_cstr()?; let replica_identity = match buf.read_i8()? { REPLICA_IDENTITY_DEFAULT_BYTE => ReplicaIdentity::Default, REPLICA_IDENTITY_NOTHING_BYTE => ReplicaIdentity::Nothing, REPLICA_IDENTITY_FULL_BYTE => ReplicaIdentity::Full, REPLICA_IDENTITY_INDEX_BYTE => ReplicaIdentity::Index, byte => { return Err(ConversionError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!("unknown replica identity byte `{}`", byte), ))); } }; let num_of_column = buf.read_i16::()?; let mut columns = Vec::with_capacity(num_of_column as usize); for _ in 0..num_of_column { let flags = buf.read_i8()?; let name = buf.read_cstr()?; let o_id = buf.read_u32::()?; let type_modifier = buf.read_i32::()?; let type_o_id = Type::from_oid(o_id); let column = Column::new(flags, name, type_o_id, type_modifier); columns.push(column); } LogicalReplicationMessage::Relation(RelationBody::new( transaction_id, o_id, namespace, name, replica_identity, columns, )) } TYPE_BYTE => { buf.read_u32::()?; buf.read_cstr()?; buf.read_cstr()?; LogicalReplicationMessage::Type } INSERT_BYTE => { let transaction_id = match logical_replication_settings.streaming { true => Some(buf.read_i32::()?), false => None, }; let o_id = buf.read_u32::()?; let byte = buf.read_u8()?; let tuple = match byte { TUPLE_NEW_BYTE => TupleData::parse(&mut buf)?, byte => { return Err(ConversionError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!("unexpected tuple byte `{}`", byte), ))); } }; LogicalReplicationMessage::Insert(InsertBody::new(transaction_id, o_id, tuple)) } UPDATE_BYTE => { let transaction_id = match logical_replication_settings.streaming { true => Some(buf.read_i32::()?), false => None, }; let o_id = buf.read_u32::()?; let byte = buf.read_u8()?; let mut key_tuple = None; let mut old_tuple = None; let new_tuple = match byte { TUPLE_NEW_BYTE => TupleData::parse(&mut buf)?, TUPLE_OLD_BYTE | TUPLE_KEY_BYTE => { if byte == TUPLE_OLD_BYTE { old_tuple = Some(TupleData::parse(&mut buf)?); } else { key_tuple = Some(TupleData::parse(&mut buf)?); } match buf.read_u8()? { TUPLE_NEW_BYTE => TupleData::parse(&mut buf)?, byte => { return Err(ConversionError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!("unexpected tuple byte `{}`", byte), ))); } } } byte => { return Err(ConversionError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!("unknown tuple byte `{}`", byte), ))); } }; LogicalReplicationMessage::Update(UpdateBody::new( transaction_id, o_id, old_tuple, key_tuple, new_tuple, )) } DELETE_BYTE => { let transaction_id = match logical_replication_settings.streaming { true => Some(buf.read_i32::()?), false => None, }; let o_id = buf.read_u32::()?; let tag = buf.read_u8()?; let mut key_tuple = None; let mut old_tuple = None; match tag { TUPLE_OLD_BYTE => old_tuple = Some(TupleData::parse(&mut buf)?), TUPLE_KEY_BYTE => key_tuple = Some(TupleData::parse(&mut buf)?), tag => { return Err(ConversionError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!("unknown tuple tag `{}`", tag), ))); } } LogicalReplicationMessage::Delete(DeleteBody::new( transaction_id, o_id, old_tuple, key_tuple, )) } byte => { return Err(ConversionError::Io(io::Error::new( io::ErrorKind::InvalidInput, format!("unknown replication message tag `{}`", byte), ))); } }; Ok(logical_replication_message) } } #[non_exhaustive] #[derive(Debug)] pub enum ReplicationMessage { XLogData(XLogDataBody), PrimaryKeepAlive(PrimaryKeepAliveBody), } impl ReplicationMessage { pub fn parse(buf: Bytes) -> io::Result { let (byte, mut message) = buf.split_first().unwrap(); let replication_message = match *byte { X_LOG_DATA_BYTE => { let len = buf.len(); let wal_start = message.read_u64::()?; let wal_end = message.read_u64::()?; let timestamp = message.read_i64::()?; let len = len - message.len(); let data = buf.slice(len..); ReplicationMessage::XLogData(XLogDataBody::new(wal_start, wal_end, timestamp, data)) } PRIMARY_KEEPALIVE_BYTE => { let wal_end = message.read_u64::()?; let timestamp = message.read_i64::()?; let reply = message.read_u8()?; ReplicationMessage::PrimaryKeepAlive(PrimaryKeepAliveBody::new( wal_end, timestamp, reply == 1, )) } byte => { return Err(io::Error::new( io::ErrorKind::InvalidInput, format!("unknown replication message byte `{}`", byte), )); } }; Ok(replication_message) } }