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,366 @@
+use crate::pg_wal_parser::{ColumnValue, LogicalMessage, MessageParser, RelationInfo, TupleData};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use tokio::sync::broadcast;
+
+/// High-level cache invalidation event emitted at transaction commit
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CacheInvalidationEvent {
+    /// Transaction ID from PostgreSQL
+    pub xid: u32,
+    /// Commit LSN/watermark for checkpoint
+    pub commit_lsn: u64,
+    /// End LSN of the transaction
+    pub end_lsn: u64,
+    /// Timestamp of commit (microseconds since PG epoch)
+    pub timestamp: i64,
+    /// All operations in this transaction, emitted atomically
+    pub changes: Vec<Change>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Change {
+    /// Table name that changed
+    pub table: String,
+    /// Schema/namespace
+    pub schema: String,
+    /// Operation type
+    pub operation: Operation,
+    /// Row data if available
+    pub row_data: Option<RowData>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum Operation {
+    Insert,
+    Update,
+    Delete,
+    Truncate,
+}
+
+/// Simplified row representation for broadcasting
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RowData {
+    /// Mapping of column name -> value
+    pub columns: HashMap<String, CacheValue>,
+    /// Primary key values if available
+    pub key_values: Option<HashMap<String, CacheValue>>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum CacheValue {
+    Null,
+    String(String),
+    Binary(Vec<u8>),
+    UnchangedToasted,
+}
+
+/// Transaction buffer for a single transaction
+struct TransactionBuffer {
+    xid: u32,
+    changes: Vec<Change>,
+}
+
+pub struct CacheInvalidationEmitter {
+    parser: MessageParser,
+    tx: broadcast::Sender<CacheInvalidationEvent>,
+    /// Buffers for active transactions, keyed by XID
+    txn_buffers: HashMap<u32, TransactionBuffer>,
+    /// Track the current transaction being processed
+    current_xid: Option<u32>,
+}
+
+impl CacheInvalidationEmitter {
+    pub fn new(tx: broadcast::Sender<CacheInvalidationEvent>) -> Self {
+        Self {
+            parser: MessageParser::new(),
+            tx,
+            txn_buffers: HashMap::new(),
+            current_xid: None,
+        }
+    }
+
+    /// Process a single logical replication message
+    /// Events are buffered until Commit, then emitted atomically
+    pub async fn process_message(&mut self, msg: LogicalMessage) -> Result<(), String> {
+        match msg {
+            LogicalMessage::Begin { xid, .. } => {
+                // Start buffering for this transaction
+                self.current_xid = Some(xid);
+                self.txn_buffers.insert(xid, TransactionBuffer {
+                    xid,
+                    changes: Vec::new(),
+                });
+                Ok(())
+            }
+            LogicalMessage::Relation(_) => {
+                // Relation updates internal state via parser
+                Ok(())
+            }
+            LogicalMessage::Insert {
+                xid,
+                rel_oid,
+                new_tuple,
+            } => {
+                let change = self.build_insert_change(rel_oid, new_tuple)?;
+                self.buffer_change(xid, change);
+                Ok(())
+            }
+            LogicalMessage::Update {
+                xid,
+                rel_oid,
+                key_tuple,
+                old_tuple: _,
+                new_tuple,
+            } => {
+                let change = self.build_update_change(rel_oid, new_tuple, key_tuple)?;
+                self.buffer_change(xid, change);
+                Ok(())
+            }
+            LogicalMessage::Delete {
+                xid,
+                rel_oid,
+                key_tuple,
+                old_tuple,
+            } => {
+                let change = self.build_delete_change(rel_oid, key_tuple.or(old_tuple))?;
+                self.buffer_change(xid, change);
+                Ok(())
+            }
+            LogicalMessage::Commit {
+                commit_lsn,
+                end_lsn,
+                timestamp,
+                ..
+            } => {
+                // Emit the buffered transaction atomically
+                if let Some(xid) = self.current_xid.take() {
+                    self.emit_transaction(xid, commit_lsn, end_lsn, timestamp)?;
+                }
+                Ok(())
+            }
+            LogicalMessage::Truncate {
+                xid,
+                rel_oids,
+                ..
+            } => {
+                for oid in rel_oids {
+                    let change = self.build_truncate_change(oid)?;
+                    self.buffer_change(xid, change);
+                }
+                Ok(())
+            }
+            LogicalMessage::Type { .. } => {
+                // Type information, no action needed
+                Ok(())
+            }
+            LogicalMessage::Message { .. } => {
+                // Custom messages, can be ignored for cache invalidation
+                Ok(())
+            }
+        }
+    }
+
+    fn buffer_change(&mut self, xid: u32, change: Change) {
+        if let Some(buf) = self.txn_buffers.get_mut(&xid) {
+            buf.changes.push(change);
+        }
+    }
+
+    fn build_insert_change(&self, rel_oid: u32, new_tuple: TupleData) -> Result<Change, String> {
+        let relation = self
+            .parser
+            .get_relation(rel_oid)
+            .ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
+
+        let row_data = self.build_row_data(relation, &new_tuple)?;
+
+        Ok(Change {
+            table: relation.name.clone(),
+            schema: relation.namespace.clone(),
+            operation: Operation::Insert,
+            row_data: Some(row_data),
+        })
+    }
+
+    fn build_update_change(
+        &self,
+        rel_oid: u32,
+        new_tuple: TupleData,
+        key_tuple: Option<TupleData>,
+    ) -> Result<Change, String> {
+        let relation = self
+            .parser
+            .get_relation(rel_oid)
+            .ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
+
+        let mut row_data = self.build_row_data(relation, &new_tuple)?;
+
+        // If we have a key tuple, extract key values
+        if let Some(key) = key_tuple {
+            let key_values = self.extract_key_values(relation, &key)?;
+            row_data.key_values = Some(key_values);
+        }
+
+        Ok(Change {
+            table: relation.name.clone(),
+            schema: relation.namespace.clone(),
+            operation: Operation::Update,
+            row_data: Some(row_data),
+        })
+    }
+
+    fn build_delete_change(
+        &self,
+        rel_oid: u32,
+        key_or_old: Option<TupleData>,
+    ) -> Result<Change, String> {
+        let relation = self
+            .parser
+            .get_relation(rel_oid)
+            .ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
+
+        let mut row_data = None;
+        if let Some(tuple) = key_or_old {
+            let mut data = self.build_row_data(relation, &tuple)?;
+            let key_values = self.extract_key_values(relation, &tuple)?;
+            data.key_values = Some(key_values);
+            row_data = Some(data);
+        }
+
+        Ok(Change {
+            table: relation.name.clone(),
+            schema: relation.namespace.clone(),
+            operation: Operation::Delete,
+            row_data,
+        })
+    }
+
+    fn build_truncate_change(&self, rel_oid: u32) -> Result<Change, String> {
+        let relation = self
+            .parser
+            .get_relation(rel_oid)
+            .ok_or_else(|| format!("Relation OID {} not found", rel_oid))?;
+
+        Ok(Change {
+            table: relation.name.clone(),
+            schema: relation.namespace.clone(),
+            operation: Operation::Truncate,
+            row_data: None,
+        })
+    }
+
+    fn build_row_data(&self, relation: &RelationInfo, tuple: &TupleData) -> Result<RowData, String> {
+        let mut columns = HashMap::new();
+
+        for (i, value) in tuple.values.iter().enumerate() {
+            if i >= relation.columns.len() {
+                return Err(format!(
+                    "Tuple has more values than relation has columns: {} vs {}",
+                    tuple.values.len(),
+                    relation.columns.len()
+                ));
+            }
+
+            let col = &relation.columns[i];
+            let cache_value = match value {
+                ColumnValue::Null => CacheValue::Null,
+                ColumnValue::UnchangedToasted => CacheValue::UnchangedToasted,
+                ColumnValue::Text(s) => CacheValue::String(s.clone()),
+                ColumnValue::Binary(b) => CacheValue::Binary(b.clone()),
+            };
+
+            columns.insert(col.name.clone(), cache_value);
+        }
+
+        Ok(RowData {
+            columns,
+            key_values: None,
+        })
+    }
+
+    fn extract_key_values(
+        &self,
+        relation: &RelationInfo,
+        key_tuple: &TupleData,
+    ) -> Result<HashMap<String, CacheValue>, String> {
+        let mut key_values = HashMap::new();
+
+        for (i, value) in key_tuple.values.iter().enumerate() {
+            if i >= relation.columns.len() {
+                break; // Key might have fewer columns
+            }
+
+            let col = &relation.columns[i];
+            if col.flags & 1 != 0 {
+                // This column is part of the key
+                let cache_value = match value {
+                    ColumnValue::Null => CacheValue::Null,
+                    ColumnValue::UnchangedToasted => CacheValue::UnchangedToasted,
+                    ColumnValue::Text(s) => CacheValue::String(s.clone()),
+                    ColumnValue::Binary(b) => CacheValue::Binary(b.clone()),
+                };
+
+                key_values.insert(col.name.clone(), cache_value);
+            }
+        }
+
+        Ok(key_values)
+    }
+
+    /// Emit a buffered transaction atomically
+    fn emit_transaction(&mut self, xid: u32, commit_lsn: u64, end_lsn: u64, timestamp: i64) -> Result<(), String> {
+        if let Some(buf) = self.txn_buffers.remove(&xid) {
+            let event = CacheInvalidationEvent {
+                xid,
+                commit_lsn,
+                end_lsn,
+                timestamp,
+                changes: buf.changes,
+            };
+
+            self.tx.send(event).ok(); // ignore if no receivers
+        }
+
+        Ok(())
+    }
+
+    /// Parse raw replication data bytes and process messages
+    /// This is the main entry point from pgwire-replication
+    pub async fn process_raw(&mut self, data: &[u8]) -> Result<(), String> {
+        let mut buf = bytes::BytesMut::from(&data[..]);
+
+        while !buf.is_empty() {
+            if let Some(msg) = self.parser.parse_message(&mut buf)? {
+                self.process_message(msg).await?;
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Get statistics about buffered transactions
+    pub fn buffered_transaction_count(&self) -> usize {
+        self.txn_buffers.len()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_emitter_creation() {
+        let (tx, _rx) = broadcast::channel(100);
+        let _emitter = CacheInvalidationEmitter::new(tx);
+    }
+
+    #[test]
+    fn test_transaction_buffering() {
+        let (tx, _rx) = broadcast::channel(100);
+        let emitter = CacheInvalidationEmitter::new(tx);
+        
+        // Changes should be buffered per transaction until commit
+        assert_eq!(emitter.buffered_transaction_count(), 0);
+    }
+}