dbx/crates/dbx-core/tests/live_postgres_transfer.rs

use dbx_core::connection::{AppState, PoolKind};
use dbx_core::db::postgres;
use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
use dbx_core::storage::Storage;
use dbx_core::transfer::{
    get_db_type, transfer_postgres_schema_dependencies, transfer_postgres_schema_objects, transfer_table, TransferMode,
    TransferRequest,
};
use serde_json::json;

fn postgres_test_config(id: &str, database: &str) -> ConnectionConfig {
    ConnectionConfig {
        id: id.to_string(),
        name: id.to_string(),
        db_type: DatabaseType::Postgres,
        driver_profile: None,
        driver_label: None,
        url_params: None,
        host: "127.0.0.1".to_string(),
        port: 5432,
        username: "postgres".to_string(),
        password: String::new(),
        database: Some(database.to_string()),
        visible_databases: None,
        attached_databases: Vec::new(),
        color: None,
        transport_layers: Vec::new(),
        connect_timeout_secs: 5,
        query_timeout_secs: 30,
        ssl: false,
        ca_cert_path: String::new(),
        sysdba: false,
        oracle_connection_type: None,
        connection_string: None,
        redis_connection_mode: None,
        redis_sentinel_master: String::new(),
        redis_sentinel_nodes: String::new(),
        redis_sentinel_username: String::new(),
        redis_sentinel_password: String::new(),
        redis_sentinel_tls: false,
        redis_cluster_nodes: String::new(),
        external_config: None,
        jdbc_driver_class: None,
        jdbc_driver_paths: Vec::new(),
        one_time: false,
    }
}

async fn query_scalar(pool: &deadpool_postgres::Pool, sql: &str) -> serde_json::Value {
    postgres::execute_query(pool, sql).await.unwrap().rows[0][0].clone()
}

#[tokio::test]
#[ignore = "requires source/target PostgreSQL URLs via DBX_LIVE_PG_TRANSFER_SOURCE_URL and DBX_LIVE_PG_TRANSFER_TARGET_URL"]
async fn live_postgres_transfer_preserves_data_and_schema_objects() {
    let source_url = std::env::var("DBX_LIVE_PG_TRANSFER_SOURCE_URL").expect("DBX_LIVE_PG_TRANSFER_SOURCE_URL");
    let target_url = std::env::var("DBX_LIVE_PG_TRANSFER_TARGET_URL").unwrap_or_else(|_| source_url.clone());

    let source_pool = postgres::connect(&source_url, std::time::Duration::from_secs(5)).await.unwrap();
    let target_pool = postgres::connect(&target_url, std::time::Duration::from_secs(5)).await.unwrap();

    let source_database = query_scalar(&source_pool, "SELECT current_database()").await.as_str().unwrap().to_string();
    let target_database = query_scalar(&target_pool, "SELECT current_database()").await.as_str().unwrap().to_string();

    let suffix = uuid::Uuid::new_v4().simple().to_string();
    let source_schema = format!("dbx_src_{}", &suffix[..8]);
    let target_schema = format!("dbx_dst_{}", &suffix[..8]);

    let cleanup_sql = [
        format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", source_schema),
        format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", target_schema),
    ];
    let _ = postgres::execute_batch(&source_pool, &[cleanup_sql[0].clone()]).await;
    let _ = postgres::execute_batch(&target_pool, &[cleanup_sql[1].clone()]).await;

    let setup_sql = vec![
        format!("CREATE SCHEMA \"{}\"", source_schema),
        format!("CREATE TYPE \"{}\".\"user_status\" AS ENUM ('active', 'disabled')", source_schema),
        format!(
            "CREATE DOMAIN \"{}\".\"email_text\" AS text CHECK (position('@' in VALUE) > 1)",
            source_schema
        ),
        format!(
            "CREATE TABLE \"{}\".\"users\" (\
                \"id\" integer GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,\
                \"email\" \"{}\".\"email_text\" NOT NULL,\
                \"status\" \"{}\".\"user_status\" NOT NULL DEFAULT 'active',\
                \"created_at\" timestamptz NOT NULL DEFAULT now(),\
                \"active\" boolean NOT NULL DEFAULT true,\
                \"display_name\" text NOT NULL\
            )",
            source_schema, source_schema, source_schema
        ),
        format!(
            "CREATE TABLE \"{}\".\"audit_logs\" (\
                \"id\" integer GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,\
                \"user_id\" integer NOT NULL REFERENCES \"{}\".\"users\"(\"id\"),\
                \"action\" text NOT NULL,\
                \"created_at\" timestamptz NOT NULL DEFAULT now()\
            )",
            source_schema, source_schema
        ),
        format!(
            "CREATE INDEX \"users_display_name_idx\" ON \"{}\".\"users\" USING btree (lower(display_name))",
            source_schema
        ),
        format!("COMMENT ON INDEX \"{}\".\"users_display_name_idx\" IS 'lookup index'", source_schema),
        format!(
            "CREATE OR REPLACE FUNCTION \"{}\".\"log_user_insert\"() RETURNS trigger LANGUAGE plpgsql AS $$ \
             BEGIN \
                 INSERT INTO \"{}\".\"audit_logs\" (\"user_id\", \"action\") VALUES (NEW.\"id\", 'insert'); \
                 RETURN NEW; \
             END; \
             $$",
            source_schema, source_schema
        ),
        format!(
            "CREATE TRIGGER \"users_insert_audit\" AFTER INSERT ON \"{}\".\"users\" \
             FOR EACH ROW EXECUTE FUNCTION \"{}\".\"log_user_insert\"()",
            source_schema, source_schema
        ),
        format!(
            "INSERT INTO \"{}\".\"users\" (\"email\", \"status\", \"active\", \"display_name\") VALUES \
             ('alpha@example.com', 'active', true, 'Alpha'), \
             ('beta@example.com', 'disabled', false, 'Beta')",
            source_schema
        ),
        format!(
            "CREATE VIEW \"{}\".\"active_users\" AS \
             SELECT \"id\", \"email\", \"display_name\" FROM \"{}\".\"users\" WHERE \"active\"",
            source_schema, source_schema
        ),
        format!(
            "CREATE MATERIALIZED VIEW \"{}\".\"user_stats\" AS \
             SELECT \"status\", count(*)::bigint AS \"total\" FROM \"{}\".\"users\" GROUP BY \"status\"",
            source_schema, source_schema
        ),
        format!("ALTER TABLE \"{}\".\"users\" ENABLE ROW LEVEL SECURITY", source_schema),
        format!(
            "CREATE POLICY \"users_public_read\" ON \"{}\".\"users\" AS PERMISSIVE FOR SELECT TO PUBLIC USING (\"active\")",
            source_schema
        ),
        format!("GRANT USAGE ON SCHEMA \"{}\" TO PUBLIC", source_schema),
        format!("GRANT SELECT ON TABLE \"{}\".\"users\" TO PUBLIC", source_schema),
        format!("GRANT SELECT ON TABLE \"{}\".\"active_users\" TO PUBLIC", source_schema),
        format!("GRANT EXECUTE ON FUNCTION \"{}\".\"log_user_insert\"() TO PUBLIC", source_schema),
    ];
    postgres::execute_batch(&source_pool, &setup_sql).await.unwrap();

    let dir = std::env::temp_dir().join(format!("dbx-live-transfer-{}", uuid::Uuid::new_v4()));
    std::fs::create_dir_all(&dir).unwrap();
    let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
    let state = AppState::new(storage);

    let source_connection_id = "live-source";
    let target_connection_id = "live-target";
    let source_pool_key = format!("{source_connection_id}:{source_database}");
    let target_pool_key = format!("{target_connection_id}:{target_database}");

    state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone()));
    state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone()));
    state
        .configs
        .write()
        .await
        .insert(source_connection_id.to_string(), postgres_test_config(source_connection_id, &source_database));
    state
        .configs
        .write()
        .await
        .insert(target_connection_id.to_string(), postgres_test_config(target_connection_id, &target_database));

    let request = TransferRequest {
        transfer_id: format!("live-transfer-{suffix}"),
        source_connection_id: source_connection_id.to_string(),
        source_database: source_database.clone(),
        source_schema: source_schema.clone(),
        target_connection_id: target_connection_id.to_string(),
        target_database: target_database.clone(),
        target_schema: target_schema.clone(),
        tables: vec!["users".to_string(), "audit_logs".to_string()],
        create_table: true,
        mode: TransferMode::Append,
        batch_size: 100,
    };

    transfer_postgres_schema_dependencies(&state, &request, &source_pool_key, &target_pool_key, |_| {}).await.unwrap();

    let source_db_type = get_db_type(&state, source_connection_id).await.unwrap();
    let target_db_type = get_db_type(&state, target_connection_id).await.unwrap();
    for (index, table) in request.tables.iter().enumerate() {
        transfer_table(
            &state,
            &request,
            table,
            index,
            &source_db_type,
            &target_db_type,
            &source_pool_key,
            &target_pool_key,
            |_| {},
        )
        .await
        .unwrap();
    }

    transfer_postgres_schema_objects(&state, &request, &source_pool_key, &target_pool_key, |_| {}).await.unwrap();

    assert_eq!(
        query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"users\"", target_schema)).await,
        json!(2)
    );
    assert_eq!(
        query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"audit_logs\"", target_schema)).await,
        json!(2)
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!(
                "SELECT column_default NOT LIKE '%{}%' AND column_default LIKE '%user_status%' \
                 FROM information_schema.columns \
                 WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'status'",
                source_schema, target_schema
            )
        )
        .await,
        json!(true)
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!(
                "SELECT is_identity FROM information_schema.columns \
                 WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'id'",
                target_schema
            )
        )
        .await,
        json!("YES")
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!(
                "SELECT udt_name FROM information_schema.columns \
                 WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'status'",
                target_schema
            )
        )
        .await,
        json!("user_status")
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!(
                "SELECT domain_name FROM information_schema.columns \
                 WHERE table_schema = '{}' AND table_name = 'users' AND column_name = 'email'",
                target_schema
            )
        )
        .await,
        json!("email_text")
    );
    assert_eq!(
        query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"active_users\"", target_schema)).await,
        json!(1)
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!("SELECT count(*) FROM \"{}\".\"user_stats\" WHERE \"status\" = 'active'", target_schema)
        )
        .await,
        json!(1)
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!(
                "SELECT relrowsecurity FROM pg_catalog.pg_class c \
                 JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
                 WHERE n.nspname = '{}' AND c.relname = 'users'",
                target_schema
            )
        )
        .await,
        json!(true)
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!(
                "SELECT count(*) FROM pg_catalog.pg_policy p \
                 JOIN pg_catalog.pg_class c ON c.oid = p.polrelid \
                 JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \
                 WHERE n.nspname = '{}' AND c.relname = 'users' AND p.polname = 'users_public_read'",
                target_schema
            )
        )
        .await,
        json!(1)
    );
    assert_eq!(
        query_scalar(
            &target_pool,
            &format!(
                "SELECT count(*) \
                 FROM pg_catalog.pg_namespace n \
                 JOIN LATERAL aclexplode(n.nspacl) a ON true \
                 WHERE n.nspname = '{}' AND a.grantee = 0 AND a.privilege_type = 'USAGE'",
                target_schema
            )
        )
        .await,
        json!(1)
    );

    postgres::execute_query(
        &target_pool,
        &format!(
            "INSERT INTO \"{}\".\"users\" (\"email\", \"status\", \"active\", \"display_name\") \
             VALUES ('gamma@example.com', 'active', true, 'Gamma')",
            target_schema
        ),
    )
    .await
    .unwrap();

    assert_eq!(
        query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"audit_logs\"", target_schema)).await,
        json!(3)
    );

    let _ = postgres::execute_batch(&source_pool, &[cleanup_sql[0].clone()]).await;
    let _ = postgres::execute_batch(&target_pool, &[cleanup_sql[1].clone()]).await;
    let _ = std::fs::remove_dir_all(dir);
}