database_replicator/replication/

monitor.rs

// ABOUTME: Replication monitoring utilities
// ABOUTME: Queries replication status and lag from source and target databases

use anyhow::{Context, Result};
use tokio_postgres::Client;

/// Replication statistics from the source database (publisher)
#[derive(Debug, Clone)]
pub struct SourceReplicationStats {
    pub application_name: String,
    pub state: String,
    pub sent_lsn: String,
    pub write_lsn: String,
    pub flush_lsn: String,
    pub replay_lsn: String,
    pub write_lag_ms: Option<i64>,
    pub flush_lag_ms: Option<i64>,
    pub replay_lag_ms: Option<i64>,
}

/// Subscription statistics from the target database (subscriber)
#[derive(Debug, Clone)]
pub struct SubscriptionStats {
    pub subscription_name: String,
    pub pid: Option<i32>,
    pub received_lsn: Option<String>,
    pub latest_end_lsn: Option<String>,
    pub state: String,
}

/// Get replication statistics from the source database
/// Queries pg_stat_replication to see what's being replicated to subscribers
pub async fn get_replication_lag(
    client: &Client,
    subscription_name: Option<&str>,
) -> Result<Vec<SourceReplicationStats>> {
    if let Some(name) = subscription_name {
        crate::utils::validate_postgres_identifier(name).context("Invalid subscription name")?;
    }

    let rows = if let Some(sub_name) = subscription_name {
        client
            .query(
                "SELECT
                    application_name,
                    state,
                    sent_lsn::text,
                    write_lsn::text,
                    flush_lsn::text,
                    replay_lsn::text,
                    EXTRACT(EPOCH FROM write_lag) * 1000 as write_lag_ms,
                    EXTRACT(EPOCH FROM flush_lag) * 1000 as flush_lag_ms,
                    EXTRACT(EPOCH FROM replay_lag) * 1000 as replay_lag_ms
                FROM pg_stat_replication
                WHERE application_name = $1",
                &[&sub_name],
            )
            .await
    } else {
        client
            .query(
                "SELECT
                    application_name,
                    state,
                    sent_lsn::text,
                    write_lsn::text,
                    flush_lsn::text,
                    replay_lsn::text,
                    EXTRACT(EPOCH FROM write_lag) * 1000 as write_lag_ms,
                    EXTRACT(EPOCH FROM flush_lag) * 1000 as flush_lag_ms,
                    EXTRACT(EPOCH FROM replay_lag) * 1000 as replay_lag_ms
                FROM pg_stat_replication",
                &[],
            )
            .await
    }
    .context("Failed to query replication statistics")?;

    let mut stats = Vec::new();
    for row in rows {
        stats.push(SourceReplicationStats {
            application_name: row.get(0),
            state: row.get(1),
            sent_lsn: row.get(2),
            write_lsn: row.get(3),
            flush_lsn: row.get(4),
            replay_lsn: row.get(5),
            write_lag_ms: row.get(6),
            flush_lag_ms: row.get(7),
            replay_lag_ms: row.get(8),
        });
    }

    Ok(stats)
}

/// Get subscription status from the target database
/// Queries pg_stat_subscription to see subscription state and progress
pub async fn get_subscription_status(
    client: &Client,
    subscription_name: Option<&str>,
) -> Result<Vec<SubscriptionStats>> {
    if let Some(name) = subscription_name {
        crate::utils::validate_postgres_identifier(name).context("Invalid subscription name")?;
    }

    let rows = if let Some(sub_name) = subscription_name {
        client
            .query(
                "SELECT
                    subname,
                    pid,
                    received_lsn::text,
                    latest_end_lsn::text,
                    srsubstate
                FROM pg_stat_subscription
                WHERE subname = $1",
                &[&sub_name],
            )
            .await
    } else {
        client
            .query(
                "SELECT
                    subname,
                    pid,
                    received_lsn::text,
                    latest_end_lsn::text,
                    srsubstate
                FROM pg_stat_subscription",
                &[],
            )
            .await
    }
    .context("Failed to query subscription statistics")?;

    let mut stats = Vec::new();
    for row in rows {
        stats.push(SubscriptionStats {
            subscription_name: row.get(0),
            pid: row.get(1),
            received_lsn: row.get(2),
            latest_end_lsn: row.get(3),
            state: row.get(4),
        });
    }

    Ok(stats)
}

/// Check if replication is caught up (no lag)
/// Returns true if all replication slots have < 1 second of replay lag
pub async fn is_replication_caught_up(
    client: &Client,
    subscription_name: Option<&str>,
) -> Result<bool> {
    let stats = get_replication_lag(client, subscription_name).await?;

    if stats.is_empty() {
        // No active replication
        return Ok(false);
    }

    for stat in stats {
        // Check if replay lag is > 1000ms (1 second)
        if let Some(lag_ms) = stat.replay_lag_ms {
            if lag_ms > 1000 {
                return Ok(false);
            }
        } else {
            // If lag is NULL, it might be too far behind or not streaming yet
            return Ok(false);
        }
    }

    Ok(true)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::postgres::connect;

    #[tokio::test]
    #[ignore]
    async fn test_get_replication_lag() {
        // This test requires a source database with active replication
        let source_url = std::env::var("TEST_SOURCE_URL").unwrap();
        let client = connect(&source_url).await.unwrap();

        let result = get_replication_lag(&client, None).await;
        match &result {
            Ok(stats) => {
                println!("✓ Replication lag query succeeded");
                println!("Found {} replication slots", stats.len());
                for stat in stats {
                    println!(
                        "  - {}: {} (replay lag: {:?}ms)",
                        stat.application_name, stat.state, stat.replay_lag_ms
                    );
                }
            }
            Err(e) => {
                println!("Error querying replication lag: {:?}", e);
                // It's okay if no replication is active
                if !e.to_string().contains("relation") && !e.to_string().contains("permission") {
                    panic!("Unexpected error: {:?}", e);
                }
            }
        }
        assert!(result.is_ok());
    }

    #[tokio::test]
    #[ignore]
    async fn test_get_subscription_status() {
        // This test requires a target database with active subscription
        let target_url = std::env::var("TEST_TARGET_URL").unwrap();
        let client = connect(&target_url).await.unwrap();

        let result = get_subscription_status(&client, None).await;
        match &result {
            Ok(stats) => {
                println!("✓ Subscription status query succeeded");
                println!("Found {} subscriptions", stats.len());
                for stat in stats {
                    println!(
                        "  - {}: state={} (pid: {:?})",
                        stat.subscription_name, stat.state, stat.pid
                    );
                }
            }
            Err(e) => {
                println!("Error querying subscription status: {:?}", e);
                // It's okay if no subscriptions exist
                if !e.to_string().contains("relation") && !e.to_string().contains("permission") {
                    panic!("Unexpected error: {:?}", e);
                }
            }
        }
        assert!(result.is_ok());
    }

    #[tokio::test]
    #[ignore]
    async fn test_is_replication_caught_up() {
        let source_url = std::env::var("TEST_SOURCE_URL").unwrap();
        let client = connect(&source_url).await.unwrap();

        let result = is_replication_caught_up(&client, None).await;
        match &result {
            Ok(caught_up) => {
                println!("✓ Caught up check succeeded: {}", caught_up);
            }
            Err(e) => {
                println!("Error checking if caught up: {:?}", e);
                // It's okay if no replication is active
                if !e.to_string().contains("relation") && !e.to_string().contains("permission") {
                    panic!("Unexpected error: {:?}", e);
                }
            }
        }
        assert!(result.is_ok());
    }

    #[tokio::test]
    #[ignore]
    async fn test_get_replication_lag_with_name() {
        let source_url = std::env::var("TEST_SOURCE_URL").unwrap();
        let client = connect(&source_url).await.unwrap();

        // Query for a specific subscription name
        let result = get_replication_lag(&client, Some("seren_migration_sub")).await;
        match &result {
            Ok(stats) => {
                println!("✓ Named replication lag query succeeded");
                println!("Found {} matching replication slots", stats.len());
            }
            Err(e) => {
                println!("Error querying named replication lag: {:?}", e);
            }
        }
        assert!(result.is_ok());
    }
}