hypercall/

db_handler.rs

use anyhow::{Context, Result};
use diesel::pg::PgConnection;
use diesel::r2d2;
use std::sync::Arc;
use tracing::info;

use hypercall_types::EngineMessage;

use crate::journal::JournalFillSideEffect;

pub use hypercall_db_diesel::DbAuthConfig;
pub use hypercall_db_diesel::DbPool;
pub use hypercall_db_diesel::DynConnectionManager;

/// Redact credentials from a database URL for safe logging.
/// Turns `postgres://user:pass@host:5432/db` into `postgres://***@host:5432/db`.
fn redact_database_url(url: &str) -> String {
    url.find('@')
        .map(|at| format!("postgres://***@{}", &url[at + 1..]))
        .unwrap_or_else(|| "postgres://***".to_string())
}

pub use hypercall_db::IntegrityQueryResults;
pub use hypercall_db::LiquidationBonusApplyResult;
pub use hypercall_db::PortfolioReplayEvent;
pub use hypercall_db::SettlementResult;

pub struct DieselEventHandler {
    pool: Arc<DbPool>,
    /// Delegate for event handling / batch cancel / fill persistence.
    db_handler: hypercall_db_diesel::DatabaseHandler,
}

impl DieselEventHandler {
    pub fn new(database_url: &str) -> Result<Self> {
        Self::new_with_pool_size(database_url, 3)
    }

    pub fn new_with_pool_size(database_url: &str, pool_max: u32) -> Result<Self> {
        Self::new_with_auth(DbAuthConfig::password(database_url), pool_max)
    }

    pub fn new_with_auth(auth: DbAuthConfig, pool_max: u32) -> Result<Self> {
        let handler = Self::new_pool_only(&auth, pool_max)?;
        handler.run_migrations().with_context(|| {
            format!(
                "Failed to run migrations for database (auth={:?}): {}",
                auth,
                redact_database_url(&auth.current_url())
            )
        })?;
        Ok(handler)
    }

    /// Create a DieselEventHandler with a connection pool but skip migrations.
    pub fn new_readonly(database_url: &str, pool_max: u32) -> Result<Self> {
        Self::new_readonly_auth(DbAuthConfig::password(database_url), pool_max)
    }

    pub fn new_readonly_auth(auth: DbAuthConfig, pool_max: u32) -> Result<Self> {
        let redacted = redact_database_url(&auth.current_url());
        info!(
            "Initializing readonly DieselEventHandler (migrations skipped): {}",
            redacted
        );
        Self::new_pool_only(&auth, pool_max)
    }

    /// Internal: create the pool without running migrations.
    fn new_pool_only(auth: &DbAuthConfig, pool_max: u32) -> Result<Self> {
        let redacted = redact_database_url(&auth.current_url());
        info!(
            "Initializing DieselEventHandler for PostgreSQL: {}",
            redacted
        );
        let pool = Arc::new(
            hypercall_db_diesel::build_db_pool(auth, pool_max, 30_000, 10_000).with_context(
                || {
                    format!(
                        "Failed to create connection pool for database: {}",
                        redacted
                    )
                },
            )?,
        );
        let db_handler =
            hypercall_db_diesel::DatabaseHandler::with_pool_no_migrations(pool.clone());

        Ok(Self { pool, db_handler })
    }

    /// Get a connection from the pool
    pub async fn get_connection(&self) -> Result<r2d2::PooledConnection<DynConnectionManager>> {
        self.pool
            .get()
            .with_context(|| "Failed to get connection from pool")
    }

    /// Get the connection pool.
    pub fn pool(&self) -> Arc<DbPool> {
        self.pool.clone()
    }

    /// Get a clone of the connection pool (not the Arc).
    pub fn get_pool(&self) -> DbPool {
        (*self.pool).clone()
    }

    /// Create DieselEventHandler with an existing pool
    pub fn with_pool(pool: Arc<DbPool>) -> Result<Self> {
        info!("Creating DieselEventHandler with existing pool");

        let db_handler =
            hypercall_db_diesel::DatabaseHandler::with_pool_no_migrations(pool.clone());
        let handler = Self { pool, db_handler };

        handler
            .run_migrations()
            .with_context(|| "Failed to run migrations")?;

        Ok(handler)
    }

    /// Create DieselEventHandler with an existing pool, skipping migrations.
    pub fn with_pool_no_migrations(pool: Arc<DbPool>) -> Self {
        info!("Creating DieselEventHandler with existing pool (migrations skipped)");
        let db_handler =
            hypercall_db_diesel::DatabaseHandler::with_pool_no_migrations(pool.clone());
        Self { pool, db_handler }
    }

    /// Get a reference to the underlying `DatabaseHandler`.
    pub fn db_handler(&self) -> &hypercall_db_diesel::DatabaseHandler {
        &self.db_handler
    }

    pub fn run_migrations(&self) -> Result<()> {
        self.db_handler.run_migrations()
    }

    // ===== Event handling: delegates to DatabaseHandler =====

    pub async fn handle_event(&self, event: &EngineMessage) -> Result<()> {
        let mut conn = self.get_connection().await?;
        self.db_handler.handle_event_with_conn(&mut conn, event)
    }

    pub fn handle_event_with_conn(
        &self,
        conn: &mut PgConnection,
        event: &EngineMessage,
    ) -> Result<()> {
        self.db_handler.handle_event_with_conn(conn, event)
    }

    pub async fn handle_event_batch(&self, events: Vec<EngineMessage>) -> Result<()> {
        self.db_handler.handle_event_batch_sync(&events)
    }

    // ===== Batch cancel: thin wrapper converting StartupExpiredOrderCancel =====

    /// Batch-insert CANCELED rows into order_updates for orders on expired instruments
    /// discovered during replay. Delegates to `DatabaseHandler::batch_cancel_expired_orders_sync`.
    pub(crate) fn batch_cancel_expired_orders_sync(
        &self,
        cancels: &[crate::rsm::unified_engine::recovery::StartupExpiredOrderCancel],
        reason: &str,
    ) -> Result<usize> {
        let portable: Vec<hypercall_db_diesel::event_handler::ExpiredOrderCancel> = cancels
            .iter()
            .map(|c| hypercall_db_diesel::event_handler::ExpiredOrderCancel {
                order_id: c.order_id,
                wallet: c.wallet,
                symbol: c.symbol.clone(),
                price: c.price,
                size: c.size,
                side: c.side,
                tif: c.tif,
                is_perp: c.is_perp,
                underlying: c.underlying.clone(),
                reduce_only: c.reduce_only,
                nonce: c.nonce,
                signature: c.signature.clone(),
                mmp_enabled: c.mmp_enabled,
                filled_size: c.filled_size,
                client_id: c.client_id.clone(),
                timestamp: c.timestamp,
            })
            .collect();
        self.db_handler
            .batch_cancel_expired_orders_sync(&portable, reason)
    }

    // ===== Fill persistence: delegates to free function in event_handler =====

    pub(crate) fn persist_fill_with_side_effects_in_tx(
        conn: &mut PgConnection,
        fill: &hypercall_types::Fill,
        side_effects: &JournalFillSideEffect,
    ) -> Result<(bool, bool)> {
        let portable_side_effects = hypercall_db_diesel::event_handler::FillSideEffect {
            trade_id: side_effects.trade_id,
            taker_ledger_delta: side_effects.taker_ledger_delta,
            maker_ledger_delta: side_effects.maker_ledger_delta,
            taker_premium_delta: side_effects.taker_premium_delta,
            maker_premium_delta: side_effects.maker_premium_delta,
            underlying_notional: side_effects.underlying_notional,
        };
        hypercall_db_diesel::event_handler::persist_legacy_replay_fill_with_side_effects_in_tx(
            conn,
            fill,
            &portable_side_effects,
        )
    }
}

#[async_trait::async_trait]
impl hypercall_runtime_api::EngineEventPersistence for DieselEventHandler {
    async fn handle_event(&self, event: &hypercall_types::EngineMessage) -> anyhow::Result<()> {
        self.handle_event(event).await
    }
}