saddle-db 0.3.3

use std::{path::PathBuf, 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, name_mapping_bypass, result_limit_exceeded,
        transaction_begin_failed,
    },
    name_mapping::{
        MappingStartupConfig, OperationSql, PhysicalOperationPlans, StaticLogicalTable, freeze,
    },
    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;

/// Frozen deployment input for the sole process-wide database.
///
/// It contains the resolved secret value and mapping directory but exposes
/// neither. Generated framework assembly consumes it exactly once to finish
/// registering its static tables and operations before pool startup.
pub struct DatabaseStartupInjection {
    url: String,
    mapping_directory: PathBuf,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DatabaseStartupInjectionError {
    InvalidConnectionEnvironment,
    MissingConnectionSecret,
    InvalidConnectionSecret,
    InvalidMappingDirectory,
    MappingDirectoryUnavailable,
}

impl DatabaseStartupInjectionError {
    pub const fn code(self) -> &'static str {
        match self {
            Self::InvalidConnectionEnvironment => "db.startup.connection_env_invalid",
            Self::MissingConnectionSecret => "db.startup.connection_secret_missing",
            Self::InvalidConnectionSecret => "db.startup.connection_secret_invalid",
            Self::InvalidMappingDirectory => "db.startup.mapping_dir_invalid",
            Self::MappingDirectoryUnavailable => "db.startup.mapping_dir_unavailable",
        }
    }
}

impl DatabaseStartupInjection {
    /// Resolves the database portion of the single `saddle.toml` deployment
    /// input. No database connection or metadata query is performed here.
    #[doc(hidden)]
    pub fn load(
        config_directory: &std::path::Path,
        connection_environment: &str,
        mapping_directory: &std::path::Path,
    ) -> std::result::Result<Self, DatabaseStartupInjectionError> {
        if !valid_environment_name(connection_environment) {
            return Err(DatabaseStartupInjectionError::InvalidConnectionEnvironment);
        }
        let url = std::env::var(connection_environment).map_err(|error| match error {
            std::env::VarError::NotPresent => {
                DatabaseStartupInjectionError::MissingConnectionSecret
            }
            std::env::VarError::NotUnicode(_) => {
                DatabaseStartupInjectionError::InvalidConnectionSecret
            }
        })?;
        if url.trim().is_empty() {
            return Err(DatabaseStartupInjectionError::InvalidConnectionSecret);
        }
        if mapping_directory.as_os_str().is_empty()
            || mapping_directory.is_absolute()
            || mapping_directory
                .components()
                .any(|component| !matches!(component, std::path::Component::Normal(_)))
        {
            return Err(DatabaseStartupInjectionError::InvalidMappingDirectory);
        }
        let mapping_directory = config_directory.join(mapping_directory);
        if !mapping_directory.is_dir() {
            return Err(DatabaseStartupInjectionError::MappingDirectoryUnavailable);
        }
        Ok(Self {
            url,
            mapping_directory,
        })
    }

    #[doc(hidden)]
    pub fn into_database_config(self) -> DatabaseConfig {
        DatabaseConfig::new(self.url).name_mapping_directory(self.mapping_directory)
    }
}

fn valid_environment_name(value: &str) -> bool {
    let mut bytes = value.bytes();
    let Some(first) = bytes.next() else {
        return false;
    };
    (first.is_ascii_alphabetic() || first == b'_')
        && bytes.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
}

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

impl DatabaseConfig {
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            max_connections: 16,
            acquire_timeout: Duration::from_secs(5),
            name_mappings: None,
        }
    }
    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
    }

    /// Enables deployment-provided logical-to-physical name mappings.
    ///
    /// The directory must contain exactly one JSON file per registered table.
    pub fn name_mapping_directory(mut self, directory: impl Into<PathBuf>) -> Self {
        self.name_mappings
            .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
            .set_directory(directory.into());
        self
    }

    /// Registers one generated logical table whose complete mapping is required.
    pub fn register_logical_table<T: StaticLogicalTable>(mut self) -> Self {
        self.name_mappings
            .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
            .register_table::<T>();
        self
    }

    #[doc(hidden)]
    pub fn register_query_operation<O: crate::internal::StaticQueryOptionalOperation>(
        mut self,
    ) -> Self {
        self.name_mappings
            .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
            .register_query::<O>(O::OPERATION, O::LOGICAL_TABLE, O::LOGICAL_COLUMNS, O::SQL);
        self
    }

    #[doc(hidden)]
    pub fn register_write_operation<O: crate::internal::StaticWriteOperation>(mut self) -> Self {
        self.name_mappings
            .get_or_insert_with(|| MappingStartupConfig::new(PathBuf::new()))
            .register_write::<O>(O::OPERATION, O::LOGICAL_TABLE, O::LOGICAL_COLUMNS, O::SQL);
        self
    }

    /// Performs the same local mapping validation used by startup, without
    /// connecting to or inspecting the database server.
    pub fn validate_name_mappings(
        &self,
    ) -> std::result::Result<(), crate::DatabaseNameMappingError> {
        freeze(self.name_mappings.clone()).map(|_| ())
    }

    pub(crate) fn verified_connections(mut self, value: u32) -> Self {
        self.max_connections = value;
        self
    }

    pub(crate) fn deployment_url(&self) -> &str {
        &self.url
    }

    pub(crate) 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 {
    pub(crate) pool: MySqlPool,
    observer: Observer,
    cleanup: Arc<CleanupCoordinator>,
    physical_plans: Arc<PhysicalOperationPlans>,
}

impl Database {
    /// Creates the single managed pool and verifies that it can connect.
    pub async fn connect(config: DatabaseConfig, observer: Observer) -> Result<Self> {
        let physical_plans = freeze(config.name_mappings.clone())
            .map_err(|_| invalid_config("database name mapping is invalid"))?;
        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)
            .idle_timeout(None)
            .max_lifetime(None)
            .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)?;
        let mut preopened = Vec::new();
        preopened
            .try_reserve_exact(config.max_connections as usize)
            .map_err(|_| invalid_config("database connection profile is too large"))?;
        while preopened.len() < config.max_connections as usize {
            preopened.push(pool.acquire().await.map_err(map_operation_error)?);
        }
        for connection in &mut preopened {
            connection.return_to_pool().await;
        }
        Ok(Self {
            pool,
            observer,
            cleanup: CleanupCoordinator::start(),
            physical_plans: Arc::new(physical_plans),
        })
    }

    pub(crate) fn query_sql<O: crate::internal::StaticQueryOptionalOperation>(
        &self,
    ) -> OperationSql {
        self.physical_plans.query::<O>()
    }

    pub(crate) fn write_sql<O: crate::internal::StaticWriteOperation>(&self) -> OperationSql {
        self.physical_plans.write::<O>()
    }

    pub async fn query_all(
        &self,
        parent: &CallContext,
        statement: Statement,
    ) -> Result<Vec<DbRow>> {
        if self.physical_plans.enabled() {
            return Err(name_mapping_bypass());
        }
        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>> {
        if self.physical_plans.enabled() {
            return Err(name_mapping_bypass());
        }
        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> {
        if self.physical_plans.enabled() {
            return Err(name_mapping_bypass());
        }
        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,
    {
        if self.physical_plans.enabled() {
            return Err(name_mapping_bypass());
        }
        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),
        }
    }

    pub(crate) 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 physical_plans = freeze(config.name_mappings.clone())
            .map_err(|_| invalid_config("database name mapping is invalid"))?;
        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(),
            physical_plans: Arc::new(physical_plans),
        })
    }
}

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;

    struct MappedTable;

    impl crate::StaticLogicalTable for MappedTable {
        const TABLE: &'static str = "逻辑表";
        const COLUMNS: &'static [&'static str] = &["逻辑列"];
    }

    #[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"));
    }

    #[test]
    fn unified_startup_injection_resolves_secret_and_relative_mapping_directory() {
        let directory =
            std::env::temp_dir().join(format!("saddle-db-alpha10-config-{}", std::process::id()));
        let mappings = directory.join("mappings");
        let _ = std::fs::remove_dir_all(&directory);
        std::fs::create_dir_all(&mappings).unwrap();
        let executable = std::env::current_exe().unwrap();
        for mode in ["happy", "missing", "bad-path"] {
            let mut child = std::process::Command::new(&executable);
            child
                .args([
                    "--ignored",
                    "--exact",
                    "database::tests::unified_startup_injection_child",
                ])
                .env("SADDLE_ALPHA10_INJECTION_MODE", mode)
                .env("SADDLE_ALPHA10_CONFIG_ROOT", &directory)
                .env_remove("SADDLE_ALPHA10_DATABASE_TEST");
            if mode != "missing" {
                child.env(
                    "SADDLE_ALPHA10_DATABASE_TEST",
                    "mysql://deployment-secret@localhost/database",
                );
            }
            assert!(
                child.status().unwrap().success(),
                "child mode {mode} failed"
            );
        }
        std::fs::remove_dir_all(directory).unwrap();
    }

    #[test]
    #[ignore = "executed in isolated child processes by the parent test"]
    fn unified_startup_injection_child() {
        let root = PathBuf::from(std::env::var_os("SADDLE_ALPHA10_CONFIG_ROOT").unwrap());
        let mode = std::env::var("SADDLE_ALPHA10_INJECTION_MODE").unwrap();
        let mapping = if mode == "bad-path" {
            std::path::Path::new("../mappings")
        } else {
            std::path::Path::new("mappings")
        };
        let result = DatabaseStartupInjection::load(&root, "SADDLE_ALPHA10_DATABASE_TEST", mapping);
        match mode.as_str() {
            "happy" => {
                let injection = result.unwrap();
                assert_eq!(injection.mapping_directory, root.join("mappings"));
                assert!(injection.url.contains("deployment-secret"));
            }
            "missing" => assert_eq!(
                result.err(),
                Some(DatabaseStartupInjectionError::MissingConnectionSecret)
            ),
            "bad-path" => assert_eq!(
                result.err(),
                Some(DatabaseStartupInjectionError::InvalidMappingDirectory)
            ),
            _ => panic!("unknown child mode"),
        }
    }

    #[tokio::test]
    async fn mapping_enabled_rejects_legacy_raw_statement_before_database_io() {
        let directory =
            std::env::temp_dir().join(format!("saddle-db-alpha7-bypass-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&directory);
        std::fs::create_dir(&directory).unwrap();
        std::fs::write(
            directory.join("table.json"),
            r#"{"table":{"from":"逻辑表","to":"physical_table"},"columns":[{"from":"逻辑列","to":"physical_column"}]}"#,
        )
        .unwrap();
        let observer = Observer::with_writer(ObserverConfig::default(), io::sink()).unwrap();
        let database = Database::connect_lazy(
            DatabaseConfig::new("mysql://localhost/unused")
                .name_mapping_directory(&directory)
                .register_logical_table::<MappedTable>(),
            observer,
        )
        .unwrap();
        let error = database
            .write(
                &context(),
                Statement::new("legacy.write", "SELECT 1").unwrap(),
            )
            .await
            .unwrap_err();
        assert_eq!(error.code(), "db.name_mapping_required");
        database.shutdown().await.unwrap();
        std::fs::remove_dir_all(directory).unwrap();
    }

    #[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
    }
}