espresso_node/api/

sql.rs

use std::{
    collections::{HashMap, HashSet, VecDeque},
    sync::Arc,
};

use anyhow::{Context, bail, ensure};
use async_trait::async_trait;
use committable::{Commitment, Committable};
use espresso_types::{
    BlockMerkleCommitment, BlockMerkleTree, ChainConfig, FeeAccount, FeeMerkleCommitment,
    FeeMerkleTree, Leaf2, NodeState, ValidatedState, get_l1_deposits,
    v0::traits::PersistenceOptions as _,
    v0_1::IterableFeeInfo,
    v0_3::{
        REWARD_MERKLE_TREE_V1_HEIGHT, RewardAccountProofV1, RewardAccountQueryDataV1,
        RewardAccountV1, RewardMerkleTreeV1,
    },
    v0_4::{PermittedRewardMerkleTreeV2, RewardAccountV2, RewardMerkleTreeV2},
    v0_6::RewardAccountQueryDataV2,
};
use futures::future::Future;
use hotshot_query_service::{
    Resolvable,
    availability::LeafId,
    data_source::{
        VersionedDataSource,
        storage::{
            AvailabilityStorage, MerklizedStateStorage, NodeStorage,
            pruning::{PrunerCfg, PrunerConfig},
        },
    },
    merklized_state::Snapshot,
    status::HasMetrics as _,
};
use hotshot_types::{
    data::{EpochNumber, QuorumProposalWrapper, ViewNumber},
    message::Proposal,
    traits::{ValidatedState as _, election::MembershipSnapshot, metrics::Metrics},
    utils::{epoch_from_block_number, is_last_block},
    vote::HasViewNumber,
};
use jf_merkle_tree_compat::{
    ForgetableMerkleTreeScheme, ForgetableUniversalMerkleTreeScheme, LookupResult,
    MerkleTreeScheme, prelude::MerkleNode,
};
use light_client::storage::LightClientDbOptions;
use process_metrics::ProcessMetrics;
use sqlx::{Encode, Row, Type};
use tokio::time::sleep;
use vbs::version::Version;
use versions::{
    DRB_AND_HEADER_UPGRADE_VERSION, EPOCH_REWARD_VERSION, EPOCH_VERSION, NEW_PROTOCOL_VERSION,
};

use super::{
    persistence::{BackendOptions as _, Options},
    qs::{
        Backend as _, Config, Db, DbKind, SqlDataSource, SqlStorage, Transaction, TransactionMode,
        Write, query_as,
    },
};
use crate::{
    SeqTypes,
    api::{
        ApiState, BlocksFrontier, LightClientProvider, RewardMerkleTreeDataSource,
        RewardMerkleTreeV2Data,
        data_source::{DataSourceOptions, NodeStateDataSource as _, Provider, SequencerDataSource},
        options::{self, Query, get_public_env_vars, init_query_data_source},
        state::NodeApiStateImpl,
        update::{ApiEventConsumer, ApiSink},
    },
    catchup::{CatchupStorage, NullStateCatchup},
    context::TaskList,
    persistence::{ChainConfigPersistence, sql::Options as SqlOptions},
    request_response::data_source::Storage as RequestResponseStorage,
    state::{EmptyStateStart, compute_state_update, update_state_storage_loop},
    util::BoundedJoinSet,
};

pub type DataSource = SqlDataSource<SeqTypes, Provider>;

impl DataSourceOptions for Options {
    type DataSource = DataSource;

    fn enable_query_module(&self, opt: options::Options, query: Query) -> options::Options {
        opt.query_sql(query, self.clone().into_inner())
    }

    fn light_client_db(&self, query: &Query) -> anyhow::Result<LightClientDbOptions> {
        Ok(Db::light_client_db(&Config::try_from(&**self)?, query))
    }
}

#[async_trait]
impl SequencerDataSource for DataSource {
    type Options = Options;

    async fn create(opt: Self::Options, provider: Provider, reset: bool) -> anyhow::Result<Self> {
        let fetch_limit = opt.fetch_rate_limit;
        let active_fetch_delay = opt.active_fetch_delay;
        let chunk_fetch_delay = opt.chunk_fetch_delay;
        let mut cfg = Config::try_from(&*opt)?;

        if reset {
            cfg = cfg.reset_schema();
        }

        let mut builder = cfg.builder(provider).await?;

        if let Some(limit) = fetch_limit {
            builder = builder.with_rate_limit(limit);
        }

        if opt.lightweight {
            tracing::warn!("enabling light weight mode..");
            builder = builder.leaf_only();
        }

        // Transaction and payload totals count from genesis, and a pruned node never has those
        // blocks.
        if opt.prune {
            builder = builder.disable_aggregator();
        }

        if let Some(delay) = active_fetch_delay {
            builder = builder.with_active_fetch_delay(delay);
        }
        if let Some(delay) = chunk_fetch_delay {
            builder = builder.with_chunk_fetch_delay(delay);
        }
        if let Some(chunk_size) = opt.sync_status_chunk_size {
            builder = builder.with_sync_status_chunk_size(chunk_size);
        }
        if let Some(ttl) = opt.sync_status_ttl {
            builder = builder.with_sync_status_ttl(ttl);
        }
        if let Some(chunk_size) = opt.proactive_scan_chunk_size {
            builder = builder.with_proactive_range_chunk_size(chunk_size);
        }
        if let Some(timeout) = opt.proactive_fetch_timeout {
            builder = builder.with_proactive_fetch_timeout(timeout);
        }
        if let Some(interval) = opt.proactive_scan_interval {
            builder = builder.with_proactive_interval(interval);
        }
        if opt.disable_proactive_fetching {
            builder = builder.disable_proactive_fetching();
        }

        builder.build().await
    }
}

/// Start the query module of `api` over this backend, returning what consensus needs from it.
pub(in crate::api) async fn init_query_module(
    api: &options::Options,
    query_opt: Query,
    mod_opt: Options,
    state: ApiState,
    tasks: &mut TaskList,
) -> anyhow::Result<(
    Box<dyn Metrics>,
    Box<dyn ApiSink>,
    Option<RequestResponseStorage>,
)> {
    let mut provider = Provider::default();

    // Use the database itself as a fetching provider: sometimes we can fetch data that is
    // missing from the query service from ephemeral consensus storage.
    let db_provider = mod_opt.clone().create().await?;
    provider = provider
        .with_block_provider(db_provider.clone())
        .with_vid_common_provider(db_provider);
    // If that fails, fetch missing data from peers.
    let light_client_db = mod_opt.light_client_db(&query_opt)?;
    let light_client = Arc::new(LightClientProvider::new(
        query_opt.peers,
        state.clone(),
        query_opt.light_client,
        light_client_db,
    )?);
    provider = provider.with_provider(Arc::clone(&light_client));

    let ranges_concurrency = mod_opt.ranges_concurrency;
    let ds = DataSource::create(mod_opt.clone(), provider, false).await?;
    let inner_storage = ds.inner();
    tasks.spawn("process_metrics", ProcessMetrics::new(ds.metrics()).run());
    let (metrics, ds) = init_query_data_source(ds, state.clone());

    let get_node_state = {
        let state = state.clone();
        async move { state.node_state().await }
    };
    let empty_state_start = if mod_opt.prune {
        EmptyStateStart::RecentBoundary
    } else {
        EmptyStateStart::Genesis
    };
    tasks.spawn(
        "merklized state storage update loop",
        update_state_storage_loop(ds.clone(), get_node_state, light_client, empty_state_start),
    );

    // Archive mode disables the pruner, so nothing else bounds the merklized state tables.
    if mod_opt.archive && !mod_opt.archive_full_state {
        let get_node_state = {
            let state = state.clone();
            async move { state.node_state().await }
        };
        tasks.spawn(
            "archive state garbage collector",
            ArchiveStateGc::new(&mod_opt).run(inner_storage.clone(), get_node_state),
        );
    }

    let port = api.http.port;
    let env_vars = get_public_env_vars().unwrap_or_default();
    let node_cfg = api.public_node_config.as_deref().cloned();
    let modules = espresso_api::OptionalModules {
        submit: api.submit.is_some(),
        config: api.config.is_some(),
        explorer: api.explorer.is_some(),
        ..Default::default()
    };
    let max_connections = api.http.max_connections;
    // Both transports serve the same state; cloning shares the env vars and node config
    // rather than copying the genesis they embed.
    let mut api_state = NodeApiStateImpl::new(ds.clone())
        .with_env_vars(env_vars)
        .with_public_node_config(node_cfg);
    if let Some(ranges_concurrency) = ranges_concurrency {
        api_state = api_state.with_ranges_concurrency(ranges_concurrency);
    }
    let tonic_state = api_state.clone();
    tasks.spawn("API server", async move {
        if let Err(e) = espresso_api::serve_axum(port, api_state, modules, max_connections).await {
            tracing::error!("Axum server error: {}", e);
        }
        anyhow::Ok(())
    });

    if let Some(tonic_port) = api.http.tonic_port {
        tasks.spawn("Tonic gRPC server", async move {
            if let Err(e) = espresso_api::serve_tonic(tonic_port, tonic_state, modules).await {
                tracing::error!("Tonic gRPC server error: {}", e);
            }
        });
    }

    Ok((
        metrics,
        Box::new(ApiEventConsumer::from(ds)),
        Some(RequestResponseStorage::Sql(inner_storage)),
    ))
}

/// Prune archive state older than the light client contract's retained history. Block merkle paths
/// below the marker are rebuilt on request from newer snapshots; fee and reward state is gone.
///
/// TODO: the `hash` rows these deletes orphan are never collected.
#[derive(Clone, Debug)]
pub(crate) struct ArchiveStateGc {
    /// Heights retained regardless of the light client.
    min_retention: u64,
    cfg: PrunerCfg,
}

impl ArchiveStateGc {
    pub(crate) fn new(opt: &SqlOptions) -> Self {
        Self {
            min_retention: opt.archive_state_min_retention,
            cfg: PrunerCfg::from(opt.pruning),
        }
    }

    pub(crate) async fn run(
        self,
        storage: Arc<SqlStorage>,
        instance: impl Future<Output = NodeState>,
    ) {
        let node_state = instance.await;
        loop {
            sleep(self.cfg.interval()).await;
            if let Err(err) = self.collect(&storage, &node_state).await {
                tracing::warn!(
                    err = %format!("{err:#}"),
                    "archive state garbage collection failed"
                );
            }
        }
    }

    async fn collect(&self, storage: &SqlStorage, node_state: &NodeState) -> anyhow::Result<()> {
        let Some(history_start) = node_state
            .light_client_history_start()
            .await
            .context("reading the oldest light client root")?
        else {
            tracing::info!("light client reports no retained history; retaining all state");
            return Ok(());
        };
        storage
            .prune_state_below(history_start, self.min_retention, &self.cfg)
            .await
    }
}

impl RewardMerkleTreeDataSource for SqlStorage {
    async fn load_v1_reward_account_proof(
        &self,
        height: u64,
        account: RewardAccountV1,
    ) -> anyhow::Result<RewardAccountQueryDataV1> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch v1 reward account {account:?}; height {height}"
        ))?;

        let block_height = NodeStorage::<SeqTypes>::block_height(&mut tx)
            .await
            .context("getting block height")? as u64;
        ensure!(
            block_height > 0,
            "cannot get accounts for height {height}: no blocks available"
        );

        // Check if we have the desired state snapshot. If so, we can load the desired accounts
        // directly.
        if height < block_height {
            let (tree, _) = load_v1_reward_accounts(self, height, &[account])
                .await
                .with_context(|| {
                    format!("failed to load v1 reward account {account:?} at height {height}")
                })?;

            let (proof, balance) = RewardAccountProofV1::prove(&tree, account.into())
                .with_context(|| {
                    format!("reward account {account:?} not available at height {height}")
                })?;

            Ok(RewardAccountQueryDataV1 { balance, proof })
        } else {
            bail!(
                "requested height {height} is not yet available (latest block height: \
                 {block_height})"
            );
        }
    }

    fn persist_tree(
        &self,
        height: u64,
        merkle_tree: Vec<u8>,
    ) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move {
            let mut tx = self
                .write()
                .await
                .context("opening transaction for reward merkle tree v2")?;

            tx.upsert(
                "reward_merkle_tree_v2_data",
                ["height", "balances"],
                ["height"],
                [(height as i64, merkle_tree)],
            )
            .await?;

            hotshot_query_service::data_source::Transaction::commit(tx)
                .await
                .context("Transaction to store reward merkle tree v2 failed.")
        }
    }

    fn load_tree(&self, height: u64) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move {
            let mut tx = self
                .read()
                .await
                .context("opening transaction for state update")?;

            let row = sqlx::query(
                r#"
                SELECT balances
                FROM reward_merkle_tree_v2_data
                WHERE height = $1
                "#,
            )
            .bind(height as i64)
            .fetch_optional(tx.as_mut())
            .await?
            .context(format!(
                "No reward merkle tree for height {} in storage",
                height
            ))?;

            row.try_get::<Vec<u8>, _>("balances")
                .context("Missing field balances from row; this should never happen")
        }
    }

    fn load_latest_tree(
        &self,
        height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move {
            let mut tx = self
                .read()
                .await
                .context("opening transaction for load_latest_tree")?;

            let row = sqlx::query(
                r#"
                SELECT balances
                FROM reward_merkle_tree_v2_data
                WHERE height <= $1
                ORDER BY height DESC
                LIMIT 1
                "#,
            )
            .bind(height as i64)
            .fetch_optional(tx.as_mut())
            .await?
            .context(format!(
                "No reward merkle tree at or below height {} in storage",
                height
            ))?;

            row.try_get::<Vec<u8>, _>("balances")
                .context("Missing field balances from row; this should never happen")
        }
    }

    fn persist_proofs(
        &self,
        height: u64,
        proofs: impl Iterator<Item = (Vec<u8>, Vec<u8>)> + Send,
    ) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move {
            let mut iter = proofs.map(|(account, proof)| (height as i64, account, proof));

            loop {
                let mut chunk = Vec::with_capacity(20);

                for _ in 0..20 {
                    let Some(row) = iter.next() else {
                        continue;
                    };
                    chunk.push(row);
                }

                if chunk.is_empty() {
                    break;
                }

                let mut tx = self
                    .write()
                    .await
                    .context("opening transaction for state update")?;

                tokio::spawn(async move {
                    tx.upsert(
                        "reward_merkle_tree_v2_proofs",
                        ["height", "account", "proof"],
                        ["height", "account"],
                        chunk,
                    )
                    .await?;

                    hotshot_query_service::data_source::Transaction::commit(tx)
                        .await
                        .context("Transaction to store reward merkle tree failed.")?;

                    Ok::<_, anyhow::Error>(())
                });
            }
            Ok(())
        }
    }

    /// Load a reward proof for a given height and account.
    ///
    /// For V5+ (epoch rewards), if the requested height is not an epoch boundary,
    /// resolves to the previous epoch's last block. Verifies the boundary block is
    /// V5+ to handle V4→V5 upgrades.
    /// For V4 (per-block rewards), proofs are stored at every block height.
    fn load_proof(
        &self,
        height: u64,
        account: Vec<u8>,
        epoch_height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move {
            let mut tx = self
                .read()
                .await
                .context("opening transaction for load_proof")?;

            let leaf = tx
                .get_leaf(LeafId::<SeqTypes>::from(height as usize))
                .await
                .context(format!("leaf {height} not available"))?;

            // For V5+, resolve to the last epoch boundary where proofs were stored.
            // For V4, proofs exist at every height.
            let proof_height = if leaf.header().version() < EPOCH_REWARD_VERSION
                || is_last_block(height, epoch_height)
            {
                height
            } else {
                let prev_epoch_last_block =
                    epoch_from_block_number(height, epoch_height).saturating_sub(1) * epoch_height;
                let boundary_leaf = tx
                    .get_leaf(LeafId::<SeqTypes>::from(prev_epoch_last_block as usize))
                    .await
                    .context(format!(
                        "leaf at epoch boundary {prev_epoch_last_block} not available"
                    ))?;
                ensure!(
                    boundary_leaf.header().version() >= EPOCH_REWARD_VERSION,
                    "no epoch reward proofs available at boundary {prev_epoch_last_block}"
                );
                prev_epoch_last_block
            };

            sqlx::query_scalar(
                "SELECT proof FROM reward_merkle_tree_v2_proofs WHERE height = $1 AND account = $2",
            )
            .bind(proof_height as i64)
            .bind(account)
            .fetch_optional(tx.as_mut())
            .await?
            .context(format!(
                "Missing proofs at height {proof_height} (resolved from {height})"
            ))
        }
    }

    fn load_latest_proof(
        &self,
        account: Vec<u8>,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move {
            let mut tx = self
                .read()
                .await
                .context("opening transaction for state update")?;

            let row = sqlx::query(
                r#"
                SELECT proof
                FROM reward_merkle_tree_v2_proofs
                WHERE account = $1
                ORDER BY height DESC
                LIMIT 1
                "#,
            )
            .bind(account)
            .fetch_optional(tx.as_mut())
            .await?
            .context("Missing proofs")?;

            row.try_get::<Vec<u8>, _>("proof")
                .context("Missing field proof from row; this should never happen")
        }
    }

    fn garbage_collect(&self, height: u64) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move {
            let batch_size = self.get_pruning_config().unwrap_or_default().batch_size();

            // Postgres supports improved performance by using `FOR UPDATE` in a sub-select to
            // acquire exclusive locks on visited rows, which are later deleted. SQLite does not
            // support this syntax.
            let for_update = match Db::KIND {
                DbKind::Sqlite => "",
                DbKind::Postgres => "FOR UPDATE",
            };

            // Delete batches from the merkle data table until there is nothing left to delete.
            loop {
                let mut tx = self
                    .write()
                    .await
                    .context("opening transaction for state update")?;

                let res = sqlx::query(&format!(
                    "
                    WITH delete_batch AS (
                        SELECT d.height FROM reward_merkle_tree_v2_data AS d
                            WHERE d.height < $1
                            ORDER BY d.height DESC
                            LIMIT $2
                            {for_update}
                    )
                    DELETE FROM reward_merkle_tree_v2_data AS del
                    WHERE del.height IN (SELECT * FROM delete_batch)
                "
                ))
                .bind(height as i64)
                .bind(batch_size as i64)
                .execute(tx.as_mut())
                .await?;

                hotshot_query_service::data_source::Transaction::commit(tx)
                    .await
                    .context(
                        "Transaction to garbage collect reward merkle trees from storage failed.",
                    )?;

                if res.rows_affected() == 0 {
                    break;
                } else {
                    tracing::debug!(
                        "deleted {} rows from reward_merkle_tree_v2_data",
                        res.rows_affected()
                    );
                }
            }

            // Delete batches from the proofs table until there is nothing left to delete.
            loop {
                let mut tx = self
                    .write()
                    .await
                    .context("opening transaction for state update")?;

                let res = sqlx::query(&format!(
                    "
                    WITH delete_batch AS (
                        SELECT d.height, d.account FROM reward_merkle_tree_v2_proofs AS d
                            WHERE d.height < $1
                            ORDER BY d.height, d.account DESC
                            LIMIT $2
                            {for_update}
                    )
                    DELETE FROM reward_merkle_tree_v2_proofs AS del
                    WHERE (del.height, del.account) IN (SELECT * FROM delete_batch)
                ",
                ))
                .bind(height as i64)
                .bind(batch_size as i64)
                .execute(tx.as_mut())
                .await?;

                hotshot_query_service::data_source::Transaction::commit(tx)
                    .await
                    .context("Transaction to garbage collect proofs from storage failed.")?;

                if res.rows_affected() == 0 {
                    break;
                } else {
                    tracing::debug!(
                        "deleted {} rows from reward_merkle_tree_v2_proofs",
                        res.rows_affected()
                    );
                }
            }

            Ok(())
        }
    }

    fn proof_exists(&self, height: u64) -> impl Send + Future<Output = bool> {
        async move {
            let Ok(mut tx) = self.write().await else {
                return false;
            };

            sqlx::query_as(
                r#"
                SELECT EXISTS(
                SELECT 1 FROM reward_merkle_tree_v2_proofs
                WHERE height = $1
                )
                "#,
            )
            .bind(height as i64)
            .fetch_one(tx.as_mut())
            .await
            .ok()
            .unwrap_or((false,))
            .0
        }
    }
    /// Generate and persist reward proofs for the current L1-finalized height.
    ///
    /// For V5+ (epoch rewards), the reward tree is only stored at epoch boundaries.
    /// We resolve to the nearest epoch boundary to load the tree, verify it's V5+
    /// (for V4→V5 upgrades), and store the generated proofs at the same boundary height.
    /// For V4 (per-block rewards), the tree exists at every block height.
    fn persist_reward_proofs(
        &self,
        node_state: &NodeState,
        height: u64,
        version: Version,
    ) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move {
            // Only attempt to persist proofs every 30th block.
            //
            // In tests, we persist proofs at every block height so
            // reward claim tests can query proofs at arbitrary heights. This can be
            // removed once all reward claim tests are updated
            if !cfg!(any(test, feature = "testing"))
                && !(height + node_state.node_id).is_multiple_of(30)
            {
                return Ok(());
            }

            // In tests, we use the current block height directly instead
            // of querying the L1 finalized HotShot height, because test environments don't
            // deploy a light client contract
            let finalized_hotshot_height = if cfg!(any(test, feature = "testing")) {
                height
            } else {
                match node_state.finalized_hotshot_height().await {
                    Ok(h) => h,
                    Err(err) => {
                        tracing::warn!("failed to get finalized hotshot height: {err:#}");
                        return Ok(());
                    },
                }
            };

            // Resolve which height to load the reward tree from.
            // V4: tree exists at every height. V5+: only at epoch boundaries.
            let mut tree_height = finalized_hotshot_height;
            let mut proof_height = finalized_hotshot_height;
            if version >= EPOCH_REWARD_VERSION {
                let epoch_height = node_state
                    .epoch_height
                    .context("epoch_height not set in node state")?;
                if !is_last_block(finalized_hotshot_height, epoch_height) {
                    tree_height = epoch_from_block_number(finalized_hotshot_height, epoch_height)
                        .saturating_sub(1)
                        * epoch_height;
                    if tree_height == 0 {
                        return Ok(());
                    }
                    // A loop replaying history below the boundary may not have its leaf yet, and a
                    // failed read would fail the whole state update.
                    if height < tree_height {
                        return Ok(());
                    }
                    // During V4→V5 upgrades the previous epoch boundary may be V4. The data pruner
                    // may have deleted the boundary leaf, in which case the tree load below decides.
                    let mut tx = self.read().await.context("opening read transaction")?;
                    if let Ok(leaf) = tx
                        .get_leaf(LeafId::<SeqTypes>::from(tree_height as usize))
                        .await
                        && leaf.header().version() < EPOCH_REWARD_VERSION
                    {
                        return Ok(());
                    }
                }
                proof_height = tree_height;
            }

            // The state loop stores the tree when it processes `tree_height`, so a loop replaying
            // history below the L1-finalized height has nothing to load yet.
            if height < tree_height {
                return Ok(());
            }

            if self.proof_exists(proof_height).await {
                return Ok(());
            }

            let permitted_tree = match self.load_reward_merkle_tree_v2(tree_height).await {
                Ok(tree) => tree,
                Err(err) => {
                    tracing::warn!(tree_height, "failed to load reward merkle tree: {err:#}");
                    return Ok(());
                },
            };

            let tree = permitted_tree.tree;
            let iter = tree
                .iter()
                .filter_map(|(account, balance): (&RewardAccountV2, _)| {
                    let proof = espresso_types::v0_6::RewardAccountProofV2::prove(
                        &tree,
                        (*account).into(),
                    )?;
                    let proof = RewardAccountQueryDataV2 {
                        balance: (*balance).into(),
                        proof: proof.0,
                    };
                    Some((
                        bincode::serialize(&account).ok()?,
                        bincode::serialize(&proof).ok()?,
                    ))
                });

            if let Err(err) = self.persist_proofs(proof_height, iter).await {
                tracing::warn!(proof_height, "failed to persist proofs: {err:#}");
            }

            Ok(())
        }
    }
}

impl CatchupStorage for SqlStorage {
    async fn get_reward_accounts_v1(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV1],
    ) -> anyhow::Result<(RewardMerkleTreeV1, Leaf2)> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch v1 reward account {accounts:?}; height {height}"
        ))?;

        let block_height = NodeStorage::<SeqTypes>::block_height(&mut tx)
            .await
            .context("getting block height")? as u64;
        ensure!(
            block_height > 0,
            "cannot get accounts for height {height}: no blocks available"
        );

        // Check if we have the desired state snapshot. If so, we can load the desired accounts
        // directly.
        if height < block_height {
            load_v1_reward_accounts(self, height, accounts).await
        } else {
            let accounts: Vec<_> = accounts
                .iter()
                .map(|acct| RewardAccountV2::from(*acct))
                .collect();
            // If we do not have the exact snapshot we need, we can try going back to the last
            // snapshot we _do_ have and replaying subsequent blocks to compute the desired state.
            let (state, leaf) = reconstruct_state(
                instance,
                self,
                &mut tx,
                block_height - 1,
                view,
                &[],
                &accounts,
            )
            .await?;
            Ok((state.reward_merkle_tree_v1, leaf))
        }
    }

    async fn get_reward_accounts_v2(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV2],
    ) -> anyhow::Result<(RewardMerkleTreeV2, Leaf2)> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch reward account {accounts:?}; height {height}"
        ))?;

        let block_height = NodeStorage::<SeqTypes>::block_height(&mut tx)
            .await
            .context("getting block height")? as u64;
        ensure!(
            block_height > 0,
            "cannot get accounts for height {height}: no blocks available"
        );

        // Check if we have the desired state snapshot. If so, we can load the desired accounts
        // directly.
        if height < block_height {
            load_reward_merkle_tree_v2(self, height)
                .await
                .map(|(permitted_tree, leaf)| (permitted_tree.tree, leaf))
        } else {
            // If we do not have the exact snapshot we need, we can try going back to the last
            // snapshot we _do_ have and replaying subsequent blocks to compute the desired state.
            let (state, leaf) = reconstruct_state(
                instance,
                self,
                &mut tx,
                block_height - 1,
                view,
                &[],
                accounts,
            )
            .await?;
            Ok((state.reward_merkle_tree_v2, leaf))
        }
    }

    async fn get_accounts(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<(FeeMerkleTree, Leaf2)> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch account {accounts:?}; height {height}"
        ))?;

        let block_height = NodeStorage::<SeqTypes>::block_height(&mut tx)
            .await
            .context("getting block height")? as u64;
        ensure!(
            block_height > 0,
            "cannot get accounts for height {height}: no blocks available"
        );

        // Check if we have the desired state snapshot. If so, we can load the desired accounts
        // directly.
        if height < block_height {
            load_accounts(&mut tx, height, accounts).await
        } else {
            // If we do not have the exact snapshot we need, we can try going back to the last
            // snapshot we _do_ have and replaying subsequent blocks to compute the desired state.
            let (state, leaf) = reconstruct_state(
                instance,
                self,
                &mut tx,
                block_height - 1,
                view,
                accounts,
                &[],
            )
            .await?;
            Ok((state.fee_merkle_tree, leaf))
        }
    }

    async fn get_accounts_at(
        &self,
        height: u64,
        root: FeeMerkleCommitment,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<FeeMerkleTree> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch account {accounts:?}; height {height}"
        ))?;
        load_accounts_at(&mut tx, height, root, accounts).await
    }

    async fn get_frontier(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
    ) -> anyhow::Result<BlocksFrontier> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch frontier at height {height}"
        ))?;

        let block_height = NodeStorage::<SeqTypes>::block_height(&mut tx)
            .await
            .context("getting block height")? as u64;
        ensure!(
            block_height > 0,
            "cannot get frontier for height {height}: no blocks available"
        );

        // Check if we have the desired state snapshot. If so, we can load the desired frontier
        // directly.
        if height < block_height {
            load_frontier(&mut tx, height).await
        } else {
            // If we do not have the exact snapshot we need, we can try going back to the last
            // snapshot we _do_ have and replaying subsequent blocks to compute the desired state.
            let (state, _) =
                reconstruct_state(instance, self, &mut tx, block_height - 1, view, &[], &[])
                    .await?;
            match state.block_merkle_tree.lookup(height - 1) {
                LookupResult::Ok(_, proof) => Ok(proof),
                _ => {
                    bail!(
                        "state snapshot {view:?},{height} was found but does not contain frontier \
                         at height {}; this should not be possible",
                        height - 1
                    );
                },
            }
        }
    }

    async fn get_frontier_at(
        &self,
        height: u64,
        root: BlockMerkleCommitment,
    ) -> anyhow::Result<BlocksFrontier> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch frontier at height {height}"
        ))?;
        load_frontier_at(&mut tx, height, root).await
    }

    async fn get_chain_config(
        &self,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        let mut tx = self.read().await.context(format!(
            "opening transaction to fetch chain config {commitment}"
        ))?;
        load_chain_config(&mut tx, commitment).await
    }

    async fn get_leaf_chain(&self, height: u64) -> anyhow::Result<Vec<Leaf2>> {
        let mut tx = self
            .read()
            .await
            .context(format!("opening transaction to fetch leaf at {height}"))?;
        let leaf: Leaf2 = tx
            .get_leaf((height as usize).into())
            .await
            .context(format!("leaf {height} not available"))?
            .leaf()
            .clone();

        // New protocol: cert2 alone proves finality. Return the leaf range up to the finalizing
        // cert2's height.
        if leaf.block_header().version() >= NEW_PROTOCOL_VERSION {
            let cert2: espresso_types::Certificate2<SeqTypes> =
                tx.load_earliest_cert2(height).await?.context(format!(
                    "no cert2 available for new-protocol leaf at {height}"
                ))?;
            let cert2_height = cert2.data.block_number;
            let mut leaves = vec![leaf];
            if height < cert2_height {
                let descendants = tx
                    .get_leaf_range((height as usize + 1)..=cert2_height as usize)
                    .await?;
                leaves.extend(descendants.into_iter().flatten().map(|l| l.leaf().clone()));
            }
            return Ok(leaves);
        }

        // Legacy protocol: build a 3-chain that decides `height`.
        let mut last_leaf = leaf;
        let mut chain = vec![last_leaf.clone()];
        let mut h = height + 1;

        loop {
            let lqd = tx.get_leaf((h as usize).into()).await?;
            let leaf = lqd.leaf();

            if leaf.justify_qc().view_number() == last_leaf.view_number() {
                chain.push(leaf.clone());
            } else {
                h += 1;
                continue;
            }

            // just one away from deciding
            if leaf.view_number() == last_leaf.view_number() + 1 {
                last_leaf = leaf.clone();
                h += 1;
                break;
            }
            h += 1;
            last_leaf = leaf.clone();
        }

        loop {
            let lqd = tx.get_leaf((h as usize).into()).await?;
            let leaf = lqd.leaf();
            if leaf.justify_qc().view_number() == last_leaf.view_number() {
                chain.push(leaf.clone());
                break;
            }
            h += 1;
        }

        Ok(chain)
    }

    async fn load_cert2(
        &self,
        height: u64,
    ) -> anyhow::Result<Option<espresso_types::Certificate2<SeqTypes>>> {
        let mut tx = self
            .read()
            .await
            .context("opening transaction to fetch cert2")?;
        Ok(tx.load_cert2(height).await?)
    }

    async fn get_leaf(&self, height: u64) -> anyhow::Result<Leaf2> {
        let mut tx = self
            .read()
            .await
            .context(format!("opening transaction to fetch leaf at {height}"))?;
        let lqd = tx
            .get_leaf((height as usize).into())
            .await
            .context(format!("leaf {height} not available"))?;
        Ok(lqd.leaf().clone())
    }

    /// The tree is stored only at the heights it changed, so the latest row at or below `height`
    /// may be an older epoch's: it is served only if it matches the header at `height`.
    async fn load_serialized_reward_merkle_tree_v2(&self, height: u64) -> anyhow::Result<Vec<u8>> {
        let tree_bytes = self.load_latest_tree(height).await?;
        let tree_data = bincode::deserialize::<RewardMerkleTreeV2Data>(&tree_bytes)
            .context("Failed to deserialize RewardMerkleTreeV2 from storage")?;
        let tree = PermittedRewardMerkleTreeV2::try_from_kv_set(tree_data.balances)
            .await
            .context("Failed to reconstruct reward merkle tree from storage")?;
        let leaf = CatchupStorage::get_leaf(self, height).await?;
        let either::Either::Right(expected_root) = leaf.block_header().reward_merkle_tree_root()
        else {
            bail!("header {height} has no V2 reward merkle tree");
        };
        ensure!(
            tree.tree.commitment() == expected_root,
            "the reward merkle tree in storage at or below height {height} is from an older epoch"
        );
        Ok(tree_bytes)
    }
}

impl RewardMerkleTreeDataSource for DataSource {
    async fn load_v1_reward_account_proof(
        &self,
        height: u64,
        account: RewardAccountV1,
    ) -> anyhow::Result<RewardAccountQueryDataV1> {
        self.as_ref()
            .load_v1_reward_account_proof(height, account)
            .await
    }

    fn persist_tree(
        &self,
        height: u64,
        merkle_tree: Vec<u8>,
    ) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move { self.as_ref().persist_tree(height, merkle_tree).await }
    }

    fn load_tree(&self, height: u64) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move { self.as_ref().load_tree(height).await }
    }

    fn load_latest_tree(
        &self,
        height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move { self.as_ref().load_latest_tree(height).await }
    }

    fn garbage_collect(&self, height: u64) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move { self.as_ref().garbage_collect(height).await }
    }

    fn persist_proofs(
        &self,
        height: u64,
        proofs: impl Iterator<Item = (Vec<u8>, Vec<u8>)> + Send,
    ) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move { self.as_ref().persist_proofs(height, proofs).await }
    }

    fn load_proof(
        &self,
        height: u64,
        account: Vec<u8>,
        epoch_height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move {
            self.as_ref()
                .load_proof(height, account, epoch_height)
                .await
        }
    }

    fn load_latest_proof(
        &self,
        account: Vec<u8>,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<u8>>> {
        async move { self.as_ref().load_latest_proof(account).await }
    }

    fn proof_exists(&self, height: u64) -> impl Send + Future<Output = bool> {
        async move { self.as_ref().proof_exists(height).await }
    }

    fn persist_reward_proofs(
        &self,
        node_state: &NodeState,
        height: u64,
        version: Version,
    ) -> impl Send + Future<Output = anyhow::Result<()>> {
        async move {
            self.as_ref()
                .persist_reward_proofs(node_state, height, version)
                .await
        }
    }
}

// A module of its own so `RequestStorage` is not in scope where `CatchupStorage` methods are
// called by name on this type.
mod request_storage {
    use async_trait::async_trait;
    use espresso_types::{
        Certificate2, ChainConfig, FeeAccount, FeeMerkleTree, Leaf2, NodeState,
        v0_3::{RewardAccountV1, RewardMerkleTreeV1},
        v0_4::{RewardAccountV2, RewardMerkleTreeV2},
    };
    use hotshot_types::data::{VidShare, ViewNumber};

    use super::SqlStorage;
    use crate::{
        SeqTypes,
        api::{BlocksFrontier, RewardMerkleTreeDataSource},
        catchup::CatchupStorage,
        request_response::data_source::RequestStorage,
    };

    #[async_trait]
    impl RequestStorage for SqlStorage {
        async fn get_accounts(
            &self,
            instance: &NodeState,
            height: u64,
            view: ViewNumber,
            accounts: &[FeeAccount],
        ) -> anyhow::Result<(FeeMerkleTree, Leaf2)> {
            CatchupStorage::get_accounts(self, instance, height, view, accounts).await
        }

        async fn get_reward_accounts_v1(
            &self,
            instance: &NodeState,
            height: u64,
            view: ViewNumber,
            accounts: &[RewardAccountV1],
        ) -> anyhow::Result<(RewardMerkleTreeV1, Leaf2)> {
            CatchupStorage::get_reward_accounts_v1(self, instance, height, view, accounts).await
        }

        async fn get_reward_accounts_v2(
            &self,
            instance: &NodeState,
            height: u64,
            view: ViewNumber,
            accounts: &[RewardAccountV2],
        ) -> anyhow::Result<(RewardMerkleTreeV2, Leaf2)> {
            CatchupStorage::get_reward_accounts_v2(self, instance, height, view, accounts).await
        }

        async fn get_frontier(
            &self,
            instance: &NodeState,
            height: u64,
            view: ViewNumber,
        ) -> anyhow::Result<BlocksFrontier> {
            CatchupStorage::get_frontier(self, instance, height, view).await
        }

        async fn get_chain_config(
            &self,
            commitment: committable::Commitment<ChainConfig>,
        ) -> anyhow::Result<ChainConfig> {
            CatchupStorage::get_chain_config(self, commitment).await
        }

        async fn get_leaf_chain(&self, height: u64) -> anyhow::Result<Vec<Leaf2>> {
            CatchupStorage::get_leaf_chain(self, height).await
        }

        async fn load_cert2(&self, height: u64) -> anyhow::Result<Option<Certificate2<SeqTypes>>> {
            CatchupStorage::load_cert2(self, height).await
        }

        async fn load_tree(&self, height: u64) -> anyhow::Result<Vec<u8>> {
            RewardMerkleTreeDataSource::load_tree(self, height).await
        }

        async fn get_vid_share(&self, block_number: u64) -> anyhow::Result<VidShare> {
            let id = hotshot_query_service::node::BlockId::Number(block_number as usize);
            Ok(SqlStorage::get_vid_share::<SeqTypes>(self, id).await?)
        }
    }
}

impl CatchupStorage for DataSource {
    async fn get_accounts(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<(FeeMerkleTree, Leaf2)> {
        self.as_ref()
            .get_accounts(instance, height, view, accounts)
            .await
    }

    async fn get_accounts_at(
        &self,
        height: u64,
        root: FeeMerkleCommitment,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<FeeMerkleTree> {
        self.as_ref().get_accounts_at(height, root, accounts).await
    }

    async fn get_reward_accounts_v2(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV2],
    ) -> anyhow::Result<(RewardMerkleTreeV2, Leaf2)> {
        self.as_ref()
            .get_reward_accounts_v2(instance, height, view, accounts)
            .await
    }

    async fn get_reward_accounts_v1(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV1],
    ) -> anyhow::Result<(RewardMerkleTreeV1, Leaf2)> {
        self.as_ref()
            .get_reward_accounts_v1(instance, height, view, accounts)
            .await
    }

    async fn get_frontier(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
    ) -> anyhow::Result<BlocksFrontier> {
        self.as_ref().get_frontier(instance, height, view).await
    }

    async fn get_frontier_at(
        &self,
        height: u64,
        root: BlockMerkleCommitment,
    ) -> anyhow::Result<BlocksFrontier> {
        self.as_ref().get_frontier_at(height, root).await
    }

    async fn get_chain_config(
        &self,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        self.as_ref().get_chain_config(commitment).await
    }
    async fn get_leaf_chain(&self, height: u64) -> anyhow::Result<Vec<Leaf2>> {
        self.as_ref().get_leaf_chain(height).await
    }

    async fn load_cert2(
        &self,
        height: u64,
    ) -> anyhow::Result<Option<espresso_types::Certificate2<SeqTypes>>> {
        self.as_ref().load_cert2(height).await
    }

    async fn get_leaf(&self, height: u64) -> anyhow::Result<Leaf2> {
        self.as_ref().get_leaf(height).await
    }

    async fn load_serialized_reward_merkle_tree_v2(&self, height: u64) -> anyhow::Result<Vec<u8>> {
        self.as_ref()
            .load_serialized_reward_merkle_tree_v2(height)
            .await
    }
}

#[async_trait]
impl ChainConfigPersistence for Transaction<Write> {
    async fn insert_chain_config(&mut self, chain_config: ChainConfig) -> anyhow::Result<()> {
        let commitment = chain_config.commitment();
        let data = bincode::serialize(&chain_config)?;
        self.upsert(
            "chain_config",
            ["commitment", "data"],
            ["commitment"],
            [(commitment.to_string(), data)],
        )
        .await
    }
}

impl crate::api::data_source::PruningDataSource for SqlStorage {
    async fn get_oldest_block(
        &self,
    ) -> anyhow::Result<Option<hotshot_query_service::availability::BlockQueryData<SeqTypes>>> {
        let mut tx = self
            .read()
            .await
            .context("opening transaction to fetch oldest block")?;
        let row = sqlx::query("SELECT MIN(height) AS height FROM header")
            .fetch_one(tx.as_mut())
            .await
            .context("failed to query oldest block height")?;
        let height: Option<i64> = row.try_get("height")?;
        match height {
            None => Ok(None),
            Some(h) => {
                let h = usize::try_from(h).context("block height out of range")?;
                Ok(Some(
                    tx.get_block(hotshot_query_service::availability::BlockId::<SeqTypes>::from(h))
                        .await
                        .context(format!("block {h} not available"))?,
                ))
            },
        }
    }

    async fn get_oldest_leaf(
        &self,
    ) -> anyhow::Result<Option<hotshot_query_service::availability::LeafQueryData<SeqTypes>>> {
        let mut tx = self
            .read()
            .await
            .context("opening transaction to fetch oldest leaf")?;
        let row = sqlx::query("SELECT MIN(height) AS height FROM leaf2")
            .fetch_one(tx.as_mut())
            .await
            .context("failed to query oldest leaf height")?;
        let height: Option<i64> = row.try_get("height")?;
        match height {
            None => Ok(None),
            Some(h) => {
                let h = usize::try_from(h).context("leaf height out of range")?;
                Ok(Some(
                    tx.get_leaf(LeafId::<SeqTypes>::from(h))
                        .await
                        .context(format!("leaf {h} not available"))?,
                ))
            },
        }
    }
}

impl crate::api::data_source::DatabaseMetadataSource for SqlStorage {
    async fn get_table_sizes(&self) -> anyhow::Result<Vec<espresso_api::v1::database::TableSize>> {
        let mut tx = self
            .read()
            .await
            .context("opening transaction to fetch table sizes")?;

        if Db::KIND == DbKind::Postgres {
            let query = r#"
                SELECT
                    schemaname || '.' || relname AS table_name,
                    n_live_tup AS row_count,
                    pg_total_relation_size(relid) AS total_size_bytes
                FROM pg_stat_user_tables
                ORDER BY pg_total_relation_size(relid) DESC
            "#;

            let rows = sqlx::query(query)
                .fetch_all(tx.as_mut())
                .await
                .context("failed to query table sizes")?;

            let mut table_sizes = Vec::new();
            for row in rows {
                let table_name: String = row.try_get("table_name")?;
                let row_count: i64 = row.try_get("row_count").unwrap_or(-1);
                let total_size_bytes: Option<i64> = row.try_get("total_size_bytes").ok();

                table_sizes.push(espresso_api::v1::database::TableSize {
                    table_name,
                    row_count,
                    total_size_bytes,
                });
            }

            Ok(table_sizes)
        } else {
            // First, get all table names from sqlite_master
            let table_names_query = r#"
                SELECT name
                FROM sqlite_master
                WHERE type = 'table'
                AND name NOT LIKE 'sqlite_%'
                ORDER BY name
            "#;

            let table_rows = sqlx::query(table_names_query)
                .fetch_all(tx.as_mut())
                .await
                .context("failed to query table names")?;

            let mut table_sizes = Vec::new();

            // For each table, get the row count
            for row in table_rows {
                let table_name: String = row.try_get("name")?;

                // Run SELECT COUNT(*) for this specific table
                // Use format! here since we need to dynamically insert the table name
                let count_query = format!("SELECT COUNT(*) as count FROM \"{}\"", table_name);
                let count_row = sqlx::query(&count_query)
                    .fetch_one(tx.as_mut())
                    .await
                    .context(format!(
                        "failed to query row count for table {}",
                        table_name
                    ))?;

                let row_count: i64 = count_row.try_get("count").unwrap_or(0);

                table_sizes.push(espresso_api::v1::database::TableSize {
                    table_name,
                    row_count,
                    total_size_bytes: None,
                });
            }

            Ok(table_sizes)
        }
    }

    async fn get_migration_status(
        &self,
    ) -> anyhow::Result<Vec<espresso_api::v1::database::MigrationStatus>> {
        let mut tx = self
            .read()
            .await
            .context("opening transaction to fetch migration status")?;

        type MigrationStatusRow = (
            String,
            chrono::DateTime<chrono::Utc>,
            Option<chrono::DateTime<chrono::Utc>>,
            Option<i64>,
        );

        let rows: Vec<MigrationStatusRow> = sqlx::query_as(
            "SELECT name, started_at, completed_at, last_offset FROM deferred_migrations ORDER BY \
             started_at",
        )
        .fetch_all(tx.as_mut())
        .await
        .context("failed to query deferred_migrations")?;

        Ok(rows
            .into_iter()
            .map(|(name, started_at, completed_at, last_offset)| {
                espresso_api::v1::database::MigrationStatus {
                    name,
                    started_at,
                    completed_at,
                    last_offset,
                }
            })
            .collect())
    }
}

impl crate::api::data_source::DatabaseMetadataSource for DataSource {
    async fn get_table_sizes(&self) -> anyhow::Result<Vec<espresso_api::v1::database::TableSize>> {
        self.as_ref().get_table_sizes().await
    }

    async fn get_migration_status(
        &self,
    ) -> anyhow::Result<Vec<espresso_api::v1::database::MigrationStatus>> {
        self.as_ref().get_migration_status().await
    }
}

impl crate::api::data_source::PruningDataSource for DataSource {
    async fn get_oldest_block(
        &self,
    ) -> anyhow::Result<Option<hotshot_query_service::availability::BlockQueryData<SeqTypes>>> {
        self.as_ref().get_oldest_block().await
    }

    async fn get_oldest_leaf(
        &self,
    ) -> anyhow::Result<Option<hotshot_query_service::availability::LeafQueryData<SeqTypes>>> {
        self.as_ref().get_oldest_leaf().await
    }
}

async fn load_frontier<Mode: TransactionMode>(
    tx: &mut Transaction<Mode>,
    height: u64,
) -> anyhow::Result<BlocksFrontier> {
    tx.get_path(
        Snapshot::<SeqTypes, BlockMerkleTree, { BlockMerkleTree::ARITY }>::Index(height),
        height
            .checked_sub(1)
            .ok_or(anyhow::anyhow!("Subtract with overflow ({height})!"))?,
    )
    .await
    .context(format!("fetching frontier at height {height}"))
}

/// Load the frontier from the snapshot at `height` whose root is `root`, without reading the header
/// at `height`.
async fn load_frontier_at<Mode: TransactionMode>(
    tx: &mut Transaction<Mode>,
    height: u64,
    root: BlockMerkleCommitment,
) -> anyhow::Result<BlocksFrontier> {
    tx.get_path_at::<SeqTypes, BlockMerkleTree, { BlockMerkleTree::ARITY }>(
        height,
        root,
        height
            .checked_sub(1)
            .ok_or(anyhow::anyhow!("Subtract with overflow ({height})!"))?,
    )
    .await
    .context(format!("fetching frontier at height {height}"))
}

async fn load_v1_reward_accounts(
    db: &SqlStorage,
    height: u64,
    accounts: &[RewardAccountV1],
) -> anyhow::Result<(RewardMerkleTreeV1, Leaf2)> {
    // Open a new read transaction to get the leaf
    let mut tx = db
        .read()
        .await
        .with_context(|| "failed to open read transaction")?;

    // Get the leaf from the database
    let leaf = tx
        .get_leaf(LeafId::<SeqTypes>::from(height as usize))
        .await
        .context(format!("leaf {height} not available"))?;
    let header = leaf.header();

    if header.version() < EPOCH_VERSION || header.version() >= DRB_AND_HEADER_UPGRADE_VERSION {
        return Ok((
            RewardMerkleTreeV1::new(REWARD_MERKLE_TREE_V1_HEIGHT),
            leaf.leaf().clone(),
        ));
    }

    // Get the merkle root from the header and create a snapshot from it
    let merkle_root = header.reward_merkle_tree_root().unwrap_left();
    let mut snapshot = RewardMerkleTreeV1::from_commitment(merkle_root);

    // Create a bounded join set with 10 concurrent tasks
    let mut join_set = BoundedJoinSet::new(10);

    // Create a map from task ID to account
    let mut task_id_to_account = HashMap::new();

    // Loop through each account, spawning a task to get the path for the account
    for account in accounts {
        // Clone things we will need in the closure
        let db_clone = db.clone();
        let account_clone = *account;
        let header_height = header.height();

        // Create the closure that will get the path for the account
        let func = async move {
            // Open a new transaction
            let mut tx = db_clone
                .read()
                .await
                .with_context(|| "failed to open read transaction")?;

            // Get the path for the account
            let proof = tx
                .get_path(
                    Snapshot::<SeqTypes, RewardMerkleTreeV1, { RewardMerkleTreeV1::ARITY }>::Index(
                        header_height,
                    ),
                    account_clone,
                )
                .await
                .with_context(|| {
                    format!(
                        "failed to get path for v1 reward account {account_clone:?}; height \
                         {height}"
                    )
                })?;

            Ok::<_, anyhow::Error>(proof)
        };

        // Spawn the task
        let id = join_set.spawn(func).id();

        // Add the task ID to the account map
        task_id_to_account.insert(id, account);
    }

    // Wait for each task to complete
    while let Some(result) = join_set.join_next_with_id().await {
        // Get the inner result (past the join error)
        let (id, result) = result.with_context(|| "failed to join task")?;

        // Get the proof from the result
        let proof = result?;

        // Get the account from the task ID to account map
        let account = task_id_to_account
            .remove(&id)
            .with_context(|| "task ID for spawned task not found")?;

        match proof.proof.first().with_context(|| {
            format!("empty proof for v1 reward account {account:?}; height {height}")
        })? {
            MerkleNode::Leaf { pos, elem, .. } => {
                snapshot.remember(*pos, *elem, proof)?;
            },
            MerkleNode::Empty => {
                snapshot.non_membership_remember(*account, proof)?;
            },
            _ => {
                bail!("invalid proof for v1 reward account {account:?}; height {height}");
            },
        }
    }

    Ok((snapshot, leaf.leaf().clone()))
}

/// Loads reward accounts for new reward merkle tree (V4).
async fn load_reward_merkle_tree_v2(
    db: &SqlStorage,
    height: u64,
) -> anyhow::Result<(PermittedRewardMerkleTreeV2, Leaf2)> {
    // Open a new read transaction to get the leaf
    let mut tx = db
        .read()
        .await
        .with_context(|| "failed to open read transaction")?;

    // Get the leaf from the database
    let leaf = tx
        .get_leaf(LeafId::<SeqTypes>::from(height as usize))
        .await
        .with_context(|| format!("leaf {height} not available"))?;

    let snapshot = db.load_reward_merkle_tree_v2(height).await?;

    Ok((snapshot, leaf.leaf().clone()))
}

async fn load_accounts<Mode: TransactionMode>(
    tx: &mut Transaction<Mode>,
    height: u64,
    accounts: &[FeeAccount],
) -> anyhow::Result<(FeeMerkleTree, Leaf2)> {
    let leaf = tx
        .get_leaf(LeafId::<SeqTypes>::from(height as usize))
        .await
        .context(format!("leaf {height} not available"))?;
    let header = leaf.header();
    let snapshot =
        load_accounts_at(tx, header.height(), header.fee_merkle_tree_root(), accounts).await?;
    Ok((snapshot, leaf.leaf().clone()))
}

/// Load `accounts` from the snapshot at `height` whose root is `root`, without reading the leaf or
/// header at `height`.
async fn load_accounts_at<Mode: TransactionMode>(
    tx: &mut Transaction<Mode>,
    height: u64,
    root: FeeMerkleCommitment,
    accounts: &[FeeAccount],
) -> anyhow::Result<FeeMerkleTree> {
    let mut snapshot = FeeMerkleTree::from_commitment(root);
    for account in accounts {
        let proof = tx
            .get_path_at::<SeqTypes, FeeMerkleTree, { FeeMerkleTree::ARITY }>(
                height, root, *account,
            )
            .await
            .context(format!("fetching account {account}; height {height}"))?;
        match proof.proof.first().context(format!(
            "empty proof for account {account}; height {height}"
        ))? {
            MerkleNode::Leaf { pos, elem, .. } => {
                snapshot.remember(*pos, *elem, proof)?;
            },
            MerkleNode::Empty => {
                snapshot.non_membership_remember(*account, proof)?;
            },
            _ => {
                bail!("Invalid proof");
            },
        }
    }

    Ok(snapshot)
}

async fn load_chain_config<Mode: TransactionMode>(
    tx: &mut Transaction<Mode>,
    commitment: Commitment<ChainConfig>,
) -> anyhow::Result<ChainConfig> {
    let (data,) = query_as::<(Vec<u8>,)>("SELECT data from chain_config where commitment = $1")
        .bind(commitment.to_string())
        .fetch_one(tx.as_mut())
        .await
        .unwrap();

    bincode::deserialize(&data[..]).context("failed to deserialize")
}

/// Reconstructs the `ValidatedState` from a specific block height up to a given view.
///
/// This loads all required fee and reward accounts into the Merkle tree before applying the
/// State Transition Function (STF), preventing recursive catchup during STF replay.
///
/// Note: Even if the primary goal is to catch up the block Merkle tree,
/// fee and reward header dependencies must still be present beforehand
/// This is because reconstructing the `ValidatedState` involves replaying the STF over a
/// range of leaves, and the STF requires all associated data to be present in the `ValidatedState`;
/// otherwise, it will attempt to trigger catchup itself.
#[tracing::instrument(skip(instance, db, tx))]
pub(crate) async fn reconstruct_state<Mode: TransactionMode>(
    instance: &NodeState,
    db: &SqlStorage,
    tx: &mut Transaction<Mode>,
    from_height: u64,
    to_view: ViewNumber,
    fee_accounts: &[FeeAccount],
    reward_accounts: &[RewardAccountV2],
) -> anyhow::Result<(ValidatedState, Leaf2)> {
    tracing::info!("attempting to reconstruct fee state");
    let from_leaf = tx
        .get_leaf((from_height as usize).into())
        .await
        .context(format!("leaf {from_height} not available"))?;
    let from_leaf: Leaf2 = from_leaf.leaf().clone();
    ensure!(
        from_leaf.view_number() < to_view,
        "state reconstruction: starting state {:?} must be before ending state {to_view:?}",
        from_leaf.view_number(),
    );

    // Get the sequence of headers we will be applying to compute the latest state.
    let mut leaves = VecDeque::new();
    let to_leaf = get_leaf_from_proposal(tx, "view = $1", &(to_view.u64() as i64))
        .await
        .context(format!(
            "unable to reconstruct state because leaf {to_view:?} is not available"
        ))?;
    let mut parent = to_leaf.parent_commitment();
    tracing::debug!(?to_leaf, ?parent, view = ?to_view, "have required leaf");
    leaves.push_front(to_leaf.clone());
    while parent != Committable::commit(&from_leaf) {
        let leaf = get_leaf_from_proposal(tx, "leaf_hash = $1", &parent.to_string())
            .await
            .context(format!(
                "unable to reconstruct state because leaf {parent} is not available"
            ))?;
        parent = leaf.parent_commitment();
        tracing::debug!(?leaf, ?parent, "have required leaf");
        leaves.push_front(leaf);
    }

    // Get the initial state.
    let mut parent = from_leaf;
    let mut state = ValidatedState::from_header(parent.block_header());

    // Pre-load the state with the accounts we care about to ensure they will be present in the
    // final state.
    let mut catchup = NullStateCatchup::default();

    let mut fee_accounts = fee_accounts.iter().copied().collect::<HashSet<_>>();
    // Add in all the accounts we will need to replay any of the headers, to ensure that we don't
    // need to do catchup recursively.
    tracing::info!(
        "reconstructing fee accounts state for from height {from_height} to view {to_view}"
    );

    let dependencies =
        fee_header_dependencies(&mut catchup, tx, instance, &parent, &leaves).await?;
    fee_accounts.extend(dependencies);
    let fee_accounts = fee_accounts.into_iter().collect::<Vec<_>>();
    state.fee_merkle_tree = load_accounts(tx, from_height, &fee_accounts)
        .await
        .context("unable to reconstruct state because accounts are not available at origin")?
        .0;
    ensure!(
        state.fee_merkle_tree.commitment() == parent.block_header().fee_merkle_tree_root(),
        "loaded fee state does not match parent header"
    );

    tracing::info!(
        "reconstructing reward accounts for from height {from_height} to view {to_view}"
    );

    let mut reward_accounts = reward_accounts.iter().copied().collect::<HashSet<_>>();

    // Collect all reward account dependencies needed for replaying the STF.
    // These accounts must be preloaded into the reward Merkle tree to prevent recursive catchups.
    let dependencies = reward_header_dependencies(instance, &leaves).await?;
    reward_accounts.extend(dependencies);
    let reward_accounts = reward_accounts.into_iter().collect::<Vec<_>>();

    // Load all required reward accounts and update the reward Merkle tree.
    match parent.block_header().reward_merkle_tree_root() {
        either::Either::Left(expected_root) => {
            let accts = reward_accounts
                .into_iter()
                .map(RewardAccountV1::from)
                .collect::<Vec<_>>();
            state.reward_merkle_tree_v1 = load_v1_reward_accounts(db, from_height, &accts)
                .await
                .context(
                    "unable to reconstruct state because v1 reward accounts are not available at \
                     origin",
                )?
                .0;
            ensure!(
                state.reward_merkle_tree_v1.commitment() == expected_root,
                "loaded v1 reward state does not match parent header"
            );
        },
        either::Either::Right(expected_root) => {
            let version = parent.block_header().version();
            let epoch_height = instance
                .epoch_height
                .context("epoch_height not set but parent has V2 reward tree")?;

            // V5+ only stores the reward merkle tree at epoch's last block, so if we're
            // reconstructing at a non boundary height there is no row at from_height.
            // Load the latest tree instead
            // the commitment check below will catch it if
            // we ended up with a tree from an older epoch.
            // But this should never happen as we don't garbage collect the latest tree
            if version >= EPOCH_REWARD_VERSION && !is_last_block(from_height, epoch_height) {
                let tree = db
                    .load_latest_reward_merkle_tree_v2(from_height)
                    .await
                    .context("RewardMerkleTreeV2 not available at or below origin")?;
                state.reward_merkle_tree_v2 = tree.tree;
            } else {
                state.reward_merkle_tree_v2 = load_reward_merkle_tree_v2(db, from_height)
                    .await
                    .context("RewardMerkleTreeV2 not available at origin")?
                    .0
                    .tree;
            }
            ensure!(
                state.reward_merkle_tree_v2.commitment() == expected_root,
                "loaded reward state does not match parent header"
            );
        },
    }

    // We need the blocks frontier as well, to apply the STF.
    let frontier = load_frontier(tx, from_height)
        .await
        .context("unable to reconstruct state because frontier is not available at origin")?;
    match frontier
        .proof
        .first()
        .context("empty proof for frontier at origin")?
    {
        MerkleNode::Leaf { pos, elem, .. } => state
            .block_merkle_tree
            .remember(*pos, *elem, frontier)
            .context("failed to remember frontier")?,
        _ => bail!("invalid frontier proof"),
    }

    // Apply subsequent headers to compute the later state.
    for proposal in leaves {
        state = compute_state_update(&state, instance, &catchup, &parent, &proposal)
            .await
            .context(format!(
                "unable to reconstruct state because state update {} failed",
                proposal.height(),
            ))?
            .0;
        parent = proposal;
    }

    tracing::info!(from_height, ?to_view, "successfully reconstructed state");
    Ok((state, to_leaf))
}

/// Get the dependencies needed to apply the STF to the given list of headers.
///
/// Returns
/// * A state catchup implementation seeded with all the chain configs required to apply the headers
///   in `leaves`
/// * The set of accounts that must be preloaded to apply these headers
async fn fee_header_dependencies<Mode: TransactionMode>(
    catchup: &mut NullStateCatchup,
    tx: &mut Transaction<Mode>,
    instance: &NodeState,
    mut parent: &Leaf2,
    leaves: impl IntoIterator<Item = &Leaf2>,
) -> anyhow::Result<HashSet<FeeAccount>> {
    let mut accounts = HashSet::default();

    for proposal in leaves {
        let header = proposal.block_header();
        let height = header.height();
        let view = proposal.view_number();
        tracing::debug!(height, ?view, "fetching dependencies for proposal");

        let header_cf = header.chain_config();
        let chain_config = if header_cf.commit() == instance.chain_config.commit() {
            instance.chain_config
        } else {
            match header_cf.resolve() {
                Some(cf) => cf,
                None => {
                    tracing::info!(
                        height,
                        ?view,
                        commit = %header_cf.commit(),
                        "chain config not available, attempting to load from storage",
                    );
                    let cf = load_chain_config(tx, header_cf.commit())
                        .await
                        .context(format!(
                            "loading chain config {} for header {},{:?}",
                            header_cf.commit(),
                            header.height(),
                            proposal.view_number()
                        ))?;

                    // If we had to fetch a chain config now, store it in the catchup implementation
                    // so the STF will be able to look it up later.
                    catchup.add_chain_config(cf);
                    cf
                },
            }
        };

        accounts.insert(chain_config.fee_recipient);
        accounts.extend(
            get_l1_deposits(instance, header, parent, chain_config.fee_contract)
                .await
                .into_iter()
                .map(|fee| fee.account()),
        );
        accounts.extend(header.fee_info().accounts());
        parent = proposal;
    }
    Ok(accounts)
}

/// Identifies all reward accounts required to replay the State Transition Function
/// for the given leaf proposals. These accounts should be present in the Merkle tree
/// *before* applying the STF to avoid recursive catchup (i.e., STF triggering another catchup).
async fn reward_header_dependencies(
    instance: &NodeState,
    leaves: impl IntoIterator<Item = &Leaf2>,
) -> anyhow::Result<HashSet<RewardAccountV2>> {
    let mut reward_accounts = HashSet::default();
    let epoch_height = instance.epoch_height;

    let Some(epoch_height) = epoch_height else {
        tracing::info!("epoch height is not set. returning empty reward_header_dependencies");
        return Ok(HashSet::new());
    };

    let coordinator = instance.coordinator.clone();
    let first_epoch = coordinator.membership().first_epoch();
    // add all the chain configs needed to apply STF to headers to the catchup
    for proposal in leaves {
        let header = proposal.block_header();

        let height = header.height();
        let view = proposal.view_number();
        tracing::debug!(height, ?view, "fetching dependencies for proposal");

        let version = header.version();
        // Skip if version is less than epoch version
        if version < EPOCH_VERSION || version >= EPOCH_REWARD_VERSION {
            continue;
        }

        let proposal_epoch = EpochNumber::new(epoch_from_block_number(height, epoch_height));

        // reward distribution starts third epoch onwards
        if proposal_epoch <= first_epoch + 1 {
            continue;
        }

        let epoch_membership = match coordinator.membership_for_epoch(Some(proposal_epoch)) {
            Ok(e) => e,
            Err(err) => {
                tracing::info!(
                    "failed to get membership for epoch={proposal_epoch:?}. err={err:#}"
                );

                coordinator
                    .wait_for_catchup(proposal_epoch)
                    .await
                    .context(format!("failed to catchup for epoch={proposal_epoch}"))?
            },
        };

        let snapshot = epoch_membership
            .snapshot()
            .with_context(|| format!("no committee for epoch={proposal_epoch}"))?;
        let leader = snapshot.lookup_leader(proposal.view_number())?;
        let validator = snapshot.validator_config(&leader)?;

        reward_accounts.insert(RewardAccountV2(validator.account));

        let delegators: Vec<RewardAccountV2> = validator
            .delegators
            .keys()
            .map(|d| RewardAccountV2(*d))
            .collect();

        reward_accounts.extend(delegators);
    }
    Ok(reward_accounts)
}

async fn get_leaf_from_proposal<Mode, P>(
    tx: &mut Transaction<Mode>,
    where_clause: &str,
    param: P,
) -> anyhow::Result<Leaf2>
where
    P: Type<Db> + for<'q> Encode<'q, Db>,
{
    let (data,) = query_as::<(Vec<u8>,)>(&format!(
        "SELECT data FROM quorum_proposals2 WHERE {where_clause} LIMIT 1",
    ))
    .bind(param)
    .fetch_one(tx.as_mut())
    .await?;
    let proposal: Proposal<SeqTypes, QuorumProposalWrapper<SeqTypes>> =
        bincode::deserialize(&data)?;
    Ok(Leaf2::from_quorum_proposal(&proposal.data))
}

#[cfg(any(test, feature = "testing"))]
pub(crate) mod impl_testable_data_source {

    use hotshot_query_service::{
        availability::LeafQueryData,
        data_source::{
            Transaction as _,
            storage::{
                MerklizedStateHeightStorage as _, SerializableRetry as _,
                UpdateAvailabilityStorage as _, pruning::PrunedHeightStorage as _,
                sql::StorageConnectionType,
            },
        },
        merklized_state::UpdateStateData,
        serializable_retry,
    };

    pub(crate) use super::super::persistence::tmp_options;
    use super::{super::qs::testing::TmpDb, *};
    use crate::api::{
        self,
        data_source::testing::{
            DeferredMigration, TestableSequencerDataSource, TestableSqlDataSource,
        },
    };

    #[async_trait]
    impl TestableSequencerDataSource for DataSource {
        type Storage = TmpDb;

        async fn create_storage() -> Self::Storage {
            TmpDb::init().await
        }

        fn persistence_options(storage: &Self::Storage) -> Self::Options {
            Options::try_from(tmp_options(storage)).unwrap()
        }

        fn leaf_only_ds_options(
            storage: &Self::Storage,
            opt: api::Options,
        ) -> anyhow::Result<api::Options> {
            let mut ds_opts = tmp_options(storage);
            ds_opts.lightweight = true;
            Ok(opt.query_sql(Default::default(), ds_opts))
        }

        fn options(storage: &Self::Storage, opt: api::Options) -> api::Options {
            opt.query_sql(Query::default(), tmp_options(storage))
        }
    }

    #[async_trait]
    impl TestableSqlDataSource for DataSource {
        type SqlStorage = SqlStorage;

        fn tmp_options(storage: &Self::Storage) -> SqlOptions {
            tmp_options(storage)
        }

        fn sql_persistence_options(opt: SqlOptions) -> Self::Options {
            Options::try_from(opt).unwrap()
        }

        async fn on_request_data_source(
            storage: &Self::Storage,
            provider: Provider,
        ) -> anyhow::Result<Self> {
            Config::try_from(&tmp_options(storage))?
                .builder(provider)
                .await?
                .disable_proactive_fetching()
                .disable_aggregator()
                .build()
                .await
        }

        async fn prune_state_below(ds: &Self, height: u64, cfg: &PrunerCfg) -> anyhow::Result<()> {
            let sql: &SqlStorage = ds.as_ref();
            sql.prune_state_below(height, 0, cfg).await
        }

        async fn connect(
            opt: &SqlOptions,
            connection_type: StorageConnectionType,
        ) -> anyhow::Result<SqlStorage> {
            Ok(SqlStorage::connect(Config::try_from(opt)?, connection_type).await?)
        }

        async fn local_block_height(db: &SqlStorage) -> anyhow::Result<u64> {
            let mut tx = db.read().await?;
            Ok(NodeStorage::<SeqTypes>::block_height(&mut tx).await? as u64)
        }

        // Read in this order because every move of the pruned height, up by a batch or down by
        // the startup repair, lands below the head, and the head only grows, so a head read
        // after the pruned height is always above it.
        async fn state_heights(db: &SqlStorage) -> anyhow::Result<(Option<u64>, u64)> {
            let mut tx = db.read().await?;
            let pruned = tx.load_state_pruned_height().await?;
            let head = tx.get_last_state_height().await? as u64;
            Ok((pruned, head))
        }

        async fn state_start_height(db: &SqlStorage) -> anyhow::Result<Option<u64>> {
            Ok(db.read().await?.load_state_start_height().await?)
        }

        async fn lowest_block_tree_node(db: &SqlStorage) -> anyhow::Result<u64> {
            let mut tx = db.read().await?;
            let lowest: Option<i64> =
                sqlx::query_scalar("SELECT min(created) FROM block_merkle_tree")
                    .fetch_one(tx.as_mut())
                    .await?;
            lowest
                .map(|height| height as u64)
                .context("no block tree nodes")
        }

        async fn insert_leaf(
            db: &SqlStorage,
            leaf: &LeafQueryData<SeqTypes>,
        ) -> anyhow::Result<()> {
            let mut tx = db.write().await?;
            tx.insert_leaf(leaf).await?;
            tx.commit().await
        }

        async fn set_last_state_height(db: &SqlStorage, height: u64) -> anyhow::Result<()> {
            let mut tx = db.write().await?;
            UpdateStateData::<SeqTypes, BlockMerkleTree, { BlockMerkleTree::ARITY }>::set_last_state_height(
                &mut tx,
                height as usize,
            )
            .await?;
            tx.commit().await
        }

        async fn stamp_state_pruned_height(db: &SqlStorage, height: u64) -> anyhow::Result<()> {
            let mut tx = db.write().await?;
            tx.upsert(
                "pruned_height",
                ["id", "last_height"],
                ["id"],
                [(Transaction::<Write>::STATE_PRUNED_HEIGHT_ID, height as i64)],
            )
            .await?;
            tx.commit().await
        }

        // The node may be writing new blocks meanwhile, so a serialization conflict is retried.
        async fn prune_data_to(db: &SqlStorage, pruned: u64) -> anyhow::Result<()> {
            serializable_retry!(db, || async {
                let mut tx = db.write().await?;
                tx.upsert(
                    "pruned_height",
                    ["id", "last_height"],
                    ["id"],
                    [(Transaction::<Write>::PRUNED_HEIGHT_ID, pruned as i64)],
                )
                .await?;
                for statement in [
                    "DELETE FROM transactions WHERE block_height <= $1",
                    "DELETE FROM leaf2 WHERE height <= $1",
                    "DELETE FROM header WHERE height <= $1",
                ] {
                    sqlx::query(statement)
                        .bind(pruned as i64)
                        .execute(tx.as_mut())
                        .await?;
                }
                tx.commit().await
            })
            .await
        }

        async fn delete_data_between(db: &SqlStorage, from: u64, to: u64) -> anyhow::Result<()> {
            let mut tx = db.write().await?;
            for statement in [
                "DELETE FROM transactions WHERE block_height > $1 AND block_height <= $2",
                "DELETE FROM leaf2 WHERE height > $1 AND height <= $2",
                "DELETE FROM header WHERE height > $1 AND height <= $2",
            ] {
                sqlx::query(statement)
                    .bind(from as i64)
                    .bind(to as i64)
                    .execute(tx.as_mut())
                    .await?;
            }
            tx.commit().await
        }

        async fn count_leaves_up_to(db: &SqlStorage, height: u64) -> anyhow::Result<i64> {
            let mut tx = db.read().await?;
            Ok(
                sqlx::query_scalar("SELECT count(*) FROM leaf2 WHERE height <= $1")
                    .bind(height as i64)
                    .fetch_one(tx.as_mut())
                    .await?,
            )
        }

        async fn data_heights(db: &SqlStorage) -> anyhow::Result<(Option<u64>, Option<u64>)> {
            let mut tx = db.read().await?;
            let pruned = tx.load_pruned_height().await?;
            let lowest: Option<i64> =
                sqlx::query_scalar("SELECT min(height) FROM leaf2 WHERE height > 0")
                    .fetch_one(tx.as_mut())
                    .await?;
            Ok((pruned, lowest.map(|height| height as u64)))
        }

        async fn insert_deferred_migration(
            db: &SqlStorage,
            migration: DeferredMigration,
        ) -> anyhow::Result<()> {
            let DeferredMigration {
                name,
                started_at,
                completed_at,
                last_offset,
            } = migration;
            let mut tx = db.write().await?;
            sqlx::query(
                "INSERT INTO deferred_migrations (name, started_at, completed_at, last_offset)
                 VALUES ($1, $2, $3, $4)",
            )
            .bind(name)
            .bind(started_at)
            .bind(completed_at)
            .bind(last_offset)
            .execute(tx.as_mut())
            .await?;
            tx.commit().await
        }

        async fn reconstruct_state(
            db: &SqlStorage,
            instance: &NodeState,
            from_height: u64,
            to_view: ViewNumber,
            fee_accounts: &[FeeAccount],
            reward_accounts: &[RewardAccountV2],
        ) -> anyhow::Result<(ValidatedState, Leaf2)> {
            let mut tx = db.read().await?;
            reconstruct_state(
                instance,
                db,
                &mut tx,
                from_height,
                to_view,
                fee_accounts,
                reward_accounts,
            )
            .await
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use alloy::primitives::{Address, U256};
    use espresso_types::{
        BLOCK_MERKLE_TREE_HEIGHT, BackoffParams, BlockMerkleTree, FEE_MERKLE_TREE_HEIGHT,
        FeeAccount, FeeAmount, FeeMerkleTree, NodeState,
        v0::traits::StateCatchup,
        v0_4::{REWARD_MERKLE_TREE_V2_HEIGHT, RewardMerkleTreeV2},
    };
    use hotshot_query_service::{
        QueryError,
        data_source::{
            Transaction, VersionedDataSource,
            storage::{
                UpdateAvailabilityStorage, pruning::PrunedHeightStorage, sql::StorageConnectionType,
            },
        },
        merklized_state::{MerklizedState, UpdateStateData},
        node::NodeDataSource,
    };
    use hotshot_types::data::ViewNumber;
    use jf_merkle_tree_compat::{
        AppendableMerkleTreeScheme, ForgetableMerkleTreeScheme, LookupResult, MerkleTreeScheme,
        ToTraversalPath, UniversalMerkleTreeScheme,
    };
    use light_client::testing::{leaf_chain, leaf_chain_with_upgrade};
    use versions::{DRB_AND_HEADER_UPGRADE_VERSION, EPOCH_REWARD_VERSION, Upgrade};

    use super::{
        super::qs::{Config, SqlStorage, Transaction as SqlTransaction, Write, testing::TmpDb},
        ArchiveStateGc, SeqTypes,
        impl_testable_data_source::tmp_options,
        query_as,
    };
    use crate::{
        api::{RewardMerkleTreeDataSource, data_source::SequencerDataSource},
        catchup::{CatchupStorage, SqlStateCatchup},
    };

    async fn write_fee_state(storage: &SqlStorage, account: FeeAccount, balance: u64, height: u64) {
        let mut tree = FeeMerkleTree::new(FEE_MERKLE_TREE_HEIGHT);
        tree.update(account, FeeAmount::from(balance)).unwrap();
        let proof = match tree.universal_lookup(account) {
            LookupResult::Ok(_, proof) => proof,
            _ => panic!("account not in tree"),
        };
        let path = <FeeAccount as ToTraversalPath<{ FeeMerkleTree::ARITY }>>::to_traversal_path(
            &account,
            tree.height(),
        );

        let mut tx = storage.write().await.unwrap();
        UpdateStateData::<SeqTypes, FeeMerkleTree, { FeeMerkleTree::ARITY }>::insert_merkle_nodes(
            &mut tx, proof, path, height,
        )
        .await
        .unwrap();
        UpdateStateData::<SeqTypes, FeeMerkleTree, { FeeMerkleTree::ARITY }>::set_last_state_height(
            &mut tx,
            height as usize,
        )
        .await
        .unwrap();
        Transaction::commit(tx).await.unwrap();
    }

    async fn fee_state_heights(storage: &SqlStorage) -> Vec<u64> {
        let mut tx = storage.read().await.unwrap();
        query_as::<(i64,)>(&format!(
            "SELECT DISTINCT created FROM {} ORDER BY created",
            FeeMerkleTree::state_type()
        ))
        .fetch_all(tx.as_mut())
        .await
        .unwrap()
        .into_iter()
        .map(|(height,)| height as u64)
        .collect()
    }

    /// The collector's options reach storage: the fee tree is a state table and the floor applies.
    #[tokio::test]
    #[test_log::test]
    async fn test_archive_state_gc_prunes_fee_state() {
        let db = TmpDb::init().await;
        let mut opt = tmp_options(&db);
        opt.archive_state_min_retention = 2;
        let storage = SqlStorage::connect(
            Config::try_from(&opt).unwrap(),
            StorageConnectionType::Query,
        )
        .await
        .unwrap();

        let account = FeeAccount::from(Address::repeat_byte(0x42));
        for height in 1..=5 {
            write_fee_state(&storage, account, 100 * height, height).await;
        }

        // Light client history at 4, floor 2 under head 5: target 3, marker 2.
        let gc = ArchiveStateGc::new(&opt);
        storage
            .prune_state_below(4, gc.min_retention, &gc.cfg)
            .await
            .unwrap();
        assert_eq!(fee_state_heights(&storage).await, [2, 3, 4, 5]);
        let mut tx = storage.read().await.unwrap();
        assert_eq!(tx.load_state_pruned_height().await.unwrap(), Some(2));
        assert_eq!(tx.load_pruned_height().await.unwrap(), None);
    }

    /// The data pruner deletes the leaf and header at the state head and keeps the ones above it,
    /// while the state pruner keeps the head's snapshot. The local catchup still reads that
    /// snapshot, by the root it is given.
    #[tokio::test]
    #[test_log::test]
    async fn test_catchup_reads_snapshot_without_header() {
        let db = TmpDb::init().await;
        let storage = Arc::new(
            SqlStorage::connect(
                Config::try_from(&tmp_options(&db)).unwrap(),
                StorageConnectionType::Query,
            )
            .await
            .unwrap(),
        );

        // A head snapshot whose leaf and header the data pruner deleted, below a stored leaf.
        let height = 3;
        let leaves = leaf_chain(0..=height + 1, DRB_AND_HEADER_UPGRADE_VERSION).await;
        let mut tx = storage.write().await.unwrap();
        tx.insert_leaf(&leaves[height as usize + 1]).await.unwrap();
        Transaction::commit(tx).await.unwrap();

        let account = FeeAccount::from(Address::repeat_byte(0x42));
        let absent = FeeAccount::from(Address::repeat_byte(0x43));
        write_fee_state(&storage, account, 100, height).await;
        let mut fee_tree = FeeMerkleTree::new(FEE_MERKLE_TREE_HEIGHT);
        fee_tree.update(account, FeeAmount::from(100)).unwrap();
        let fee_root = fee_tree.commitment();

        let mut block_tree = BlockMerkleTree::new(BLOCK_MERKLE_TREE_HEIGHT);
        for leaf in &leaves[..height as usize] {
            block_tree.push(leaf.block_hash()).unwrap();
        }
        let frontier = block_tree.lookup(height - 1).expect_ok().unwrap().1;
        let mut tx = storage.write().await.unwrap();
        UpdateStateData::<SeqTypes, BlockMerkleTree, { BlockMerkleTree::ARITY }>::insert_merkle_nodes(
            &mut tx,
            frontier,
            ToTraversalPath::<{ BlockMerkleTree::ARITY }>::to_traversal_path(
                &(height - 1),
                BLOCK_MERKLE_TREE_HEIGHT,
            ),
            height,
        )
        .await
        .unwrap();
        Transaction::commit(tx).await.unwrap();

        let instance = NodeState::mock();
        let view = ViewNumber::new(height);

        // The reads through the leaf and header fail, so the catchup has to read by root.
        let err = storage
            .get_accounts(&instance, height, view, &[account])
            .await
            .unwrap_err();
        assert!(
            format!("{err:#}").contains(&format!("leaf {height} not available")),
            "{err:#}"
        );
        storage
            .get_frontier(&instance, height, view)
            .await
            .unwrap_err();

        let catchup = SqlStateCatchup::new(storage, BackoffParams::disabled());

        let proofs = catchup
            .try_fetch_accounts(0, &instance, height, view, fee_root, &[account, absent])
            .await
            .unwrap();
        let balances = proofs
            .iter()
            .map(|proof| proof.verify(&fee_root).unwrap())
            .collect::<Vec<_>>();
        assert_eq!(balances, [U256::from(100), U256::ZERO]);

        let mut block_tree_from_root = BlockMerkleTree::from_commitment(block_tree.commitment());
        catchup
            .try_remember_blocks_merkle_tree(0, &instance, height, view, &mut block_tree_from_root)
            .await
            .unwrap();
        assert!(block_tree_from_root.lookup(height - 1).expect_ok().is_ok());

        // A root the snapshot does not have is rejected, not trusted.
        let wrong_root = FeeMerkleTree::new(FEE_MERKLE_TREE_HEIGHT).commitment();
        catchup
            .try_fetch_accounts(0, &instance, height, view, wrong_root, &[account])
            .await
            .unwrap_err();
    }

    async fn insert_test_header(
        tx: &mut SqlTransaction<Write>,
        block_height: u64,
        reward_tree: &RewardMerkleTreeV2,
    ) {
        let reward_commitment = serde_json::to_value(reward_tree.commitment()).unwrap();
        let test_data = serde_json::json!({
            "block_merkle_tree_root": format!("block_root_{}", block_height),
            "fee_merkle_tree_root": format!("fee_root_{}", block_height),
            "fields": {
                RewardMerkleTreeV2::header_state_commitment_field(): reward_commitment
            }
        });
        tx.upsert(
            "header",
            [
                "height",
                "hash",
                "payload_hash",
                "timestamp",
                "data",
                "ns_table",
            ],
            ["height"],
            [(
                block_height as i64,
                format!("hash_{}", block_height),
                format!("payload_{}", block_height),
                block_height as i64,
                test_data,
                "ns_table".to_string(),
            )],
        )
        .await
        .unwrap();
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_merkle_proof_gc() {
        let db = TmpDb::init().await;
        let opt = tmp_options(&db);
        let cfg = Config::try_from(&opt).expect("failed to create config from options");
        let storage = SqlStorage::connect(cfg, StorageConnectionType::Query)
            .await
            .expect("failed to connect to storage");

        let account = vec![0; 32];

        // Insert mock leaves at heights 0-2 so load_proof can check leaf version.
        let leaves = leaf_chain(0..=2, DRB_AND_HEADER_UPGRADE_VERSION).await;
        let mut tx = storage.write().await.unwrap();
        for leaf in &leaves {
            tx.insert_leaf(leaf).await.unwrap();
        }
        Transaction::commit(tx).await.unwrap();

        // Insert some mock proofs at heights 0-2.
        for h in 0..=2 {
            storage
                .persist_proofs(h, [(account.clone(), h.to_le_bytes().to_vec())].into_iter())
                .await
                .unwrap();
        }

        // Test garbage collection.
        storage.garbage_collect(2).await.unwrap();

        // Make sure the proof at height 2 is still available.
        assert_eq!(
            storage.load_proof(2, account.clone(), 0).await.unwrap(),
            2u64.to_le_bytes()
        );

        // Meanwhile, the proofs at heights 0-1 have been garbage collected.
        for h in 0..2 {
            let err = storage.load_proof(h, account.clone(), 0).await.unwrap_err();
            assert!(err.to_string().contains("Missing proof"), "{err:#}");
        }

        // Garbage collect the remaining proof.
        storage.garbage_collect(3).await.unwrap();
        let err = storage.load_proof(2, account, 0).await.unwrap_err();
        assert!(err.to_string().contains("Missing proof"), "{err:#}");
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_load_proof_v5_epoch_boundary() {
        let db = TmpDb::init().await;
        let opt = tmp_options(&db);
        let cfg = Config::try_from(&opt).expect("failed to create config from options");
        let storage = SqlStorage::connect(cfg, StorageConnectionType::Query)
            .await
            .expect("failed to connect to storage");

        let epoch_height = 10u64;
        let account = vec![0; 32];

        // Create V5 leaves at heights 0..=15.
        let leaves = leaf_chain(0..=15, EPOCH_REWARD_VERSION).await;
        let mut tx = storage.write().await.unwrap();
        for leaf in &leaves {
            tx.insert_leaf(leaf).await.unwrap();
        }
        Transaction::commit(tx).await.unwrap();

        // Store proofs only at epoch boundaries (height 10).
        let boundary_proof = b"proof_at_10".to_vec();
        {
            let mut tx = storage.write().await.unwrap();
            tx.upsert(
                "reward_merkle_tree_v2_proofs",
                ["height", "account", "proof"],
                ["height", "account"],
                [(10i64, account.clone(), boundary_proof.clone())],
            )
            .await
            .unwrap();
            Transaction::commit(tx).await.unwrap();
        }

        // Querying at the epoch boundary itself should return the proof directly.
        assert_eq!(
            storage
                .load_proof(10, account.clone(), epoch_height)
                .await
                .unwrap(),
            boundary_proof,
        );

        // Querying at a non-boundary V5 height (e.g. 15) should resolve to the
        // previous epoch boundary (10) and return its proof.
        assert_eq!(
            storage
                .load_proof(15, account.clone(), epoch_height)
                .await
                .unwrap(),
            boundary_proof,
        );
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_load_proof_v4_to_v5_upgrade_boundary() {
        let db = TmpDb::init().await;
        let opt = tmp_options(&db);
        let cfg = Config::try_from(&opt).expect("failed to create config from options");
        let storage = SqlStorage::connect(cfg, StorageConnectionType::Query)
            .await
            .expect("failed to connect to storage");

        let epoch_height = 10u64;
        let account = vec![0; 32];

        // V4 leaves at heights 0..=10, V5 from height 11 onward.
        // The upgrade happens at height 11.
        let upgrade = Upgrade::new(DRB_AND_HEADER_UPGRADE_VERSION, EPOCH_REWARD_VERSION);
        let leaves = leaf_chain_with_upgrade(0..=15, 11, upgrade).await;
        {
            let mut tx = storage.write().await.unwrap();
            for leaf in &leaves {
                tx.insert_leaf(leaf).await.unwrap();
            }
            Transaction::commit(tx).await.unwrap();
        }

        // Querying at V5 height 15 resolves to prev epoch boundary
        // But leaf at height 10 is V4, so this should fail.
        storage
            .load_proof(15, account.clone(), epoch_height)
            .await
            .unwrap_err();

        let v4_proof = b"v4_proof_at_5".to_vec();
        {
            let mut tx = storage.write().await.unwrap();
            tx.upsert(
                "reward_merkle_tree_v2_proofs",
                ["height", "account", "proof"],
                ["height", "account"],
                [(5i64, account.clone(), v4_proof.clone())],
            )
            .await
            .unwrap();
            Transaction::commit(tx).await.unwrap();
        }
        assert_eq!(
            storage.load_proof(5, account.clone(), 0).await.unwrap(),
            v4_proof,
        );
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_get_table_sizes() {
        use crate::api::data_source::DatabaseMetadataSource;

        let db = TmpDb::init().await;
        let opt = tmp_options(&db);
        let cfg = Config::try_from(&opt).expect("failed to create config from options");
        let storage = SqlStorage::connect(cfg, StorageConnectionType::Query)
            .await
            .expect("failed to connect to storage");

        // Insert some test data to ensure tables have rows
        let mut tx = storage.write().await.unwrap();

        // Insert a test header
        let reward_tree = RewardMerkleTreeV2::new(REWARD_MERKLE_TREE_V2_HEIGHT);
        insert_test_header(&mut tx, 1, &reward_tree).await;

        tx.commit().await.unwrap();

        // Call get_table_sizes and verify it doesn't error
        let table_sizes = storage
            .get_table_sizes()
            .await
            .expect("get_table_sizes should succeed");

        // Verify we got some tables back
        assert!(
            !table_sizes.is_empty(),
            "should have at least one table in the database"
        );
    }

    /// A pruned node runs without an aggregator, so it has no transaction or payload totals.
    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_pruned_node_has_no_totals() {
        for prune in [false, true] {
            let storage = TmpDb::init().await;
            let mut opt = tmp_options(&storage);
            opt.prune = prune;
            let ds = super::DataSource::create(
                opt.try_into().unwrap(),
                super::Provider::default(),
                false,
            )
            .await
            .unwrap();

            let res = NodeDataSource::count_transactions_in_range(&ds, .., None).await;
            if prune {
                assert!(matches!(res, Err(QueryError::NotFound)), "{res:?}");
            } else {
                assert_eq!(res.unwrap(), 0);
            }
        }
    }
}