saddle-db 0.1.1

use std::{str::FromStr, sync::Arc, time::Duration};

use futures_util::TryStreamExt;
use saddle_core::{ComponentLifecycle, LifecycleFuture, OperationId, Result};
use saddle_observability::{CallKind, Observer};
use sqlx::{
    Connection, MySqlPool,
    mysql::{MySqlConnectOptions, MySqlConnection, MySqlPoolOptions},
};

use crate::{
    CallContext, DbRow, MAX_RESULT_BYTES, Statement, Transaction, TransactionFuture, WriteResult,
    cleanup::CleanupCoordinator,
    error::{invalid_config, map_operation_error, result_limit_exceeded, transaction_begin_failed},
    row::row_payload_bytes,
};

pub const MAX_QUERY_ROWS: usize = 10_000;
/// Maximum MySQL protocol packet accepted by the V1 deployment contract.
///
/// Pool startup rejects servers configured with a larger `max_allowed_packet`,
/// and every physical pool connection repeats the check. This value remains
/// below MySQL's `0xFF_FF_FF` continuation threshold, so sqlx receives one
/// bounded buffer and cannot enter its two-fragment aggregate preallocation.
pub const MAX_INBOUND_PACKET_BYTES: u64 = 8_388_608;

/// Configuration for the single V1 MySQL/MariaDB data source.
#[derive(Clone)]
pub struct DatabaseConfig {
    url: String,
    max_connections: u32,
    acquire_timeout: Duration,
}

impl DatabaseConfig {
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            max_connections: 16,
            acquire_timeout: Duration::from_secs(5),
        }
    }
    pub fn max_connections(mut self, value: u32) -> Self {
        self.max_connections = value;
        self
    }
    pub fn acquire_timeout(mut self, value: Duration) -> Self {
        self.acquire_timeout = value;
        self
    }

    fn options(&self) -> Result<MySqlConnectOptions> {
        if self.max_connections == 0 {
            return Err(invalid_config("max connections must be greater than zero"));
        }
        if self.acquire_timeout.is_zero() {
            return Err(invalid_config("acquire timeout must be greater than zero"));
        }
        MySqlConnectOptions::from_str(&self.url)
            .map_err(|_| invalid_config("database URL is not a valid MySQL/MariaDB URL"))
    }
}

/// The process-wide managed database capability.
#[derive(Clone)]
pub struct Database {
    pool: MySqlPool,
    observer: Observer,
    cleanup: Arc<CleanupCoordinator>,
}

impl Database {
    /// Creates the single managed pool and verifies that it can connect.
    pub async fn connect(config: DatabaseConfig, observer: Observer) -> Result<Self> {
        let options = config.options()?;
        let mut preflight = MySqlConnection::connect_with(&options)
            .await
            .map_err(map_operation_error)?;
        let server_packet_limit = sqlx::query_scalar::<_, u64>("SELECT @@max_allowed_packet")
            .fetch_one(&mut preflight)
            .await
            .map_err(map_operation_error)?;
        preflight.close().await.map_err(map_operation_error)?;
        if server_packet_limit > MAX_INBOUND_PACKET_BYTES {
            return Err(invalid_config(
                "server max_allowed_packet exceeds the V1 inbound allocation limit",
            ));
        }
        let pool = MySqlPoolOptions::new()
            .max_connections(config.max_connections)
            .acquire_timeout(config.acquire_timeout)
            .after_connect(|connection, _metadata| {
                Box::pin(async move {
                    let packet_limit = sqlx::query_scalar::<_, u64>("SELECT @@max_allowed_packet")
                        .fetch_one(connection)
                        .await?;
                    if packet_limit > MAX_INBOUND_PACKET_BYTES {
                        return Err(sqlx::Error::Protocol(
                            "server packet limit exceeds Saddle V1 allocation boundary".to_owned(),
                        ));
                    }
                    Ok(())
                })
            })
            .connect_with(options)
            .await
            .map_err(map_operation_error)?;
        Ok(Self {
            pool,
            observer,
            cleanup: CleanupCoordinator::start(),
        })
    }

    pub async fn query_all(
        &self,
        parent: &CallContext,
        statement: Statement,
    ) -> Result<Vec<DbRow>> {
        statement.validate()?;
        let operation = statement.operation().to_owned();
        let call = self.observer.start_child_call(
            parent,
            CallKind::Database,
            "database",
            "database",
            OperationId::from(operation),
        );
        let result = async {
            let mut stream = statement.query().fetch(&self.pool);
            let mut rows = Vec::new();
            let mut result_bytes = 0_usize;
            while let Some(row) = stream.try_next().await.map_err(map_operation_error)? {
                if rows.len() == MAX_QUERY_ROWS {
                    return Err(result_limit_exceeded());
                }
                result_bytes = result_bytes.saturating_add(row_payload_bytes(&row)?);
                if result_bytes > MAX_RESULT_BYTES {
                    return Err(result_limit_exceeded());
                }
                rows.push(DbRow(row));
            }
            Ok(rows)
        }
        .await;
        finish_call(call, &result);
        result
    }

    pub async fn query_optional(
        &self,
        parent: &CallContext,
        statement: Statement,
    ) -> Result<Option<DbRow>> {
        statement.validate()?;
        let operation = statement.operation().to_owned();
        let call = self.observer.start_child_call(
            parent,
            CallKind::Database,
            "database",
            "database",
            OperationId::from(operation),
        );
        let result = statement
            .query()
            .fetch_optional(&self.pool)
            .await
            .map_err(map_operation_error)
            .and_then(|row| {
                row.map(|row| {
                    row_payload_bytes(&row)?;
                    Ok(DbRow(row))
                })
                .transpose()
            });
        finish_call(call, &result);
        result
    }

    pub async fn write(&self, parent: &CallContext, statement: Statement) -> Result<WriteResult> {
        statement.validate()?;
        let operation = statement.operation().to_owned();
        let call = self.observer.start_child_call(
            parent,
            CallKind::Database,
            "database",
            "database",
            OperationId::from(operation),
        );
        let result = statement
            .query()
            .execute(&self.pool)
            .await
            .map(|result| WriteResult::new(result.rows_affected(), result.last_insert_id()))
            .map_err(map_operation_error);
        finish_call(call, &result);
        result
    }

    /// Runs one explicit, single-level transaction.
    ///
    /// The callback returns a boxed async block because that is the minimal
    /// Rust API that safely ties all operations to the borrowed transaction.
    pub async fn transaction<T, F>(
        &self,
        parent: &CallContext,
        operation: impl Into<OperationId>,
        work: F,
    ) -> Result<T>
    where
        T: Send,
        F: for<'a> FnOnce(&'a mut Transaction) -> TransactionFuture<'a, T> + Send,
    {
        let operation = operation.into();
        crate::statement::validate_operation(operation.as_str())?;
        let call = self.observer.start_child_call(
            parent,
            CallKind::Transaction,
            "database",
            "database",
            operation,
        );
        let cleanup = match self.cleanup.transaction_sender() {
            Some(cleanup) => cleanup,
            None => {
                let error = transaction_begin_failed();
                call.fail(&error);
                return Err(error);
            }
        };
        let begin = self.observer.start_child_call(
            call.context(),
            CallKind::Transaction,
            "database",
            "database",
            "begin",
        );
        let raw = match self.pool.begin().await {
            Ok(transaction) => {
                begin.succeed();
                transaction
            }
            Err(_) => {
                let error = transaction_begin_failed();
                begin.fail(&error);
                call.fail(&error);
                return Err(error);
            }
        };
        let mut transaction = Transaction::new(raw, self.observer.clone(), call, cleanup);
        match work(&mut transaction).await {
            Ok(value) => transaction.commit().await.map(|()| value),
            Err(work_error) => Err(transaction.rollback(work_error).await),
        }
    }

    async fn close(&self) -> Result<()> {
        let cleanup = self.cleanup.shutdown().await;
        self.pool.close().await;
        cleanup
    }

    #[cfg(test)]
    pub(crate) fn connect_lazy(config: DatabaseConfig, observer: Observer) -> Result<Self> {
        let options = config.options()?;
        let pool = MySqlPoolOptions::new()
            .max_connections(config.max_connections)
            .acquire_timeout(config.acquire_timeout)
            .connect_lazy_with(options);
        Ok(Self {
            pool,
            observer,
            cleanup: CleanupCoordinator::start(),
        })
    }
}

impl ComponentLifecycle for Database {
    fn name(&self) -> &'static str {
        "database"
    }
    fn start(&self) -> LifecycleFuture<'_> {
        Box::pin(async { Ok(()) })
    }
    fn shutdown(&self) -> LifecycleFuture<'_> {
        Box::pin(async move { self.close().await })
    }
}

fn finish_call<T>(call: saddle_observability::ActiveCall, result: &Result<T>) {
    match result {
        Ok(_) => call.succeed(),
        Err(error) => call.fail(error),
    }
}

#[cfg(test)]
mod tests {
    use saddle_core::{ApplicationId, ErrorKind, ModuleId, ServiceId, SpanId, TraceId};
    use saddle_observability::ObserverConfig;
    use serde_json::Value;
    use std::{
        io,
        sync::{Arc, Mutex},
    };

    use super::*;
    use crate::SaddleError;

    #[derive(Clone, Default)]
    struct Capture(Arc<Mutex<Vec<u8>>>);
    impl io::Write for Capture {
        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
            self.0.lock().unwrap().extend_from_slice(bytes);
            Ok(bytes.len())
        }
        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }
    fn context() -> CallContext {
        CallContext::new(
            ApplicationId::from("shop"),
            ModuleId::from("orders"),
            ServiceId::from("orders"),
            OperationId::from("create"),
            TraceId::from_u128(1),
            SpanId::from_u64(2),
        )
    }

    #[test]
    fn configuration_rejects_invalid_bounds_without_exposing_url() {
        let observer = Observer::with_writer(ObserverConfig::default(), io::sink()).unwrap();
        let error = Database::connect_lazy(
            DatabaseConfig::new("mysql://secret@localhost/db").max_connections(0),
            observer,
        )
        .err()
        .unwrap();
        assert_eq!(error.code(), "db.invalid_config");
        assert!(!error.to_string().contains("secret"));
    }

    #[tokio::test]
    async fn closed_pool_query_has_stable_error_and_trace_record() {
        let capture = Capture::default();
        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
        let database = Database::connect_lazy(
            DatabaseConfig::new("mysql://localhost/db"),
            observer.clone(),
        )
        .unwrap();
        database.close().await.unwrap();
        let error = match database
            .query_all(
                &context(),
                Statement::new("orders.list", "SELECT 1").unwrap(),
            )
            .await
        {
            Ok(_) => panic!("closed pool query unexpectedly succeeded"),
            Err(error) => error,
        };
        assert_eq!(error.kind(), ErrorKind::Unavailable);
        assert_eq!(error.code(), "db.connection_unavailable");
        observer.flush().await.unwrap();
        let output = String::from_utf8(capture.0.lock().unwrap().clone()).unwrap();
        let records: Vec<Value> = output
            .lines()
            .map(|line| serde_json::from_str(line).unwrap())
            .collect();
        assert_eq!(records[0]["call_kind"], "database");
        assert_eq!(records[1]["error_code"], "db.connection_unavailable");
        assert_eq!(records[0]["trace_id"], context().trace_id().to_string());
        assert!(!output.contains("SELECT 1"));
    }

    #[tokio::test]
    async fn transaction_begin_failure_records_phase_and_stable_error() {
        let capture = Capture::default();
        let observer = Observer::with_writer(ObserverConfig::default(), capture.clone()).unwrap();
        let database = Database::connect_lazy(
            DatabaseConfig::new("mysql://localhost/db"),
            observer.clone(),
        )
        .unwrap();
        database.pool.close().await;
        let error = database
            .transaction(&context(), "orders.create", |_transaction| {
                Box::pin(async { Ok::<_, SaddleError>(()) })
            })
            .await
            .unwrap_err();
        assert_eq!(error.code(), "db.transaction_begin_failed");
        observer.flush().await.unwrap();
        let output = String::from_utf8(capture.0.lock().unwrap().clone()).unwrap();
        let records: Vec<Value> = output
            .lines()
            .map(|line| serde_json::from_str(line).unwrap())
            .collect();
        assert_eq!(records.len(), 4);
        assert_eq!(records[0]["operation"], "orders.create");
        assert_eq!(records[1]["operation"], "begin");
        assert_eq!(records[2]["error_code"], "db.transaction_begin_failed");
        assert_eq!(records[3]["error_code"], "db.transaction_begin_failed");
        assert!(
            records
                .iter()
                .all(|record| record["trace_id"] == context().trace_id().to_string())
        );
    }

    #[allow(dead_code)]
    async fn transaction_usage_compiles(database: &Database, context: &CallContext) -> Result<()> {
        database
            .transaction(context, "orders.create", |transaction| {
                Box::pin(async move {
                    transaction
                        .write(
                            Statement::new("orders.insert", "INSERT INTO orders(id) VALUES (?)")?
                                .bind(1_u64)?,
                        )
                        .await?;
                    transaction
                        .query_optional(
                            Statement::new("orders.find", "SELECT id FROM orders WHERE id = ?")?
                                .bind(1_u64)?,
                        )
                        .await?;
                    Ok(())
                })
            })
            .await
    }
}