hypercall_db_diesel/

snapshots.rs

//! Snapshot persistence implementations for DatabaseHandler.
//!
//! Implements InstrumentsSnapshotWriter, InstrumentsSnapshotReader,
//! PortfolioSnapshotWriter, PortfolioSnapshotReader, and EngineCommandOffsetReader.

use anyhow::{Context, Result};
use diesel::prelude::*;
use diesel::sql_types::{BigInt, Bytea, Int4, Text};
use diesel::RunQueryDsl;
use tracing::{debug, info};

use crate::database_handler::DatabaseHandler;
use hypercall_db::types::snapshots::{
    InstrumentSnapshotEntry, InstrumentsSnapshotData, InstrumentsSnapshotInput,
    PortfolioSnapshotAccountEntry, PortfolioSnapshotData, PortfolioSnapshotInput,
    SnapshotOffsetEntry,
};

// ---- Diesel QueryableByName helper structs ----

#[derive(QueryableByName)]
struct SnapshotIdResult {
    #[diesel(sql_type = BigInt)]
    id: i64,
}

#[derive(QueryableByName)]
struct InstrumentEntryRow {
    #[diesel(sql_type = BigInt)]
    #[allow(dead_code)]
    snapshot_id: i64,
    #[diesel(sql_type = Text)]
    symbol: String,
    #[diesel(sql_type = Bytea)]
    data: Vec<u8>,
}

#[derive(QueryableByName)]
struct PortfolioAccountRow {
    #[diesel(sql_type = Text)]
    wallet: String,
    #[diesel(sql_type = Bytea)]
    data: Vec<u8>,
}

#[derive(QueryableByName)]
struct OffsetRow {
    #[diesel(sql_type = Text)]
    stream: String,
    #[diesel(sql_type = Int4)]
    partition: i32,
    #[diesel(sql_type = BigInt)]
    offset: i64,
}

// ---- InstrumentsSnapshotWriter ----

impl hypercall_db::InstrumentsSnapshotWriter for DatabaseHandler {
    fn write_instruments_snapshot_sync(&self, input: &InstrumentsSnapshotInput) -> Result<i64> {
        let mut conn = self.pool().get().context("Failed to get DB connection")?;

        conn.transaction(|conn| {
            // Insert snapshot header
            let snapshot_id: i64 = diesel::sql_query(
                "INSERT INTO engine_snapshots (snapshot_type, engine_version) VALUES ($1, NULL) RETURNING id",
            )
            .bind::<Text, _>(&input.snapshot_type)
            .get_result::<SnapshotIdResult>(conn)
            .context("Failed to insert snapshot header")?
            .id;

            debug!(
                "Created instruments snapshot header id={} entries={}",
                snapshot_id,
                input.entries.len()
            );

            // Insert entries
            for entry in &input.entries {
                diesel::sql_query(
                    "INSERT INTO instruments_snapshot_entries (snapshot_id, symbol, data) VALUES ($1, $2, $3)",
                )
                .bind::<BigInt, _>(snapshot_id)
                .bind::<Text, _>(&entry.symbol)
                .bind::<Bytea, _>(&entry.data)
                .execute(conn)
                .with_context(|| format!("Failed to insert instrument {}", entry.symbol))?;
            }

            // Insert offsets
            let mut offset_rows = 0;
            for (stream, partitions) in &input.offsets {
                for (partition, offset) in partitions {
                    diesel::sql_query(
                        "INSERT INTO instruments_snapshot_offsets (snapshot_id, stream, partition, \"offset\") VALUES ($1, $2, $3, $4)",
                    )
                    .bind::<BigInt, _>(snapshot_id)
                    .bind::<Text, _>(stream)
                    .bind::<Int4, _>(*partition)
                    .bind::<BigInt, _>(*offset)
                    .execute(conn)
                    .with_context(|| format!("Failed to insert offset for {}", stream))?;
                    offset_rows += 1;
                }
            }

            // Cleanup old snapshots
            let deleted = diesel::sql_query(
                "DELETE FROM engine_snapshots WHERE snapshot_type = $1 AND id NOT IN (
                    SELECT id FROM engine_snapshots WHERE snapshot_type = $1 ORDER BY created_at DESC LIMIT $2
                )",
            )
            .bind::<Text, _>(&input.snapshot_type)
            .bind::<BigInt, _>(input.retention_count)
            .execute(conn)
            .context("Failed to cleanup snapshots")?;

            if deleted > 0 {
                info!("Deleted {} old instruments snapshot(s)", deleted);
            }

            info!(
                "Stored instruments snapshot id={} offset_rows={}",
                snapshot_id, offset_rows
            );
            Ok(snapshot_id)
        })
    }
}

// ---- InstrumentsSnapshotReader ----

impl hypercall_db::InstrumentsSnapshotReader for DatabaseHandler {
    fn get_latest_instruments_snapshot_id_sync(&self) -> Result<Option<i64>> {
        let mut conn = self.pool().get().context("Failed to get DB connection")?;

        let result = diesel::sql_query(
            "SELECT id FROM engine_snapshots WHERE snapshot_type = 'instruments' ORDER BY created_at DESC LIMIT 1",
        )
        .get_result::<SnapshotIdResult>(&mut conn)
        .optional()
        .context("Failed to fetch snapshot id")?;

        Ok(result.map(|r| r.id))
    }

    fn load_instruments_snapshot_sync(&self, snapshot_id: i64) -> Result<InstrumentsSnapshotData> {
        let mut conn = self.pool().get().context("Failed to get DB connection")?;

        let entries = diesel::sql_query(
            "SELECT snapshot_id, symbol, data FROM instruments_snapshot_entries WHERE snapshot_id = $1",
        )
        .bind::<BigInt, _>(snapshot_id)
        .load::<InstrumentEntryRow>(&mut conn)
        .context("Failed to load instrument rows")?;

        let entry_data: Vec<InstrumentSnapshotEntry> = entries
            .into_iter()
            .map(|row| InstrumentSnapshotEntry {
                symbol: row.symbol,
                data: row.data,
            })
            .collect();

        let offsets = diesel::sql_query(
            "SELECT stream, partition, \"offset\" FROM instruments_snapshot_offsets WHERE snapshot_id = $1",
        )
        .bind::<BigInt, _>(snapshot_id)
        .load::<OffsetRow>(&mut conn)
        .context("Failed to load offsets")?;

        let offset_data: Vec<SnapshotOffsetEntry> = offsets
            .into_iter()
            .map(|row| SnapshotOffsetEntry {
                stream: row.stream,
                partition: row.partition,
                offset: row.offset,
            })
            .collect();

        Ok(InstrumentsSnapshotData {
            entries: entry_data,
            offsets: offset_data,
        })
    }
}

// ---- PortfolioSnapshotWriter ----

impl hypercall_db::PortfolioSnapshotWriter for DatabaseHandler {
    fn write_portfolio_snapshot_sync(&self, input: &PortfolioSnapshotInput) -> Result<i64> {
        let mut conn = self.pool().get().context("Failed to get DB connection")?;

        conn.transaction(|conn| {
            // Insert snapshot header
            let snapshot_id: i64 = diesel::sql_query(
                "INSERT INTO engine_snapshots (snapshot_type, engine_version) VALUES ('portfolio', NULL) RETURNING id",
            )
            .get_result::<SnapshotIdResult>(conn)
            .context("Failed to insert snapshot header")?
            .id;

            debug!("Created portfolio snapshot header with id={}", snapshot_id);

            // Insert accounts
            let mut account_count = 0;
            for account in &input.accounts {
                diesel::sql_query(
                    "INSERT INTO portfolio_snapshot_accounts (snapshot_id, wallet, data) VALUES ($1, $2, $3)",
                )
                .bind::<BigInt, _>(snapshot_id)
                .bind::<Text, _>(&account.wallet)
                .bind::<Bytea, _>(&account.data)
                .execute(conn)
                .with_context(|| format!("Failed to insert account {}", account.wallet))?;
                account_count += 1;
            }

            debug!("Inserted {} portfolio accounts", account_count);

            // Insert offsets
            let mut offset_count = 0;
            for (stream, partitions) in &input.offsets {
                for (partition, offset) in partitions {
                    diesel::sql_query(
                        "INSERT INTO portfolio_snapshot_offsets (snapshot_id, stream, partition, \"offset\") VALUES ($1, $2, $3, $4)",
                    )
                    .bind::<BigInt, _>(snapshot_id)
                    .bind::<Text, _>(stream)
                    .bind::<Int4, _>(*partition)
                    .bind::<BigInt, _>(*offset)
                    .execute(conn)
                    .with_context(|| format!("Failed to insert offset for {}", stream))?;
                    offset_count += 1;
                }
            }

            debug!("Inserted {} offset entries", offset_count);

            // Cleanup old snapshots beyond retention limit
            let deleted = diesel::sql_query(
                "DELETE FROM engine_snapshots WHERE snapshot_type = 'portfolio' AND id NOT IN (
                    SELECT id FROM engine_snapshots
                    WHERE snapshot_type = 'portfolio'
                    ORDER BY created_at DESC
                    LIMIT $1
                )",
            )
            .bind::<BigInt, _>(input.retention_count)
            .execute(conn)
            .context("Failed to cleanup old snapshots")?;

            if deleted > 0 {
                info!("Cleaned up {} old portfolio snapshot(s)", deleted);
            }

            info!(
                "Created portfolio snapshot id={} with {} accounts and {} offsets (retaining {})",
                snapshot_id, account_count, offset_count, input.retention_count
            );

            Ok(snapshot_id)
        })
    }
}

// ---- PortfolioSnapshotReader ----

impl hypercall_db::PortfolioSnapshotReader for DatabaseHandler {
    fn get_latest_portfolio_snapshot_id_sync(&self) -> Result<Option<i64>> {
        let mut conn = self.pool().get().context("Failed to get DB connection")?;

        let result = diesel::sql_query(
            "SELECT id FROM engine_snapshots WHERE snapshot_type = 'portfolio' ORDER BY created_at DESC LIMIT 1",
        )
        .get_result::<SnapshotIdResult>(&mut conn)
        .optional()
        .context("Failed to query latest snapshot")?;

        Ok(result.map(|r| r.id))
    }

    fn portfolio_snapshot_exists_sync(&self, snapshot_id: i64) -> Result<bool> {
        let mut conn = self.pool().get().context("Failed to get DB connection")?;

        let exists: Option<SnapshotIdResult> = diesel::sql_query(
            "SELECT id FROM engine_snapshots WHERE id = $1 AND snapshot_type = 'portfolio'",
        )
        .bind::<BigInt, _>(snapshot_id)
        .get_result(&mut conn)
        .optional()
        .context("Failed to verify snapshot")?;

        Ok(exists.is_some())
    }

    fn load_portfolio_snapshot_sync(&self, snapshot_id: i64) -> Result<PortfolioSnapshotData> {
        let mut conn = self.pool().get().context("Failed to get DB connection")?;

        // Load accounts
        let account_rows: Vec<PortfolioAccountRow> = diesel::sql_query(
            "SELECT wallet, data FROM portfolio_snapshot_accounts WHERE snapshot_id = $1",
        )
        .bind::<BigInt, _>(snapshot_id)
        .load(&mut conn)
        .context("Failed to load accounts")?;

        let accounts: Vec<PortfolioSnapshotAccountEntry> = account_rows
            .into_iter()
            .map(|row| PortfolioSnapshotAccountEntry {
                wallet: row.wallet,
                data: row.data,
            })
            .collect();

        // Load offsets
        let offset_rows: Vec<OffsetRow> = diesel::sql_query(
            "SELECT stream, partition, \"offset\" FROM portfolio_snapshot_offsets WHERE snapshot_id = $1",
        )
        .bind::<BigInt, _>(snapshot_id)
        .load(&mut conn)
        .context("Failed to load offsets")?;

        let offsets: Vec<SnapshotOffsetEntry> = offset_rows
            .into_iter()
            .map(|row| SnapshotOffsetEntry {
                stream: row.stream,
                partition: row.partition,
                offset: row.offset,
            })
            .collect();

        Ok(PortfolioSnapshotData { accounts, offsets })
    }
}

// Note: get_next_engine_command_id_sync is implemented via JournalReplayReader
// in hypercall-db-diesel/src/replay.rs. No separate EngineCommandOffsetReader needed.