axum-postgres-tx 0.3.0

#![cfg(feature = "integration")]

use axum::{
    Json,
    body::Body,
    http::Request,
    http::StatusCode,
    routing::{get, post},
};
use serde_json::json;
use tower::{Layer, ServiceExt};

use axum_postgres_tx::{layer::Layer as TxLayer, tx::Tx};

#[cfg(feature = "bb8")]
type TestPool = bb8_postgres::bb8::Pool<
    bb8_postgres::PostgresConnectionManager<bb8_postgres::tokio_postgres::NoTls>,
>;

#[cfg(feature = "deadpool")]
type TestPool = deadpool_postgres::Pool;

fn db_url() -> String {
    let user = std::env::var("PGUSER").unwrap_or_else(|_| "test_user".into());
    let db = std::env::var("PGDATABASE").unwrap_or_else(|_| "test_db".into());
    let host = std::env::var("PGHOST").unwrap_or_default();
    if host.starts_with('/') {
        format!("host={host} user={user} dbname={db}")
    } else if host.is_empty() {
        format!("postgres://{user}@/{db}")
    } else {
        format!("postgres://{user}@{host}/{db}")
    }
}

#[cfg(feature = "bb8")]
async fn create_pool() -> TestPool {
    use bb8_postgres::{PostgresConnectionManager, tokio_postgres::NoTls};
    let manager = PostgresConnectionManager::new_from_stringlike(db_url(), NoTls).unwrap();
    bb8_postgres::bb8::Pool::builder()
        .max_size(5)
        .build(manager)
        .await
        .unwrap()
}

#[cfg(feature = "deadpool")]
async fn create_pool() -> TestPool {
    use deadpool_postgres::{Config, Runtime, tokio_postgres::NoTls};
    let mut cfg = Config::new();
    let url = db_url();
    let parts: Vec<&str> = url.split_whitespace().collect();
    for part in &parts {
        if let Some(val) = part.strip_prefix("host=") {
            cfg.host = Some(val.to_string());
        } else if let Some(val) = part.strip_prefix("user=") {
            cfg.user = Some(val.to_string());
        } else if let Some(val) = part.strip_prefix("dbname=") {
            cfg.dbname = Some(val.to_string());
        } else if let Some(val) = part.strip_prefix("password=") {
            cfg.password = Some(val.to_string());
        }
    }
    if cfg.host.is_none() {
        cfg.host = Some("localhost".to_string());
    }
    cfg.create_pool(Some(Runtime::Tokio1), NoTls).unwrap()
}

async fn commit_handler(tx: Tx) -> Json<serde_json::Value> {
    tx.execute("SELECT 1", &[]).await.unwrap();
    Json(json!({"ok": true}))
}

async fn error_handler(tx: Tx) -> StatusCode {
    let _ = tx
        .execute(
            "INSERT INTO rollback_test (val) VALUES ('should_not_persist')",
            &[],
        )
        .await;
    StatusCode::INTERNAL_SERVER_ERROR
}

async fn insert_handler(tx: Tx) -> Json<serde_json::Value> {
    tx.execute("INSERT INTO tx_test (val) VALUES ('hello')", &[])
        .await
        .unwrap();
    Json(json!({"inserted": true}))
}

async fn read_handler(tx: Tx) -> Json<serde_json::Value> {
    tx.execute("SELECT 1", &[]).await.unwrap();
    Json(json!({"ok": true}))
}

#[tokio::test]
async fn layer_inserts_extension_and_commits_on_success() {
    let pool = create_pool().await;
    let layer = TxLayer::from(pool);
    let svc = layer.layer(get(commit_handler));

    let response = svc
        .oneshot(Request::get("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn layer_rolls_back_on_error_response() {
    let pool = create_pool().await;

    #[cfg(feature = "bb8")]
    {
        pool.get()
            .await
            .unwrap()
            .execute(
                "CREATE TABLE IF NOT EXISTS rollback_test (id SERIAL PRIMARY KEY, val TEXT NOT NULL)",
                &[],
            )
            .await
            .unwrap();
    }

    #[cfg(feature = "deadpool")]
    {
        pool.get()
            .await
            .unwrap()
            .execute(
                "CREATE TABLE IF NOT EXISTS rollback_test (id SERIAL PRIMARY KEY, val TEXT NOT NULL)",
                &[],
            )
            .await
            .unwrap();
    }

    let layer = TxLayer::from(pool.clone());
    let svc = layer.layer(post(error_handler));

    let response = svc
        .oneshot(Request::post("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);

    #[cfg(feature = "bb8")]
    {
        let rows = pool
            .get()
            .await
            .unwrap()
            .query(
                "SELECT val FROM rollback_test WHERE val = 'should_not_persist'",
                &[],
            )
            .await
            .unwrap();
        assert!(rows.is_empty(), "transaction should have been rolled back");
    }

    #[cfg(feature = "deadpool")]
    {
        let rows = pool
            .get()
            .await
            .unwrap()
            .query(
                "SELECT val FROM rollback_test WHERE val = 'should_not_persist'",
                &[],
            )
            .await
            .unwrap();
        assert!(rows.is_empty(), "transaction should have been rolled back");
    }
}

#[tokio::test]
async fn tx_extractor_commit_for_post() {
    let pool = create_pool().await;

    #[cfg(feature = "bb8")]
    {
        pool.get()
            .await
            .unwrap()
            .execute(
                "CREATE TABLE IF NOT EXISTS tx_test (id SERIAL PRIMARY KEY, val TEXT NOT NULL)",
                &[],
            )
            .await
            .unwrap();
    }

    #[cfg(feature = "deadpool")]
    {
        pool.get()
            .await
            .unwrap()
            .execute(
                "CREATE TABLE IF NOT EXISTS tx_test (id SERIAL PRIMARY KEY, val TEXT NOT NULL)",
                &[],
            )
            .await
            .unwrap();
    }

    let layer = TxLayer::from(pool.clone());
    let svc = layer.layer(post(insert_handler));

    let response = svc
        .oneshot(Request::post("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);

    #[cfg(feature = "bb8")]
    {
        let row = pool
            .get()
            .await
            .unwrap()
            .query_one("SELECT val FROM tx_test WHERE val = 'hello'", &[])
            .await
            .unwrap();
        let val: &str = row.get(0);
        assert_eq!(val, "hello");
    }

    #[cfg(feature = "deadpool")]
    {
        let row = pool
            .get()
            .await
            .unwrap()
            .query_one("SELECT val FROM tx_test WHERE val = 'hello'", &[])
            .await
            .unwrap();
        let val: &str = row.get(0);
        assert_eq!(val, "hello");
    }
}

#[tokio::test]
async fn tx_extractor_read_only_for_get() {
    let pool = create_pool().await;

    #[cfg(feature = "bb8")]
    {
        pool.get()
            .await
            .unwrap()
            .execute(
                "CREATE TABLE IF NOT EXISTS ro_test (id SERIAL PRIMARY KEY, val TEXT NOT NULL)",
                &[],
            )
            .await
            .unwrap();
    }

    #[cfg(feature = "deadpool")]
    {
        pool.get()
            .await
            .unwrap()
            .execute(
                "CREATE TABLE IF NOT EXISTS ro_test (id SERIAL PRIMARY KEY, val TEXT NOT NULL)",
                &[],
            )
            .await
            .unwrap();
    }

    let layer = TxLayer::from(pool);
    let svc = layer.layer(get(read_handler));

    let response = svc
        .oneshot(Request::get("/").body(Body::empty()).unwrap())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
}