database_replicator/migration/

dump.rs

// ABOUTME: Wrapper for pg_dump command to export database objects
// ABOUTME: Handles global objects, schema, and data export

use crate::filters::ReplicationFilter;
use anyhow::{Context, Result};
use std::collections::BTreeSet;
use std::fs;
use std::process::{Command, Stdio};
use std::time::Duration;

/// Dump global objects (roles, tablespaces) using pg_dumpall
pub async fn dump_globals(source_url: &str, output_path: &str) -> Result<()> {
    tracing::info!("Dumping global objects to {}", output_path);

    // Parse URL and create .pgpass file for secure authentication
    let parts = crate::utils::parse_postgres_url(source_url)
        .with_context(|| format!("Failed to parse source URL: {}", source_url))?;
    let pgpass = crate::utils::PgPassFile::new(&parts)
        .context("Failed to create .pgpass file for authentication")?;

    let env_vars = parts.to_pg_env_vars();
    let output_path_owned = output_path.to_string();

    // Wrap subprocess execution with retry logic
    crate::utils::retry_subprocess_with_backoff(
        || {
            let mut cmd = Command::new("pg_dumpall");
            cmd.arg("--globals-only")
                .arg("--no-role-passwords") // Don't dump passwords
                .arg("--verbose") // Show progress
                .arg("--host")
                .arg(&parts.host)
                .arg("--port")
                .arg(parts.port.to_string())
                .arg("--database")
                .arg(&parts.database)
                .arg(format!("--file={}", output_path_owned))
                .env("PGPASSFILE", pgpass.path())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit());

            // Add username if specified
            if let Some(user) = &parts.user {
                cmd.arg("--username").arg(user);
            }

            // Apply query parameters as environment variables (SSL, channel_binding, etc.)
            for (env_var, value) in &env_vars {
                cmd.env(env_var, value);
            }

            // Apply TCP keepalive parameters to prevent idle connection timeouts
            for (env_var, value) in crate::utils::get_keepalive_env_vars() {
                cmd.env(env_var, value);
            }

            cmd.status().context(
                "Failed to execute pg_dumpall. Is PostgreSQL client installed?\n\
                 Install with:\n\
                 - Ubuntu/Debian: sudo apt-get install postgresql-client\n\
                 - macOS: brew install postgresql\n\
                 - RHEL/CentOS: sudo yum install postgresql",
            )
        },
        3,                      // Max 3 retries
        Duration::from_secs(1), // Start with 1 second delay
        "pg_dumpall (dump globals)",
    )
    .await
    .context(
        "pg_dumpall failed to dump global objects.\n\
         \n\
         Common causes:\n\
         - Connection authentication failed\n\
         - User lacks sufficient privileges (need SUPERUSER or pg_read_all_settings role)\n\
         - Network connectivity issues\n\
         - Invalid connection string\n\
         - Connection timeout or network issues",
    )?;

    tracing::info!("✓ Global objects dumped successfully");
    Ok(())
}

/// Update a globals dump so duplicate role creation errors become harmless notices.
///
/// `pg_dumpall --globals-only` emits `CREATE ROLE` statements that fail if the
/// role already exists on the target cluster. When an operator reruns
/// replication against the same target, those statements cause `psql` to exit
/// with status 3 which previously triggered noisy retries and prevented the rest
/// of the globals from being applied. By wrapping each `CREATE ROLE` statement
/// in a `DO $$ ... EXCEPTION WHEN duplicate_object` block, we allow Postgres to
/// skip recreating existing roles while still applying subsequent `ALTER ROLE`
/// and `GRANT` statements.
pub fn sanitize_globals_dump(path: &str) -> Result<()> {
    let content = fs::read_to_string(path)
        .with_context(|| format!("Failed to read globals dump at {}", path))?;

    if let Some(updated) = rewrite_create_role_statements(&content) {
        fs::write(path, updated)
            .with_context(|| format!("Failed to update globals dump at {}", path))?;
    }

    Ok(())
}

fn rewrite_create_role_statements(sql: &str) -> Option<String> {
    if sql.is_empty() {
        return None;
    }

    let mut output = String::with_capacity(sql.len() + 1024);
    let mut modified = false;
    let mut cursor = 0;

    while cursor < sql.len() {
        if let Some(rel_pos) = sql[cursor..].find('\n') {
            let end = cursor + rel_pos + 1;
            let chunk = &sql[cursor..end];
            if let Some(transformed) = wrap_create_role_line(chunk) {
                output.push_str(&transformed);
                modified = true;
            } else {
                output.push_str(chunk);
            }
            cursor = end;
        } else {
            let chunk = &sql[cursor..];
            if let Some(transformed) = wrap_create_role_line(chunk) {
                output.push_str(&transformed);
                modified = true;
            } else {
                output.push_str(chunk);
            }
            break;
        }
    }

    if modified {
        Some(output)
    } else {
        None
    }
}

fn wrap_create_role_line(chunk: &str) -> Option<String> {
    let trimmed = chunk.trim_start();
    if !trimmed.starts_with("CREATE ROLE ") {
        return None;
    }

    let statement = trimmed.trim_end();
    let statement_body = statement.trim_end_matches(';').trim_end();
    let leading_ws_len = chunk.len() - trimmed.len();
    let leading_ws = &chunk[..leading_ws_len];
    let newline = if chunk.ends_with("\r\n") {
        "\r\n"
    } else if chunk.ends_with('\n') {
        "\n"
    } else {
        ""
    };

    let role_token = extract_role_token(statement_body)?;

    let notice_name = escape_single_quotes(&unquote_role_name(&role_token));

    let mut block = String::with_capacity(chunk.len() + 128);
    block.push_str(leading_ws);
    block.push_str("DO $$\n");
    block.push_str(leading_ws);
    block.push_str("BEGIN\n");
    block.push_str(leading_ws);
    block.push_str("    ");
    block.push_str(statement_body);
    block.push_str(";\n");
    block.push_str(leading_ws);
    block.push_str("EXCEPTION\n");
    block.push_str(leading_ws);
    block.push_str("    WHEN duplicate_object THEN\n");
    block.push_str(leading_ws);
    block.push_str("        RAISE NOTICE 'Role ");
    block.push_str(&notice_name);
    block.push_str(" already exists on target, skipping CREATE ROLE';\n");
    block.push_str(leading_ws);
    block.push_str("END $$;");

    if !newline.is_empty() {
        block.push_str(newline);
    }

    Some(block)
}

fn extract_role_token(statement: &str) -> Option<String> {
    let remainder = statement.strip_prefix("CREATE ROLE")?.trim_start();

    if remainder.starts_with('"') {
        let mut idx = 1;
        let bytes = remainder.as_bytes();
        while idx < bytes.len() {
            if bytes[idx] == b'"' {
                if idx + 1 < bytes.len() && bytes[idx + 1] == b'"' {
                    idx += 2;
                    continue;
                } else {
                    idx += 1;
                    break;
                }
            }
            idx += 1;
        }
        if idx <= remainder.len() {
            return Some(remainder[..idx].to_string());
        }
        None
    } else {
        let mut end = remainder.len();
        for (i, ch) in remainder.char_indices() {
            if ch.is_whitespace() || ch == ';' {
                end = i;
                break;
            }
        }
        if end == 0 {
            None
        } else {
            Some(remainder[..end].to_string())
        }
    }
}

fn unquote_role_name(token: &str) -> String {
    if token.starts_with('"') && token.ends_with('"') && token.len() >= 2 {
        let inner = &token[1..token.len() - 1];
        inner.replace("\"\"", "\"")
    } else {
        token.to_string()
    }
}

fn escape_single_quotes(value: &str) -> String {
    value.replace('\'', "''")
}

/// Dump schema (DDL) for a specific database
pub async fn dump_schema(
    source_url: &str,
    database: &str,
    output_path: &str,
    filter: &ReplicationFilter,
) -> Result<()> {
    tracing::info!(
        "Dumping schema for database '{}' to {}",
        database,
        output_path
    );

    // Parse URL and create .pgpass file for secure authentication
    let parts = crate::utils::parse_postgres_url(source_url)
        .with_context(|| format!("Failed to parse source URL: {}", source_url))?;
    let pgpass = crate::utils::PgPassFile::new(&parts)
        .context("Failed to create .pgpass file for authentication")?;

    let env_vars = parts.to_pg_env_vars();
    let output_path_owned = output_path.to_string();

    // Collect filter options
    let exclude_tables = get_schema_excluded_tables_for_db(filter, database);
    let include_tables = get_included_tables_for_db(filter, database);

    // Wrap subprocess execution with retry logic
    crate::utils::retry_subprocess_with_backoff(
        || {
            let mut cmd = Command::new("pg_dump");
            cmd.arg("--schema-only")
                .arg("--no-owner") // Don't include ownership commands
                .arg("--no-privileges") // We'll handle privileges separately
                .arg("--verbose"); // Show progress

            // Add table filtering if specified
            // Only exclude explicit exclude_tables from schema dump (NOT schema_only or predicate tables)
            if let Some(ref exclude) = exclude_tables {
                if !exclude.is_empty() {
                    for table in exclude {
                        cmd.arg("--exclude-table").arg(table);
                    }
                }
            }

            // If include_tables is specified, only dump those tables
            if let Some(ref include) = include_tables {
                if !include.is_empty() {
                    for table in include {
                        cmd.arg("--table").arg(table);
                    }
                }
            }

            cmd.arg("--host")
                .arg(&parts.host)
                .arg("--port")
                .arg(parts.port.to_string())
                .arg("--dbname")
                .arg(&parts.database)
                .arg(format!("--file={}", output_path_owned))
                .env("PGPASSFILE", pgpass.path())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit());

            // Add username if specified
            if let Some(user) = &parts.user {
                cmd.arg("--username").arg(user);
            }

            // Apply query parameters as environment variables (SSL, channel_binding, etc.)
            for (env_var, value) in &env_vars {
                cmd.env(env_var, value);
            }

            // Apply TCP keepalive parameters to prevent idle connection timeouts
            for (env_var, value) in crate::utils::get_keepalive_env_vars() {
                cmd.env(env_var, value);
            }

            cmd.status().context(
                "Failed to execute pg_dump. Is PostgreSQL client installed?\n\
                 Install with:\n\
                 - Ubuntu/Debian: sudo apt-get install postgresql-client\n\
                 - macOS: brew install postgresql\n\
                 - RHEL/CentOS: sudo yum install postgresql",
            )
        },
        3,                      // Max 3 retries
        Duration::from_secs(1), // Start with 1 second delay
        "pg_dump (dump schema)",
    )
    .await
    .with_context(|| {
        format!(
            "pg_dump failed to dump schema for database '{}'.\n\
             \n\
             Common causes:\n\
             - Database does not exist\n\
             - Connection authentication failed\n\
             - User lacks privileges to read database schema\n\
             - Network connectivity issues\n\
             - Connection timeout or network issues",
            database
        )
    })?;

    tracing::info!("✓ Schema dumped successfully");
    Ok(())
}

/// Dump data for a specific database using optimized directory format
///
/// Uses PostgreSQL directory format dump with:
/// - Parallel dumps for faster performance
/// - Maximum compression (level 9)
/// - Large object (blob) support
/// - Directory output for efficient parallel restore
///
/// The number of parallel jobs is automatically determined based on available CPU cores.
pub async fn dump_data(
    source_url: &str,
    database: &str,
    output_path: &str,
    filter: &ReplicationFilter,
) -> Result<()> {
    // Determine optimal number of parallel jobs (number of CPUs, capped at 8)
    let num_cpus = std::thread::available_parallelism()
        .map(|n| n.get().min(8))
        .unwrap_or(4);

    tracing::info!(
        "Dumping data for database '{}' to {} (parallel={}, compression=9, format=directory)",
        database,
        output_path,
        num_cpus
    );

    // Parse URL and create .pgpass file for secure authentication
    let parts = crate::utils::parse_postgres_url(source_url)
        .with_context(|| format!("Failed to parse source URL: {}", source_url))?;
    let pgpass = crate::utils::PgPassFile::new(&parts)
        .context("Failed to create .pgpass file for authentication")?;

    let env_vars = parts.to_pg_env_vars();
    let output_path_owned = output_path.to_string();

    // Collect filter options
    let exclude_tables = get_data_excluded_tables_for_db(filter, database);
    let include_tables = get_included_tables_for_db(filter, database);

    // Wrap subprocess execution with retry logic
    crate::utils::retry_subprocess_with_backoff(
        || {
            let mut cmd = Command::new("pg_dump");
            cmd.arg("--data-only")
                .arg("--no-owner")
                .arg("--format=directory") // Directory format enables parallel operations
                .arg("--blobs") // Include large objects (blobs)
                .arg("--compress=9") // Maximum compression for smaller dump size
                .arg(format!("--jobs={}", num_cpus)) // Parallel dump jobs
                .arg("--verbose"); // Show progress

            // Add table filtering if specified
            // Exclude explicit excludes, schema_only tables, and predicate tables from data dump
            if let Some(ref exclude) = exclude_tables {
                if !exclude.is_empty() {
                    for table in exclude {
                        cmd.arg("--exclude-table-data").arg(table);
                    }
                }
            }

            // If include_tables is specified, only dump data for those tables
            if let Some(ref include) = include_tables {
                if !include.is_empty() {
                    for table in include {
                        cmd.arg("--table").arg(table);
                    }
                }
            }

            cmd.arg("--host")
                .arg(&parts.host)
                .arg("--port")
                .arg(parts.port.to_string())
                .arg("--dbname")
                .arg(&parts.database)
                .arg(format!("--file={}", output_path_owned))
                .env("PGPASSFILE", pgpass.path())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit());

            // Add username if specified
            if let Some(user) = &parts.user {
                cmd.arg("--username").arg(user);
            }

            // Apply query parameters as environment variables (SSL, channel_binding, etc.)
            for (env_var, value) in &env_vars {
                cmd.env(env_var, value);
            }

            // Apply TCP keepalive parameters to prevent idle connection timeouts
            for (env_var, value) in crate::utils::get_keepalive_env_vars() {
                cmd.env(env_var, value);
            }

            cmd.status().context(
                "Failed to execute pg_dump. Is PostgreSQL client installed?\n\
                 Install with:\n\
                 - Ubuntu/Debian: sudo apt-get install postgresql-client\n\
                 - macOS: brew install postgresql\n\
                 - RHEL/CentOS: sudo yum install postgresql",
            )
        },
        3,                      // Max 3 retries
        Duration::from_secs(1), // Start with 1 second delay
        "pg_dump (dump data)",
    )
    .await
    .with_context(|| {
        format!(
            "pg_dump failed to dump data for database '{}'.\n\
             \n\
             Common causes:\n\
             - Database does not exist\n\
             - Connection authentication failed\n\
             - User lacks privileges to read table data\n\
             - Network connectivity issues\n\
             - Insufficient disk space for dump directory\n\
             - Output directory already exists (pg_dump requires non-existent path)\n\
             - Connection timeout or network issues",
            database
        )
    })?;

    tracing::info!(
        "✓ Data dumped successfully using {} parallel jobs",
        num_cpus
    );
    Ok(())
}

/// Extract table names to exclude from SCHEMA dumps (--exclude-table flag)
/// Only excludes explicit exclude_tables - NOT schema_only or predicate tables
/// (those need their schema created, just not bulk data copied)
/// Returns schema-qualified names in format: "schema"."table"
fn get_schema_excluded_tables_for_db(
    filter: &ReplicationFilter,
    db_name: &str,
) -> Option<Vec<String>> {
    let mut tables = BTreeSet::new();

    // Handle explicit exclude_tables (format: "database.table")
    // These tables are completely excluded (no schema, no data)
    if let Some(explicit) = filter.exclude_tables() {
        for full_name in explicit {
            let parts: Vec<&str> = full_name.split('.').collect();
            if parts.len() == 2 && parts[0] == db_name {
                // Format as "public"."table" for consistency
                tables.insert(format!("\"public\".\"{}\"", parts[1]));
            }
        }
    }

    if tables.is_empty() {
        None
    } else {
        Some(tables.into_iter().collect())
    }
}

/// Extract table names to exclude from DATA dumps (--exclude-table-data flag)
/// Excludes explicit excludes, schema_only tables, and predicate tables
/// (predicate tables will be copied separately with filtering)
/// Returns schema-qualified names in format: "schema"."table"
fn get_data_excluded_tables_for_db(
    filter: &ReplicationFilter,
    db_name: &str,
) -> Option<Vec<String>> {
    let mut tables = BTreeSet::new();

    // Handle explicit exclude_tables (format: "database.table")
    // Default to public schema for backward compatibility
    if let Some(explicit) = filter.exclude_tables() {
        for full_name in explicit {
            let parts: Vec<&str> = full_name.split('.').collect();
            if parts.len() == 2 && parts[0] == db_name {
                // Format as "public"."table" for consistency
                tables.insert(format!("\"public\".\"{}\"", parts[1]));
            }
        }
    }

    // schema_only_tables and predicate_tables already return schema-qualified names
    for table in filter.schema_only_tables(db_name) {
        tables.insert(table);
    }

    for (table, _) in filter.predicate_tables(db_name) {
        tables.insert(table);
    }

    if tables.is_empty() {
        None
    } else {
        Some(tables.into_iter().collect())
    }
}

/// Extract table names for a specific database from include_tables filter
/// Returns schema-qualified names in format: "schema"."table"
fn get_included_tables_for_db(filter: &ReplicationFilter, db_name: &str) -> Option<Vec<String>> {
    filter.include_tables().map(|tables| {
        tables
            .iter()
            .filter_map(|full_name| {
                let parts: Vec<&str> = full_name.split('.').collect();
                if parts.len() == 2 && parts[0] == db_name {
                    // Format as "public"."table" for consistency
                    Some(format!("\"public\".\"{}\"", parts[1]))
                } else {
                    None
                }
            })
            .collect()
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[tokio::test]
    #[ignore]
    async fn test_dump_globals() {
        let url = std::env::var("TEST_SOURCE_URL").unwrap();
        let dir = tempdir().unwrap();
        let output = dir.path().join("globals.sql");

        let result = dump_globals(&url, output.to_str().unwrap()).await;

        assert!(result.is_ok());
        assert!(output.exists());

        // Verify file contains SQL
        let content = std::fs::read_to_string(&output).unwrap();
        assert!(content.contains("CREATE ROLE") || !content.is_empty());
    }

    #[tokio::test]
    #[ignore]
    async fn test_dump_schema() {
        let url = std::env::var("TEST_SOURCE_URL").unwrap();
        let dir = tempdir().unwrap();
        let output = dir.path().join("schema.sql");

        // Extract database name from URL
        let db = url.split('/').next_back().unwrap_or("postgres");

        let filter = crate::filters::ReplicationFilter::empty();
        let result = dump_schema(&url, db, output.to_str().unwrap(), &filter).await;

        assert!(result.is_ok());
        assert!(output.exists());
    }

    #[test]
    fn test_get_schema_excluded_tables_for_db() {
        let filter = crate::filters::ReplicationFilter::new(
            None,
            None,
            None,
            Some(vec![
                "db1.table1".to_string(),
                "db1.table2".to_string(),
                "db2.table3".to_string(),
            ]),
        )
        .unwrap();

        // Schema exclusion only includes explicit exclude_tables
        let tables = get_schema_excluded_tables_for_db(&filter, "db1").unwrap();
        // Should return schema-qualified names
        assert_eq!(
            tables,
            vec!["\"public\".\"table1\"", "\"public\".\"table2\""]
        );

        let tables = get_schema_excluded_tables_for_db(&filter, "db2").unwrap();
        assert_eq!(tables, vec!["\"public\".\"table3\""]);

        let tables = get_schema_excluded_tables_for_db(&filter, "db3");
        assert!(tables.is_none() || tables.unwrap().is_empty());
    }

    #[test]
    fn test_get_data_excluded_tables_for_db() {
        let filter = crate::filters::ReplicationFilter::new(
            None,
            None,
            None,
            Some(vec![
                "db1.table1".to_string(),
                "db1.table2".to_string(),
                "db2.table3".to_string(),
            ]),
        )
        .unwrap();

        // Data exclusion includes explicit exclude_tables, schema_only, and predicate tables
        let tables = get_data_excluded_tables_for_db(&filter, "db1").unwrap();
        // Should return schema-qualified names
        assert_eq!(
            tables,
            vec!["\"public\".\"table1\"", "\"public\".\"table2\""]
        );

        let tables = get_data_excluded_tables_for_db(&filter, "db2").unwrap();
        assert_eq!(tables, vec!["\"public\".\"table3\""]);

        let tables = get_data_excluded_tables_for_db(&filter, "db3");
        assert!(tables.is_none() || tables.unwrap().is_empty());
    }

    #[test]
    fn test_get_included_tables_for_db() {
        let filter = crate::filters::ReplicationFilter::new(
            None,
            None,
            Some(vec![
                "db1.users".to_string(),
                "db1.orders".to_string(),
                "db2.products".to_string(),
            ]),
            None,
        )
        .unwrap();

        let tables = get_included_tables_for_db(&filter, "db1").unwrap();
        // Should return schema-qualified names in original order
        assert_eq!(
            tables,
            vec!["\"public\".\"users\"", "\"public\".\"orders\""]
        );

        let tables = get_included_tables_for_db(&filter, "db2").unwrap();
        assert_eq!(tables, vec!["\"public\".\"products\""]);

        let tables = get_included_tables_for_db(&filter, "db3");
        assert!(tables.is_none() || tables.unwrap().is_empty());
    }

    #[test]
    fn test_get_schema_excluded_tables_for_db_with_empty_filter() {
        let filter = crate::filters::ReplicationFilter::empty();
        let tables = get_schema_excluded_tables_for_db(&filter, "db1");
        assert!(tables.is_none());
    }

    #[test]
    fn test_get_data_excluded_tables_for_db_with_empty_filter() {
        let filter = crate::filters::ReplicationFilter::empty();
        let tables = get_data_excluded_tables_for_db(&filter, "db1");
        assert!(tables.is_none());
    }

    #[test]
    fn test_get_included_tables_for_db_with_empty_filter() {
        let filter = crate::filters::ReplicationFilter::empty();
        let tables = get_included_tables_for_db(&filter, "db1");
        assert!(tables.is_none());
    }

    #[test]
    fn test_rewrite_create_role_statements_wraps_unquoted_role() {
        let sql = "CREATE ROLE replicator WITH LOGIN;\nALTER ROLE replicator WITH LOGIN;\n";
        let rewritten = rewrite_create_role_statements(sql).expect("rewrite expected");

        assert!(rewritten.contains("DO $$"));
        assert!(rewritten.contains("Role replicator already exists"));
        assert!(rewritten.contains("CREATE ROLE replicator WITH LOGIN;"));
        assert!(rewritten.contains("ALTER ROLE replicator WITH LOGIN;"));
    }

    #[test]
    fn test_rewrite_create_role_statements_wraps_quoted_role() {
        let sql = "    CREATE ROLE \"Andre Admin\";\n";
        let rewritten = rewrite_create_role_statements(sql).expect("rewrite expected");

        assert!(rewritten.contains("DO $$"));
        assert!(rewritten.contains("Andre Admin already exists"));
        assert!(rewritten.contains("CREATE ROLE \"Andre Admin\""));
        assert!(rewritten.starts_with("    DO $$"));
    }

    #[test]
    fn test_rewrite_create_role_statements_noop_when_absent() {
        let sql = "ALTER ROLE existing WITH LOGIN;\n";
        assert!(rewrite_create_role_statements(sql).is_none());
    }
}