sui_indexer_alt_reader/
epochs.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::BTreeMap;
5use std::collections::HashMap;
6
7use async_graphql::dataloader::Loader;
8use diesel::ExpressionMethods;
9use diesel::QueryDsl;
10use diesel::sql_types::Array;
11use diesel::sql_types::BigInt;
12use sui_indexer_alt_schema::epochs::StoredEpochEnd;
13use sui_indexer_alt_schema::epochs::StoredEpochStart;
14use sui_indexer_alt_schema::schema::kv_epoch_ends;
15use sui_indexer_alt_schema::schema::kv_epoch_starts;
16
17use crate::error::Error;
18use crate::pg_reader::PgReader;
19
20/// Key for fetching information about the start of an epoch.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22pub struct EpochStartKey(pub u64);
23
24/// Key for fetching information about the latest epoch to have started as of a given checkpoint.
25#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
26pub struct CheckpointBoundedEpochStartKey(pub u64);
27
28/// Key for fetching information about the end of an epoch (which must already be finished).
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30pub struct EpochEndKey(pub u64);
31
32#[async_trait::async_trait]
33impl Loader<EpochStartKey> for PgReader {
34    type Value = StoredEpochStart;
35    type Error = Error;
36
37    async fn load(
38        &self,
39        keys: &[EpochStartKey],
40    ) -> Result<HashMap<EpochStartKey, Self::Value>, Error> {
41        use kv_epoch_starts::dsl as s;
42
43        if keys.is_empty() {
44            return Ok(HashMap::new());
45        }
46
47        let mut conn = self.connect().await?;
48
49        let ids: Vec<_> = keys.iter().map(|e| e.0 as i64).collect();
50        let epochs: Vec<StoredEpochStart> = conn
51            .results(s::kv_epoch_starts.filter(s::epoch.eq_any(ids)))
52            .await?;
53
54        Ok(epochs
55            .into_iter()
56            .map(|e| (EpochStartKey(e.epoch as u64), e))
57            .collect())
58    }
59}
60
61#[async_trait::async_trait]
62impl Loader<CheckpointBoundedEpochStartKey> for PgReader {
63    type Value = StoredEpochStart;
64    type Error = Error;
65
66    async fn load(
67        &self,
68        keys: &[CheckpointBoundedEpochStartKey],
69    ) -> Result<HashMap<CheckpointBoundedEpochStartKey, Self::Value>, Error> {
70        if keys.is_empty() {
71            return Ok(HashMap::new());
72        }
73
74        let mut conn = self.connect().await?;
75
76        let cps: Vec<_> = keys.iter().map(|e| e.0 as i64).collect();
77        let query = diesel::sql_query(
78            r#"
79                SELECT
80                    v.*
81                FROM (
82                    SELECT UNNEST($1) cp_sequence_number
83                ) k
84                CROSS JOIN LATERAL (
85                    SELECT
86                        epoch,
87                        protocol_version,
88                        cp_lo,
89                        start_timestamp_ms,
90                        reference_gas_price,
91                        system_state
92                    FROM
93                        kv_epoch_starts
94                    WHERE
95                        kv_epoch_starts.cp_lo <= k.cp_sequence_number
96                    ORDER BY
97                        kv_epoch_starts.cp_lo DESC
98                    LIMIT
99                        1
100                ) v
101            "#,
102        )
103        .bind::<Array<BigInt>, _>(cps);
104
105        let stored_epochs: Vec<StoredEpochStart> = conn.results(query).await?;
106
107        // A single data loader request may contain multiple keys for the same epoch. Store them in
108        // an ordered map, so that we can find the latest version for each key.
109        let cp_to_stored: BTreeMap<_, _> = stored_epochs
110            .into_iter()
111            .map(|epoch| (epoch.cp_lo as u64, epoch))
112            .collect();
113
114        Ok(keys
115            .iter()
116            .filter_map(|key| {
117                let stored = cp_to_stored.range(..=key.0).last()?.1;
118                Some((*key, stored.clone()))
119            })
120            .collect())
121    }
122}
123
124#[async_trait::async_trait]
125impl Loader<EpochEndKey> for PgReader {
126    type Value = StoredEpochEnd;
127    type Error = Error;
128
129    async fn load(&self, keys: &[EpochEndKey]) -> Result<HashMap<EpochEndKey, Self::Value>, Error> {
130        use kv_epoch_ends::dsl as e;
131
132        if keys.is_empty() {
133            return Ok(HashMap::new());
134        }
135
136        let mut conn = self.connect().await?;
137
138        let ids: Vec<_> = keys.iter().map(|e| e.0 as i64).collect();
139        let epochs: Vec<StoredEpochEnd> = conn
140            .results(e::kv_epoch_ends.filter(e::epoch.eq_any(ids)))
141            .await?;
142
143        Ok(epochs
144            .into_iter()
145            .map(|e| (EpochEndKey(e.epoch as u64), e))
146            .collect())
147    }
148}
sui_indexer_alt_reader/epochs.rs

sui_indexer_alt_reader/
epochs.rs
