neocrates 0.1.47

use sqlx::postgres::{PgPool, PgPoolOptions};
use thiserror::Error;
use tracing::{error, info, warn};
use url::Url;

// ── Error & Result ────────────────────────────────────────────────────────────

#[derive(Debug, Error)]
pub enum SqlxError {
    #[error("Database error: {0}")]
    Database(#[from] sqlx::Error),

    #[error("Migration error: {0}")]
    Migration(#[from] sqlx::migrate::MigrateError),

    #[error("Database URL parse error: {0}")]
    UrlParseError(#[from] url::ParseError),

    #[error("Database name missing in URL")]
    DatabaseNameMissing,

    #[error("Configuration error: {0}")]
    Config(String),

    #[error(transparent)]
    Other(#[from] anyhow::Error),
}

pub type SqlxResult<T> = Result<T, SqlxError>;

// ── Type aliases ──────────────────────────────────────────────────────────────

/// Convenience alias for a PostgreSQL transaction.
///
/// `'a` is typically `'static` when the connection comes from a [`SqlxPool`],
/// and a shorter lifetime when it comes from a single [`sqlx::PgConnection`].
///
/// # Service-layer pattern
///
/// Service methods that must run inside *or* outside a transaction should
/// accept `conn: &mut sqlx::PgConnection`. Both `PoolConnection` (via
/// `pool.acquire()`) and `SqlxTx` (via `pool.begin()`) deref to `PgConnection`,
/// so the same function works in either context:
///
/// ```rust,ignore
/// async fn insert_user(conn: &mut sqlx::PgConnection, name: &str) -> SqlxResult<i64> {
///     sqlx_fetch_scalar!(
///         &mut *conn,
///         sqlx::query_scalar::<_, i64>(
///             "INSERT INTO users (name) VALUES ($1) RETURNING id"
///         ).bind(name)
///     ).await
/// }
///
/// // Standalone (no transaction)
/// let mut conn = pool.acquire().await?;
/// insert_user(&mut conn, "alice").await?;
///
/// // Inside a transaction
/// sqlx_with_tx!(pool, |tx| {
///     insert_user(&mut *tx, "alice").await?;
///     Ok(())
/// }).await?;
/// ```
pub type SqlxTx<'a> = sqlx::Transaction<'a, sqlx::Postgres>;

// ── Internal helper ───────────────────────────────────────────────────────────

/// Connect to the system `postgres` database and create `db_name` if it does
/// not yet exist. Mirrors the behaviour of `dieselhelper::pool::ensure_database_exists`.
async fn ensure_database_exists(url: &str) -> SqlxResult<()> {
    let parsed = Url::parse(url)?;

    let db_name = parsed
        .path_segments()
        .and_then(|segs| segs.filter(|s| !s.is_empty()).last())
        .map(str::to_string)
        .filter(|s| !s.trim().is_empty())
        .ok_or(SqlxError::DatabaseNameMissing)?;

    // Build a URL that points to the system database
    let mut system_url = parsed.clone();
    system_url.set_path("/postgres");

    let system_pool = PgPoolOptions::new()
        .max_connections(1)
        .connect(system_url.as_str())
        .await?;

    let exists: bool =
        sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)")
            .bind(&db_name)
            .fetch_one(&system_pool)
            .await?;

    if exists {
        info!("Database '{}' already exists", db_name);
    } else {
        let sanitized = db_name.replace('"', "\"\"");
        sqlx::query(&format!("CREATE DATABASE \"{}\"", sanitized))
            .execute(&system_pool)
            .await?;
        info!("Database '{}' created", db_name);
    }

    system_pool.close().await;
    Ok(())
}

// ── SqlxPool ──────────────────────────────────────────────────────────────────

/// Async PostgreSQL connection pool backed by [`sqlx::PgPool`].
///
/// # Quick start
/// ```rust,ignore
/// let pool = SqlxPool::from_env().await?;
/// let row: (i64,) = sqlx::query_as("SELECT 1").fetch_one(pool.pool()).await?;
/// ```
#[derive(Clone, Debug)]
pub struct SqlxPool {
    pool: PgPool,
}

impl SqlxPool {
    /// Create a new pool, auto-creating the target database if it does not exist.
    ///
    /// # Arguments
    /// * `url` – full PostgreSQL connection string (e.g. `postgres://user:pass@host/db`)
    /// * `max_connections` – maximum number of connections in the pool
    pub async fn new(url: &str, max_connections: u32) -> SqlxResult<Self> {
        ensure_database_exists(url).await?;

        let pool = PgPoolOptions::new()
            .max_connections(max_connections)
            .after_connect(|conn, _meta| {
                Box::pin(async move {
                    sqlx::query("SET TIME ZONE 'UTC'")
                        .execute(&mut *conn)
                        .await
                        .map(|_| ())
                })
            })
            .connect(url)
            .await?;

        Ok(Self { pool })
    }

    /// Create a pool from environment variables.
    ///
    /// Reads `DATABASE_URL` (required) and `DATABASE_POOL_SIZE` (optional, default 10).
    pub async fn from_env() -> SqlxResult<Self> {
        let url = std::env::var("DATABASE_URL")
            .map_err(|_| SqlxError::Config("DATABASE_URL env var not set".into()))?;
        let max_connections = std::env::var("DATABASE_POOL_SIZE")
            .ok()
            .and_then(|s| s.parse::<u32>().ok())
            .unwrap_or(10);
        Self::new(&url, max_connections).await
    }

    /// Returns a reference to the underlying [`PgPool`].
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Begin a new database transaction.
    ///
    /// The transaction is automatically rolled back when dropped if not committed.
    pub async fn begin(&self) -> SqlxResult<sqlx::Transaction<'_, sqlx::Postgres>> {
        self.pool.begin().await.map_err(SqlxError::from)
    }

    /// Execute `f` inside a transaction that is automatically committed on
    /// success or rolled back on error.
    ///
    /// Prefer the [`sqlx_with_tx!`] macro to avoid writing `Box::pin(async move { … })`.
    ///
    /// # Example
    /// ```rust,ignore
    /// let user_id: i64 = sqlx_with_tx!(pool, |tx| {
    ///     let id = sqlx_fetch_scalar!(
    ///         &mut *tx,
    ///         sqlx::query_scalar::<_, i64>(
    ///             "INSERT INTO users (name) VALUES ($1) RETURNING id"
    ///         ).bind("alice")
    ///     ).await?;
    ///     Ok(id)
    /// }).await?;
    /// ```
    pub async fn with_transaction<F, T>(&self, f: F) -> SqlxResult<T>
    where
        F: for<'c> FnOnce(
            &'c mut SqlxTx<'static>,
        ) -> std::pin::Pin<
            Box<dyn std::future::Future<Output = SqlxResult<T>> + 'c>,
        >,
    {
        let mut tx = self.pool.begin().await?;
        match f(&mut tx).await {
            Ok(val) => {
                tx.commit().await?;
                Ok(val)
            }
            Err(e) => {
                if let Err(rb_err) = tx.rollback().await {
                    warn!("SQLx transaction rollback failed: {}", rb_err);
                }
                Err(e)
            }
        }
    }

    /// Run all pending SQLx migrations from the given migrator.
    ///
    /// # Example
    /// ```rust,ignore
    /// use sqlx::migrate::Migrator;
    /// use std::path::Path;
    ///
    /// let migrator = Migrator::new(Path::new("./migrations")).await?;
    /// pool.run_migrations(&migrator).await?;
    /// ```
    pub async fn run_migrations(
        &self,
        migrator: &sqlx::migrate::Migrator,
    ) -> SqlxResult<()> {
        migrator.run(&self.pool).await.map_err(SqlxError::from)
    }

    /// Returns the current number of connections in the pool (idle + active).
    pub fn size(&self) -> u32 {
        self.pool.size()
    }

    /// Returns the number of idle connections.
    pub fn idle(&self) -> u32 {
        self.pool.num_idle() as u32
    }

    /// Perform a lightweight connectivity check (`SELECT 1`).
    pub async fn health_check(&self) -> SqlxResult<()> {
        sqlx::query("SELECT 1")
            .execute(&self.pool)
            .await
            .map_err(|e| {
                error!("SQLx health check failed: {}", e);
                SqlxError::from(e)
            })?;
        info!("SQLx health check passed: db connection test successful");
        Ok(())
    }
}

// Allow `SqlxPool` to be used directly wherever `&PgPool` is accepted.
impl std::ops::Deref for SqlxPool {
    type Target = PgPool;
    fn deref(&self) -> &Self::Target {
        &self.pool
    }
}

// ── Macro ─────────────────────────────────────────────────────────────────────

/// Convenience wrapper around [`SqlxPool::with_transaction`].
///
/// Avoids writing `Box::pin(async move { … })` by hand.
/// The transaction is committed on success and rolled back on error.
///
/// Use `&mut *tx` to reborrow the transaction as `&mut PgConnection` so it
/// can be passed to multiple query macros inside the block.
///
/// # Example
/// ```rust,ignore
/// let user_id: i64 = sqlx_with_tx!(pool, |tx| {
///     sqlx_execute!(
///         &mut *tx,
///         sqlx::query("INSERT INTO users (name) VALUES ($1)").bind("alice")
///     ).await?;
///     let id = sqlx_fetch_scalar!(
///         &mut *tx,
///         sqlx::query_scalar::<_, i64>("SELECT lastval()")
///     ).await?;
///     Ok(id)
/// }).await?;
/// ```
#[macro_export]
macro_rules! sqlx_with_tx {
    ($pool:expr, |$tx:ident| $body:block) => {
        $pool.with_transaction(|$tx| Box::pin(async move { $body }))
    };
}