Skip to main content

sui_indexer_alt_reader/
package_resolver.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashMap;
5use std::sync::Arc;
6
7use async_graphql::dataloader::DataLoader;
8use async_graphql::dataloader::Loader;
9use diesel::prelude::QueryableByName;
10use diesel::sql_types::Array;
11use diesel::sql_types::Bytea;
12use move_core_types::account_address::AccountAddress;
13use sui_indexer_alt_schema::schema::kv_packages;
14use sui_package_resolver::Package;
15use sui_package_resolver::PackageStore;
16use sui_package_resolver::PackageStoreWithLruCache;
17use sui_package_resolver::Result;
18use sui_package_resolver::error::Error;
19use sui_types::object::Object;
20
21use crate::pg_reader::PgReader;
22
23const STORE: &str = "PostgreSQL";
24
25pub type PackageCache = PackageStoreWithLruCache<DbPackageStore>;
26pub struct DbPackageStore(Arc<DataLoader<PgReader>>);
27
28#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
29struct PackageKey(AccountAddress);
30
31impl DbPackageStore {
32    pub fn new(loader: Arc<DataLoader<PgReader>>) -> Self {
33        Self(loader)
34    }
35}
36
37#[async_trait::async_trait]
38impl PackageStore for DbPackageStore {
39    async fn fetch(&self, id: AccountAddress) -> Result<Arc<Package>> {
40        let Self(loader) = self;
41        let Some(package) = loader.load_one(PackageKey(id)).await? else {
42            return Err(Error::PackageNotFound(id));
43        };
44
45        Ok(package)
46    }
47}
48
49#[async_trait::async_trait]
50impl Loader<PackageKey> for PgReader {
51    type Value = Arc<Package>;
52    type Error = Error;
53
54    async fn load(&self, keys: &[PackageKey]) -> Result<HashMap<PackageKey, Arc<Package>>> {
55        let mut id_to_package = HashMap::new();
56        if keys.is_empty() {
57            return Ok(id_to_package);
58        }
59
60        let mut conn = self.connect().await.map_err(|e| Error::Store {
61            store: STORE,
62            error: e.to_string(),
63        })?;
64
65        #[derive(QueryableByName)]
66        #[diesel(table_name = kv_packages)]
67        struct SerializedPackage {
68            serialized_object: Vec<u8>,
69        }
70
71        let ids: Vec<_> = keys.iter().map(|PackageKey(id)| id.into_bytes()).collect();
72        let query = diesel::sql_query(
73            r#"
74                SELECT
75                    v.serialized_object
76                FROM (
77                    SELECT UNNEST($1) package_id
78                ) k
79                CROSS JOIN LATERAL (
80                    SELECT
81                        serialized_object
82                    FROM
83                        kv_packages
84                    WHERE
85                        kv_packages.package_id = k.package_id
86                    ORDER BY
87                        package_version DESC
88                    LIMIT 1
89                ) v
90            "#,
91        )
92        .bind::<Array<Bytea>, _>(ids);
93
94        let stored_packages: Vec<SerializedPackage> =
95            conn.results(query).await.map_err(|e| Error::Store {
96                store: STORE,
97                error: e.to_string(),
98            })?;
99
100        for stored in stored_packages {
101            let object: Object = bcs::from_bytes(&stored.serialized_object)?;
102            let Some(move_package) = object.data.try_as_package() else {
103                return Err(Error::NotAPackage(object.id().into()));
104            };
105
106            let package = Package::read_from_package(move_package)?;
107            id_to_package.insert(PackageKey(*move_package.id()), Arc::new(package));
108        }
109
110        Ok(id_to_package)
111    }
112}