sui_indexer_alt_reader/
objects.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashMap;
5
6use anyhow::Context;
7use prost_types::FieldMask;
8use sui_rpc::field::FieldMaskUtil;
9use sui_rpc::proto::sui::rpc::v2 as proto;
10use sui_types::base_types::ObjectID;
11use sui_types::object::Object;
12
13use crate::error::Error;
14use crate::ledger_grpc_reader::ChunkedLoader;
15use crate::ledger_grpc_reader::LedgerGrpcReader;
16
17/// Key for fetching the contents a particular version of an object.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
19pub struct VersionedObjectKey(pub ObjectID, pub u64);
20
21#[async_trait::async_trait]
22impl ChunkedLoader<VersionedObjectKey> for LedgerGrpcReader {
23    type Value = Object;
24    type Error = Error;
25
26    fn chunk_size(&self) -> usize {
27        self.max_batch_get_objects()
28    }
29
30    async fn load_chunk(
31        &self,
32        keys: &[VersionedObjectKey],
33    ) -> Result<HashMap<VersionedObjectKey, Object>, Error> {
34        let requests = keys
35            .iter()
36            .map(|key| {
37                let mut req = proto::GetObjectRequest::new(&key.0.into());
38                req.version = Some(key.1);
39                req
40            })
41            .collect();
42
43        let mut request = proto::BatchGetObjectsRequest::default();
44        request.requests = requests;
45        request.read_mask = Some(FieldMask::from_paths(["bcs"]));
46
47        let batch_response = self.batch_get_objects(request).await?;
48
49        let mut results = HashMap::new();
50        for obj_result in batch_response.objects {
51            if let Some(proto::get_object_result::Result::Object(object)) = obj_result.result {
52                let obj: Object = object
53                    .bcs
54                    .as_ref()
55                    .context("Missing bcs in object")?
56                    .deserialize()
57                    .context("Failed to deserialize object")?;
58                results.insert(VersionedObjectKey(obj.id(), obj.version().into()), obj);
59            }
60        }
61        Ok(results)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use async_graphql::dataloader::Loader;
68    use sui_sdk_types::Address;
69    use sui_types::base_types::ObjectID;
70
71    use super::*;
72    use crate::ledger_grpc_reader::test_support::mock_reader;
73
74    #[tokio::test]
75    async fn load_chunks_oversized_batches() {
76        let (reader, mock, server) = mock_reader().await;
77        let limit = reader.max_batch_get_objects();
78
79        let keys: Vec<VersionedObjectKey> = (0..limit + 50)
80            .map(|i| VersionedObjectKey(ObjectID::random(), i as u64))
81            .collect();
82
83        let result = reader.load(&keys).await.expect("load should succeed");
84        assert!(result.is_empty());
85
86        let batches = mock.object_batches();
87        assert_eq!(batches.len(), 2);
88        assert!(batches.iter().all(|batch| batch.len() <= limit));
89
90        let mut requested: Vec<(String, Option<u64>)> = batches
91            .into_iter()
92            .flatten()
93            .map(|req| (req.object_id.unwrap_or_default(), req.version))
94            .collect();
95        requested.sort();
96        let mut expected: Vec<(String, Option<u64>)> = keys
97            .iter()
98            .map(|key| {
99                let address: Address = key.0.into();
100                (address.to_string(), Some(key.1))
101            })
102            .collect();
103        expected.sort();
104        assert_eq!(requested, expected);
105
106        server.abort();
107    }
108}
sui_indexer_alt_reader/objects.rs

sui_indexer_alt_reader/
objects.rs
