integresql 0.2.0

use std::error::Error;

use integresql::{ConnectionSettings, DbManager, IntegresqlError, TemplateDb};
use tokio::sync::OnceCell;
use tokio_postgres::Client;
use tokio_postgres::NoTls;

const SCHEMA_SQL: &str = "
CREATE SCHEMA test_schema;

CREATE TABLE test_schema.people (
    id SERIAL PRIMARY KEY,
    first_name TEXT,
    last_name TEXT,
    birthdate DATE
);
";

const DATA_SQL: &str = "
INSERT INTO test_schema.people (first_name, last_name, birthdate) VALUES
('John', 'Doe', '1980-01-01'),
('Jane', 'Smith', '1990-02-02'),
('Alice', 'Johnson', '2000-03-03');
";

async fn connect(conn_str: impl AsRef<str>) -> Result<Client, tokio_postgres::Error> {
    let (client, connection) = tokio_postgres::connect(conn_str.as_ref(), NoTls).await?;
    tokio::spawn(async move {
        if let Err(e) = connection.await {
            eprintln!("Connection error: {}", e);
        }
    });
    Ok(client)
}

/// Sets up the template database with your schema and data
async fn setup_template(config: ConnectionSettings) -> Result<(), Box<dyn Error + 'static>> {
    let conn_str = config.to_libpq_url();
    let client = connect(&conn_str).await?;

    client.batch_execute(SCHEMA_SQL).await?;
    client.batch_execute(DATA_SQL).await?;

    drop(client);
    tokio::task::yield_now().await; // Yield to ensure the connection is closed
    Ok(())
}

/// Caches the template DB info in a Tokio-friendly way
static TEMPLATE_DB: OnceCell<TemplateDb> = OnceCell::const_new();

/// Returns connection settings to a fresh, writable test database cloned from
/// the template database.
///
/// If your tests won't modify the database, you can call `get_readonly_test_db()`
/// instead for even faster performance.
///
/// The underlying Postgres DB will automatically be reset when the TestDb is
/// dropped.
async fn get_test_db() -> Result<ConnectionSettings, IntegresqlError> {
    let template_db = TEMPLATE_DB
        .get_or_init(|| async {
            let db_manager = DbManager::from_env();
            let template_key = (SCHEMA_SQL, DATA_SQL);
            db_manager
                .get_template_db_async(template_key, setup_template)
                .await
                .expect("Failed to create template database")
        })
        .await;

    let test_db = template_db.get_writable_test_db()?;
    Ok(test_db)
}

async fn test_insert() {
    let test_db = get_test_db().await.expect("Failed to get test database");
    let client = connect(test_db.to_libpq_url())
        .await
        .expect("Failed to connect to test database");

    let insert_sql = "INSERT INTO test_schema.people (first_name, last_name, birthdate) VALUES ($1, $2, $3::TEXT::DATE)";
    client
        .execute(insert_sql, &[&"Bob", &"Brown", &"1995-05-05"])
        .await
        .expect("Failed to insert row");

    let new_count: i64 = client
        .query_one("SELECT count(*) FROM test_schema.people", &[])
        .await
        .expect("Query failed")
        .get(0);
    assert_eq!(new_count, 4, "Row count should be 4 after insert");
}

async fn test_query() {
    let test_db = get_test_db().await.expect("Failed to get test database");
    let client = connect(test_db.to_libpq_url())
        .await
        .expect("Failed to connect to test database");

    let query_sql = "SELECT first_name, last_name FROM test_schema.people ORDER BY id";
    let rows = client
        .query(query_sql, &[])
        .await
        .expect("Query failed")
        .into_iter()
        .map(|row| {
            let first_name: String = row.get(0);
            let last_name: String = row.get(1);
            (first_name, last_name)
        })
        .collect::<Vec<_>>();

    assert_eq!(rows.len(), 3, "Should return 3 rows");
    assert_eq!(rows[0], ("John".to_string(), "Doe".to_string()));
    assert_eq!(rows[1], ("Jane".to_string(), "Smith".to_string()));
    assert_eq!(rows[2], ("Alice".to_string(), "Johnson".to_string()));
}

#[tokio::main]
async fn main() {
    // To keep this example runnable, per Rust style, the main function will
    // drive the tests manually. Usually, you'd use the `#[tokio::test]`
    // attribute for async tests while using tokio.
    test_insert().await;
    test_query().await;
}