drasi_source_postgres/

stream.rs

// Copyright 2025 The Drasi Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use anyhow::{anyhow, Result};
use chrono::Utc;
use log::{debug, error, info, trace, warn};
use std::collections::HashMap;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{oneshot, RwLock};
use tokio::time::{interval, sleep};

use super::connection::{parse_lsn, ReplicationConnection};
use super::decoder::PgOutputDecoder;
use super::protocol::BackendMessage;
use super::types::{StandbyStatusUpdate, WalMessage};
use super::{PostgresSourceConfig, ReplayState};
use drasi_core::models::{Element, ElementMetadata, ElementReference, SourceChange};
use drasi_lib::channels::{ComponentStatus, SourceEvent, SourceEventWrapper};
use drasi_lib::component_graph::ComponentStatusHandle;
use drasi_lib::sources::base::SourceBase;

pub struct ReplicationStream {
    config: PostgresSourceConfig,
    source_id: String,
    connection: Option<ReplicationConnection>,
    decoder: PgOutputDecoder,
    #[allow(dead_code)]
    status_handle: ComponentStatusHandle,
    base: SourceBase,
    replay_state: Arc<ReplayState>,
    read_lsn: u64,
    start_lsn: Option<u64>,
    last_feedback_time: std::time::Instant,
    pending_transaction: Option<Vec<(SourceChange, u64)>>,
    relations: HashMap<u32, RelationMapping>,
    table_primary_keys: Arc<RwLock<HashMap<String, Vec<String>>>>,
}

struct RelationMapping {
    #[allow(dead_code)]
    table_name: String,
    #[allow(dead_code)]
    schema_name: String,
    label: String,
}

impl ReplicationStream {
    pub(crate) fn new(
        config: PostgresSourceConfig,
        source_id: String,
        status_handle: ComponentStatusHandle,
        base: SourceBase,
        replay_state: Arc<ReplayState>,
        start_lsn: Option<u64>,
    ) -> Self {
        Self {
            config,
            source_id,
            connection: None,
            decoder: PgOutputDecoder::new(),
            status_handle,
            base,
            replay_state,
            read_lsn: 0,
            start_lsn,
            last_feedback_time: std::time::Instant::now(),
            pending_transaction: None,
            relations: HashMap::new(),
            table_primary_keys: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    // Note: table_primary_keys is initialized empty and remains so.
    // Element IDs are generated from configured table_keys (in config.table_keys),
    // or fall back to using all column values if no keys are configured.

    pub async fn run(
        &mut self,
        startup_tx: Option<oneshot::Sender<std::result::Result<(), String>>>,
    ) -> Result<()> {
        info!("Starting replication stream for source {}", self.source_id);

        // Connect and setup replication
        if let Err(error) = self.connect_and_setup().await {
            if let Some(tx) = startup_tx {
                let _ = tx.send(Err(format!("{error:#}")));
            }
            return Err(error);
        }
        if let Some(tx) = startup_tx {
            let _ = tx.send(Ok(()));
        }

        // Start streaming loop
        let mut keepalive_interval = interval(Duration::from_secs(10));

        loop {
            // Check if we should stop
            {
                let status = self.status_handle.get_status().await;
                if status == ComponentStatus::Stopping || status == ComponentStatus::Stopped {
                    info!("Received stop signal, shutting down replication");
                    break;
                }
            }

            tokio::select! {
                // Check for replication messages
                result = self.read_next_message() => {
                    match result {
                        Ok(Some(msg)) => {
                            if let Err(e) = self.handle_message(msg).await {
                                error!("Error handling message: {e}");
                                // Attempt recovery
                                if let Err(e) = self.recover_connection().await {
                                    error!("Failed to recover connection: {e}");
                                    return Err(e);
                                }
                            }
                        }
                        Ok(None) => {
                            // No message available
                        }
                        Err(e) => {
                            error!("Error reading message: {e}");
                            // Attempt recovery
                            if let Err(e) = self.recover_connection().await {
                                error!("Failed to recover connection: {e}");
                                return Err(e);
                            }
                        }
                    }
                }

                // Send periodic keepalives
                _ = keepalive_interval.tick() => {
                    if let Err(e) = self.send_feedback(false).await {
                        warn!("Failed to send keepalive: {e}");
                    }
                }
            }
        }

        // Clean shutdown
        self.shutdown().await?;
        Ok(())
    }

    async fn connect_and_setup(&mut self) -> Result<()> {
        info!("Connecting to PostgreSQL for replication");

        // Create connection
        let mut conn = ReplicationConnection::connect(
            &self.config.host,
            self.config.port,
            &self.config.database,
            &self.config.user,
            &self.config.password,
        )
        .await?;

        // Identify system
        let system_info = conn.identify_system().await?;
        info!("Connected to PostgreSQL system: {system_info:?}");

        // Create or verify replication slot
        let slot_info = conn
            .create_replication_slot(&self.config.slot_name, false)
            .await?;
        info!("Using replication slot: {slot_info:?}");

        // Parse starting LSN — use explicit start_lsn if provided (replay),
        // otherwise use the slot's consistent_point.
        let slot_lsn =
            if !slot_info.consistent_point.is_empty() && slot_info.consistent_point != "0/0" {
                parse_lsn(&slot_info.consistent_point)?
            } else {
                0
            };
        self.read_lsn = self.start_lsn.unwrap_or(slot_lsn);
        self.replay_state
            .read_lsn
            .store(self.read_lsn, Ordering::Release);

        // Build replication options
        let mut options = HashMap::new();
        options.insert("proto_version".to_string(), "1".to_string());
        options.insert(
            "publication_names".to_string(),
            self.config.publication_name.clone(),
        );

        // Start replication
        conn.start_replication(&self.config.slot_name, Some(self.read_lsn), options)
            .await?;

        self.connection = Some(conn);
        info!(
            "Replication started from read LSN {:x} (slot watermark {:x})",
            self.read_lsn, slot_lsn
        );

        Ok(())
    }

    async fn read_next_message(&mut self) -> Result<Option<BackendMessage>> {
        if let Some(conn) = &mut self.connection {
            // Try to read with a short timeout to avoid blocking forever
            match tokio::time::timeout(Duration::from_millis(100), conn.read_replication_message())
                .await
            {
                Ok(Ok(msg)) => Ok(Some(msg)),
                Ok(Err(e)) => Err(e),
                Err(_) => Ok(None), // Timeout, no message available
            }
        } else {
            Err(anyhow!("No connection available"))
        }
    }

    async fn handle_message(&mut self, msg: BackendMessage) -> Result<()> {
        match msg {
            BackendMessage::CopyData(data) => {
                self.handle_copy_data(&data).await?;
            }
            BackendMessage::PrimaryKeepaliveMessage {
                wal_end,
                timestamp: _,
                reply,
            } => {
                self.read_lsn = wal_end;
                self.replay_state
                    .read_lsn
                    .store(self.read_lsn, Ordering::Release);
                if reply == 1 {
                    self.send_feedback(true).await?;
                }
            }
            BackendMessage::ErrorResponse(err) => {
                error!("Server error: {}", err.message);
                return Err(anyhow!("Server error: {}", err.message));
            }
            _ => {
                trace!("Ignoring message: {msg:?}");
            }
        }
        Ok(())
    }

    async fn handle_copy_data(&mut self, data: &[u8]) -> Result<()> {
        if data.is_empty() {
            return Ok(());
        }

        // First byte indicates the message type
        let msg_type = data[0];

        match msg_type {
            b'w' => {
                // XLogData message
                self.handle_xlog_data(&data[1..]).await?;
            }
            b'k' => {
                // Primary keepalive
                self.handle_keepalive(&data[1..]).await?;
            }
            _ => {
                warn!("Unknown copy data message type: 0x{msg_type:02x}");
            }
        }

        Ok(())
    }

    async fn handle_xlog_data(&mut self, data: &[u8]) -> Result<()> {
        if data.len() < 24 {
            return Err(anyhow!("XLogData message too short: {} bytes", data.len()));
        }

        // Parse XLogData header
        let _start_lsn = u64::from_be_bytes([
            data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
        ]);
        let end_lsn = u64::from_be_bytes([
            data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15],
        ]);
        let _timestamp = i64::from_be_bytes([
            data[16], data[17], data[18], data[19], data[20], data[21], data[22], data[23],
        ]);

        // Update read LSN
        self.read_lsn = end_lsn;
        self.replay_state
            .read_lsn
            .store(self.read_lsn, Ordering::Release);

        // Decode the actual WAL message
        let wal_data = &data[24..];

        // Try to decode, but don't fail on partial messages
        if !wal_data.is_empty() {
            let msg_type = wal_data[0];
            debug!(
                "Attempting to decode WAL message type: {} ({}), data length: {}",
                msg_type as char,
                msg_type,
                wal_data.len()
            );
        }

        match self.decoder.decode_message(wal_data) {
            Ok(Some(wal_msg)) => {
                self.process_wal_message(wal_msg).await?;
            }
            Ok(None) => {
                // No message or skipped message type
            }
            Err(e) => {
                // Log but don't fail on decode errors - might be partial message
                if !wal_data.is_empty() {
                    debug!(
                        "Failed to decode WAL message type {} ({}): {}, data length: {}",
                        wal_data[0] as char,
                        wal_data[0],
                        e,
                        wal_data.len()
                    );
                }
                // We'll get the rest in the next message
            }
        }

        // Send feedback periodically
        if self.last_feedback_time.elapsed() > Duration::from_secs(5) {
            self.send_feedback(false).await?;
        }

        Ok(())
    }

    async fn handle_keepalive(&mut self, data: &[u8]) -> Result<()> {
        if data.len() < 17 {
            return Err(anyhow!("Keepalive message too short"));
        }

        let wal_end = u64::from_be_bytes([
            data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
        ]);
        let reply = data[16];

        self.read_lsn = wal_end;
        self.replay_state
            .read_lsn
            .store(self.read_lsn, Ordering::Release);

        if reply == 1 {
            self.send_feedback(true).await?;
        }

        Ok(())
    }

    async fn process_wal_message(&mut self, msg: WalMessage) -> Result<()> {
        match msg {
            WalMessage::Begin(_) => {
                // Start a new transaction — track changes with their LSN
                self.pending_transaction = Some(Vec::new());
            }
            WalMessage::Commit(tx_info) => {
                // Commit the transaction — stamp all changes with the commit LSN
                // for atomic checkpoint semantics
                if let Some(changes) = self.pending_transaction.take() {
                    for (change, _) in changes {
                        self.dispatch_change(change, tx_info.commit_lsn).await;
                    }
                    debug!(
                        "Committed transaction {} with LSN {:x}",
                        tx_info.xid, tx_info.commit_lsn
                    );
                }
            }
            WalMessage::Relation(relation) => {
                // Store relation mapping - use table name as-is for label (no uppercase)
                // This ensures consistency with bootstrap which uses the actual table name case
                let label = relation.name.clone();
                self.relations.insert(
                    relation.id,
                    RelationMapping {
                        table_name: relation.name.clone(),
                        schema_name: relation.namespace.clone(),
                        label,
                    },
                );

                // Update decoder's relation info
                // The decoder already handles this internally
            }
            WalMessage::Insert { relation_id, tuple } => {
                if let Some(change) = self.convert_insert(relation_id, tuple).await? {
                    if let Some(tx) = &mut self.pending_transaction {
                        tx.push((change, self.read_lsn));
                    } else {
                        self.dispatch_change(change, self.read_lsn).await;
                    }
                }
            }
            WalMessage::Update {
                relation_id,
                old_tuple,
                new_tuple,
            } => {
                if let Some(change) = self
                    .convert_update(relation_id, old_tuple, new_tuple)
                    .await?
                {
                    if let Some(tx) = &mut self.pending_transaction {
                        tx.push((change, self.read_lsn));
                    } else {
                        self.dispatch_change(change, self.read_lsn).await;
                    }
                }
            }
            WalMessage::Delete {
                relation_id,
                old_tuple,
            } => {
                if let Some(change) = self.convert_delete(relation_id, old_tuple).await? {
                    if let Some(tx) = &mut self.pending_transaction {
                        tx.push((change, self.read_lsn));
                    } else {
                        self.dispatch_change(change, self.read_lsn).await;
                    }
                }
            }
            WalMessage::Truncate { relation_ids } => {
                warn!("Truncate not yet implemented for relations: {relation_ids:?}");
            }
        }
        Ok(())
    }

    async fn convert_insert(
        &self,
        relation_id: u32,
        tuple: Vec<super::types::PostgresValue>,
    ) -> Result<Option<SourceChange>> {
        // Get relation info
        let relation = self
            .decoder
            .get_relation(relation_id)
            .ok_or_else(|| anyhow!("Unknown relation {relation_id}"))?;

        let mapping = self
            .relations
            .get(&relation_id)
            .ok_or_else(|| anyhow!("No mapping for relation {relation_id}"))?;

        // Convert tuple to properties
        let mut properties = drasi_core::models::ElementPropertyMap::new();
        for (i, value) in tuple.iter().enumerate() {
            if let Some(column) = relation.columns.get(i) {
                let json_value = value.to_json();
                if !json_value.is_null() {
                    properties.insert(
                        &column.name,
                        drasi_lib::sources::manager::convert_json_to_element_value(&json_value),
                    );
                }
            }
        }

        // Create element ID (use primary key if available, otherwise generate)
        let element_id = self.generate_element_id(relation, &tuple).await?;

        // Create the element
        let element = Element::Node {
            metadata: ElementMetadata {
                reference: ElementReference::new(&self.source_id, &element_id),
                labels: Arc::from([Arc::from(mapping.label.as_str())]),
                effective_from: Utc::now().timestamp_millis() as u64,
            },
            properties,
        };

        Ok(Some(SourceChange::Insert { element }))
    }

    async fn convert_update(
        &self,
        relation_id: u32,
        old_tuple: Option<Vec<super::types::PostgresValue>>,
        new_tuple: Vec<super::types::PostgresValue>,
    ) -> Result<Option<SourceChange>> {
        let relation = self
            .decoder
            .get_relation(relation_id)
            .ok_or_else(|| anyhow!("Unknown relation {relation_id}"))?;

        let mapping = self
            .relations
            .get(&relation_id)
            .ok_or_else(|| anyhow!("No mapping for relation {relation_id}"))?;

        // Generate element ID (should be the same for both old and new tuples)
        let element_id = self.generate_element_id(relation, &new_tuple).await?;

        if old_tuple.is_none() {
            warn!("UPDATE without old tuple for relation {relation_id}, preserving UPDATE");
        }

        // Create properties for after state
        // Note: We allow UPDATE without old_tuple to avoid converting to INSERT.
        let mut after_properties = drasi_core::models::ElementPropertyMap::new();

        // Process new tuple (after state)
        for (i, column) in relation.columns.iter().enumerate() {
            if let Some(value) = new_tuple.get(i) {
                let json_value = value.to_json();
                if !json_value.is_null() {
                    after_properties.insert(
                        &column.name,
                        drasi_lib::sources::manager::convert_json_to_element_value(&json_value),
                    );
                }
            }
        }

        let after_element = Element::Node {
            metadata: ElementMetadata {
                reference: ElementReference::new(&self.source_id, &element_id),
                labels: Arc::from([Arc::from(mapping.label.as_str())]),
                effective_from: Utc::now().timestamp_millis() as u64,
            },
            properties: after_properties,
        };

        Ok(Some(SourceChange::Update {
            element: after_element,
        }))
    }

    async fn convert_delete(
        &self,
        relation_id: u32,
        old_tuple: Vec<super::types::PostgresValue>,
    ) -> Result<Option<SourceChange>> {
        let relation = self
            .decoder
            .get_relation(relation_id)
            .ok_or_else(|| anyhow!("Unknown relation {relation_id}"))?;

        let mapping = self
            .relations
            .get(&relation_id)
            .ok_or_else(|| anyhow!("No mapping for relation {relation_id}"))?;

        let element_id = self.generate_element_id(relation, &old_tuple).await?;

        Ok(Some(SourceChange::Delete {
            metadata: ElementMetadata {
                reference: ElementReference::new(&self.source_id, &element_id),
                labels: Arc::from([Arc::from(mapping.label.as_str())]),
                effective_from: Utc::now().timestamp_millis() as u64,
            },
        }))
    }

    /// Generate a stable element ID for a tuple based on primary key values.
    ///
    /// Priority order:
    /// 1. User-configured key columns (from table_keys config)
    /// 2. Automatically detected primary keys from PostgreSQL system catalogs
    /// 3. UUID fallback if no keys are available
    ///
    /// Element ID format:
    /// - Single key: Table name prefix with value (e.g., "stocks:AAPL")
    /// - Composite key: Table name prefix with values joined (e.g., "portfolio:tenant1_user2")
    /// - No key: Table name prefix with UUID (e.g., "orders:550e8400-e29b-41d4-a716-446655440000")
    async fn generate_element_id(
        &self,
        relation: &super::types::RelationInfo,
        tuple: &[super::types::PostgresValue],
    ) -> Result<String> {
        // Get the table name (use schema-qualified if not in public schema)
        let table_name = if relation.namespace == "public" {
            relation.name.clone()
        } else {
            format!("{}.{}", relation.namespace, relation.name)
        };

        // Get primary key columns for this table
        let primary_keys = self.table_primary_keys.read().await;
        let pk_columns = primary_keys.get(&table_name);

        // Check configured table_keys first
        let configured_keys = self
            .config
            .table_keys
            .iter()
            .find(|tk| tk.table == table_name)
            .map(|tk| &tk.key_columns);

        // Use configured keys if available, otherwise use detected primary keys
        let key_columns = configured_keys.or(pk_columns);

        if let Some(keys) = key_columns {
            let mut key_parts = Vec::new();

            for (i, column) in relation.columns.iter().enumerate() {
                if keys.contains(&column.name) {
                    if let Some(value) = tuple.get(i) {
                        let json_val = value.to_json();
                        if !json_val.is_null() {
                            // Remove quotes from JSON string representation
                            let val_str = json_val.to_string();
                            let cleaned = val_str.trim_matches('"');
                            key_parts.push(cleaned.to_string());
                        }
                    }
                }
            }

            if !key_parts.is_empty() {
                // Include table name as prefix to ensure uniqueness across tables
                return Ok(format!("{}:{}", table_name, key_parts.join("_")));
            }
        }

        // No primary key found or all key values are NULL
        warn!("No primary key value found for table '{table_name}'. Consider adding 'table_keys' configuration.");
        // Still include table name prefix for consistency
        Ok(format!("{}:{}", table_name, uuid::Uuid::new_v4()))
    }

    async fn send_feedback(&mut self, reply_requested: bool) -> Result<()> {
        if let Some(conn) = &mut self.connection {
            // Use the confirmed source position (min LSN across all query
            // position handles) for flush_lsn / apply_lsn.  This tells
            // Postgres that WAL up to this LSN has been durably processed
            // by all subscribers.  Using read_lsn here would advance the
            // slot watermark past un-checkpointed query positions, causing
            // PositionUnavailable on crash+restart.
            let confirmed_lsn = match self.base.compute_confirmed_source_position().await {
                Some(bytes) if bytes.len() == 8 => {
                    let arr: [u8; 8] = bytes[..8].try_into().expect("length already checked");
                    u64::from_be_bytes(arr)
                }
                Some(bytes) => {
                    warn!(
                        "[{}] Confirmed source position has unexpected length {} (expected 8); \
                             not advancing flush_lsn",
                        self.source_id,
                        bytes.len()
                    );
                    0
                }
                None => 0, // No confirmed position yet — don't advance
            };

            // Apply the flush fence: during the subscription window after a
            // replay restart, cap flush_lsn to prevent the slot's restart_lsn
            // from advancing past positions that pending subscribers still need.
            let fence = self.replay_state.effective_flush_fence();
            let (effective_lsn, was_clamped) = if fence < u64::MAX && confirmed_lsn > fence {
                (fence, true)
            } else {
                (confirmed_lsn, false)
            };

            let status = StandbyStatusUpdate {
                write_lsn: self.read_lsn,
                flush_lsn: effective_lsn,
                apply_lsn: effective_lsn,
                reply_requested,
            };

            conn.send_standby_status(status).await?;
            self.last_feedback_time = std::time::Instant::now();

            // Prune the sequence→position map up to the confirmed sequence
            // now that feedback was successfully sent. Skip pruning when the
            // fence clamped flush_lsn — Postgres hasn't actually acknowledged
            // the full confirmed position, so we must retain map entries for
            // accurate re-computation after the fence lifts.
            if !was_clamped && effective_lsn > 0 {
                if let Some(confirmed_seq) = self.base.compute_confirmed_position().await {
                    self.base.prune_position_map(confirmed_seq).await;
                }
            }

            trace!(
                "[{}] Sent feedback: write_lsn={:x}, flush_lsn={:x}{}",
                self.source_id,
                self.read_lsn,
                effective_lsn,
                if was_clamped { " (fenced)" } else { "" }
            );
        }

        Ok(())
    }

    /// Dispatch a single change event through the framework's `dispatch_event()`.
    ///
    /// Sets `source_position` to the 8-byte big-endian WAL LSN so the checkpoint
    /// framework can persist the position for recovery/replay.
    async fn dispatch_change(&self, change: SourceChange, lsn: u64) {
        let mut profiling = drasi_lib::profiling::ProfilingMetadata::new();
        profiling.source_send_ns = Some(drasi_lib::profiling::timestamp_ns());

        let mut wrapper = SourceEventWrapper::with_profiling(
            self.source_id.clone(),
            SourceEvent::Change(change),
            chrono::Utc::now(),
            profiling,
        );

        // Attach the WAL LSN as opaque source_position bytes for checkpoint/recovery
        wrapper.set_source_position(super::connection::lsn_to_position_bytes(lsn));

        // Use dispatch_event() which stamps the monotonic sequence
        if let Err(e) = self.base.dispatch_event(wrapper).await {
            debug!(
                "[{}] Failed to dispatch change (no subscribers): {}",
                self.source_id, e
            );
        }
    }

    #[allow(dead_code)]
    async fn check_stop_signal(&self) -> bool {
        let status = self.status_handle.get_status().await;
        status == ComponentStatus::Stopping || status == ComponentStatus::Stopped
    }

    async fn recover_connection(&mut self) -> Result<()> {
        warn!("Attempting to recover connection");

        // Close existing connection if any
        if let Some(conn) = self.connection.take() {
            let _ = conn.close().await;
        }

        // Wait a bit before reconnecting
        sleep(Duration::from_secs(5)).await;

        // Try to reconnect
        self.connect_and_setup().await?;

        info!("Connection recovered successfully");
        Ok(())
    }

    async fn shutdown(&mut self) -> Result<()> {
        info!("Shutting down replication stream");

        // Send final feedback
        let _ = self.send_feedback(false).await;

        // Close connection
        if let Some(conn) = self.connection.take() {
            conn.close().await?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use chrono::Utc;
    use drasi_core::models::validate_effective_from;

    /// Validates that the timestamp pattern used in convert_insert/convert_update/convert_delete
    /// produces a value in the millisecond range.
    #[test]
    fn effective_from_uses_milliseconds() {
        let effective_from = Utc::now().timestamp_millis() as u64;
        assert!(
            validate_effective_from(effective_from).is_ok(),
            "Postgres CDC effective_from ({effective_from}) should be in millisecond range"
        );
    }

    /// Verifies that using nanoseconds would be caught by the validator.
    #[test]
    fn effective_from_rejects_nanoseconds_pattern() {
        let bad_effective_from = Utc::now().timestamp_nanos_opt().unwrap() as u64;
        assert!(
            validate_effective_from(bad_effective_from).is_err(),
            "Nanosecond timestamp ({bad_effective_from}) should be rejected"
        );
    }
}