espresso_node/persistence/

sql.rs

use std::{
    collections::BTreeMap,
    path::PathBuf,
    str::FromStr,
    sync::Arc,
    time::{Duration, Instant},
};

use alloy::primitives::Address;
use anyhow::{Context, bail, ensure};
use async_trait::async_trait;
use clap::Parser;
use committable::Committable;
use derivative::Derivative;
use derive_more::derive::{From, Into};
use either::Either;
use espresso_types::{
    AuthenticatedValidatorMap, BackoffParams, BlockMerkleTree, FeeMerkleTree, Leaf, Leaf2,
    NetworkConfig, Payload, PubKey, Ratio, RegisteredValidatorMap, StakeTableHash, parse_duration,
    parse_size,
    traits::{EventsPersistenceRead, MembershipPersistence, StakeTuple},
    v0::traits::{EventConsumer, PersistenceOptions, SequencerPersistence, StateCatchup},
    v0_3::{
        AuthenticatedValidator, EventKey, IndexedStake, RegisteredValidator, RewardAmount,
        StakeTableEvent,
    },
    v0_4::{REWARD_MERKLE_TREE_V2_HEIGHT, RewardAccountV2, RewardMerkleTreeV2},
};
use futures::stream::StreamExt;
use hotshot::InitializerEpochInfo;
use hotshot_libp2p_networking::network::behaviours::dht::store::persistent::{
    DhtPersistentStorage, SerializableRecord,
};
use hotshot_new_protocol::message::Certificate2;
use hotshot_query_service::{
    availability::BlockId,
    data_source::{
        Transaction as _, VersionedDataSource,
        storage::{
            AvailabilityStorage,
            pruning::PrunerCfg,
            sql::{
                Config, Db, Read, SqlStorage, StorageConnectionType, Transaction, Write,
                include_migrations, query_as, syntax_helpers::MAX_FN,
            },
        },
    },
    fetching::{
        Provider,
        request::{PayloadRequest, VidCommonRequest},
    },
    merklized_state::MerklizedState,
};
use hotshot_types::{
    data::{
        DaProposal, DaProposal2, EpochNumber, QuorumProposal, QuorumProposalWrapper,
        QuorumProposalWrapperLegacy, VidCommitment, VidCommon, VidDisperseShare, VidDisperseShare0,
    },
    drb::{DrbInput, DrbResult},
    event::{Event, EventType, HotShotAction, LeafInfo},
    message::{Proposal, convert_proposal},
    new_protocol::CoordinatorEvent,
    simple_certificate::{
        CertificatePair, LightClientStateUpdateCertificateV1, LightClientStateUpdateCertificateV2,
        NextEpochQuorumCertificate2, QuorumCertificate, QuorumCertificate2, UpgradeCertificate,
    },
    traits::{
        block_contents::{BlockHeader, BlockPayload},
        metrics::Metrics,
    },
    vote::HasViewNumber,
};
use indexmap::IndexMap;
use itertools::Itertools;
use jf_merkle_tree_compat::MerkleTreeScheme;
use sqlx::{Executor, QueryBuilder, Row, query};

use crate::{
    NodeType, RECENT_STAKE_TABLES_LIMIT, SeqTypes, ViewNumber,
    api::RewardMerkleTreeV2Data,
    catchup::SqlStateCatchup,
    persistence::{migrate_network_config, persistence_metrics::PersistenceMetricsValue},
};

/// Options for Postgres-backed persistence.
#[derive(Parser, Clone, Derivative)]
#[derivative(Debug)]
pub struct PostgresOptions {
    /// Hostname for the remote Postgres database server.
    #[clap(long, env = "ESPRESSO_NODE_POSTGRES_HOST")]
    pub(crate) host: Option<String>,

    /// Port for the remote Postgres database server.
    #[clap(long, env = "ESPRESSO_NODE_POSTGRES_PORT")]
    pub(crate) port: Option<u16>,

    /// Name of database to connect to.
    #[clap(long, env = "ESPRESSO_NODE_POSTGRES_DATABASE")]
    pub(crate) database: Option<String>,

    /// Postgres user to connect as.
    #[clap(long, env = "ESPRESSO_NODE_POSTGRES_USER")]
    pub(crate) user: Option<String>,

    /// Password for Postgres user.
    #[clap(long, env = "ESPRESSO_NODE_POSTGRES_PASSWORD")]
    // Hide from debug output since may contain sensitive data.
    #[derivative(Debug = "ignore")]
    pub(crate) password: Option<String>,

    /// Use TLS for an encrypted connection to the database.
    #[clap(long, env = "ESPRESSO_NODE_POSTGRES_USE_TLS")]
    pub(crate) use_tls: bool,

    /// Disable `DEFERRABLE` on read transactions for the query service.
    ///
    /// When true, read transactions on Postgres start with `SERIALIZABLE READ ONLY` (no
    /// `DEFERRABLE`), so they begin immediately rather than waiting for a safe serializable
    /// snapshot. This trades start-up latency for the chance of a serialization-error retry,
    /// and is opt-in.
    #[clap(
        long,
        env = "ESPRESSO_NODE_POSTGRES_NO_DEFERRABLE",
        default_value_t = false
    )]
    pub(crate) no_deferrable: bool,
}

impl Default for PostgresOptions {
    fn default() -> Self {
        Self::parse_from(std::iter::empty::<String>())
    }
}

#[derive(Parser, Clone, Derivative, Default, From, Into)]
#[derivative(Debug)]
pub struct SqliteOptions {
    /// Base directory for the SQLite database.
    /// The SQLite file will be created in the `sqlite` subdirectory with filename as `database`.
    #[clap(
        long,
        env = "ESPRESSO_NODE_STORAGE_PATH",
        value_parser = build_sqlite_path
    )]
    pub(crate) path: PathBuf,
}

pub fn build_sqlite_path(path: &str) -> anyhow::Result<PathBuf> {
    let sub_dir = PathBuf::from_str(path)?.join("sqlite");

    // if `sqlite` sub dir does not exist then create it
    if !sub_dir.exists() {
        std::fs::create_dir_all(&sub_dir)
            .with_context(|| format!("failed to create directory: {sub_dir:?}"))?;
    }

    Ok(sub_dir.join("database"))
}

/// Options for database-backed persistence, supporting both Postgres and SQLite.
#[derive(Parser, Clone, Derivative, From, Into)]
#[derivative(Debug)]
pub struct Options {
    #[cfg(not(feature = "embedded-db"))]
    #[clap(flatten)]
    pub(crate) postgres_options: PostgresOptions,

    #[cfg(feature = "embedded-db")]
    #[clap(flatten)]
    pub(crate) sqlite_options: SqliteOptions,

    /// Database URI for Postgres or SQLite.
    ///
    /// This is a shorthand for setting a number of other options all at once. The URI has the
    /// following format ([brackets] indicate optional segments):
    ///
    /// - **Postgres:** `postgres[ql]://[username[:password]@][host[:port],]/database[?parameter_list]`
    /// - **SQLite:** `sqlite://path/to/db.sqlite`
    ///
    /// Options set explicitly via other env vars or flags will take precedence, so you can use this
    /// URI to set a baseline and then use other parameters to override or add configuration. In
    /// addition, there are some parameters which cannot be set via the URI, such as TLS.
    // Hide from debug output since may contain sensitive data.
    #[derivative(Debug = "ignore")]
    pub(crate) uri: Option<String>,

    /// This will enable the pruner and set the default pruning parameters unless provided.
    /// Default parameters:
    /// - pruning_threshold: 3 TB
    /// - minimum_retention: 1 day
    /// - target_retention: 7 days
    /// - batch_size: 1000
    /// - max_usage: 80%
    /// - interval: 1 hour
    #[clap(long, env = "ESPRESSO_NODE_DATABASE_PRUNE")]
    pub(crate) prune: bool,

    /// Pruning parameters.
    #[clap(flatten)]
    pub(crate) pruning: PruningOptions,

    /// Pruning parameters for ephemeral consensus storage.
    #[clap(flatten)]
    pub(crate) consensus_pruning: ConsensusPruningOptions,

    /// Specifies the maximum number of concurrent fetch requests allowed from peers.
    #[clap(long, env = "ESPRESSO_NODE_FETCH_RATE_LIMIT")]
    pub(crate) fetch_rate_limit: Option<usize>,

    /// The minimum delay between active fetches in a stream.
    #[clap(long, env = "ESPRESSO_NODE_ACTIVE_FETCH_DELAY", value_parser = parse_duration)]
    pub(crate) active_fetch_delay: Option<Duration>,

    /// The minimum delay between loading chunks in a stream.
    #[clap(long, env = "ESPRESSO_NODE_CHUNK_FETCH_DELAY", value_parser = parse_duration)]
    pub(crate) chunk_fetch_delay: Option<Duration>,

    /// The number of items to process in a single transaction when scanning the database for
    /// missing objects.
    #[clap(long, env = "ESPRESSO_NODE_SYNC_STATUS_CHUNK_SIZE")]
    pub(crate) sync_status_chunk_size: Option<usize>,

    /// Duration to cache sync status results for.
    #[clap(long, env = "ESPRESSO_NODE_SYNC_STATUS_TTL", value_parser = parse_duration)]
    pub(crate) sync_status_ttl: Option<Duration>,

    /// The number of items to process at a time when scanning for proactive fetching.
    #[clap(long, env = "ESPRESSO_NODE_PROACTIVE_SCAN_CHUNK_SIZE")]
    pub(crate) proactive_scan_chunk_size: Option<usize>,

    /// The time interval between proactive fetching scans.
    #[clap(long, env = "ESPRESSO_NODE_PROACTIVE_SCAN_INTERVAL", value_parser = parse_duration)]
    pub(crate) proactive_scan_interval: Option<Duration>,

    /// Disable the proactive scanner task.
    #[clap(long, env = "ESPRESSO_NODE_DISABLE_PROACTIVE_FETCHING")]
    pub(crate) disable_proactive_fetching: bool,

    /// Disable pruning and reconstruct previously pruned data.
    ///
    /// While running without pruning is the default behavior, the default will not try to
    /// reconstruct data that was pruned in a previous run where pruning was enabled. This option
    /// instructs the service to run without pruning _and_ reconstruct all previously pruned data by
    /// fetching from peers.
    #[clap(long, env = "ESPRESSO_NODE_ARCHIVE", conflicts_with = "prune")]
    pub(crate) archive: bool,

    /// Turns on leaf only data storage
    #[clap(
        long,
        env = "ESPRESSO_NODE_LIGHTWEIGHT",
        default_value_t = false,
        conflicts_with = "archive"
    )]
    pub(crate) lightweight: bool,

    /// The maximum idle time of a database connection.
    ///
    /// Any connection which has been open and unused longer than this duration will be
    /// automatically closed to reduce load on the server.
    #[clap(long, env = "ESPRESSO_NODE_DATABASE_IDLE_CONNECTION_TIMEOUT", value_parser = parse_duration, default_value = "10m")]
    pub(crate) idle_connection_timeout: Duration,

    /// The maximum lifetime of a database connection.
    ///
    /// Any connection which has been open longer than this duration will be automatically closed
    /// (and, if needed, replaced), even if it is otherwise healthy. It is good practice to refresh
    /// even healthy connections once in a while (e.g. daily) in case of resource leaks in the
    /// server implementation.
    #[clap(long, env = "ESPRESSO_NODE_DATABASE_CONNECTION_TIMEOUT", value_parser = parse_duration, default_value = "30m")]
    pub(crate) connection_timeout: Duration,

    #[clap(long, env = "ESPRESSO_NODE_DATABASE_SLOW_STATEMENT_THRESHOLD", value_parser = parse_duration, default_value = "1s")]
    pub(crate) slow_statement_threshold: Duration,

    /// The maximum time a single SQL statement is allowed to run before being canceled.
    ///
    /// This helps prevent queries from running indefinitely and consuming resources.
    /// Set to 10 minutes by default
    #[clap(long, env = "ESPRESSO_NODE_DATABASE_STATEMENT_TIMEOUT", value_parser = parse_duration, default_value = "10m")]
    pub(crate) statement_timeout: Duration,

    /// The minimum number of database connections to maintain at any time.
    ///
    /// The database client will, to the best of its ability, maintain at least `min` open
    /// connections at all times. This can be used to reduce the latency hit of opening new
    /// connections when at least this many simultaneous connections are frequently needed.
    #[clap(
        long,
        env = "ESPRESSO_NODE_DATABASE_MIN_CONNECTIONS",
        default_value = "0"
    )]
    pub(crate) min_connections: u32,

    /// Allows setting a different maximum number of connections for query operations.
    /// Default value of None implies using the min_connections value.
    #[cfg(not(feature = "embedded-db"))]
    #[clap(long, env = "ESPRESSO_NODE_DATABASE_QUERY_MIN_CONNECTIONS", default_value = None)]
    pub(crate) query_min_connections: Option<u32>,

    /// The maximum number of database connections to maintain at any time.
    ///
    /// Once `max` connections are in use simultaneously, further attempts to acquire a connection
    /// (or begin a transaction) will block until one of the existing connections is released.
    #[clap(
        long,
        env = "ESPRESSO_NODE_DATABASE_MAX_CONNECTIONS",
        default_value = "25"
    )]
    pub(crate) max_connections: u32,

    /// Allows setting a different maximum number of connections for query operations.
    /// Default value of None implies using the max_connections value.
    #[cfg(not(feature = "embedded-db"))]
    #[clap(long, env = "ESPRESSO_NODE_DATABASE_QUERY_MAX_CONNECTIONS", default_value = None)]
    pub(crate) query_max_connections: Option<u32>,

    // Keep the database connection pool when persistence is created,
    // allowing it to be reused across multiple instances instead of creating
    // a new pool each time such as for API, consensus storage etc
    // This also ensures all storage instances adhere to the MAX_CONNECTIONS limit if set
    //
    // Note: Cloning the `Pool` is lightweight and efficient because it simply
    // creates a new reference-counted handle to the underlying pool state.
    #[clap(skip)]
    pub(crate) pool: Option<sqlx::Pool<Db>>,
}

impl Default for Options {
    fn default() -> Self {
        Self::parse_from(std::iter::empty::<String>())
    }
}

#[cfg(not(feature = "embedded-db"))]
impl From<PostgresOptions> for Config {
    fn from(opt: PostgresOptions) -> Self {
        let mut cfg = Config::default();

        if let Some(host) = opt.host {
            cfg = cfg.host(host);
        }

        if let Some(port) = opt.port {
            cfg = cfg.port(port);
        }

        if let Some(database) = &opt.database {
            cfg = cfg.database(database);
        }

        if let Some(user) = &opt.user {
            cfg = cfg.user(user);
        }

        if let Some(password) = &opt.password {
            cfg = cfg.password(password);
        }

        if opt.use_tls {
            cfg = cfg.tls();
        }

        cfg = cfg.max_connections(20);
        cfg = cfg.idle_connection_timeout(Duration::from_secs(120));
        cfg = cfg.connection_timeout(Duration::from_secs(10240));
        cfg = cfg.slow_statement_threshold(Duration::from_secs(1));
        cfg = cfg.statement_timeout(Duration::from_secs(600)); // 10 minutes default

        hotshot_query_service::data_source::storage::sql::set_no_deferrable_on_read(
            opt.no_deferrable,
        );

        cfg
    }
}

#[cfg(feature = "embedded-db")]
impl From<SqliteOptions> for Config {
    fn from(opt: SqliteOptions) -> Self {
        let mut cfg = Config::default();

        cfg = cfg.db_path(opt.path);

        cfg = cfg.max_connections(20);
        cfg = cfg.idle_connection_timeout(Duration::from_secs(120));
        cfg = cfg.connection_timeout(Duration::from_secs(10240));
        cfg = cfg.slow_statement_threshold(Duration::from_secs(2));
        cfg = cfg.statement_timeout(Duration::from_secs(600));
        cfg
    }
}

#[cfg(not(feature = "embedded-db"))]
impl From<PostgresOptions> for Options {
    fn from(opt: PostgresOptions) -> Self {
        Options {
            postgres_options: opt,
            max_connections: 20,
            idle_connection_timeout: Duration::from_secs(120),
            connection_timeout: Duration::from_secs(10240),
            slow_statement_threshold: Duration::from_secs(1),
            statement_timeout: Duration::from_secs(600),
            ..Default::default()
        }
    }
}

#[cfg(feature = "embedded-db")]
impl From<SqliteOptions> for Options {
    fn from(opt: SqliteOptions) -> Self {
        Options {
            sqlite_options: opt,
            max_connections: 5,
            idle_connection_timeout: Duration::from_secs(120),
            connection_timeout: Duration::from_secs(10240),
            slow_statement_threshold: Duration::from_secs(1),
            uri: None,
            statement_timeout: Duration::from_secs(600),
            prune: false,
            pruning: Default::default(),
            consensus_pruning: Default::default(),
            fetch_rate_limit: None,
            active_fetch_delay: None,
            chunk_fetch_delay: None,
            sync_status_chunk_size: None,
            sync_status_ttl: None,
            proactive_scan_chunk_size: None,
            proactive_scan_interval: None,
            disable_proactive_fetching: false,
            archive: false,
            lightweight: false,
            min_connections: 0,
            pool: None,
        }
    }
}
impl TryFrom<&Options> for Config {
    type Error = anyhow::Error;

    fn try_from(opt: &Options) -> Result<Self, Self::Error> {
        let mut cfg = match &opt.uri {
            Some(uri) => uri.parse()?,
            None => Self::default(),
        };

        if let Some(pool) = &opt.pool {
            cfg = cfg.pool(pool.clone());
        }

        cfg = cfg.max_connections(opt.max_connections);
        cfg = cfg.idle_connection_timeout(opt.idle_connection_timeout);
        cfg = cfg.min_connections(opt.min_connections);

        #[cfg(not(feature = "embedded-db"))]
        {
            cfg =
                cfg.query_max_connections(opt.query_max_connections.unwrap_or(opt.max_connections));
            cfg =
                cfg.query_min_connections(opt.query_min_connections.unwrap_or(opt.min_connections));

            hotshot_query_service::data_source::storage::sql::set_no_deferrable_on_read(
                opt.postgres_options.no_deferrable,
            );
        }

        cfg = cfg.connection_timeout(opt.connection_timeout);
        cfg = cfg.slow_statement_threshold(opt.slow_statement_threshold);
        cfg = cfg.statement_timeout(opt.statement_timeout);

        #[cfg(not(feature = "embedded-db"))]
        {
            cfg = cfg.migrations(include_migrations!(
                "$CARGO_MANIFEST_DIR/api/migrations/postgres"
            ));

            let pg_options = &opt.postgres_options;

            if let Some(host) = &pg_options.host {
                cfg = cfg.host(host.clone());
            }

            if let Some(port) = pg_options.port {
                cfg = cfg.port(port);
            }

            if let Some(database) = &pg_options.database {
                cfg = cfg.database(database);
            }

            if let Some(user) = &pg_options.user {
                cfg = cfg.user(user);
            }

            if let Some(password) = &pg_options.password {
                cfg = cfg.password(password);
            }

            if pg_options.use_tls {
                cfg = cfg.tls();
            }
        }

        #[cfg(feature = "embedded-db")]
        {
            cfg = cfg.migrations(include_migrations!(
                "$CARGO_MANIFEST_DIR/api/migrations/sqlite"
            ));

            cfg = cfg.db_path(opt.sqlite_options.path.clone());
        }

        if opt.prune {
            cfg = cfg.pruner_cfg(PrunerCfg::from(opt.pruning))?;
        }
        if opt.archive {
            cfg = cfg.archive();
        }

        Ok(cfg)
    }
}

/// Pruning parameters.
#[derive(Parser, Clone, Copy, Debug)]
pub struct PruningOptions {
    /// Threshold for pruning, specified in bytes.
    /// If the disk usage surpasses this threshold, pruning is initiated for data older than the specified minimum retention period.
    /// Pruning continues until the disk usage drops below the MAX USAGE.
    #[clap(long, env = "ESPRESSO_NODE_PRUNER_PRUNING_THRESHOLD", value_parser = parse_size)]
    pub(crate) pruning_threshold: Option<u64>,

    /// Minimum retention period.
    /// Data is retained for at least this duration, even if there's no free disk space.
    #[clap(
        long,
        env = "ESPRESSO_NODE_PRUNER_MINIMUM_RETENTION",
        value_parser = parse_duration,
    )]
    pub(crate) minimum_retention: Option<Duration>,

    /// Target retention period.
    /// Data older than this is pruned to free up space.
    #[clap(
        long,
        env = "ESPRESSO_NODE_PRUNER_TARGET_RETENTION",
        value_parser = parse_duration,
    )]
    pub(crate) target_retention: Option<Duration>,

    /// Batch size for pruning.
    /// This is the number of blocks data to delete in a single transaction.
    #[clap(long, env = "ESPRESSO_NODE_PRUNER_BATCH_SIZE")]
    pub(crate) batch_size: Option<u64>,

    /// Maximum disk usage (in basis points).
    ///
    /// Pruning stops once the disk usage falls below this value, even if
    /// some data older than the `MINIMUM_RETENTION` remains. Values range
    /// from 0 (0%) to 10000 (100%).
    #[clap(long, env = "ESPRESSO_NODE_PRUNER_MAX_USAGE")]
    pub(crate) max_usage: Option<u16>,

    /// Interval for running the pruner.
    #[clap(
        long,
        env = "ESPRESSO_NODE_PRUNER_INTERVAL",
        value_parser = parse_duration,
    )]
    pub(crate) interval: Option<Duration>,

    /// Number of SQLite pages to vacuum from the freelist
    /// during each pruner cycle.
    /// This value corresponds to `N` in the SQLite PRAGMA `incremental_vacuum(N)`,
    #[clap(long, env = "ESPRESSO_NODE_PRUNER_INCREMENTAL_VACUUM_PAGES")]
    pub(crate) pages: Option<u64>,
}

impl Default for PruningOptions {
    fn default() -> Self {
        Self::parse_from(std::iter::empty::<String>())
    }
}

impl From<PruningOptions> for PrunerCfg {
    fn from(opt: PruningOptions) -> Self {
        let mut cfg = PrunerCfg::new();
        if let Some(threshold) = opt.pruning_threshold {
            cfg = cfg.with_pruning_threshold(threshold);
        }
        if let Some(min) = opt.minimum_retention {
            cfg = cfg.with_minimum_retention(min);
        }
        if let Some(target) = opt.target_retention {
            cfg = cfg.with_target_retention(target);
        }
        if let Some(batch) = opt.batch_size {
            cfg = cfg.with_batch_size(batch);
        }
        if let Some(max) = opt.max_usage {
            cfg = cfg.with_max_usage(max);
        }
        if let Some(interval) = opt.interval {
            cfg = cfg.with_interval(interval);
        }

        if let Some(pages) = opt.pages {
            cfg = cfg.with_incremental_vacuum_pages(pages)
        }

        cfg = cfg.with_state_tables(vec![
            BlockMerkleTree::state_type().to_string(),
            FeeMerkleTree::state_type().to_string(),
        ]);

        cfg
    }
}

/// Pruning parameters for ephemeral consensus storage.
#[derive(Parser, Clone, Copy, Debug)]
pub struct ConsensusPruningOptions {
    /// Number of views to try to retain in consensus storage before data that hasn't been archived
    /// is garbage collected.
    ///
    /// The longer this is, the more certain that all data will eventually be archived, even if
    /// there are temporary problems with archive storage or partially missing data. This can be set
    /// very large, as most data is garbage collected as soon as it is finalized by consensus. This
    /// setting only applies to views which never get decided (ie forks in consensus) and views for
    /// which this node is partially offline. These should be exceptionally rare.
    ///
    /// Note that in extreme scenarios, data may be garbage collected even before TARGET_RETENTION
    /// views, if consensus storage exceeds TARGET_USAGE. For a hard lower bound on how long
    /// consensus data will be retained, see MINIMUM_RETENTION.
    ///
    /// The default of 302000 views equates to approximately 1 week (604800 seconds) at an average
    /// view time of 2s.
    #[clap(
        name = "TARGET_RETENTION",
        long = "consensus-storage-target-retention",
        env = "ESPRESSO_NODE_CONSENSUS_STORAGE_TARGET_RETENTION",
        default_value = "302000"
    )]
    pub(crate) target_retention: u64,

    /// Minimum number of views to try to retain in consensus storage before data that hasn't been
    /// archived is garbage collected.
    ///
    /// This bound allows data to be retained even if consensus storage occupies more than
    /// TARGET_USAGE. This can be used to ensure sufficient time to move consensus data to archival
    /// storage as necessary, even under extreme circumstances where otherwise garbage collection
    /// would kick in based on TARGET_RETENTION.
    ///
    /// The default of 130000 views equates to approximately 3 days (259200 seconds) at an average
    /// view time of 2s.
    #[clap(
        name = "MINIMUM_RETENTION",
        long = "consensus-storage-minimum-retention",
        env = "ESPRESSO_NODE_CONSENSUS_STORAGE_MINIMUM_RETENTION",
        default_value = "130000"
    )]
    pub(crate) minimum_retention: u64,

    /// Amount (in bytes) of data to retain in consensus storage before garbage collecting more
    /// aggressively.
    ///
    /// See also TARGET_RETENTION and MINIMUM_RETENTION.
    #[clap(
        name = "TARGET_USAGE",
        long = "consensus-storage-target-usage",
        env = "ESPRESSO_NODE_CONSENSUS_STORAGE_TARGET_USAGE",
        default_value = "1000000000"
    )]
    pub(crate) target_usage: u64,
}

impl Default for ConsensusPruningOptions {
    fn default() -> Self {
        Self::parse_from(std::iter::empty::<String>())
    }
}

#[async_trait]
impl PersistenceOptions for Options {
    type Persistence = Persistence;

    fn set_view_retention(&mut self, view_retention: u64) {
        self.consensus_pruning.target_retention = view_retention;
        self.consensus_pruning.minimum_retention = view_retention;
    }

    async fn create(&mut self) -> anyhow::Result<Self::Persistence> {
        let config = (&*self).try_into()?;
        let persistence = Persistence {
            db: SqlStorage::connect(config, StorageConnectionType::Sequencer).await?,
            gc_opt: self.consensus_pruning,
            internal_metrics: PersistenceMetricsValue::default(),
        };
        persistence.migrate_quorum_proposal_leaf_hashes().await?;
        self.pool = Some(persistence.db.pool());
        Ok(persistence)
    }

    async fn reset(self) -> anyhow::Result<()> {
        SqlStorage::connect(
            Config::try_from(&self)?.reset_schema(),
            StorageConnectionType::Sequencer,
        )
        .await?;
        Ok(())
    }
}

#[derive(Debug, Clone, Copy)]
pub enum DataMigration {
    X25519Keys,
}

impl DataMigration {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::X25519Keys => "x25519_keys",
        }
    }
}

/// Postgres-backed persistence.
#[derive(Clone, Debug)]
pub struct Persistence {
    db: SqlStorage,
    gc_opt: ConsensusPruningOptions,
    /// A reference to the internal metrics
    internal_metrics: PersistenceMetricsValue,
}

/// PostgreSQL error code for serialization failures under SERIALIZABLE isolation.
/// Transactions that fail with this code are safe to retry from scratch.
const PG_SERIALIZATION_FAILURE_CODE: &str = "40001";

#[derive(Debug)]
struct DecidedLeaf {
    info: LeafInfo<SeqTypes>,
    cert: CertificatePair<SeqTypes>,
}

fn decide_events_from_chain(
    mut chain: Vec<DecidedLeaf>,
    cert2: Option<Certificate2<SeqTypes>>,
    deciding_qc: Option<Arc<CertificatePair<SeqTypes>>>,
) -> Vec<CoordinatorEvent<SeqTypes>> {
    let split_idx = chain
        .iter()
        .position(|leaf| leaf.info.leaf.block_header().version() < versions::NEW_PROTOCOL_VERSION)
        .unwrap_or(chain.len());
    let legacy_leaves = chain.split_off(split_idx);
    let new_leaves = chain;

    let mut events = Vec::with_capacity(2);
    if !legacy_leaves.is_empty() {
        let committing_qc = legacy_leaves[0].cert.clone();
        let deciding_qc = new_leaves
            .is_empty()
            .then_some(deciding_qc)
            .flatten()
            .filter(|qc| qc.view_number() == committing_qc.view_number() + 1);
        let view_number = legacy_leaves[0].info.leaf.view_number();
        let leaf_chain = legacy_leaves
            .into_iter()
            .map(|leaf| leaf.info)
            .collect::<Vec<_>>();

        events.push(CoordinatorEvent::LegacyEvent(Event {
            view_number,
            event: EventType::Decide {
                leaf_chain: Arc::new(leaf_chain),
                committing_qc: Arc::new(committing_qc),
                deciding_qc,
                block_size: None,
            },
        }));
    }

    if new_leaves.is_empty() && cert2.is_some() {
        tracing::warn!(
            "decide_events_from_chain called with cert2 but no new-protocol leaves; cert2 will be \
             dropped"
        );
    }

    if !new_leaves.is_empty() {
        // cert1 is the QC for the newest leaf
        // ancestors are certified by
        // their successor's justify_qc. cert2 finalizes the newest leaf.
        // update() uses cert1 to build LeafQueryData for
        // the newest leaf and only attaches cert2 to it.
        let cert1 = new_leaves[0].cert.qc().clone();
        let leaf_infos = new_leaves.into_iter().map(|leaf| leaf.info).collect();

        events.push(CoordinatorEvent::NewDecide {
            leaf_infos,
            cert1,
            cert2,
        });
    }

    events
}

impl Persistence {
    /// Ensure the `leaf_hash` column is populated for all existing quorum proposals.
    ///
    /// This column was added in a migration, but because it requires computing a commitment of the
    /// existing data, it is not easy to populate in the SQL migration itself. Thus, on startup, we
    /// check if there are any just-migrated quorum proposals with a `NULL` value for this column,
    /// and if so we populate the column manually.
    async fn migrate_quorum_proposal_leaf_hashes(&self) -> anyhow::Result<()> {
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;

                let mut proposals = tx.fetch("SELECT * FROM quorum_proposals");

                let mut updates = vec![];
                while let Some(row) = proposals.next().await {
                    let row = row?;

                    let hash: Option<String> = row.try_get("leaf_hash")?;
                    if hash.is_none() {
                        let view: i64 = row.try_get("view")?;
                        let data: Vec<u8> = row.try_get("data")?;
                        let proposal: Proposal<SeqTypes, QuorumProposal<SeqTypes>> =
                            bincode::deserialize(&data)?;
                        let leaf = Leaf::from_quorum_proposal(&proposal.data);
                        let leaf_hash = Committable::commit(&leaf);
                        tracing::info!(view, %leaf_hash, "populating quorum proposal leaf hash");
                        updates.push((view, leaf_hash.to_string()));
                    }
                }
                drop(proposals);

                tx.upsert("quorum_proposals", ["view", "leaf_hash"], ["view"], updates)
                    .await?;

                tx.commit().await
            })
            .await
    }

    async fn is_migration_complete(&self, name: &str, table_name: &str) -> anyhow::Result<bool> {
        let mut tx = self.db.read().await?;
        let (completed,): (bool,) =
            query_as("SELECT completed FROM data_migrations WHERE name = $1 AND table_name = $2")
                .bind(name)
                .bind(table_name)
                .fetch_one(tx.as_mut())
                .await
                .context("migration tracking row missing - schema may be out of sync")?;
        Ok(completed)
    }

    async fn mark_migration_complete(
        tx: &mut Transaction<Write>,
        name: &str,
        table_name: &str,
        migrated_rows: usize,
    ) -> anyhow::Result<()> {
        tx.execute(
            query(
                "UPDATE data_migrations SET completed = true, migrated_rows = $1 WHERE name = $2 \
                 AND table_name = $3",
            )
            .bind(migrated_rows as i64)
            .bind(name)
            .bind(table_name),
        )
        .await?;
        Ok(())
    }

    async fn generate_decide_events(
        &self,
        deciding_qc: Option<Arc<CertificatePair<SeqTypes>>>,
        consumer: &impl EventConsumer,
    ) -> anyhow::Result<()> {
        let mut last_processed_view: Option<i64> = self
            .db
            .read()
            .await?
            .fetch_optional("SELECT last_processed_view FROM event_stream WHERE id = 1 LIMIT 1")
            .await?
            .map(|row| row.get("last_processed_view"));
        loop {
            // In SQLite, overlapping read and write transactions can lead to database errors. To
            // avoid this:
            // - start a read transaction to query and collect all the necessary data.
            // - Commit (or implicitly drop) the read transaction once the data is fetched.
            // - use the collected data to generate a "decide" event for the consumer.
            // - begin a write transaction to delete the data and update the event stream.
            let mut tx = self.db.read().await?;

            // Collect a chain of consecutive leaves, starting from the first view after the last
            // decide. This will correspond to a decide event, and defines a range of views which
            // can be garbage collected. This may even include views for which there was no leaf,
            // for which we might still have artifacts like proposals that never finalized.
            let from_view = match last_processed_view {
                Some(v) => v + 1,
                None => 0,
            };
            tracing::debug!(?from_view, "generate decide event");

            let mut parent = None;
            let mut rows = query(
                "SELECT leaf, qc, next_epoch_qc FROM anchor_leaf2 WHERE view >= $1 ORDER BY view",
            )
            .bind(from_view)
            .fetch(tx.as_mut());
            let mut leaves: Vec<(Leaf2, CertificatePair<SeqTypes>)> = vec![];
            let mut final_qc = None;
            while let Some(row) = rows.next().await {
                let row = match row {
                    Ok(row) => row,
                    Err(err) => {
                        // If there's an error getting a row, try generating an event with the rows
                        // we do have.
                        tracing::warn!("error loading row: {err:#}");
                        break;
                    },
                };

                let leaf_data: Vec<u8> = row.get("leaf");
                let leaf = bincode::deserialize::<Leaf2>(&leaf_data)?;
                let qc_data: Vec<u8> = row.get("qc");
                let qc = bincode::deserialize::<QuorumCertificate2<SeqTypes>>(&qc_data)?;
                let next_epoch_qc = match row.get::<Option<Vec<u8>>, _>("next_epoch_qc") {
                    Some(bytes) => {
                        Some(bincode::deserialize::<NextEpochQuorumCertificate2<SeqTypes>>(&bytes)?)
                    },
                    None => None,
                };
                let height = leaf.block_header().block_number();

                // Ensure we are only dealing with a consecutive chain of leaves. We don't want to
                // garbage collect any views for which we missed a leaf or decide event; at least
                // not right away, in case we need to recover that data later.
                if let Some(parent) = parent
                    && height != parent + 1
                {
                    tracing::debug!(
                        height,
                        parent,
                        "ending decide event at non-consecutive leaf"
                    );
                    break;
                }
                parent = Some(height);
                let cert = CertificatePair::new(qc, next_epoch_qc);
                final_qc = Some(cert.clone());
                leaves.push((leaf, cert));
            }
            drop(rows);

            let Some(final_qc) = final_qc else {
                // End event processing when there are no more decided views.
                tracing::debug!(from_view, "no new leaves at decide");
                return Ok(());
            };

            // Find the range of views encompassed by this leaf chain. All data in this range can be
            // processed by the consumer and then deleted.
            let from_view = leaves[0].0.view_number();
            let to_view = leaves[leaves.len() - 1].0.view_number();

            // Collect VID shares for the decide event.
            let mut vid_shares = tx
                .fetch_all(
                    query("SELECT view, data FROM vid_share2 where view >= $1 AND view <= $2")
                        .bind(from_view.u64() as i64)
                        .bind(to_view.u64() as i64),
                )
                .await?
                .into_iter()
                .map(|row| {
                    let view: i64 = row.get("view");
                    let data: Vec<u8> = row.get("data");
                    let vid_proposal = bincode::deserialize::<
                        Proposal<SeqTypes, VidDisperseShare<SeqTypes>>,
                    >(&data)?;
                    Ok((view as u64, vid_proposal))
                })
                .collect::<anyhow::Result<BTreeMap<_, _>>>()?;

            // Collect DA proposals for the decide event.
            let mut da_proposals = tx
                .fetch_all(
                    query("SELECT view, data FROM da_proposal2 where view >= $1 AND view <= $2")
                        .bind(from_view.u64() as i64)
                        .bind(to_view.u64() as i64),
                )
                .await?
                .into_iter()
                .map(|row| {
                    let view: i64 = row.get("view");
                    let data: Vec<u8> = row.get("data");
                    let da_proposal =
                        bincode::deserialize::<Proposal<SeqTypes, DaProposal2<SeqTypes>>>(&data)?;
                    Ok((view as u64, da_proposal.data))
                })
                .collect::<anyhow::Result<BTreeMap<_, _>>>()?;

            // Collect state certs for the decide event.
            let state_certs = Self::load_state_certs(&mut tx, from_view, to_view)
                .await
                .inspect_err(|err| {
                    tracing::error!(
                        ?from_view,
                        ?to_view,
                        "failed to load state certificates. error={err:#}"
                    );
                })?;

            let cert2 = tx
                .fetch_optional(
                    query("SELECT data FROM decided_cert2 WHERE view = $1")
                        .bind(to_view.u64() as i64),
                )
                .await?
                .map(|row| {
                    let bytes: Vec<u8> = row.get("data");
                    bincode::deserialize::<Certificate2<SeqTypes>>(&bytes)
                        .context("deserializing decided cert2")
                })
                .transpose()?;
            drop(tx);

            // Collate all the information by view number and construct a chain of leaves.
            let chain = leaves
                .into_iter()
                // Go in reverse chronological order, as expected by Decide events.
                .rev()
                .map(|(mut leaf, cert)| {
                    let view = leaf.view_number();

                    // Include the VID share if available.
                    let vid_proposal = vid_shares.remove(&view);
                    if vid_proposal.is_none() {
                        tracing::debug!(?view, "VID share not available at decide");
                    }
                    let vid_share = vid_proposal.as_ref().map(|proposal| proposal.data.clone());

                    // Fill in the full block payload using the DA proposals we had persisted.
                    if let Some(proposal) = da_proposals.remove(&view) {
                        let payload =
                            Payload::from_bytes(&proposal.encoded_transactions, &proposal.metadata);
                        leaf.fill_block_payload_unchecked(payload);
                    } else if view == ViewNumber::genesis() {
                        // We don't get a DA proposal for the genesis view, but we know what the
                        // payload always is.
                        leaf.fill_block_payload_unchecked(Payload::empty().0);
                    } else {
                        tracing::debug!(?view, "DA proposal not available at decide");
                    }

                    let state_cert = state_certs.get(&view).cloned();

                    let info = LeafInfo {
                        leaf,
                        vid_share,
                        state_cert,
                        // Note: the following fields are not used in Decide event processing,
                        // and should be removed. For now, we just default them.
                        state: Default::default(),
                        delta: Default::default(),
                    };
                    DecidedLeaf { info, cert }
                })
                .collect();

            tracing::debug!(
                ?from_view,
                ?to_view,
                ?final_qc,
                ?chain,
                "generating decide event"
            );

            for event in decide_events_from_chain(chain, cert2, deciding_qc.clone()) {
                consumer.handle_event(&event).await?;
            }

            let from_view_i64 = from_view.u64() as i64;
            let to_view_i64 = to_view.u64() as i64;
            let serialized_state_certs = state_certs
                .into_iter()
                .map(|(epoch, cert)| Ok((epoch as i64, bincode::serialize(&cert)?)))
                .collect::<anyhow::Result<Vec<(i64, Vec<u8>)>>>()?;

            // Now that we have definitely processed leaves up to `to_view`, we can update
            // `last_processed_view` so we don't process these leaves again. We may still fail at
            // this point, or shut down, and fail to complete this update. At worst this will lead
            // to us sending a duplicate decide event the next time we are called; this is fine as
            // the event consumer is required to be idempotent.
            WRITE_BACKOFF
                .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                    let mut tx = self.db.write().await?;
                    tx.upsert(
                        "event_stream",
                        ["id", "last_processed_view"],
                        ["id"],
                        [(1i32, to_view_i64)],
                    )
                    .await?;

                    // Store all the finalized state certs
                    for (epoch, state_cert_bytes) in &serialized_state_certs {
                        tx.upsert(
                            "finalized_state_cert",
                            ["epoch", "state_cert"],
                            ["epoch"],
                            [(*epoch, state_cert_bytes.clone())],
                        )
                        .await?;
                    }

                    // Delete the data that has been fully processed.
                    tx.execute(
                        query("DELETE FROM vid_share2 where view >= $1 AND view <= $2")
                            .bind(from_view_i64)
                            .bind(to_view_i64),
                    )
                    .await?;
                    tx.execute(
                        query("DELETE FROM da_proposal2 where view >= $1 AND view <= $2")
                            .bind(from_view_i64)
                            .bind(to_view_i64),
                    )
                    .await?;
                    tx.execute(
                        query("DELETE FROM quorum_proposals2 where view >= $1 AND view <= $2")
                            .bind(from_view_i64)
                            .bind(to_view_i64),
                    )
                    .await?;
                    tx.execute(
                        query("DELETE FROM quorum_certificate2 where view >= $1 AND view <= $2")
                            .bind(from_view_i64)
                            .bind(to_view_i64),
                    )
                    .await?;
                    tx.execute(
                        query("DELETE FROM state_cert where view >= $1 AND view <= $2")
                            .bind(from_view_i64)
                            .bind(to_view_i64),
                    )
                    .await?;
                    tx.execute(
                        query("DELETE FROM decided_cert2 where view >= $1 AND view <= $2")
                            .bind(from_view_i64)
                            .bind(to_view_i64),
                    )
                    .await?;

                    // Clean up leaves, but do not delete the most recent one (all leaves with a view
                    // number less than the given value). This is necessary to ensure that, in case of
                    // a restart, we can resume from the last decided leaf.
                    tx.execute(
                        query("DELETE FROM anchor_leaf2 WHERE view >= $1 AND view < $2")
                            .bind(from_view_i64)
                            .bind(to_view_i64),
                    )
                    .await?;

                    tx.commit().await?;
                    Ok(())
                })
                .await?;
            last_processed_view = Some(to_view_i64);
        }
    }

    async fn load_state_certs(
        tx: &mut Transaction<Read>,
        from_view: ViewNumber,
        to_view: ViewNumber,
    ) -> anyhow::Result<BTreeMap<u64, LightClientStateUpdateCertificateV2<SeqTypes>>> {
        let rows = tx
            .fetch_all(
                query("SELECT view, state_cert FROM state_cert WHERE view >= $1 AND view <= $2")
                    .bind(from_view.u64() as i64)
                    .bind(to_view.u64() as i64),
            )
            .await?;

        let mut result = BTreeMap::new();

        for row in rows {
            let data: Vec<u8> = row.get("state_cert");

            let cert: LightClientStateUpdateCertificateV2<SeqTypes> = bincode::deserialize(&data)
                .or_else(|err_v2| {
                bincode::deserialize::<LightClientStateUpdateCertificateV1<SeqTypes>>(&data)
                    .map(Into::into)
                    .context(format!(
                        "Failed to deserialize LightClientStateUpdateCertificate: with v1 and v2. \
                         error: {err_v2}"
                    ))
            })?;

            result.insert(cert.epoch.u64(), cert);
        }

        Ok(result)
    }

    #[tracing::instrument(skip(self))]
    async fn prune(&self, cur_view: ViewNumber) -> anyhow::Result<()> {
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;

                // Prune everything older than the target retention period.
                prune_to_view(
                    &mut tx,
                    cur_view.u64().saturating_sub(self.gc_opt.target_retention),
                )
                .await?;

                // Check our storage usage; if necessary we will prune more aggressively (up to the
                // minimum retention) to get below the target usage.
                #[cfg(feature = "embedded-db")]
                let usage_query = format!(
                    "SELECT sum(pgsize) FROM dbstat WHERE name IN ({})",
                    PRUNE_TABLES
                        .iter()
                        .map(|table| format!("'{table}'"))
                        .join(",")
                );

                #[cfg(not(feature = "embedded-db"))]
                let usage_query = {
                    let table_sizes = PRUNE_TABLES
                        .iter()
                        .map(|table| format!("pg_table_size('{table}')"))
                        .join(" + ");
                    format!("SELECT {table_sizes}")
                };

                let (usage,): (i64,) = query_as(&usage_query).fetch_one(tx.as_mut()).await?;
                tracing::debug!(usage, "consensus storage usage after pruning");

                if (usage as u64) > self.gc_opt.target_usage {
                    tracing::warn!(
                        usage,
                        gc_opt = ?self.gc_opt,
                        "consensus storage is running out of space, pruning to minimum retention"
                    );
                    prune_to_view(
                        &mut tx,
                        cur_view.u64().saturating_sub(self.gc_opt.minimum_retention),
                    )
                    .await?;
                }

                tx.commit().await
            })
            .await
    }
}

/// Maximum number of retries on PostgreSQL serialization conflicts (error 40001).
const WRITE_RETRY_MAX: u32 = 5;

/// Backoff parameters for write-transaction retries.
const WRITE_BACKOFF: BackoffParams = BackoffParams::new(
    Duration::from_millis(10),
    Duration::from_millis(500),
    2,
    Ratio {
        numerator: 5,
        denominator: 10,
    },
);

fn is_serialization_error(err: &anyhow::Error) -> bool {
    err.chain()
        .filter_map(|e| e.downcast_ref::<sqlx::Error>())
        .filter_map(|e| e.as_database_error())
        .any(|e| e.code().as_deref() == Some(PG_SERIALIZATION_FAILURE_CODE))
}

const PRUNE_TABLES: &[&str] = &[
    "anchor_leaf2",
    "vid_share2",
    "da_proposal2",
    "quorum_proposals2",
    "quorum_certificate2",
];

async fn prune_to_view(tx: &mut Transaction<Write>, view: u64) -> anyhow::Result<()> {
    if view == 0 {
        // Nothing to prune, the entire chain is younger than the retention period.
        return Ok(());
    }
    tracing::debug!(view, "pruning consensus storage");

    for table in PRUNE_TABLES {
        let res = query(&format!("DELETE FROM {table} WHERE view < $1"))
            .bind(view as i64)
            .execute(tx.as_mut())
            .await
            .context(format!("pruning {table}"))?;
        if res.rows_affected() > 0 {
            tracing::info!(
                "garbage collected {} rows from {table}",
                res.rows_affected()
            );
        }
    }

    Ok(())
}

#[async_trait]
impl SequencerPersistence for Persistence {
    async fn migrate_reward_merkle_tree_v2(&self) -> anyhow::Result<()> {
        let batch_size: i64 = 1000;

        let result = {
            let mut tx = self.db.read().await?;
            query_as::<(bool, i64)>(
                "SELECT completed, migrated_rows FROM epoch_migration WHERE table_name = \
                 'reward_merkle_tree_v2_data'",
            )
            .fetch_optional(tx.as_mut())
            .await?
        };

        let (is_completed, mut offset) = result.unwrap_or((false, 0));

        if is_completed {
            tracing::info!("reward_merkle_tree_v2 migration already done");
            return Ok(());
        }

        let max_height: Option<i64> = {
            let mut tx = self.db.read().await?;
            query_as::<(Option<i64>,)>("SELECT MAX(created) FROM reward_merkle_tree_v2")
                .fetch_one(tx.as_mut())
                .await?
                .0
        };

        let max_height = match max_height {
            Some(h) => h,
            None => {
                tracing::info!("no reward data found in reward_merkle_tree_v2, skipping migration");
                return Ok(());
            },
        };

        tracing::warn!(
            "migrating reward_merkle_tree_v2 to reward_merkle_tree_v2_data at height \
             {max_height}..."
        );

        let mut balances: Vec<(RewardAccountV2, RewardAmount)> = Vec::new();

        loop {
            let mut tx = self.db.read().await?;

            #[cfg(not(feature = "embedded-db"))]
            let rows = query_as::<(serde_json::Value, serde_json::Value)>(
                "SELECT DISTINCT ON (idx) idx, entry
                   FROM reward_merkle_tree_v2
                  WHERE idx IS NOT NULL AND entry IS NOT NULL
                  ORDER BY idx, created DESC
                  LIMIT $1 OFFSET $2",
            )
            .bind(batch_size)
            .bind(offset)
            .fetch_all(tx.as_mut())
            .await
            .context("loading reward accounts from reward_merkle_tree_v2")?;

            #[cfg(feature = "embedded-db")]
            let rows = query_as::<(serde_json::Value, serde_json::Value)>(
                "SELECT idx, entry FROM (
                     SELECT idx, entry, ROW_NUMBER() OVER (PARTITION BY idx ORDER BY created DESC) \
                 as rn
                       FROM reward_merkle_tree_v2
                      WHERE idx IS NOT NULL AND entry IS NOT NULL
                 ) sub
                 WHERE rn = 1 ORDER BY idx
                 LIMIT $1 OFFSET $2",
            )
            .bind(batch_size)
            .bind(offset)
            .fetch_all(tx.as_mut())
            .await
            .context("loading reward accounts from reward_merkle_tree_v2")?;

            drop(tx);

            if rows.is_empty() {
                break;
            }

            let rows_count = rows.len();

            for (idx, entry) in rows {
                let account: RewardAccountV2 =
                    serde_json::from_value(idx).context("deserializing reward account")?;
                let balance: RewardAmount = serde_json::from_value(entry).context(format!(
                    "deserializing reward balance for account {account}"
                ))?;
                balances.push((account, balance));
            }

            offset += rows_count as i64;
            let mut tx = self.db.write().await?;
            tx.upsert(
                "epoch_migration",
                ["table_name", "completed", "migrated_rows"],
                ["table_name"],
                [("reward_merkle_tree_v2_data".to_string(), false, offset)],
            )
            .await?;
            tx.commit().await?;

            tracing::info!(
                "reward_merkle_tree_v2 progress: rows={} offset={}",
                rows_count,
                offset
            );

            if rows_count < batch_size as usize {
                break;
            }
        }

        if balances.is_empty() {
            tracing::info!("no reward accounts found, skipping tree rebuild");
            return Ok(());
        }

        tracing::info!(
            "rebuilding RewardMerkleTreeV2 from {} accounts",
            balances.len()
        );

        let tree = RewardMerkleTreeV2::from_kv_set(REWARD_MERKLE_TREE_V2_HEIGHT, balances)
            .context("failed to rebuild RewardMerkleTreeV2 from balances")?;

        let mut tx = self.db.read().await?;
        let header = tx
            .get_header(BlockId::<SeqTypes>::from(max_height as usize))
            .await
            .context(format!("header {max_height} not available"))?;
        drop(tx);

        match header.reward_merkle_tree_root() {
            Either::Right(expected_root) => {
                ensure!(
                    tree.commitment() == expected_root,
                    "rebuilt RewardMerkleTreeV2 commitment {} does not match header commitment {} \
                     at height {max_height}",
                    tree.commitment(),
                    expected_root,
                );
            },
            Either::Left(_) => {
                bail!(
                    "header at height {max_height} has a v1 reward merkle tree root, expected v2"
                );
            },
        }

        let tree_data: RewardMerkleTreeV2Data = (&tree)
            .try_into()
            .context("failed to convert RewardMerkleTreeV2 to RewardMerkleTreeV2Data")?;
        let serialized =
            bincode::serialize(&tree_data).context("failed to serialize RewardMerkleTreeV2Data")?;

        let mut tx = self.db.write().await?;
        tx.upsert(
            "reward_merkle_tree_v2_data",
            ["height", "balances"],
            ["height"],
            [(max_height, serialized)],
        )
        .await?;
        tx.commit().await?;

        // Mark migration as complete, and clean up old tables.
        let mut tx = self.db.write().await?;
        tx.upsert(
            "epoch_migration",
            ["table_name", "completed", "migrated_rows"],
            ["table_name"],
            [("reward_merkle_tree_v2_data".to_string(), true, offset)],
        )
        .await?;
        let truncate = if cfg!(feature = "embedded-db") {
            "DELETE FROM"
        } else {
            "TRUNCATE"
        };
        query(&format!("{truncate} reward_merkle_tree_v2"))
            .execute(tx.as_mut())
            .await?;
        query(&format!("{truncate} reward_merkle_tree"))
            .execute(tx.as_mut())
            .await?;
        tx.commit().await?;

        tracing::warn!("migrated reward_merkle_tree_v2 at height {max_height}");

        Ok(())
    }

    fn into_catchup_provider(
        self,
        backoff: BackoffParams,
    ) -> anyhow::Result<Arc<dyn StateCatchup>> {
        Ok(Arc::new(SqlStateCatchup::new(Arc::new(self.db), backoff)))
    }

    async fn load_config(&self) -> anyhow::Result<Option<NetworkConfig>> {
        tracing::info!("loading config from Postgres");

        // Select the most recent config (although there should only be one).
        let Some(row) = self
            .db
            .read()
            .await?
            .fetch_optional("SELECT config FROM network_config ORDER BY id DESC LIMIT 1")
            .await?
        else {
            tracing::info!("config not found");
            return Ok(None);
        };
        let json = row.try_get("config")?;

        let json = migrate_network_config(json).context("migration of network config failed")?;
        let config = serde_json::from_value(json).context("malformed config file")?;

        Ok(Some(config))
    }

    async fn save_config(&self, cfg: &NetworkConfig) -> anyhow::Result<()> {
        tracing::info!("saving config to database");
        let json = serde_json::to_value(cfg)?;

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.execute(
                    query("INSERT INTO network_config (config) VALUES ($1)").bind(json.clone()),
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn append_decided_leaves(
        &self,
        view: ViewNumber,
        leaf_chain: impl IntoIterator<Item = (&LeafInfo<SeqTypes>, CertificatePair<SeqTypes>)> + Send,
        deciding_qc: Option<Arc<CertificatePair<SeqTypes>>>,
        consumer: &(impl EventConsumer + 'static),
    ) -> anyhow::Result<()> {
        let values = leaf_chain
            .into_iter()
            .map(|(info, cert)| {
                // The leaf may come with a large payload attached. We don't care about this payload
                // because we already store it separately, as part of the DA proposal. Storing it
                // here contributes to load on the DB for no reason, so we remove it before
                // serializing the leaf.
                let mut leaf = info.leaf.clone();
                leaf.unfill_block_payload();

                let view = cert.view_number().u64() as i64;
                let leaf_bytes = bincode::serialize(&leaf)?;
                let qc_bytes = bincode::serialize(cert.qc())?;
                let next_epoch_qc_bytes = match cert.next_epoch_qc() {
                    Some(qc) => Some(bincode::serialize(qc)?),
                    None => None,
                };
                Ok((view, leaf_bytes, qc_bytes, next_epoch_qc_bytes))
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        // First, append the new leaves. We do this in its own transaction because even if GC or the
        // event consumer later fails, there is no need to abort the storage of the leaves.
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "anchor_leaf2",
                    ["view", "leaf", "qc", "next_epoch_qc"],
                    ["view"],
                    values.clone(),
                )
                .await?;
                tx.commit().await
            })
            .await?;

        // Generate an event for the new leaves and, only if it succeeds, clean up data we no longer
        // need.
        if let Err(err) = self.generate_decide_events(deciding_qc, consumer).await {
            // GC/event processing failure is not an error, since by this point we have at least
            // managed to persist the decided leaves successfully, and GC will just run again at the
            // next decide. Log an error but do not return it.
            tracing::warn!(?view, "event processing failed: {err:#}");
            return Ok(());
        }

        // Garbage collect data which was not included in any decide event, but which at this point
        // is old enough to just forget about.
        if let Err(err) = self.prune(view).await {
            tracing::warn!(?view, "pruning failed: {err:#}");
        }

        Ok(())
    }

    async fn load_latest_acted_view(&self) -> anyhow::Result<Option<ViewNumber>> {
        Ok(self
            .db
            .read()
            .await?
            .fetch_optional(query("SELECT view FROM highest_voted_view WHERE id = 0"))
            .await?
            .map(|row| {
                let view: i64 = row.get("view");
                ViewNumber::new(view as u64)
            }))
    }

    async fn load_restart_view(&self) -> anyhow::Result<Option<ViewNumber>> {
        Ok(self
            .db
            .read()
            .await?
            .fetch_optional(query("SELECT view FROM restart_view WHERE id = 0"))
            .await?
            .map(|row| {
                let view: i64 = row.get("view");
                ViewNumber::new(view as u64)
            }))
    }

    async fn load_anchor_leaf(
        &self,
    ) -> anyhow::Result<Option<(Leaf2, QuorumCertificate2<SeqTypes>)>> {
        let Some(row) = self
            .db
            .read()
            .await?
            .fetch_optional("SELECT leaf, qc FROM anchor_leaf2 ORDER BY view DESC LIMIT 1")
            .await?
        else {
            return Ok(None);
        };

        let leaf_bytes: Vec<u8> = row.get("leaf");
        let leaf2: Leaf2 = bincode::deserialize(&leaf_bytes)?;

        let qc_bytes: Vec<u8> = row.get("qc");
        let qc2: QuorumCertificate2<SeqTypes> = bincode::deserialize(&qc_bytes)?;

        Ok(Some((leaf2, qc2)))
    }

    async fn load_anchor_view(&self) -> anyhow::Result<ViewNumber> {
        let mut tx = self.db.read().await?;
        let (view,) = query_as::<(i64,)>("SELECT coalesce(max(view), 0) FROM anchor_leaf2")
            .fetch_one(tx.as_mut())
            .await?;
        Ok(ViewNumber::new(view as u64))
    }

    async fn load_da_proposal(
        &self,
        view: ViewNumber,
    ) -> anyhow::Result<Option<Proposal<SeqTypes, DaProposal2<SeqTypes>>>> {
        let result = self
            .db
            .read()
            .await?
            .fetch_optional(
                query("SELECT data FROM da_proposal2 where view = $1").bind(view.u64() as i64),
            )
            .await?;

        result
            .map(|row| {
                let bytes: Vec<u8> = row.get("data");
                anyhow::Result::<_>::Ok(bincode::deserialize(&bytes)?)
            })
            .transpose()
    }

    async fn load_vid_share(
        &self,
        view: ViewNumber,
    ) -> anyhow::Result<Option<Proposal<SeqTypes, VidDisperseShare<SeqTypes>>>> {
        let result = self
            .db
            .read()
            .await?
            .fetch_optional(
                query("SELECT data FROM vid_share2 where view = $1").bind(view.u64() as i64),
            )
            .await?;

        result
            .map(|row| {
                let bytes: Vec<u8> = row.get("data");
                anyhow::Result::<_>::Ok(bincode::deserialize(&bytes)?)
            })
            .transpose()
    }

    async fn load_quorum_proposals(
        &self,
    ) -> anyhow::Result<BTreeMap<ViewNumber, Proposal<SeqTypes, QuorumProposalWrapper<SeqTypes>>>>
    {
        let rows = self
            .db
            .read()
            .await?
            .fetch_all("SELECT * FROM quorum_proposals2")
            .await?;

        Ok(BTreeMap::from_iter(
            rows.into_iter()
                .map(|row| {
                    let view: i64 = row.get("view");
                    let view_number: ViewNumber = ViewNumber::new(view.try_into()?);
                    let bytes: Vec<u8> = row.get("data");
                    let proposal: Proposal<SeqTypes, QuorumProposalWrapper<SeqTypes>> =
                        bincode::deserialize(&bytes).or_else(|error| {
                            bincode::deserialize::<
                                Proposal<SeqTypes, QuorumProposalWrapperLegacy<SeqTypes>>,
                            >(&bytes)
                            .map(convert_proposal)
                            .inspect_err(|err_v3| {
                                tracing::warn!(
                                    ?view_number,
                                    %error,
                                    %err_v3,
                                    "ignoring malformed quorum proposal DB row"
                                );
                            })
                        })?;
                    Ok((view_number, proposal))
                })
                .collect::<anyhow::Result<Vec<_>>>()?,
        ))
    }

    async fn load_quorum_proposal(
        &self,
        view: ViewNumber,
    ) -> anyhow::Result<Proposal<SeqTypes, QuorumProposalWrapper<SeqTypes>>> {
        let mut tx = self.db.read().await?;
        let (data,) =
            query_as::<(Vec<u8>,)>("SELECT data FROM quorum_proposals2 WHERE view = $1 LIMIT 1")
                .bind(view.u64() as i64)
                .fetch_one(tx.as_mut())
                .await?;
        let proposal: Proposal<SeqTypes, QuorumProposalWrapper<SeqTypes>> =
            bincode::deserialize(&data).or_else(|error| {
                bincode::deserialize::<Proposal<SeqTypes, QuorumProposalWrapperLegacy<SeqTypes>>>(
                    &data,
                )
                .map(convert_proposal)
                .context(format!(
                    "Failed to deserialize quorum proposal for view {view}. error={error}"
                ))
            })?;

        Ok(proposal)
    }

    async fn append_vid(
        &self,
        proposal: &Proposal<SeqTypes, VidDisperseShare<SeqTypes>>,
    ) -> anyhow::Result<()> {
        let view = proposal.data.view_number().u64();
        let payload_hash = proposal.data.payload_commitment();
        let data_bytes = bincode::serialize(proposal).unwrap();

        let now = Instant::now();
        let res = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "vid_share2",
                    ["view", "data", "payload_hash"],
                    ["view"],
                    [(view as i64, data_bytes.clone(), payload_hash.to_string())],
                )
                .await?;
                tx.commit().await
            })
            .await;
        self.internal_metrics
            .internal_append_vid_duration
            .add_point(now.elapsed().as_secs_f64());
        res
    }

    async fn append_da(
        &self,
        proposal: &Proposal<SeqTypes, DaProposal<SeqTypes>>,
        vid_commit: VidCommitment,
    ) -> anyhow::Result<()> {
        let data = &proposal.data;
        let view = data.view_number().u64();
        let data_bytes = bincode::serialize(proposal).unwrap();

        let now = Instant::now();
        let res = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "da_proposal",
                    ["view", "data", "payload_hash"],
                    ["view"],
                    [(view as i64, data_bytes.clone(), vid_commit.to_string())],
                )
                .await?;
                tx.commit().await
            })
            .await;
        self.internal_metrics
            .internal_append_da_duration
            .add_point(now.elapsed().as_secs_f64());
        res
    }

    async fn record_action(
        &self,
        view: ViewNumber,
        _epoch: Option<EpochNumber>,
        action: HotShotAction,
    ) -> anyhow::Result<()> {
        // Todo Remove this after https://github.com/EspressoSystems/espresso-network/issues/1931
        if !matches!(action, HotShotAction::Propose | HotShotAction::Vote) {
            return Ok(());
        }

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let stmt = format!(
                    "INSERT INTO highest_voted_view (id, view) VALUES (0, $1)
                ON CONFLICT (id) DO UPDATE SET view = {MAX_FN}(highest_voted_view.view, \
                     excluded.view)"
                );

                let mut tx = self.db.write().await?;
                tx.execute(query(&stmt).bind(view.u64() as i64)).await?;

                if matches!(action, HotShotAction::Vote) {
                    let restart_view = view + 1;
                    let stmt = format!(
                        "INSERT INTO restart_view (id, view) VALUES (0, $1)
                    ON CONFLICT (id) DO UPDATE SET view = {MAX_FN}(restart_view.view, \
                         excluded.view)"
                    );
                    tx.execute(query(&stmt).bind(restart_view.u64() as i64))
                        .await?;
                }

                tx.commit().await
            })
            .await
    }

    async fn append_quorum_proposal2(
        &self,
        proposal: &Proposal<SeqTypes, QuorumProposalWrapper<SeqTypes>>,
    ) -> anyhow::Result<()> {
        let view_number = proposal.data.view_number().u64();

        let proposal_bytes = bincode::serialize(&proposal).context("serializing proposal")?;
        let leaf_hash = Committable::commit(&Leaf2::from_quorum_proposal(&proposal.data));

        // We also keep track of any QC we see in case we need it to recover our archival storage.
        let justify_qc = proposal.data.justify_qc();
        let justify_qc_bytes = bincode::serialize(&justify_qc).context("serializing QC")?;
        let justify_qc_view = justify_qc.view_number.u64() as i64;
        let justify_qc_leaf_commit = justify_qc.data.leaf_commit.to_string();

        let now = Instant::now();
        let res = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "quorum_proposals2",
                    ["view", "leaf_hash", "data"],
                    ["view"],
                    [(
                        view_number as i64,
                        leaf_hash.to_string(),
                        proposal_bytes.clone(),
                    )],
                )
                .await?;
                tx.upsert(
                    "quorum_certificate2",
                    ["view", "leaf_hash", "data"],
                    ["view"],
                    [(
                        justify_qc_view,
                        justify_qc_leaf_commit.clone(),
                        justify_qc_bytes.clone(),
                    )],
                )
                .await?;
                tx.commit().await
            })
            .await;
        self.internal_metrics
            .internal_append_quorum2_duration
            .add_point(now.elapsed().as_secs_f64());
        res
    }

    async fn append_cert2(
        &self,
        view: ViewNumber,
        cert2: Certificate2<SeqTypes>,
    ) -> anyhow::Result<()> {
        let data = bincode::serialize(&cert2).context("serializing cert2")?;
        let view_i64 = view.u64() as i64;
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "decided_cert2",
                    ["view", "data"],
                    ["view"],
                    [(view_i64, data.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn load_cert2(&self, view: ViewNumber) -> anyhow::Result<Option<Certificate2<SeqTypes>>> {
        let row = self
            .db
            .read()
            .await?
            .fetch_optional(
                query("SELECT data FROM decided_cert2 WHERE view = $1").bind(view.u64() as i64),
            )
            .await?;
        row.map(|row| {
            let bytes: Vec<u8> = row.get("data");
            bincode::deserialize::<Certificate2<SeqTypes>>(&bytes).context("deserializing cert2")
        })
        .transpose()
    }

    async fn load_upgrade_certificate(
        &self,
    ) -> anyhow::Result<Option<UpgradeCertificate<SeqTypes>>> {
        let result = self
            .db
            .read()
            .await?
            .fetch_optional("SELECT * FROM upgrade_certificate where id = true")
            .await?;

        result
            .map(|row| {
                let bytes: Vec<u8> = row.get("data");
                anyhow::Result::<_>::Ok(bincode::deserialize(&bytes)?)
            })
            .transpose()
    }

    async fn store_upgrade_certificate(
        &self,
        decided_upgrade_certificate: Option<UpgradeCertificate<SeqTypes>>,
    ) -> anyhow::Result<()> {
        let certificate = match decided_upgrade_certificate {
            Some(cert) => cert,
            None => return Ok(()),
        };
        let upgrade_certificate_bytes =
            bincode::serialize(&certificate).context("serializing upgrade certificate")?;
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "upgrade_certificate",
                    ["id", "data"],
                    ["id"],
                    [(true, upgrade_certificate_bytes.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn migrate_anchor_leaf(&self) -> anyhow::Result<()> {
        let batch_size: i64 = 10000;
        let mut tx = self.db.read().await?;

        // The SQL migration populates the table name and sets a default value of 0 for migrated rows.
        // so, fetch_one() would always return a row
        // The number of migrated rows is updated after each batch insert.
        // This allows the types migration to resume from where it left off.
        let (is_completed, mut offset) = query_as::<(bool, i64)>(
            "SELECT completed, migrated_rows from epoch_migration WHERE table_name = 'anchor_leaf'",
        )
        .fetch_one(tx.as_mut())
        .await?;

        if is_completed {
            tracing::info!("decided leaves already migrated");
            return Ok(());
        }

        tracing::warn!("migrating decided leaves..");
        loop {
            let mut tx = self.db.read().await?;
            let rows = query(
                "SELECT view, leaf, qc FROM anchor_leaf WHERE view >= $1 ORDER BY view LIMIT $2",
            )
            .bind(offset)
            .bind(batch_size)
            .fetch_all(tx.as_mut())
            .await?;

            drop(tx);
            if rows.is_empty() {
                break;
            }
            let mut values = Vec::new();

            for row in rows.iter() {
                let leaf: Vec<u8> = row.try_get("leaf")?;
                let qc: Vec<u8> = row.try_get("qc")?;
                let leaf1: Leaf = bincode::deserialize(&leaf)?;
                let qc1: QuorumCertificate<SeqTypes> = bincode::deserialize(&qc)?;
                let view: i64 = row.try_get("view")?;

                let leaf2: Leaf2 = leaf1.into();
                let qc2: QuorumCertificate2<SeqTypes> = qc1.to_qc2();

                let leaf2_bytes = bincode::serialize(&leaf2)?;
                let qc2_bytes = bincode::serialize(&qc2)?;

                values.push((view, leaf2_bytes, qc2_bytes));
            }

            let mut query_builder: sqlx::QueryBuilder<Db> =
                sqlx::QueryBuilder::new("INSERT INTO anchor_leaf2 (view, leaf, qc) ");

            offset = values.last().context("last row")?.0;

            query_builder.push_values(values, |mut b, (view, leaf, qc)| {
                b.push_bind(view).push_bind(leaf).push_bind(qc);
            });

            // Offset tracking prevents duplicate inserts
            // Added as a safeguard.
            query_builder.push(" ON CONFLICT DO NOTHING");

            let query = query_builder.build();

            let mut tx = self.db.write().await?;
            query.execute(tx.as_mut()).await?;

            tx.upsert(
                "epoch_migration",
                ["table_name", "completed", "migrated_rows"],
                ["table_name"],
                [("anchor_leaf".to_string(), false, offset)],
            )
            .await?;
            tx.commit().await?;

            tracing::info!(
                "anchor leaf migration progress: rows={} offset={}",
                rows.len(),
                offset
            );

            if rows.len() < batch_size as usize {
                break;
            }
        }

        tracing::warn!("migrated decided leaves");

        let mut tx = self.db.write().await?;
        tx.upsert(
            "epoch_migration",
            ["table_name", "completed", "migrated_rows"],
            ["table_name"],
            [("anchor_leaf".to_string(), true, offset)],
        )
        .await?;
        tx.commit().await?;

        tracing::info!("updated epoch_migration table for anchor_leaf");

        Ok(())
    }

    async fn migrate_da_proposals(&self) -> anyhow::Result<()> {
        let batch_size: i64 = 10000;
        let mut tx = self.db.read().await?;

        let (is_completed, mut offset) = query_as::<(bool, i64)>(
            "SELECT completed, migrated_rows from epoch_migration WHERE table_name = 'da_proposal'",
        )
        .fetch_one(tx.as_mut())
        .await?;

        if is_completed {
            tracing::info!("da proposals migration already done");
            return Ok(());
        }

        tracing::warn!("migrating da proposals..");

        loop {
            let mut tx = self.db.read().await?;
            let rows = query(
                "SELECT payload_hash, data FROM da_proposal WHERE view >= $1 ORDER BY view LIMIT \
                 $2",
            )
            .bind(offset)
            .bind(batch_size)
            .fetch_all(tx.as_mut())
            .await?;

            drop(tx);
            if rows.is_empty() {
                break;
            }
            let mut values = Vec::new();

            for row in rows.iter() {
                let data: Vec<u8> = row.try_get("data")?;
                let payload_hash: String = row.try_get("payload_hash")?;

                let da_proposal: Proposal<SeqTypes, DaProposal<SeqTypes>> =
                    bincode::deserialize(&data)?;
                let da_proposal2: Proposal<SeqTypes, DaProposal2<SeqTypes>> =
                    convert_proposal(da_proposal);

                let view = da_proposal2.data.view_number.u64() as i64;
                let data = bincode::serialize(&da_proposal2)?;

                values.push((view, payload_hash, data));
            }

            let mut query_builder: sqlx::QueryBuilder<Db> =
                sqlx::QueryBuilder::new("INSERT INTO da_proposal2 (view, payload_hash, data) ");

            offset = values.last().context("last row")?.0;
            query_builder.push_values(values, |mut b, (view, payload_hash, data)| {
                b.push_bind(view).push_bind(payload_hash).push_bind(data);
            });
            query_builder.push(" ON CONFLICT DO NOTHING");
            let query = query_builder.build();

            let mut tx = self.db.write().await?;
            query.execute(tx.as_mut()).await?;

            tx.upsert(
                "epoch_migration",
                ["table_name", "completed", "migrated_rows"],
                ["table_name"],
                [("da_proposal".to_string(), false, offset)],
            )
            .await?;
            tx.commit().await?;

            tracing::info!(
                "DA proposals migration progress: rows={} offset={}",
                rows.len(),
                offset
            );
            if rows.len() < batch_size as usize {
                break;
            }
        }

        tracing::warn!("migrated da proposals");

        let mut tx = self.db.write().await?;
        tx.upsert(
            "epoch_migration",
            ["table_name", "completed", "migrated_rows"],
            ["table_name"],
            [("da_proposal".to_string(), true, offset)],
        )
        .await?;
        tx.commit().await?;

        tracing::info!("updated epoch_migration table for da_proposal");

        Ok(())
    }

    async fn migrate_vid_shares(&self) -> anyhow::Result<()> {
        let batch_size: i64 = 10000;

        let mut tx = self.db.read().await?;

        let (is_completed, mut offset) = query_as::<(bool, i64)>(
            "SELECT completed, migrated_rows from epoch_migration WHERE table_name = 'vid_share'",
        )
        .fetch_one(tx.as_mut())
        .await?;

        if is_completed {
            tracing::info!("vid_share migration already done");
            return Ok(());
        }

        tracing::warn!("migrating vid shares..");
        loop {
            let mut tx = self.db.read().await?;
            let rows = query(
                "SELECT payload_hash, data FROM vid_share WHERE view >= $1 ORDER BY view LIMIT $2",
            )
            .bind(offset)
            .bind(batch_size)
            .fetch_all(tx.as_mut())
            .await?;

            drop(tx);
            if rows.is_empty() {
                break;
            }
            let mut values = Vec::new();

            for row in rows.iter() {
                let data: Vec<u8> = row.try_get("data")?;
                let payload_hash: String = row.try_get("payload_hash")?;

                let vid_share: Proposal<SeqTypes, VidDisperseShare0<SeqTypes>> =
                    bincode::deserialize(&data)?;
                let vid_share2: Proposal<SeqTypes, VidDisperseShare<SeqTypes>> =
                    convert_proposal(vid_share);

                let view = vid_share2.data.view_number().u64() as i64;
                let data = bincode::serialize(&vid_share2)?;

                values.push((view, payload_hash, data));
            }

            let mut query_builder: sqlx::QueryBuilder<Db> =
                sqlx::QueryBuilder::new("INSERT INTO vid_share2 (view, payload_hash, data) ");

            offset = values.last().context("last row")?.0;

            query_builder.push_values(values, |mut b, (view, payload_hash, data)| {
                b.push_bind(view).push_bind(payload_hash).push_bind(data);
            });

            let query = query_builder.build();

            let mut tx = self.db.write().await?;
            query.execute(tx.as_mut()).await?;

            tx.upsert(
                "epoch_migration",
                ["table_name", "completed", "migrated_rows"],
                ["table_name"],
                [("vid_share".to_string(), false, offset)],
            )
            .await?;
            tx.commit().await?;

            tracing::info!(
                "VID shares migration progress: rows={} offset={}",
                rows.len(),
                offset
            );
            if rows.len() < batch_size as usize {
                break;
            }
        }

        tracing::warn!("migrated vid shares");

        let mut tx = self.db.write().await?;
        tx.upsert(
            "epoch_migration",
            ["table_name", "completed", "migrated_rows"],
            ["table_name"],
            [("vid_share".to_string(), true, offset)],
        )
        .await?;
        tx.commit().await?;

        tracing::info!("updated epoch_migration table for vid_share");

        Ok(())
    }

    async fn migrate_quorum_proposals(&self) -> anyhow::Result<()> {
        let batch_size: i64 = 10000;
        let mut tx = self.db.read().await?;

        let (is_completed, mut offset) = query_as::<(bool, i64)>(
            "SELECT completed, migrated_rows from epoch_migration WHERE table_name = \
             'quorum_proposals'",
        )
        .fetch_one(tx.as_mut())
        .await?;

        if is_completed {
            tracing::info!("quorum proposals migration already done");
            return Ok(());
        }

        tracing::warn!("migrating quorum proposals..");

        loop {
            let mut tx = self.db.read().await?;
            let rows = query(
                "SELECT view, leaf_hash, data FROM quorum_proposals WHERE view >= $1 ORDER BY \
                 view LIMIT $2",
            )
            .bind(offset)
            .bind(batch_size)
            .fetch_all(tx.as_mut())
            .await?;

            drop(tx);

            if rows.is_empty() {
                break;
            }

            let mut values = Vec::new();

            for row in rows.iter() {
                let leaf_hash: String = row.try_get("leaf_hash")?;
                let data: Vec<u8> = row.try_get("data")?;

                let quorum_proposal: Proposal<SeqTypes, QuorumProposal<SeqTypes>> =
                    bincode::deserialize(&data)?;
                let quorum_proposal2: Proposal<SeqTypes, QuorumProposalWrapper<SeqTypes>> =
                    convert_proposal(quorum_proposal);

                let view = quorum_proposal2.data.view_number().u64() as i64;
                let data = bincode::serialize(&quorum_proposal2)?;

                values.push((view, leaf_hash, data));
            }

            let mut query_builder: sqlx::QueryBuilder<Db> =
                sqlx::QueryBuilder::new("INSERT INTO quorum_proposals2 (view, leaf_hash, data) ");

            offset = values.last().context("last row")?.0;
            query_builder.push_values(values, |mut b, (view, leaf_hash, data)| {
                b.push_bind(view).push_bind(leaf_hash).push_bind(data);
            });

            query_builder.push(" ON CONFLICT DO NOTHING");

            let query = query_builder.build();

            let mut tx = self.db.write().await?;
            query.execute(tx.as_mut()).await?;

            tx.upsert(
                "epoch_migration",
                ["table_name", "completed", "migrated_rows"],
                ["table_name"],
                [("quorum_proposals".to_string(), false, offset)],
            )
            .await?;
            tx.commit().await?;

            tracing::info!(
                "quorum proposals migration progress: rows={} offset={}",
                rows.len(),
                offset
            );

            if rows.len() < batch_size as usize {
                break;
            }
        }

        tracing::warn!("migrated quorum proposals");

        let mut tx = self.db.write().await?;
        tx.upsert(
            "epoch_migration",
            ["table_name", "completed", "migrated_rows"],
            ["table_name"],
            [("quorum_proposals".to_string(), true, offset)],
        )
        .await?;
        tx.commit().await?;

        tracing::info!("updated epoch_migration table for quorum_proposals");

        Ok(())
    }

    async fn migrate_quorum_certificates(&self) -> anyhow::Result<()> {
        let batch_size: i64 = 10000;
        let mut tx = self.db.read().await?;

        let (is_completed, mut offset) = query_as::<(bool, i64)>(
            "SELECT completed, migrated_rows from epoch_migration WHERE table_name = \
             'quorum_certificate'",
        )
        .fetch_one(tx.as_mut())
        .await?;

        if is_completed {
            tracing::info!("quorum certificates migration already done");
            return Ok(());
        }

        tracing::warn!("migrating quorum certificates..");
        loop {
            let mut tx = self.db.read().await?;
            let rows = query(
                "SELECT view, leaf_hash, data FROM quorum_certificate WHERE view >= $1 ORDER BY \
                 view LIMIT $2",
            )
            .bind(offset)
            .bind(batch_size)
            .fetch_all(tx.as_mut())
            .await?;

            drop(tx);
            if rows.is_empty() {
                break;
            }
            let mut values = Vec::new();

            for row in rows.iter() {
                let leaf_hash: String = row.try_get("leaf_hash")?;
                let data: Vec<u8> = row.try_get("data")?;

                let qc: QuorumCertificate<SeqTypes> = bincode::deserialize(&data)?;
                let qc2: QuorumCertificate2<SeqTypes> = qc.to_qc2();

                let view = qc2.view_number().u64() as i64;
                let data = bincode::serialize(&qc2)?;

                values.push((view, leaf_hash, data));
            }

            let mut query_builder: sqlx::QueryBuilder<Db> =
                sqlx::QueryBuilder::new("INSERT INTO quorum_certificate2 (view, leaf_hash, data) ");

            offset = values.last().context("last row")?.0;

            query_builder.push_values(values, |mut b, (view, leaf_hash, data)| {
                b.push_bind(view).push_bind(leaf_hash).push_bind(data);
            });

            query_builder.push(" ON CONFLICT DO NOTHING");
            let query = query_builder.build();

            let mut tx = self.db.write().await?;
            query.execute(tx.as_mut()).await?;

            tx.upsert(
                "epoch_migration",
                ["table_name", "completed", "migrated_rows"],
                ["table_name"],
                [("quorum_certificate".to_string(), false, offset)],
            )
            .await?;
            tx.commit().await?;

            tracing::info!(
                "Quorum certificates migration progress: rows={} offset={}",
                rows.len(),
                offset
            );

            if rows.len() < batch_size as usize {
                break;
            }
        }

        tracing::warn!("migrated quorum certificates");

        let mut tx = self.db.write().await?;
        tx.upsert(
            "epoch_migration",
            ["table_name", "completed", "migrated_rows"],
            ["table_name"],
            [("quorum_certificate".to_string(), true, offset)],
        )
        .await?;
        tx.commit().await?;
        tracing::info!("updated epoch_migration table for quorum_certificate");

        Ok(())
    }

    /// Migrate stake table data to include x25519_key and p2p_addr fields.
    ///
    /// Data written before x25519 support lacks these fields. This migration
    /// deserializes legacy records and re-serializes them with the new fields set to None.
    async fn migrate_x25519_keys(&self) -> anyhow::Result<()> {
        use super::RegisteredValidatorNoX25519;

        let name = DataMigration::X25519Keys.as_str();

        // Migrate bincode storage (epoch_drb_and_root.stake).
        if !self
            .is_migration_complete(name, "epoch_drb_and_root")
            .await?
        {
            let rows: Vec<(i64, Vec<u8>)> = {
                let mut tx = self.db.read().await?;
                query_as("SELECT epoch, stake FROM epoch_drb_and_root WHERE stake IS NOT NULL")
                    .fetch_all(tx.as_mut())
                    .await?
            };

            let num_rows = rows.len();
            let mut tx = self.db.write().await?;
            for (epoch, stake_bytes) in rows {
                // Try current format first
                if bincode::deserialize::<AuthenticatedValidatorMap>(&stake_bytes).is_ok() {
                    continue;
                }

                // Legacy format without x25519 fields
                let old_validators: IndexMap<Address, RegisteredValidatorNoX25519> =
                    bincode::deserialize(&stake_bytes)
                        .context("deserializing legacy stake table")?;
                let validators: AuthenticatedValidatorMap = old_validators
                    .into_iter()
                    .map(|(addr, v)| {
                        let registered = v.migrate();
                        (
                            addr,
                            AuthenticatedValidator::try_from(registered)
                                .expect("stake tables only contain authenticated validators"),
                        )
                    })
                    .collect();

                let new_bytes =
                    bincode::serialize(&validators).context("serializing stake table")?;

                tracing::debug!(epoch, "migrating x25519 keys in stake table");
                tx.execute(
                    query("UPDATE epoch_drb_and_root SET stake = $1 WHERE epoch = $2")
                        .bind(&new_bytes)
                        .bind(epoch),
                )
                .await?;
            }
            Self::mark_migration_complete(&mut tx, name, "epoch_drb_and_root", num_rows).await?;
            tx.commit().await?;
            tracing::info!(
                num_rows,
                "x25519_keys migration completed for epoch_drb_and_root"
            );
        }

        // Migrate JSONB storage (stake_table_validators).
        if !self
            .is_migration_complete(name, "stake_table_validators")
            .await?
        {
            let rows: Vec<(i64, String, serde_json::Value)> = {
                let mut tx = self.db.read().await?;
                query_as("SELECT epoch, address, validator FROM stake_table_validators")
                    .fetch_all(tx.as_mut())
                    .await?
            };

            let num_rows = rows.len();
            let mut tx = self.db.write().await?;
            for (epoch, address, validator_json) in rows {
                // Check if JSON already has x25519 fields (can't rely on deserialization
                // since serde_json fills missing Option<T> fields with None).
                if validator_json
                    .as_object()
                    .is_some_and(|obj| obj.contains_key("x25519_key"))
                {
                    continue;
                }

                // Deserialize (serde_json fills missing Option fields with None),
                // then re-serialize to ensure x25519_key and p2p_addr are present.
                let validator: RegisteredValidator<PubKey> =
                    serde_json::from_value(validator_json).context("deserializing validator")?;

                let new_json = serde_json::to_value(&validator).context("serializing validator")?;

                tracing::debug!(epoch, %address, "migrating x25519 keys for validator");
                tx.execute(
                    query(
                        "UPDATE stake_table_validators SET validator = $1 WHERE epoch = $2 AND \
                         address = $3",
                    )
                    .bind(&new_json)
                    .bind(epoch)
                    .bind(&address),
                )
                .await?;
            }
            Self::mark_migration_complete(&mut tx, name, "stake_table_validators", num_rows)
                .await?;
            tx.commit().await?;
            tracing::info!(
                num_rows,
                "x25519_keys migration completed for stake_table_validators"
            );
        }

        Ok(())
    }

    async fn store_next_epoch_quorum_certificate(
        &self,
        high_qc: NextEpochQuorumCertificate2<SeqTypes>,
    ) -> anyhow::Result<()> {
        let qc2_bytes = bincode::serialize(&high_qc).context("serializing next epoch qc")?;
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "next_epoch_quorum_certificate",
                    ["id", "data"],
                    ["id"],
                    [(true, qc2_bytes.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn load_next_epoch_quorum_certificate(
        &self,
    ) -> anyhow::Result<Option<NextEpochQuorumCertificate2<SeqTypes>>> {
        let result = self
            .db
            .read()
            .await?
            .fetch_optional("SELECT * FROM next_epoch_quorum_certificate where id = true")
            .await?;

        result
            .map(|row| {
                let bytes: Vec<u8> = row.get("data");
                anyhow::Result::<_>::Ok(bincode::deserialize(&bytes)?)
            })
            .transpose()
    }

    async fn store_eqc(
        &self,
        high_qc: QuorumCertificate2<SeqTypes>,
        next_epoch_high_qc: NextEpochQuorumCertificate2<SeqTypes>,
    ) -> anyhow::Result<()> {
        let eqc_bytes =
            bincode::serialize(&(high_qc, next_epoch_high_qc)).context("serializing eqc")?;
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert("eqc", ["id", "data"], ["id"], [(true, eqc_bytes.clone())])
                    .await?;
                tx.commit().await
            })
            .await
    }

    async fn load_eqc(
        &self,
    ) -> Option<(
        QuorumCertificate2<SeqTypes>,
        NextEpochQuorumCertificate2<SeqTypes>,
    )> {
        let result = self
            .db
            .read()
            .await
            .ok()?
            .fetch_optional("SELECT * FROM eqc where id = true")
            .await
            .ok()?;

        result
            .map(|row| {
                let bytes: Vec<u8> = row.get("data");
                bincode::deserialize(&bytes)
            })
            .transpose()
            .ok()?
    }

    async fn append_da2(
        &self,
        proposal: &Proposal<SeqTypes, DaProposal2<SeqTypes>>,
        vid_commit: VidCommitment,
    ) -> anyhow::Result<()> {
        let data = &proposal.data;
        let view = data.view_number().u64();
        let data_bytes = bincode::serialize(proposal).unwrap();

        let now = Instant::now();
        let res = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "da_proposal2",
                    ["view", "data", "payload_hash"],
                    ["view"],
                    [(view as i64, data_bytes.clone(), vid_commit.to_string())],
                )
                .await?;
                tx.commit().await
            })
            .await;
        self.internal_metrics
            .internal_append_da2_duration
            .add_point(now.elapsed().as_secs_f64());
        res
    }

    async fn store_drb_result(
        &self,
        epoch: EpochNumber,
        drb_result: DrbResult,
    ) -> anyhow::Result<()> {
        let epoch_i64 = epoch.u64() as i64;
        let drb_result_vec = Vec::from(drb_result);
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "epoch_drb_and_root",
                    ["epoch", "drb_result"],
                    ["epoch"],
                    [(epoch_i64, drb_result_vec.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn store_epoch_root(
        &self,
        epoch: EpochNumber,
        block_header: <SeqTypes as NodeType>::BlockHeader,
    ) -> anyhow::Result<()> {
        let epoch_i64 = epoch.u64() as i64;
        let block_header_bytes =
            bincode::serialize(&block_header).context("serializing block header")?;

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "epoch_drb_and_root",
                    ["epoch", "block_header"],
                    ["epoch"],
                    [(epoch_i64, block_header_bytes.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn store_drb_input(&self, drb_input: DrbInput) -> anyhow::Result<()> {
        if let Ok(loaded_drb_input) = self.load_drb_input(drb_input.epoch).await {
            if loaded_drb_input.difficulty_level != drb_input.difficulty_level {
                tracing::error!("Overwriting {loaded_drb_input:?} in storage with {drb_input:?}");
            } else if loaded_drb_input.iteration >= drb_input.iteration {
                anyhow::bail!(
                    "DrbInput in storage {:?} is more recent than {:?}, refusing to update",
                    loaded_drb_input,
                    drb_input
                )
            }
        }

        let drb_epoch_i64 = drb_input.epoch as i64;
        let drb_input_bytes = bincode::serialize(&drb_input)
            .context("Failed to serialize DrbInput. This is not fatal, but should never happen.")?;

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "drb",
                    ["epoch", "drb_input"],
                    ["epoch"],
                    [(drb_epoch_i64, drb_input_bytes.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn load_drb_input(&self, epoch: u64) -> anyhow::Result<DrbInput> {
        let row = self
            .db
            .read()
            .await?
            .fetch_optional(query("SELECT drb_input FROM drb WHERE epoch = $1").bind(epoch as i64))
            .await?;

        match row {
            None => anyhow::bail!("No DrbInput for epoch {} in storage", epoch),
            Some(row) => {
                let drb_input_bytes: Vec<u8> = row.try_get("drb_input")?;
                let drb_input = bincode::deserialize(&drb_input_bytes)
                    .context("Failed to deserialize drb_input from storage")?;

                Ok(drb_input)
            },
        }
    }

    async fn add_state_cert(
        &self,
        state_cert: LightClientStateUpdateCertificateV2<SeqTypes>,
    ) -> anyhow::Result<()> {
        let view_number = state_cert.light_client_state.view_number as i64;
        let state_cert_bytes = bincode::serialize(&state_cert)
            .context("serializing light client state update certificate")?;

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "state_cert",
                    ["view", "state_cert"],
                    ["view"],
                    [(view_number, state_cert_bytes.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn load_state_cert(
        &self,
    ) -> anyhow::Result<Option<LightClientStateUpdateCertificateV2<SeqTypes>>> {
        let Some(row) = self
            .db
            .read()
            .await?
            .fetch_optional(
                "SELECT state_cert FROM finalized_state_cert ORDER BY epoch DESC LIMIT 1",
            )
            .await?
        else {
            return Ok(None);
        };
        let bytes: Vec<u8> = row.get("state_cert");

        let cert = match bincode::deserialize(&bytes) {
            Ok(cert) => cert,
            Err(err) => {
                tracing::info!(
                    error = %err,
                    "Failed to deserialize state certificate with v2. attempting with v1"
                );

                let v1_cert =
                    bincode::deserialize::<LightClientStateUpdateCertificateV1<SeqTypes>>(&bytes)
                        .with_context(|| {
                        format!("Failed to deserialize using both v1 and v2. error: {err}")
                    })?;

                v1_cert.into()
            },
        };

        Ok(Some(cert))
    }

    async fn get_state_cert_by_epoch(
        &self,
        epoch: u64,
    ) -> anyhow::Result<Option<LightClientStateUpdateCertificateV2<SeqTypes>>> {
        let Some(row) = self
            .db
            .read()
            .await?
            .fetch_optional(
                query("SELECT state_cert FROM finalized_state_cert WHERE epoch = $1")
                    .bind(epoch as i64),
            )
            .await?
        else {
            return Ok(None);
        };
        let bytes: Vec<u8> = row.get("state_cert");

        let cert = match bincode::deserialize(&bytes) {
            Ok(cert) => cert,
            Err(err) => {
                tracing::info!(
                    error = %err,
                    "Failed to deserialize state certificate with v2. attempting with v1"
                );

                let v1_cert =
                    bincode::deserialize::<LightClientStateUpdateCertificateV1<SeqTypes>>(&bytes)
                        .with_context(|| {
                        format!("Failed to deserialize using both v1 and v2. error: {err}")
                    })?;

                v1_cert.into()
            },
        };

        Ok(Some(cert))
    }

    async fn insert_state_cert(
        &self,
        epoch: u64,
        cert: LightClientStateUpdateCertificateV2<SeqTypes>,
    ) -> anyhow::Result<()> {
        let epoch_i64 = epoch as i64;
        let bytes = bincode::serialize(&cert)
            .with_context(|| format!("Failed to serialize state cert for epoch {epoch}"))?;

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "finalized_state_cert",
                    ["epoch", "state_cert"],
                    ["epoch"],
                    [(epoch_i64, bytes.clone())],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn load_start_epoch_info(&self) -> anyhow::Result<Vec<InitializerEpochInfo<SeqTypes>>> {
        let rows = self
            .db
            .read()
            .await?
            .fetch_all(
                query("SELECT * from epoch_drb_and_root ORDER BY epoch DESC LIMIT $1")
                    .bind(RECENT_STAKE_TABLES_LIMIT as i64),
            )
            .await?;

        // reverse the rows vector to return the most recent epochs, but in ascending order
        rows.into_iter()
            .rev()
            .map(|row| {
                let epoch: i64 = row.try_get("epoch")?;
                let drb_result: Option<Vec<u8>> = row.try_get("drb_result")?;
                let block_header: Option<Vec<u8>> = row.try_get("block_header")?;
                if let Some(drb_result) = drb_result {
                    let drb_result_array = drb_result
                        .try_into()
                        .or_else(|_| bail!("invalid drb result"))?;
                    let block_header: Option<<SeqTypes as NodeType>::BlockHeader> = block_header
                        .map(|data| bincode::deserialize(&data))
                        .transpose()?;
                    Ok(Some(InitializerEpochInfo::<SeqTypes> {
                        epoch: EpochNumber::new(epoch as u64),
                        drb_result: drb_result_array,
                        block_header,
                    }))
                } else {
                    // Right now we skip the epoch_drb_and_root row if there is no drb result.
                    // This seems reasonable based on the expected order of events, but please double check!
                    Ok(None)
                }
            })
            .filter_map(|e| match e {
                Err(v) => Some(Err(v)),
                Ok(Some(v)) => Some(Ok(v)),
                Ok(None) => None,
            })
            .collect()
    }

    fn enable_metrics(&mut self, metrics: &dyn Metrics) {
        self.internal_metrics = PersistenceMetricsValue::new(metrics);
    }
}

#[async_trait]
impl MembershipPersistence for Persistence {
    async fn load_stake(&self, epoch: EpochNumber) -> anyhow::Result<Option<StakeTuple>> {
        let result = self
            .db
            .read()
            .await?
            .fetch_optional(
                query(
                    "SELECT stake, block_reward, stake_table_hash FROM epoch_drb_and_root WHERE \
                     epoch = $1",
                )
                .bind(epoch.u64() as i64),
            )
            .await?;

        result
            .map(|row| {
                let stake_table_bytes: Vec<u8> = row.get("stake");
                let reward_bytes: Option<Vec<u8>> = row.get("block_reward");
                let stake_table_hash_bytes: Option<Vec<u8>> = row.get("stake_table_hash");
                let stake_table: AuthenticatedValidatorMap =
                    bincode::deserialize(&stake_table_bytes)
                        .context("deserializing stake table")?;
                let reward: Option<RewardAmount> = reward_bytes
                    .map(|b| bincode::deserialize(&b).context("deserializing block_reward"))
                    .transpose()?;
                let stake_table_hash: Option<StakeTableHash> = stake_table_hash_bytes
                    .map(|b| bincode::deserialize(&b).context("deserializing stake table hash"))
                    .transpose()?;

                Ok((stake_table, reward, stake_table_hash))
            })
            .transpose()
    }

    async fn load_latest_stake(&self, limit: u64) -> anyhow::Result<Option<Vec<IndexedStake>>> {
        let mut tx = self.db.read().await?;

        let rows = match query_as::<(i64, Vec<u8>, Option<Vec<u8>>, Option<Vec<u8>>)>(
            "SELECT epoch, stake, block_reward, stake_table_hash FROM epoch_drb_and_root WHERE \
             stake is NOT NULL ORDER BY epoch DESC LIMIT $1",
        )
        .bind(limit as i64)
        .fetch_all(tx.as_mut())
        .await
        {
            Ok(bytes) => bytes,
            Err(err) => {
                tracing::error!("error loading stake tables: {err:#}");
                bail!("{err:#}");
            },
        };

        let stakes: anyhow::Result<Vec<IndexedStake>> = rows
            .into_iter()
            .map(
                |(id, stake_bytes, reward_bytes_opt, stake_table_hash_bytes_opt)| {
                    let stake_table: AuthenticatedValidatorMap =
                        bincode::deserialize(&stake_bytes).context("deserializing stake table")?;

                    let block_reward: Option<RewardAmount> = reward_bytes_opt
                        .map(|b| bincode::deserialize(&b).context("deserializing block_reward"))
                        .transpose()?;

                    let stake_table_hash: Option<StakeTableHash> = stake_table_hash_bytes_opt
                        .map(|b| bincode::deserialize(&b).context("deserializing stake table hash"))
                        .transpose()?;

                    Ok((
                        EpochNumber::new(id as u64),
                        (stake_table, block_reward),
                        stake_table_hash,
                    ))
                },
            )
            .collect();

        Ok(Some(stakes?))
    }

    async fn store_stake(
        &self,
        epoch: EpochNumber,
        stake: AuthenticatedValidatorMap,
        block_reward: Option<RewardAmount>,
        stake_table_hash: Option<StakeTableHash>,
    ) -> anyhow::Result<()> {
        let epoch_i64 = epoch.u64() as i64;
        let stake_table_bytes = bincode::serialize(&stake).context("serializing stake table")?;
        let reward_bytes = block_reward
            .map(|r| bincode::serialize(&r).context("serializing block reward"))
            .transpose()?;
        let stake_table_hash_bytes = stake_table_hash
            .map(|h| bincode::serialize(&h).context("serializing stake table hash"))
            .transpose()?;
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                tx.upsert(
                    "epoch_drb_and_root",
                    ["epoch", "stake", "block_reward", "stake_table_hash"],
                    ["epoch"],
                    [(
                        epoch_i64,
                        stake_table_bytes.clone(),
                        reward_bytes.clone(),
                        stake_table_hash_bytes.clone(),
                    )],
                )
                .await?;
                tx.commit().await
            })
            .await
    }

    async fn store_events(
        &self,
        l1_finalized: u64,
        events: Vec<(EventKey, StakeTableEvent)>,
    ) -> anyhow::Result<()> {
        let l1_finalized_i64: i64 = l1_finalized.try_into()?;
        let serialized_events = events
            .into_iter()
            .map(|((block_number, index), event)| {
                Ok((
                    i64::try_from(block_number)?,
                    i64::try_from(index)?,
                    serde_json::to_value(event).context("l1 event to value")?,
                ))
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;

                // check last l1 block if there is any
                let last_processed_l1_block = query_as::<(i64,)>(
                    "SELECT last_l1_block FROM stake_table_events_l1_block where id = 0",
                )
                .fetch_optional(tx.as_mut())
                .await?
                .map(|(l1,)| l1);

                tracing::debug!("last l1 finalizes in database = {last_processed_l1_block:?}");

                // skip events storage if the database already has higher l1 block events
                let serialized_events_len = serialized_events.len();
                if last_processed_l1_block > Some(l1_finalized_i64) {
                    tracing::debug!(
                        ?last_processed_l1_block,
                        l1_finalized,
                        serialized_events_len,
                        "last l1 finalized stored is already higher"
                    );
                    return Ok(());
                }

                if !serialized_events.is_empty() {
                    let mut query_builder: sqlx::QueryBuilder<Db> = sqlx::QueryBuilder::new(
                        "INSERT INTO stake_table_events (l1_block, log_index, event) ",
                    );

                    query_builder.push_values(
                        serialized_events.iter().cloned(),
                        |mut b, (l1_block, log_index, event)| {
                            b.push_bind(l1_block).push_bind(log_index).push_bind(event);
                        },
                    );

                    query_builder.push(" ON CONFLICT DO NOTHING");
                    let query = query_builder.build();

                    query.execute(tx.as_mut()).await?;
                }

                // update l1 block
                tx.upsert(
                    "stake_table_events_l1_block",
                    ["id", "last_l1_block"],
                    ["id"],
                    [(0_i32, l1_finalized_i64)],
                )
                .await?;

                tx.commit().await?;

                Ok(())
            })
            .await
    }

    /// Loads all events from persistent storage up to the specified L1 block.
    ///
    /// # Returns
    ///
    /// Returns a tuple containing:
    /// - `Option<u64>` - The queried L1 block for which all events have been successfully fetched.
    /// - `Vec<(EventKey, StakeTableEvent)>` - A list of events, where each entry is a tuple of the event key
    /// event key is (l1 block number, log index)
    ///   and the corresponding StakeTable event.
    ///
    async fn load_events(
        &self,
        from_l1_block: u64,
        to_l1_block: u64,
    ) -> anyhow::Result<(
        Option<EventsPersistenceRead>,
        Vec<(EventKey, StakeTableEvent)>,
    )> {
        let mut tx = self.db.read().await?;

        // check last l1 block if there is any
        let res = query_as::<(i64,)>(
            "SELECT last_l1_block FROM stake_table_events_l1_block where id = 0",
        )
        .fetch_optional(tx.as_mut())
        .await?;

        let Some((last_processed_l1_block,)) = res else {
            // this just means we dont have any events stored
            return Ok((None, Vec::new()));
        };

        // Determine the L1 block for querying events.
        // If the last stored L1 block is greater than the requested block, limit the query to the requested block.
        // Otherwise, query up to the last stored block.
        let to_l1_block = to_l1_block.try_into()?;
        let query_l1_block = if last_processed_l1_block > to_l1_block {
            to_l1_block
        } else {
            last_processed_l1_block
        };

        let rows = query(
            "SELECT l1_block, log_index, event FROM stake_table_events WHERE $1 <= l1_block AND \
             l1_block <= $2 ORDER BY l1_block ASC, log_index ASC",
        )
        .bind(i64::try_from(from_l1_block)?)
        .bind(query_l1_block)
        .fetch_all(tx.as_mut())
        .await?;

        let events = rows
            .into_iter()
            .map(|row| {
                let l1_block: i64 = row.try_get("l1_block")?;
                let log_index: i64 = row.try_get("log_index")?;
                let event = serde_json::from_value(row.try_get("event")?)?;

                Ok(((l1_block.try_into()?, log_index.try_into()?), event))
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        // Determine the read state based on the queried block range.
        // - If the persistence returned events up to the requested block, the read is complete.
        // - Otherwise, indicate that the read is up to the last processed block.
        if query_l1_block == to_l1_block {
            Ok((Some(EventsPersistenceRead::Complete), events))
        } else {
            Ok((
                Some(EventsPersistenceRead::UntilL1Block(
                    query_l1_block.try_into()?,
                )),
                events,
            ))
        }
    }

    async fn delete_stake_tables(&self) -> anyhow::Result<()> {
        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;
                #[cfg(not(feature = "embedded-db"))]
                query(
                    "TRUNCATE stake_table_events, stake_table_events_l1_block, \
                     epoch_drb_and_root, stake_table_validators",
                )
                .execute(tx.as_mut())
                .await?;
                #[cfg(feature = "embedded-db")]
                {
                    query("DELETE FROM stake_table_events")
                        .execute(tx.as_mut())
                        .await?;
                    query("DELETE FROM stake_table_events_l1_block")
                        .execute(tx.as_mut())
                        .await?;
                    query("DELETE FROM epoch_drb_and_root")
                        .execute(tx.as_mut())
                        .await?;
                    query("DELETE FROM stake_table_validators")
                        .execute(tx.as_mut())
                        .await?;
                }
                tx.commit().await?;
                Ok(())
            })
            .await
    }

    async fn store_all_validators(
        &self,
        epoch: EpochNumber,
        all_validators: RegisteredValidatorMap,
    ) -> anyhow::Result<()> {
        if all_validators.is_empty() {
            return Ok(());
        }

        let epoch_i64 = epoch.u64() as i64;
        let serialized_validators = all_validators
            .into_iter()
            .map(|(address, validator)| {
                let validator_json =
                    serde_json::to_value(&validator).context("serializing validator to json")?;
                Ok((address.to_string(), validator_json))
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                let mut tx = self.db.write().await?;

                let mut query_builder = QueryBuilder::new(
                    "INSERT INTO stake_table_validators (epoch, address, validator) ",
                );

                query_builder.push_values(
                    serialized_validators.iter().cloned(),
                    |mut b, (address, validator)| {
                        b.push_bind(epoch_i64)
                            .push_bind(address)
                            .push_bind(validator);
                    },
                );

                query_builder.push(
                    " ON CONFLICT (epoch, address) DO UPDATE SET validator = EXCLUDED.validator",
                );

                let query = query_builder.build();

                query.execute(tx.as_mut()).await?;

                tx.commit().await?;
                Ok(())
            })
            .await
    }

    async fn load_all_validators(
        &self,
        epoch: EpochNumber,
        offset: u64,
        limit: u64,
    ) -> anyhow::Result<Vec<RegisteredValidator<PubKey>>> {
        let mut tx = self.db.read().await?;

        // Use LOWER(address) in ORDER BY to ensure consistent ordering for SQlite and Postgres.
        // Postgres sorts text case sensitively by default, while SQLite sorts case insensitively.
        // Applying LOWER() makes the result consistent.
        let rows = query(
            "SELECT address, validator
         FROM stake_table_validators
         WHERE epoch = $1
         ORDER BY LOWER(address) ASC
         LIMIT $2 OFFSET $3",
        )
        .bind(epoch.u64() as i64)
        .bind(limit as i64)
        .bind(offset as i64)
        .fetch_all(tx.as_mut())
        .await?;
        rows.into_iter()
            .map(|row| {
                let validator_json: serde_json::Value = row.try_get("validator")?;
                serde_json::from_value::<RegisteredValidator<PubKey>>(validator_json)
                    .map_err(Into::into)
            })
            .collect()
    }
}

#[async_trait]
impl DhtPersistentStorage for Persistence {
    /// Save the DHT to the database
    ///
    /// # Errors
    /// - If we fail to serialize the records
    /// - If we fail to write the serialized records to the DB
    async fn save(&self, records: Vec<SerializableRecord>) -> anyhow::Result<()> {
        // Bincode-serialize the records
        let to_save =
            bincode::serialize(&records).with_context(|| "failed to serialize records")?;

        // Prepare the statement
        let stmt = "INSERT INTO libp2p_dht (id, serialized_records) VALUES (0, $1) ON CONFLICT \
                    (id) DO UPDATE SET serialized_records = $1";

        WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || async {
                // Execute the query
                let mut tx = self
                    .db
                    .write()
                    .await
                    .with_context(|| "failed to start an atomic DB transaction")?;
                tx.execute(query(stmt).bind(to_save.clone()))
                    .await
                    .with_context(|| "failed to execute DB query")?;

                // Commit the state
                tx.commit().await.with_context(|| "failed to commit to DB")
            })
            .await
    }

    /// Load the DHT from the database
    ///
    /// # Errors
    /// - If we fail to read from the DB
    /// - If we fail to deserialize the records
    async fn load(&self) -> anyhow::Result<Vec<SerializableRecord>> {
        // Fetch the results from the DB
        let result = self
            .db
            .read()
            .await
            .with_context(|| "failed to start a DB read transaction")?
            .fetch_one("SELECT * FROM libp2p_dht where id = 0")
            .await
            .with_context(|| "failed to fetch from DB")?;

        // Get the `serialized_records` row
        let serialied_records: Vec<u8> = result.get("serialized_records");

        // Deserialize it
        let records: Vec<SerializableRecord> = bincode::deserialize(&serialied_records)
            .with_context(|| "Failed to deserialize records")?;

        Ok(records)
    }
}

#[async_trait]
impl Provider<SeqTypes, VidCommonRequest> for Persistence {
    #[tracing::instrument(skip(self))]
    async fn fetch(&self, req: VidCommonRequest) -> Option<VidCommon> {
        let mut tx = match self.db.read().await {
            Ok(tx) => tx,
            Err(err) => {
                tracing::warn!("could not open transaction: {err:#}");
                return None;
            },
        };

        let bytes = match query_as::<(Vec<u8>,)>(
            "SELECT data FROM vid_share2 WHERE payload_hash = $1 LIMIT 1",
        )
        .bind(req.0.to_string())
        .fetch_optional(tx.as_mut())
        .await
        {
            Ok(Some((bytes,))) => bytes,
            Ok(None) => return None,
            Err(err) => {
                tracing::error!("error loading VID share: {err:#}");
                return None;
            },
        };

        let share: Proposal<SeqTypes, VidDisperseShare<SeqTypes>> =
            match bincode::deserialize(&bytes) {
                Ok(share) => share,
                Err(err) => {
                    tracing::warn!("error decoding VID share: {err:#}");
                    return None;
                },
            };

        match share.data {
            VidDisperseShare::V0(vid) => Some(VidCommon::V0(vid.common)),
            VidDisperseShare::V1(vid) => Some(VidCommon::V1(vid.common)),
            VidDisperseShare::V2(vid) => Some(VidCommon::V2(vid.common)),
        }
    }
}

#[async_trait]
impl Provider<SeqTypes, PayloadRequest> for Persistence {
    #[tracing::instrument(skip(self))]
    async fn fetch(&self, req: PayloadRequest) -> Option<Payload> {
        let mut tx = match self.db.read().await {
            Ok(tx) => tx,
            Err(err) => {
                tracing::warn!("could not open transaction: {err:#}");
                return None;
            },
        };

        let bytes = match query_as::<(Vec<u8>,)>(
            "SELECT data FROM da_proposal2 WHERE payload_hash = $1 LIMIT 1",
        )
        .bind(req.0.to_string())
        .fetch_optional(tx.as_mut())
        .await
        {
            Ok(Some((bytes,))) => bytes,
            Ok(None) => return None,
            Err(err) => {
                tracing::warn!("error loading DA proposal: {err:#}");
                return None;
            },
        };

        let proposal: Proposal<SeqTypes, DaProposal2<SeqTypes>> = match bincode::deserialize(&bytes)
        {
            Ok(proposal) => proposal,
            Err(err) => {
                tracing::error!("error decoding DA proposal: {err:#}");
                return None;
            },
        };

        Some(Payload::from_bytes(
            &proposal.data.encoded_transactions,
            &proposal.data.metadata,
        ))
    }
}

#[cfg(test)]
mod testing {
    use hotshot_query_service::data_source::storage::sql::testing::TmpDb;

    use super::*;
    use crate::persistence::tests::TestablePersistence;

    #[async_trait]
    impl TestablePersistence for Persistence {
        type Storage = Arc<TmpDb>;

        async fn tmp_storage() -> Self::Storage {
            Arc::new(TmpDb::init().await)
        }

        #[allow(refining_impl_trait)]
        fn options(db: &Self::Storage) -> Options {
            #[cfg(not(feature = "embedded-db"))]
            {
                PostgresOptions {
                    port: Some(db.port()),
                    host: Some(db.host()),
                    user: Some("postgres".into()),
                    password: Some("password".into()),
                    ..Default::default()
                }
                .into()
            }

            #[cfg(feature = "embedded-db")]
            {
                SqliteOptions { path: db.path() }.into()
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::{AtomicU32, Ordering};

    use committable::{Commitment, CommitmentBoundsArkless};
    use espresso_types::{Header, Leaf, NodeState, ValidatedState, traits::NullEventConsumer};
    use futures::stream::TryStreamExt;
    use hotshot_example_types::node_types::TEST_VERSIONS;
    use hotshot_types::{
        data::{
            EpochNumber, QuorumProposal2, ns_table::parse_ns_table,
            vid_disperse::AvidMDisperseShare,
        },
        message::convert_proposal,
        simple_certificate::QuorumCertificate,
        simple_vote::QuorumData,
        traits::{
            EncodeBytes,
            block_contents::{BlockHeader, GENESIS_VID_NUM_STORAGE_NODES},
            signature_key::SignatureKey,
        },
        utils::EpochTransitionIndicator,
        vid::{
            advz::advz_scheme,
            avidm::{AvidMScheme, init_avidm_param},
        },
    };
    use jf_advz::VidScheme;

    use super::*;
    use crate::{BLSPubKey, PubKey, persistence::tests::TestablePersistence as _};

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_quorum_proposals_leaf_hash_migration() {
        // Create some quorum proposals to test with.
        let leaf: Leaf2 = Leaf::genesis(
            &ValidatedState::default(),
            &NodeState::mock(),
            TEST_VERSIONS.test.base,
        )
        .await
        .into();
        let privkey = BLSPubKey::generated_from_seed_indexed([0; 32], 1).1;
        let signature = PubKey::sign(&privkey, &[]).unwrap();
        let mut quorum_proposal = Proposal {
            data: QuorumProposal2::<SeqTypes> {
                epoch: None,
                block_header: leaf.block_header().clone(),
                view_number: ViewNumber::genesis(),
                justify_qc: QuorumCertificate::genesis(
                    &ValidatedState::default(),
                    &NodeState::mock(),
                    TEST_VERSIONS.test,
                )
                .await
                .to_qc2(),
                upgrade_certificate: None,
                view_change_evidence: None,
                next_drb_result: None,
                next_epoch_justify_qc: None,
                state_cert: None,
            },
            signature,
            _pd: Default::default(),
        };

        let qp1: Proposal<SeqTypes, QuorumProposal<SeqTypes>> =
            convert_proposal(quorum_proposal.clone());

        quorum_proposal.data.view_number = ViewNumber::new(1);

        let qp2: Proposal<SeqTypes, QuorumProposal<SeqTypes>> =
            convert_proposal(quorum_proposal.clone());
        let qps = [qp1, qp2];

        // Create persistence and add the quorum proposals with NULL leaf hash.
        let db = Persistence::tmp_storage().await;
        let persistence = Persistence::connect(&db).await;
        let mut tx = persistence.db.write().await.unwrap();
        let params = qps
            .iter()
            .map(|qp| {
                (
                    qp.data.view_number.u64() as i64,
                    bincode::serialize(&qp).unwrap(),
                )
            })
            .collect::<Vec<_>>();
        tx.upsert("quorum_proposals", ["view", "data"], ["view"], params)
            .await
            .unwrap();
        tx.commit().await.unwrap();

        // Create a new persistence and ensure the commitments get populated.
        let persistence = Persistence::connect(&db).await;
        let mut tx = persistence.db.read().await.unwrap();
        let rows = tx
            .fetch("SELECT * FROM quorum_proposals ORDER BY view ASC")
            .try_collect::<Vec<_>>()
            .await
            .unwrap();
        assert_eq!(rows.len(), qps.len());
        for (row, qp) in rows.into_iter().zip(qps) {
            assert_eq!(row.get::<i64, _>("view"), qp.data.view_number.u64() as i64);
            assert_eq!(
                row.get::<Vec<u8>, _>("data"),
                bincode::serialize(&qp).unwrap()
            );
            assert_eq!(
                row.get::<String, _>("leaf_hash"),
                Committable::commit(&Leaf::from_quorum_proposal(&qp.data)).to_string()
            );
        }
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_x25519_keys_migration() {
        use std::collections::HashMap;

        use crate::persistence::RegisteredValidatorNoX25519;

        let mut validator = RegisteredValidator::mock();
        validator.delegators.clear();
        validator.stake = alloy::primitives::U256::from(1000u64);

        let epoch = 1i64;
        let address = validator.account;

        // Create legacy data without x25519 fields
        let legacy = RegisteredValidatorNoX25519 {
            account: validator.account,
            stake_table_key: validator.stake_table_key,
            state_ver_key: validator.state_ver_key.clone(),
            stake: validator.stake,
            commission: validator.commission,
            delegators: HashMap::new(),
            authenticated: true,
        };

        // Bincode: serialize as legacy map
        let mut legacy_map: IndexMap<Address, RegisteredValidatorNoX25519> = IndexMap::new();
        legacy_map.insert(address, legacy);
        let stake_bytes = bincode::serialize(&legacy_map).unwrap();

        // JSON: serialize without x25519 fields
        let json_legacy = RegisteredValidatorNoX25519 {
            account: validator.account,
            stake_table_key: validator.stake_table_key,
            state_ver_key: validator.state_ver_key.clone(),
            stake: validator.stake,
            commission: validator.commission,
            delegators: HashMap::new(),
            authenticated: true,
        };
        let validator_json = serde_json::to_value(&json_legacy).unwrap();

        let db = Persistence::tmp_storage().await;
        let persistence = Persistence::connect(&db).await;
        let mut tx = persistence.db.write().await.unwrap();

        tx.execute(
            query(
                "INSERT INTO stake_table_validators (epoch, address, validator) VALUES ($1, $2, \
                 $3)",
            )
            .bind(epoch)
            .bind(format!("{:?}", address))
            .bind(&validator_json),
        )
        .await
        .unwrap();

        tx.execute(
            query("INSERT INTO epoch_drb_and_root (epoch, stake) VALUES ($1, $2)")
                .bind(epoch)
                .bind(&stake_bytes),
        )
        .await
        .unwrap();

        // Reset migration state so it runs on the newly inserted data
        tx.execute(query(
            "UPDATE data_migrations SET completed = false, migrated_rows = 0 WHERE name = \
             'x25519_keys'",
        ))
        .await
        .unwrap();

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

        // Verify JSON was inserted without x25519_key
        {
            let mut tx = persistence.db.read().await.unwrap();
            let row: (serde_json::Value,) =
                query_as("SELECT validator FROM stake_table_validators WHERE epoch = $1")
                    .bind(epoch)
                    .fetch_one(tx.as_mut())
                    .await
                    .unwrap();
            let json_obj = row.0.as_object().unwrap();
            assert!(!json_obj.contains_key("x25519_key"));
        }

        // Run migrations
        let persistence = Persistence::connect(&db).await;
        persistence.migrate_storage().await.unwrap();

        // Verify stake_table_validators now has x25519_key field
        {
            let mut tx = persistence.db.read().await.unwrap();
            let row: (serde_json::Value,) =
                query_as("SELECT validator FROM stake_table_validators WHERE epoch = $1")
                    .bind(epoch)
                    .fetch_one(tx.as_mut())
                    .await
                    .unwrap();
            let json_obj = row.0.as_object().unwrap();
            assert!(json_obj.contains_key("x25519_key"));
            assert!(json_obj.get("x25519_key").unwrap().is_null());
        }

        // Verify epoch_drb_and_root stake was migrated
        {
            let mut tx = persistence.db.read().await.unwrap();
            let row: (Vec<u8>,) = query_as("SELECT stake FROM epoch_drb_and_root WHERE epoch = $1")
                .bind(epoch)
                .fetch_one(tx.as_mut())
                .await
                .unwrap();
            let migrated_map: AuthenticatedValidatorMap = bincode::deserialize(&row.0).unwrap();
            assert!(migrated_map.contains_key(&address));
            let v = migrated_map.get(&address).unwrap();
            assert!(v.x25519_key.is_none());
        }

        // Verify migration tracking
        {
            let mut tx = persistence.db.read().await.unwrap();
            let row: (bool, i64) = query_as(
                "SELECT completed, migrated_rows FROM data_migrations WHERE name = 'x25519_keys' \
                 AND table_name = 'epoch_drb_and_root'",
            )
            .fetch_one(tx.as_mut())
            .await
            .unwrap();
            assert!(row.0);
            assert_eq!(row.1, 1);
        }
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_store_all_validators_authenticated_and_unauthenticated() {
        use std::collections::HashMap;

        use alloy::primitives::{Address, U256};
        use hotshot_types::light_client::StateVerKey;
        use indexmap::IndexMap;

        let tmp = Persistence::tmp_storage().await;
        let storage = Persistence::connect(&tmp).await;

        // Create an authenticated validator
        let authenticated_validator = RegisteredValidator {
            account: Address::random(),
            stake_table_key: BLSPubKey::generated_from_seed_indexed([0u8; 32], 0).0,
            state_ver_key: StateVerKey::default(),
            stake: U256::from(1000),
            commission: 100,
            delegators: HashMap::new(),
            authenticated: true,
            x25519_key: None,
            p2p_addr: None,
        };

        // Create an unauthenticated validator
        let unauthenticated_validator = RegisteredValidator {
            account: Address::random(),
            stake_table_key: BLSPubKey::generated_from_seed_indexed([0u8; 32], 1).0,
            state_ver_key: StateVerKey::default(),
            stake: U256::from(2000),
            commission: 200,
            delegators: HashMap::new(),
            authenticated: false,
            x25519_key: None,
            p2p_addr: None,
        };

        let mut validators: IndexMap<Address, RegisteredValidator<BLSPubKey>> = IndexMap::new();
        validators.insert(
            authenticated_validator.account,
            authenticated_validator.clone(),
        );
        validators.insert(
            unauthenticated_validator.account,
            unauthenticated_validator.clone(),
        );

        // Store both validators
        storage
            .store_all_validators(EpochNumber::new(1), validators)
            .await
            .unwrap();

        // Load and verify
        let loaded = storage
            .load_all_validators(EpochNumber::new(1), 0, 100)
            .await
            .unwrap();
        assert_eq!(loaded.len(), 2);

        // Find each validator and verify authenticated state is preserved
        let loaded_auth = loaded
            .iter()
            .find(|v| v.account == authenticated_validator.account)
            .unwrap();
        assert!(
            loaded_auth.authenticated,
            "authenticated validator should remain authenticated"
        );

        let loaded_unauth = loaded
            .iter()
            .find(|v| v.account == unauthenticated_validator.account)
            .unwrap();
        assert!(
            !loaded_unauth.authenticated,
            "unauthenticated validator should remain unauthenticated"
        );
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_fetching_providers() {
        let tmp = Persistence::tmp_storage().await;
        let storage = Persistence::connect(&tmp).await;

        // Mock up some data.
        let leaf = Leaf2::genesis(
            &ValidatedState::default(),
            &NodeState::mock(),
            TEST_VERSIONS.test.base,
        )
        .await;
        let leaf_payload = leaf.block_payload().unwrap();
        let leaf_payload_bytes_arc = leaf_payload.encode();

        let avidm_param = init_avidm_param(2).unwrap();
        let weights = vec![1u32; 2];

        let ns_table = parse_ns_table(
            leaf_payload.byte_len().as_usize(),
            &leaf_payload.ns_table().encode(),
        );
        let (payload_commitment, shares) =
            AvidMScheme::ns_disperse(&avidm_param, &weights, &leaf_payload_bytes_arc, ns_table)
                .unwrap();
        let (pubkey, privkey) = BLSPubKey::generated_from_seed_indexed([0; 32], 1);
        let vid_share = convert_proposal(
            AvidMDisperseShare::<SeqTypes> {
                view_number: ViewNumber::new(0),
                payload_commitment,
                share: shares[0].clone(),
                recipient_key: pubkey,
                epoch: None,
                target_epoch: None,
                common: avidm_param.clone(),
            }
            .to_proposal(&privkey)
            .unwrap()
            .clone(),
        );

        let quorum_proposal = QuorumProposalWrapper::<SeqTypes> {
            proposal: QuorumProposal2::<SeqTypes> {
                block_header: leaf.block_header().clone(),
                view_number: leaf.view_number(),
                justify_qc: leaf.justify_qc(),
                upgrade_certificate: None,
                view_change_evidence: None,
                next_drb_result: None,
                next_epoch_justify_qc: None,
                epoch: None,
                state_cert: None,
            },
        };
        let quorum_proposal_signature =
            BLSPubKey::sign(&privkey, &bincode::serialize(&quorum_proposal).unwrap())
                .expect("Failed to sign quorum proposal");
        let quorum_proposal = Proposal {
            data: quorum_proposal,
            signature: quorum_proposal_signature,
            _pd: Default::default(),
        };

        let block_payload_signature = BLSPubKey::sign(&privkey, &leaf_payload_bytes_arc)
            .expect("Failed to sign block payload");
        let da_proposal = Proposal {
            data: DaProposal2::<SeqTypes> {
                encoded_transactions: leaf_payload_bytes_arc,
                metadata: leaf_payload.ns_table().clone(),
                view_number: ViewNumber::new(0),
                epoch: None,
                epoch_transition_indicator: EpochTransitionIndicator::NotInTransition,
            },
            signature: block_payload_signature,
            _pd: Default::default(),
        };

        let mut next_quorum_proposal = quorum_proposal.clone();
        next_quorum_proposal.data.proposal.view_number += 1;
        next_quorum_proposal.data.proposal.justify_qc.view_number += 1;
        next_quorum_proposal
            .data
            .proposal
            .justify_qc
            .data
            .leaf_commit = Committable::commit(&leaf.clone());

        // Add to database.
        storage
            .append_da2(&da_proposal, VidCommitment::V1(payload_commitment))
            .await
            .unwrap();
        storage.append_vid(&vid_share).await.unwrap();
        storage
            .append_quorum_proposal2(&quorum_proposal)
            .await
            .unwrap();

        // Add an extra quorum proposal so we have a QC pointing back at `leaf`.
        storage
            .append_quorum_proposal2(&next_quorum_proposal)
            .await
            .unwrap();

        // Fetch it as if we were rebuilding an archive.
        assert_eq!(
            Some(VidCommon::V1(avidm_param)),
            storage
                .fetch(VidCommonRequest(vid_share.data.payload_commitment()))
                .await
        );
        assert_eq!(
            leaf_payload,
            storage
                .fetch(PayloadRequest(vid_share.data.payload_commitment()))
                .await
                .unwrap()
        );
    }

    /// Test conditions that trigger pruning.
    ///
    /// This is a configurable test that can be used to test different configurations of GC,
    /// `pruning_opt`. The test populates the database with some data for view 1, asserts that it is
    /// retained for view 2, and then asserts that it is pruned by view 3. There are various
    /// different configurations that can achieve this behavior, such that the data is retained and
    /// then pruned due to different logic and code paths.
    async fn test_pruning_helper(pruning_opt: ConsensusPruningOptions) {
        let tmp = Persistence::tmp_storage().await;
        let mut opt = Persistence::options(&tmp);
        opt.consensus_pruning = pruning_opt;
        let storage = opt.create().await.unwrap();

        let data_view = ViewNumber::new(1);

        // Populate some data.
        let leaf = Leaf2::genesis(
            &ValidatedState::default(),
            &NodeState::mock(),
            TEST_VERSIONS.test.base,
        )
        .await;
        let leaf_payload = leaf.block_payload().unwrap();
        let leaf_payload_bytes_arc = leaf_payload.encode();

        let avidm_param = init_avidm_param(2).unwrap();
        let weights = vec![1u32; 2];

        let ns_table = parse_ns_table(
            leaf_payload.byte_len().as_usize(),
            &leaf_payload.ns_table().encode(),
        );
        let (payload_commitment, shares) =
            AvidMScheme::ns_disperse(&avidm_param, &weights, &leaf_payload_bytes_arc, ns_table)
                .unwrap();

        let (pubkey, privkey) = BLSPubKey::generated_from_seed_indexed([0; 32], 1);
        let vid = convert_proposal(
            AvidMDisperseShare::<SeqTypes> {
                view_number: data_view,
                payload_commitment,
                share: shares[0].clone(),
                recipient_key: pubkey,
                epoch: None,
                target_epoch: None,
                common: avidm_param,
            }
            .to_proposal(&privkey)
            .unwrap()
            .clone(),
        );
        let quorum_proposal = QuorumProposalWrapper::<SeqTypes> {
            proposal: QuorumProposal2::<SeqTypes> {
                epoch: None,
                block_header: leaf.block_header().clone(),
                view_number: data_view,
                justify_qc: QuorumCertificate2::genesis(
                    &ValidatedState::default(),
                    &NodeState::mock(),
                    TEST_VERSIONS.test,
                )
                .await,
                upgrade_certificate: None,
                view_change_evidence: None,
                next_drb_result: None,
                next_epoch_justify_qc: None,
                state_cert: None,
            },
        };
        let quorum_proposal_signature =
            BLSPubKey::sign(&privkey, &bincode::serialize(&quorum_proposal).unwrap())
                .expect("Failed to sign quorum proposal");
        let quorum_proposal = Proposal {
            data: quorum_proposal,
            signature: quorum_proposal_signature,
            _pd: Default::default(),
        };

        let block_payload_signature = BLSPubKey::sign(&privkey, &leaf_payload_bytes_arc)
            .expect("Failed to sign block payload");
        let da_proposal = Proposal {
            data: DaProposal2::<SeqTypes> {
                encoded_transactions: leaf_payload_bytes_arc.clone(),
                metadata: leaf_payload.ns_table().clone(),
                view_number: data_view,
                epoch: Some(EpochNumber::new(0)),
                epoch_transition_indicator: EpochTransitionIndicator::NotInTransition,
            },
            signature: block_payload_signature,
            _pd: Default::default(),
        };

        tracing::info!(?vid, ?da_proposal, ?quorum_proposal, "append data");
        storage.append_vid(&vid).await.unwrap();
        storage
            .append_da2(&da_proposal, VidCommitment::V1(payload_commitment))
            .await
            .unwrap();
        storage
            .append_quorum_proposal2(&quorum_proposal)
            .await
            .unwrap();

        // The first decide doesn't trigger any garbage collection, even though our usage exceeds
        // the target, because of the minimum retention.
        tracing::info!("decide view 1");
        storage
            .append_decided_leaves(data_view + 1, [], None, &NullEventConsumer)
            .await
            .unwrap();
        assert_eq!(
            storage.load_vid_share(data_view).await.unwrap().unwrap(),
            vid
        );
        assert_eq!(
            storage.load_da_proposal(data_view).await.unwrap().unwrap(),
            da_proposal
        );
        assert_eq!(
            storage.load_quorum_proposal(data_view).await.unwrap(),
            quorum_proposal
        );

        // After another view, our data is beyond the minimum retention (though not the target
        // retention) so it gets pruned.
        tracing::info!("decide view 2");
        storage
            .append_decided_leaves(data_view + 2, [], None, &NullEventConsumer)
            .await
            .unwrap();
        assert!(storage.load_vid_share(data_view).await.unwrap().is_none(),);
        assert!(storage.load_da_proposal(data_view).await.unwrap().is_none());
        storage.load_quorum_proposal(data_view).await.unwrap_err();
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_pruning_minimum_retention() {
        test_pruning_helper(ConsensusPruningOptions {
            // Use a very low target usage, to show that we still retain data up to the minimum
            // retention even when usage is above target.
            target_usage: 0,
            minimum_retention: 1,
            // Use a very high target retention, so that pruning is only triggered by the minimum
            // retention.
            target_retention: u64::MAX,
        })
        .await
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_pruning_target_retention() {
        test_pruning_helper(ConsensusPruningOptions {
            target_retention: 1,
            // Use a very low minimum retention, so that data is only kept around due to the target
            // retention.
            minimum_retention: 0,
            // Use a very high target usage, so that pruning is only triggered by the target
            // retention.
            target_usage: u64::MAX,
        })
        .await
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_consensus_migration() {
        let tmp = Persistence::tmp_storage().await;
        let mut opt = Persistence::options(&tmp);

        let storage = opt.create().await.unwrap();

        let rows = 300;

        assert!(storage.load_state_cert().await.unwrap().is_none());

        for i in 0..rows {
            let view = ViewNumber::new(i);
            let validated_state = ValidatedState::default();
            let instance_state = NodeState::default();

            let (pubkey, privkey) = BLSPubKey::generated_from_seed_indexed([0; 32], i);
            let (payload, metadata) =
                Payload::from_transactions([], &validated_state, &instance_state)
                    .await
                    .unwrap();

            let payload_bytes = payload.encode();

            let block_header = Header::genesis(
                &instance_state,
                payload.clone(),
                &metadata,
                TEST_VERSIONS.test.base,
            );

            let null_quorum_data = QuorumData {
                leaf_commit: Commitment::<Leaf>::default_commitment_no_preimage(),
            };

            let justify_qc = QuorumCertificate::new(
                null_quorum_data.clone(),
                null_quorum_data.commit(),
                view,
                None,
                std::marker::PhantomData,
            );

            let quorum_proposal = QuorumProposal {
                block_header,
                view_number: view,
                justify_qc: justify_qc.clone(),
                upgrade_certificate: None,
                proposal_certificate: None,
            };

            let quorum_proposal_signature =
                BLSPubKey::sign(&privkey, &bincode::serialize(&quorum_proposal).unwrap())
                    .expect("Failed to sign quorum proposal");

            let proposal = Proposal {
                data: quorum_proposal.clone(),
                signature: quorum_proposal_signature,
                _pd: std::marker::PhantomData::<SeqTypes>,
            };

            let proposal_bytes = bincode::serialize(&proposal)
                .context("serializing proposal")
                .unwrap();

            let mut leaf = Leaf::from_quorum_proposal(&quorum_proposal);
            leaf.fill_block_payload(
                payload,
                GENESIS_VID_NUM_STORAGE_NODES,
                TEST_VERSIONS.test.base,
            )
            .unwrap();

            let mut tx = storage.db.write().await.unwrap();

            let qc_bytes = bincode::serialize(&justify_qc).unwrap();
            let leaf_bytes = bincode::serialize(&leaf).unwrap();

            tx.upsert(
                "anchor_leaf",
                ["view", "leaf", "qc"],
                ["view"],
                [(i as i64, leaf_bytes, qc_bytes)],
            )
            .await
            .unwrap();

            let state_cert = LightClientStateUpdateCertificateV2::<SeqTypes> {
                epoch: EpochNumber::new(i),
                light_client_state: Default::default(), // filling arbitrary value
                next_stake_table_state: Default::default(), // filling arbitrary value
                signatures: vec![],                     // filling arbitrary value
                auth_root: Default::default(),
            };
            // manually upsert the state cert to the finalized database
            let state_cert_bytes = bincode::serialize(&state_cert).unwrap();
            tx.upsert(
                "finalized_state_cert",
                ["epoch", "state_cert"],
                ["epoch"],
                [(i as i64, state_cert_bytes)],
            )
            .await
            .unwrap();

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

            let disperse = advz_scheme(GENESIS_VID_NUM_STORAGE_NODES)
                .disperse(payload_bytes.clone())
                .unwrap();

            let vid = VidDisperseShare0::<SeqTypes> {
                view_number: ViewNumber::new(i),
                payload_commitment: Default::default(),
                share: disperse.shares[0].clone(),
                common: disperse.common,
                recipient_key: pubkey,
            };

            let (payload, metadata) =
                Payload::from_transactions([], &ValidatedState::default(), &NodeState::default())
                    .await
                    .unwrap();

            let da = DaProposal::<SeqTypes> {
                encoded_transactions: payload.encode(),
                metadata,
                view_number: ViewNumber::new(i),
            };

            let block_payload_signature =
                BLSPubKey::sign(&privkey, &payload_bytes).expect("Failed to sign block payload");

            let da_proposal = Proposal {
                data: da,
                signature: block_payload_signature,
                _pd: Default::default(),
            };

            storage
                .append_vid(&convert_proposal(vid.to_proposal(&privkey).unwrap()))
                .await
                .unwrap();
            storage
                .append_da(&da_proposal, VidCommitment::V0(disperse.commit))
                .await
                .unwrap();

            let leaf_hash = Committable::commit(&leaf);
            let mut tx = storage.db.write().await.expect("failed to start write tx");
            tx.upsert(
                "quorum_proposals",
                ["view", "leaf_hash", "data"],
                ["view"],
                [(i as i64, leaf_hash.to_string(), proposal_bytes)],
            )
            .await
            .expect("failed to upsert quorum proposal");

            let justify_qc = &proposal.data.justify_qc;
            let justify_qc_bytes = bincode::serialize(&justify_qc)
                .context("serializing QC")
                .unwrap();
            tx.upsert(
                "quorum_certificate",
                ["view", "leaf_hash", "data"],
                ["view"],
                [(
                    justify_qc.view_number.u64() as i64,
                    justify_qc.data.leaf_commit.to_string(),
                    &justify_qc_bytes,
                )],
            )
            .await
            .expect("failed to upsert qc");

            tx.commit().await.expect("failed to commit");
        }

        storage.migrate_storage().await.unwrap();

        let mut tx = storage.db.read().await.unwrap();
        let (anchor_leaf2_count,) = query_as::<(i64,)>("SELECT COUNT(*) from anchor_leaf2")
            .fetch_one(tx.as_mut())
            .await
            .unwrap();
        assert_eq!(
            anchor_leaf2_count, rows as i64,
            "anchor leaf count does not match rows",
        );

        let (da_proposal_count,) = query_as::<(i64,)>("SELECT COUNT(*) from da_proposal2")
            .fetch_one(tx.as_mut())
            .await
            .unwrap();
        assert_eq!(
            da_proposal_count, rows as i64,
            "da proposal count does not match rows",
        );

        let (vid_share_count,) = query_as::<(i64,)>("SELECT COUNT(*) from vid_share2")
            .fetch_one(tx.as_mut())
            .await
            .unwrap();
        assert_eq!(
            vid_share_count, rows as i64,
            "vid share count does not match rows"
        );

        let (quorum_proposals_count,) =
            query_as::<(i64,)>("SELECT COUNT(*) from quorum_proposals2")
                .fetch_one(tx.as_mut())
                .await
                .unwrap();
        assert_eq!(
            quorum_proposals_count, rows as i64,
            "quorum proposals count does not match rows",
        );

        let (quorum_certificates_count,) =
            query_as::<(i64,)>("SELECT COUNT(*) from quorum_certificate2")
                .fetch_one(tx.as_mut())
                .await
                .unwrap();
        assert_eq!(
            quorum_certificates_count, rows as i64,
            "quorum certificates count does not match rows",
        );

        let (state_cert_count,) = query_as::<(i64,)>("SELECT COUNT(*) from finalized_state_cert")
            .fetch_one(tx.as_mut())
            .await
            .unwrap();
        assert_eq!(
            state_cert_count, rows as i64,
            "Light client state update certificates count does not match rows",
        );
        assert_eq!(
            storage.load_state_cert().await.unwrap().unwrap(),
            LightClientStateUpdateCertificateV2::<SeqTypes> {
                epoch: EpochNumber::new(rows - 1),
                light_client_state: Default::default(),
                next_stake_table_state: Default::default(),
                signatures: vec![],
                auth_root: Default::default(),
            },
            "Wrong light client state update certificate in the storage",
        );

        storage.migrate_storage().await.unwrap();
    }

    /// Regression test for an ambiguous behavior in `store_events`/`load_events`.
    ///
    /// Previously, `store_events` did nothing when given an empty events list (in fact,
    /// `fetch_and_store_stake_table_events` was not even calling it). But this means that the
    /// `stake_table_events_l1_block` column does not get updated when we enter a new epoch with no
    /// new stake table events. This makes it impossible to distinguish between two very different
    /// scenarios:
    ///
    /// 1. The node has successfully processed events through the latest L1 finalized block, but
    ///    there are no new events from the last epoch.
    /// 2. The node is lagging behind the latest L1 finalized block, and is possibly missing some
    ///    new events.
    ///
    /// In scenario 1, clients of this node should be able to treat the empty list of stake table
    /// events as authoritative, and derive the stake table for the next epoch (which will end up
    /// being the same as the previous one. However, in scenario 2, clients need to wait, because we
    /// don't yet know whether there could be any events that modify the stake table. Thus,
    /// distinguishing these two scenarios is important.
    ///
    /// This regression test ensures that even if there are no new events, at least the
    /// `stake_table_events_l1_block` column gets updated. We can then distinguish the two scenarios
    /// using the `EventsPersistenceRead`` return value from load_events.
    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_store_events_empty() {
        let tmp = Persistence::tmp_storage().await;
        let mut opt = Persistence::options(&tmp);
        let storage = opt.create().await.unwrap();

        assert_eq!(storage.load_events(0, 100).await.unwrap(), (None, vec![]));

        // Storing an empty events list still updates the latest L1 block.
        for i in 1..=2 {
            tracing::info!(i, "update l1 height");
            storage.store_events(i, vec![]).await.unwrap();
            assert_eq!(
                storage.load_events(0, 100).await.unwrap(),
                (Some(EventsPersistenceRead::UntilL1Block(i)), vec![])
            );
        }
    }

    // ---------------------------------------------------------------------------
    // Tests for retry_if / is_serialization_error
    // ---------------------------------------------------------------------------

    /// Minimal `DatabaseError` implementation that lets tests construct a
    /// `sqlx::Error::Database(...)` with an arbitrary SQLSTATE code without
    /// needing a live database connection.
    #[derive(Debug)]
    struct MockDatabaseError {
        code: &'static str,
    }

    impl std::fmt::Display for MockDatabaseError {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "mock db error (code {})", self.code)
        }
    }

    impl std::error::Error for MockDatabaseError {}

    impl sqlx::error::DatabaseError for MockDatabaseError {
        fn message(&self) -> &str {
            "mock db error"
        }

        fn code(&self) -> Option<std::borrow::Cow<'_, str>> {
            Some(std::borrow::Cow::Borrowed(self.code))
        }

        fn kind(&self) -> sqlx::error::ErrorKind {
            sqlx::error::ErrorKind::Other
        }

        fn as_error(&self) -> &(dyn std::error::Error + Send + Sync + 'static) {
            self
        }

        fn as_error_mut(&mut self) -> &mut (dyn std::error::Error + Send + Sync + 'static) {
            self
        }

        fn into_error(self: Box<Self>) -> Box<dyn std::error::Error + Send + Sync + 'static> {
            self
        }
    }

    fn mock_serialization_error() -> anyhow::Error {
        anyhow::Error::from(sqlx::Error::Database(Box::new(MockDatabaseError {
            code: "40001",
        })))
    }

    #[test]
    fn test_is_serialization_error() {
        // PostgreSQL error code 40001 must be recognised as a serialization failure.
        assert!(is_serialization_error(&mock_serialization_error()));

        // Any other database error code must NOT match.
        let unique_violation =
            anyhow::Error::from(sqlx::Error::Database(Box::new(MockDatabaseError {
                code: "23505",
            })));
        assert!(!is_serialization_error(&unique_violation));

        // Non-database errors must not match.
        assert!(!is_serialization_error(&anyhow::anyhow!("plain error")));
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_retry_if_succeeds_immediately() {
        let calls = Arc::new(AtomicU32::new(0));
        let calls_clone = calls.clone();

        let result = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || {
                let calls = calls_clone.clone();
                async move {
                    calls.fetch_add(1, Ordering::SeqCst);
                    Ok(())
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(calls.load(Ordering::SeqCst), 1);
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_retry_if_retries_on_serialization_error() {
        let calls = Arc::new(AtomicU32::new(0));
        let calls_clone = calls.clone();

        // The closure fails twice with a serialization error, then succeeds on the third attempt.
        let result = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || {
                let calls = calls_clone.clone();
                async move {
                    let n = calls.fetch_add(1, Ordering::SeqCst);
                    if n < 2 {
                        Err(mock_serialization_error())
                    } else {
                        Ok(())
                    }
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(calls.load(Ordering::SeqCst), 3);
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_retry_if_exhausts_retries() {
        let calls = Arc::new(AtomicU32::new(0));
        let calls_clone = calls.clone();

        // The closure always fails; retry must give up after WRITE_RETRY_MAX (5) retries.
        let result: anyhow::Result<()> = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || {
                let calls = calls_clone.clone();
                async move {
                    calls.fetch_add(1, Ordering::SeqCst);
                    Err(mock_serialization_error())
                }
            })
            .await;

        assert!(result.is_err());
        // 1 initial attempt + 5 retries = 6 total calls.
        assert_eq!(calls.load(Ordering::SeqCst), 6);
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_retry_if_no_retry_on_other_errors() {
        let calls = Arc::new(AtomicU32::new(0));
        let calls_clone = calls.clone();

        // Non-serialization errors must not be retried.
        let result: anyhow::Result<()> = WRITE_BACKOFF
            .retry_if(WRITE_RETRY_MAX, is_serialization_error, || {
                let calls = calls_clone.clone();
                async move {
                    calls.fetch_add(1, Ordering::SeqCst);
                    Err(anyhow::anyhow!("unrelated error"))
                }
            })
            .await;

        assert!(result.is_err());
        assert_eq!(calls.load(Ordering::SeqCst), 1);
    }
}

#[cfg(test)]
#[cfg(not(feature = "embedded-db"))]
mod postgres_tests {
    use espresso_types::{FeeAccount, Header, Leaf, NodeState, Transaction as Tx};
    use hotshot_example_types::node_types::TEST_VERSIONS;
    use hotshot_query_service::{
        availability::{BlockQueryData, LeafQueryData},
        data_source::storage::UpdateAvailabilityStorage,
    };
    use hotshot_types::{
        data::vid_commitment,
        simple_certificate::QuorumCertificate,
        traits::{
            EncodeBytes,
            block_contents::{BlockHeader, BuilderFee, GENESIS_VID_NUM_STORAGE_NODES},
            election::Membership,
            signature_key::BuilderSignatureKey,
        },
    };

    use super::*;
    use crate::persistence::tests::TestablePersistence as _;

    async fn test_postgres_read_ns_table(instance_state: NodeState) {
        instance_state
            .coordinator
            .membership()
            .set_first_epoch(EpochNumber::genesis(), Default::default());

        let tmp = Persistence::tmp_storage().await;
        let mut opt = Persistence::options(&tmp);
        let storage = opt.create().await.unwrap();

        let txs = [
            Tx::new(10001u32.into(), vec![1, 2, 3]),
            Tx::new(10001u32.into(), vec![4, 5, 6]),
            Tx::new(10009u32.into(), vec![7, 8, 9]),
        ];

        let validated_state = Default::default();
        let justify_qc =
            QuorumCertificate::genesis(&validated_state, &instance_state, TEST_VERSIONS.test).await;
        let view_number: ViewNumber = justify_qc.view_number + 1;
        let parent_leaf = Leaf::genesis(&validated_state, &instance_state, TEST_VERSIONS.test.base)
            .await
            .into();

        let (payload, ns_table) =
            Payload::from_transactions(txs.clone(), &validated_state, &instance_state)
                .await
                .unwrap();
        let payload_bytes = payload.encode();
        let payload_commitment = vid_commitment(
            &payload_bytes,
            &ns_table.encode(),
            GENESIS_VID_NUM_STORAGE_NODES,
            instance_state.current_version,
        );
        let builder_commitment = payload.builder_commitment(&ns_table);
        let (fee_account, fee_key) = FeeAccount::generated_from_seed_indexed([0; 32], 0);
        let fee_amount = 0;
        let fee_signature = FeeAccount::sign_fee(&fee_key, fee_amount, &ns_table).unwrap();
        let block_header = Header::new(
            &validated_state,
            &instance_state,
            &parent_leaf,
            payload_commitment,
            builder_commitment,
            ns_table,
            BuilderFee {
                fee_amount,
                fee_account,
                fee_signature,
            },
            instance_state.current_version,
            view_number.u64(),
        )
        .await
        .unwrap();
        let proposal = QuorumProposal {
            block_header: block_header.clone(),
            view_number,
            justify_qc: justify_qc.clone(),
            upgrade_certificate: None,
            proposal_certificate: None,
        };
        let leaf: Leaf2 = Leaf::from_quorum_proposal(&proposal).into();
        let mut qc = justify_qc.to_qc2();
        qc.data.leaf_commit = leaf.commit();
        qc.view_number = view_number;

        let mut tx = storage.db.write().await.unwrap();
        tx.insert_leaf(&LeafQueryData::new(leaf, qc).unwrap())
            .await
            .unwrap();
        tx.insert_block(&BlockQueryData::<SeqTypes>::new(block_header, payload))
            .await
            .unwrap();
        tx.commit().await.unwrap();

        let mut tx = storage.db.read().await.unwrap();
        let rows = query(
            "
            SELECT ns_id, read_ns_id(get_ns_table(h.data), t.ns_index) AS read_ns_id
              FROM header AS h
              JOIN transactions AS t ON t.block_height = h.height
              ORDER BY t.ns_index, t.position
        ",
        )
        .fetch_all(tx.as_mut())
        .await
        .unwrap();
        assert_eq!(rows.len(), txs.len());
        for (i, row) in rows.into_iter().enumerate() {
            let ns = u64::from(txs[i].namespace()) as i64;
            assert_eq!(row.get::<i64, _>("ns_id"), ns);
            assert_eq!(row.get::<i64, _>("read_ns_id"), ns);
        }
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_postgres_read_ns_table_v0_1() {
        test_postgres_read_ns_table(NodeState::mock()).await;
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_postgres_read_ns_table_v0_2() {
        test_postgres_read_ns_table(NodeState::mock_v2()).await;
    }

    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_postgres_read_ns_table_v0_3() {
        test_postgres_read_ns_table(NodeState::mock_v3().with_epoch_height(0)).await;
    }

    /// Verify that concurrent calls to `record_action` all succeed under
    /// PostgreSQL SERIALIZABLE isolation. `WRITE_BACKOFF.retry_if` handles any
    /// 40001 serialization failures that arise when many tasks race to update
    /// the same row.
    #[test_log::test(tokio::test(flavor = "multi_thread"))]
    async fn test_record_action_concurrent() {
        let tmp = Persistence::tmp_storage().await;
        let storage = Arc::new(Persistence::connect(&tmp).await);

        let handles: Vec<_> = (0u64..20)
            .map(|i| {
                let storage = Arc::clone(&storage);
                tokio::spawn(async move {
                    storage
                        .record_action(ViewNumber::new(i), None, HotShotAction::Vote)
                        .await
                })
            })
            .collect();

        for handle in handles {
            handle.await.unwrap().unwrap();
        }

        let latest = storage.load_latest_acted_view().await.unwrap();
        assert_eq!(latest, Some(ViewNumber::new(19)));
    }
}