pg-upsert 0.1.1

use pg_embed::pg_enums::PgAuthMethod;
use pg_embed::pg_fetch::{PG_V15, PgFetchSettings};
use pg_embed::postgres::{PgEmbed, PgSettings};
use pg_upsert::{UpsertOptions, fields, upsert};
use sqlx::PgPool;
use sqlx::postgres::PgPoolOptions;
use std::path::PathBuf;
use std::time::Duration;

async fn setup_test_db() -> (PgEmbed, PgPool) {
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::time::{SystemTime, UNIX_EPOCH};

    static COUNTER: AtomicU64 = AtomicU64::new(0);

    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis();

    let counter = COUNTER.fetch_add(1, Ordering::SeqCst);
    let thread_id = std::thread::current().id();

    let database_dir = PathBuf::from(format!(
        "target/pg_data_{}_{:?}_{}",
        timestamp, thread_id, counter
    ));

    let pg_settings = PgSettings {
        database_dir,
        port: portpicker::pick_unused_port().expect("no free port"),
        user: "postgres".to_string(),
        password: "postgres".to_string(),
        auth_method: PgAuthMethod::Plain,
        persistent: false,
        timeout: Some(Duration::from_secs(30)),
        migration_dir: None,
    };

    let fetch_settings = PgFetchSettings {
        version: PG_V15,
        ..Default::default()
    };

    let mut pg = PgEmbed::new(pg_settings, fetch_settings)
        .await
        .expect("failed to create pg-embed");

    pg.setup().await.expect("failed to setup pg-embed");
    pg.start_db().await.expect("failed to start pg-embed");

    let database_name = "test_db";
    pg.create_database(database_name)
        .await
        .expect("failed to create database");

    let connection_string = pg.full_db_uri(database_name);

    let pool = PgPoolOptions::new()
        .max_connections(5)
        .connect(&connection_string)
        .await
        .expect("failed to connect to database");

    (pg, pool)
}

async fn create_users_table(pool: &PgPool) {
    sqlx::query(
        r#"
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            email TEXT,
            version INTEGER DEFAULT 1
        )
        "#,
    )
    .execute(pool)
    .await
    .expect("failed to create users table");
}

async fn get_user(pool: &PgPool, id: i32) -> Option<(i32, String, Option<String>, i32)> {
    sqlx::query_as::<_, (i32, String, Option<String>, i32)>(
        "SELECT id, name, email, version FROM users WHERE id = $1",
    )
    .bind(id)
    .fetch_optional(pool)
    .await
    .expect("failed to fetch user")
}

async fn count_users(pool: &PgPool) -> i64 {
    sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM users")
        .fetch_one(pool)
        .await
        .expect("failed to count users")
}

#[tokio::test]
async fn test_upsert_insert_new_row() {
    let (_pg, pool) = setup_test_db().await;
    create_users_table(&pool).await;

    let rows_affected = upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice",
            "email" => "alice@example.com",
            "version" => 1_i32
        ]],
        UpsertOptions::default(),
    )
    .await
    .expect("upsert failed");

    assert_eq!(rows_affected, 1);

    let user = get_user(&pool, 1).await.expect("user not found");
    assert_eq!(user.0, 1);
    assert_eq!(user.1, "Alice");
    assert_eq!(user.2, Some("alice@example.com".to_string()));
}

#[tokio::test]
async fn test_upsert_update_existing_row() {
    let (_pg, pool) = setup_test_db().await;
    create_users_table(&pool).await;

    upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice",
            "email" => "alice@example.com",
            "version" => 1_i32
        ]],
        UpsertOptions::default(),
    )
    .await
    .expect("first upsert failed");

    let rows_affected = upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice Updated",
            "email" => "alice.new@example.com",
            "version" => 2_i32
        ]],
        UpsertOptions::default(),
    )
    .await
    .expect("second upsert failed");

    assert_eq!(rows_affected, 1);

    let user = get_user(&pool, 1).await.expect("user not found");
    assert_eq!(user.1, "Alice Updated");
    assert_eq!(user.2, Some("alice.new@example.com".to_string()));
}

#[tokio::test]
async fn test_upsert_do_nothing_on_conflict() {
    let (_pg, pool) = setup_test_db().await;
    create_users_table(&pool).await;

    upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice",
            "email" => "alice@example.com",
            "version" => 1_i32
        ]],
        UpsertOptions::default(),
    )
    .await
    .expect("first upsert failed");

    let rows_affected = upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice Should Not Update",
            "email" => "nope@example.com",
            "version" => 2_i32
        ]],
        UpsertOptions {
            do_nothing_on_conflict: true,
            ..Default::default()
        },
    )
    .await
    .expect("second upsert failed");

    assert_eq!(rows_affected, 0);

    let user = get_user(&pool, 1).await.expect("user not found");
    assert_eq!(user.1, "Alice");
    assert_eq!(user.2, Some("alice@example.com".to_string()));
}

#[tokio::test]
async fn test_upsert_version_field_newer() {
    let (_pg, pool) = setup_test_db().await;
    create_users_table(&pool).await;

    upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice",
            "email" => "alice@example.com",
            "version" => 1_i32
        ]],
        UpsertOptions::default(),
    )
    .await
    .expect("first upsert failed");

    let rows_affected = upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice v5",
            "email" => "alice5@example.com",
            "version" => 5_i32
        ]],
        UpsertOptions {
            version_field: Some("version".into()),
            ..Default::default()
        },
    )
    .await
    .expect("second upsert failed");

    assert_eq!(rows_affected, 1);

    let user = get_user(&pool, 1).await.expect("user not found");
    assert_eq!(user.1, "Alice v5");
    assert_eq!(user.3, 5);
}

#[tokio::test]
async fn test_upsert_version_field_older() {
    let (_pg, pool) = setup_test_db().await;
    create_users_table(&pool).await;

    upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice",
            "email" => "alice@example.com",
            "version" => 10_i32
        ]],
        UpsertOptions::default(),
    )
    .await
    .expect("first upsert failed");

    let rows_affected = upsert(
        &pool,
        "users",
        &["id"],
        vec![fields![
            "id" => 1_i32,
            "name" => "Alice Old",
            "email" => "old@example.com",
            "version" => 5_i32
        ]],
        UpsertOptions {
            version_field: Some("version".into()),
            ..Default::default()
        },
    )
    .await
    .expect("second upsert failed");

    assert_eq!(rows_affected, 0);

    let user = get_user(&pool, 1).await.expect("user not found");
    assert_eq!(user.1, "Alice");
    assert_eq!(user.3, 10);
}

#[tokio::test]
async fn test_upsert_multiple_rows() {
    let (_pg, pool) = setup_test_db().await;
    create_users_table(&pool).await;

    let rows_affected = upsert(
        &pool,
        "users",
        &["id"],
        vec![
            fields![
                "id" => 1_i32,
                "name" => "Alice",
                "email" => "alice@example.com",
                "version" => 1_i32
            ],
            fields![
                "id" => 2_i32,
                "name" => "Bob",
                "email" => "bob@example.com",
                "version" => 1_i32
            ],
            fields![
                "id" => 3_i32,
                "name" => "Charlie",
                "email" => "charlie@example.com",
                "version" => 1_i32
            ],
        ],
        UpsertOptions::default(),
    )
    .await
    .expect("upsert failed");

    assert_eq!(rows_affected, 3);
    assert_eq!(count_users(&pool).await, 3);

    let alice = get_user(&pool, 1).await.expect("alice not found");
    let bob = get_user(&pool, 2).await.expect("bob not found");
    let charlie = get_user(&pool, 3).await.expect("charlie not found");

    assert_eq!(alice.1, "Alice");
    assert_eq!(bob.1, "Bob");
    assert_eq!(charlie.1, "Charlie");
}