sui_indexer_alt_reader/
packages.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 anyhow::Context as _;
8use async_graphql::dataloader::Loader;
9use diesel::sql_types::Array;
10use diesel::sql_types::BigInt;
11use diesel::sql_types::Bytea;
12use sui_indexer_alt_schema::packages::StoredPackage;
13use sui_indexer_alt_schema::packages::StoredPackageOriginalId;
14use sui_types::base_types::ObjectID;
15
16use crate::error::Error;
17use crate::pg_reader::PgReader;
18
19/// Key for fetching the original ID of a package
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
21pub struct PackageOriginalIdKey(pub ObjectID);
22
23/// Key for fetching the latest version of a package, based on its *original ID* and a checkpoint
24/// bound.
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Ord, PartialOrd, Hash)]
26pub struct CheckpointBoundedOriginalPackageKey(pub ObjectID, pub u64);
27
28/// Key for fetching a package by its original ID and version.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30pub struct VersionedOriginalPackageKey(pub ObjectID, pub u64);
31
32#[async_trait::async_trait]
33impl Loader<PackageOriginalIdKey> for PgReader {
34    type Value = StoredPackageOriginalId;
35    type Error = Error;
36
37    async fn load(
38        &self,
39        keys: &[PackageOriginalIdKey],
40    ) -> Result<HashMap<PackageOriginalIdKey, StoredPackageOriginalId>, Error> {
41        if keys.is_empty() {
42            return Ok(HashMap::new());
43        }
44
45        let mut conn = self.connect().await?;
46
47        let ids: Vec<_> = keys.iter().map(|k| k.0.into_bytes()).collect();
48        let query = diesel::sql_query(
49            r#"
50                SELECT
51                    k.package_id,
52                    v.original_id,
53                    v.cp_sequence_number
54                FROM (
55                    SELECT UNNEST($1) package_id
56                ) k
57                CROSS JOIN LATERAL (
58                    SELECT
59                        original_id,
60                        cp_sequence_number
61                    FROM
62                        kv_packages
63                    WHERE
64                        kv_packages.package_id = k.package_id
65                    LIMIT
66                        1
67                ) v
68            "#,
69        )
70        .bind::<Array<Bytea>, _>(ids);
71
72        let stored: Vec<StoredPackageOriginalId> = conn.results(query).await?;
73        let id_to_stored: HashMap<_, _> = stored
74            .iter()
75            .map(|package| (&package.package_id[..], package))
76            .collect();
77
78        Ok(keys
79            .iter()
80            .filter_map(|key| {
81                let stored = *id_to_stored.get(key.0.into_bytes().as_ref())?;
82                Some((*key, stored.clone()))
83            })
84            .collect())
85    }
86}
87
88#[async_trait::async_trait]
89impl Loader<CheckpointBoundedOriginalPackageKey> for PgReader {
90    type Value = StoredPackage;
91    type Error = Error;
92
93    async fn load(
94        &self,
95        keys: &[CheckpointBoundedOriginalPackageKey],
96    ) -> Result<HashMap<CheckpointBoundedOriginalPackageKey, StoredPackage>, Error> {
97        if keys.is_empty() {
98            return Ok(HashMap::new());
99        }
100
101        let mut conn = self.connect().await?;
102
103        let ids: Vec<_> = keys.iter().map(|k| k.0.into_bytes()).collect();
104        let cps: Vec<_> = keys.iter().map(|k| k.1 as i64).collect();
105        let query = diesel::sql_query(
106            r#"
107                SELECT
108                    v.*
109                FROM (
110                    SELECT
111                        UNNEST($1) original_id,
112                        UNNEST($2) cp_sequence_number
113                ) k
114                CROSS JOIN LATERAL (
115                    SELECT
116                        package_id,
117                        package_version,
118                        original_id,
119                        is_system_package,
120                        serialized_object,
121                        cp_sequence_number
122                    FROM
123                        kv_packages
124                    WHERE
125                        kv_packages.original_id = k.original_id
126                    AND kv_packages.cp_sequence_number <= k.cp_sequence_number
127                    ORDER BY
128                        cp_sequence_number DESC,
129                        package_version DESC
130                    LIMIT
131                        1
132                ) v
133            "#,
134        )
135        .bind::<Array<Bytea>, _>(ids)
136        .bind::<Array<BigInt>, _>(cps);
137
138        let stored_packages: Vec<StoredPackage> = conn.results(query).await?;
139
140        // A single data loader request may contain multiple keys for the same package ID. Store
141        // them in an ordered map, so that we can find the latest version for each key.
142        let mut key_to_stored = BTreeMap::new();
143        for package in stored_packages {
144            let id = ObjectID::from_bytes(&package.original_id)
145                .context("Failed to deserialize ObjectID")?;
146
147            let cp_sequence_number = package.cp_sequence_number as u64;
148            key_to_stored.insert(
149                CheckpointBoundedOriginalPackageKey(id, cp_sequence_number),
150                package,
151            );
152        }
153
154        Ok(keys
155            .iter()
156            .filter_map(|key| {
157                let (bound, stored) = key_to_stored.range(..=key).last()?;
158                (key.0 == bound.0).then(|| (*key, stored.clone()))
159            })
160            .collect())
161    }
162}
163
164#[async_trait::async_trait]
165impl Loader<VersionedOriginalPackageKey> for PgReader {
166    type Value = StoredPackage;
167    type Error = Error;
168
169    async fn load(
170        &self,
171        keys: &[VersionedOriginalPackageKey],
172    ) -> Result<HashMap<VersionedOriginalPackageKey, StoredPackage>, Error> {
173        if keys.is_empty() {
174            return Ok(HashMap::new());
175        }
176
177        let mut conn = self.connect().await?;
178
179        let ids: Vec<_> = keys.iter().map(|k| k.0.into_bytes()).collect();
180        let versions: Vec<_> = keys.iter().map(|k| k.1 as i64).collect();
181        let query = diesel::sql_query(
182            r#"
183                SELECT
184                    v.*
185                FROM (
186                    SELECT
187                        UNNEST($1) original_id,
188                        UNNEST($2) package_version
189                ) k
190                CROSS JOIN LATERAL (
191                    SELECT
192                        package_id,
193                        package_version,
194                        original_id,
195                        is_system_package,
196                        serialized_object,
197                        cp_sequence_number
198                    FROM
199                        kv_packages
200                    WHERE
201                        kv_packages.original_id = k.original_id
202                    AND kv_packages.package_version = k.package_version
203                    LIMIT
204                        1
205                ) v
206            "#,
207        )
208        .bind::<Array<Bytea>, _>(ids)
209        .bind::<Array<BigInt>, _>(versions);
210
211        let stored_packages: Vec<StoredPackage> = conn.results(query).await?;
212        let key_to_stored: HashMap<_, _> = stored_packages
213            .iter()
214            .map(|stored| {
215                let id = &stored.original_id[..];
216                let version = stored.package_version as u64;
217                ((id, version), stored)
218            })
219            .collect();
220
221        Ok(keys
222            .iter()
223            .filter_map(|key| {
224                let slice: &[u8] = key.0.as_ref();
225                let stored = *key_to_stored.get(&(slice, key.1))?;
226                Some((*key, stored.clone()))
227            })
228            .collect())
229    }
230}
sui_indexer_alt_reader/packages.rs

sui_indexer_alt_reader/
packages.rs
