sui_gql_client/queries/
full_objects.rs

1use std::collections::HashMap;
2
3use af_sui_types::{Object, ObjectId};
4use futures::{StreamExt as _, TryStreamExt as _};
5use graphql_extract::extract;
6use itertools::{Either, Itertools as _};
7use sui_gql_schema::scalars::Base64Bcs;
8
9use super::fragments::{ObjectFilter, ObjectKey, PageInfo, PageInfoForward};
10use super::stream;
11use crate::queries::Error;
12use crate::{missing_data, schema, GraphQlClient, GraphQlResponseExt as _};
13
14pub(super) async fn query<C: GraphQlClient>(
15    client: &C,
16    objects: impl IntoIterator<Item = (ObjectId, Option<u64>)> + Send,
17    page_size: Option<u32>,
18) -> Result<HashMap<ObjectId, Object>, Error<C::Error>> {
19    // To keep track of all ids requested.
20    let mut requested = vec![];
21
22    let (object_ids, object_keys) = objects
23        .into_iter()
24        .inspect(|(id, _)| requested.push(*id))
25        .partition_map(|(id, v)| {
26            v.map_or(Either::Left(id), |n| {
27                Either::Right(ObjectKey {
28                    object_id: id,
29                    version: n,
30                })
31            })
32        });
33
34    #[expect(
35        deprecated,
36        reason = "TODO: build query from scratch with new ObjectFilter and Query.multiGetObjects"
37    )]
38    let filter = ObjectFilter {
39        object_ids: Some(object_ids),
40        object_keys: Some(object_keys),
41        ..Default::default()
42    };
43    let vars = Variables {
44        after: None,
45        first: page_size.map(|v| v.try_into().unwrap_or(i32::MAX)),
46        filter: Some(filter),
47    };
48
49    let raw_objs: HashMap<_, _> = stream::forward(client, vars, request)
50        .map(|r| -> super::Result<_, C> {
51            let (id, obj) = r?;
52            Ok((id, obj.ok_or_else(|| missing_data!("BCS for {id}"))?))
53        })
54        .try_collect()
55        .await?;
56
57    // Ensure all requested objects were returned
58    for id in requested {
59        raw_objs
60            .contains_key(&id)
61            .then_some(())
62            .ok_or(missing_data!("Object {id}"))?;
63    }
64
65    Ok(raw_objs)
66}
67
68async fn request<C: GraphQlClient>(
69    client: &C,
70    vars: Variables,
71) -> super::Result<
72    stream::Page<impl Iterator<Item = super::Result<(ObjectId, Option<Object>), C>> + 'static>,
73    C,
74> {
75    let data = client
76        .query::<Query, _>(vars)
77        .await
78        .map_err(Error::Client)?
79        .try_into_data()?;
80
81    extract!(data => {
82        objects {
83            nodes[] {
84                id
85                object
86            }
87            page_info
88        }
89    });
90    Ok(stream::Page::new(
91        page_info,
92        nodes.map(|r| -> super::Result<_, C> {
93            let (id, obj) = r?;
94            Ok((id, obj.map(Base64Bcs::into_inner)))
95        }),
96    ))
97}
98
99#[derive(cynic::QueryVariables, Clone, Debug)]
100struct Variables {
101    filter: Option<ObjectFilter>,
102    after: Option<String>,
103    first: Option<i32>,
104}
105
106impl stream::UpdatePageInfo for Variables {
107    fn update_page_info(&mut self, info: &PageInfo) {
108        self.after.clone_from(&info.end_cursor);
109    }
110}
111
112#[derive(cynic::QueryFragment, Clone, Debug)]
113#[cynic(variables = "Variables")]
114struct Query {
115    #[arguments(filter: $filter, first: $first, after: $after)]
116    objects: ObjectConnection,
117}
118
119// =============================================================================
120//  Inner query fragments
121// =============================================================================
122
123/// `ObjectConnection` where the `Object` fragment does take any parameters.
124#[derive(cynic::QueryFragment, Clone, Debug)]
125struct ObjectConnection {
126    nodes: Vec<ObjectGql>,
127    page_info: PageInfoForward,
128}
129
130#[derive(cynic::QueryFragment, Debug, Clone)]
131#[cynic(graphql_type = "Object")]
132struct ObjectGql {
133    #[cynic(rename = "address")]
134    id: ObjectId,
135    #[cynic(rename = "bcs")]
136    object: Option<Base64Bcs<Object>>,
137}
138
139#[cfg(test)]
140#[allow(clippy::unwrap_used)]
141#[test]
142fn gql_output() -> color_eyre::Result<()> {
143    use cynic::QueryBuilder as _;
144
145    let vars = Variables {
146        filter: Some(ObjectFilter {
147            object_ids: Some(vec![
148                "0x4264c07a42f9d002c1244e43a1f0fa21c49e4a25c7202c597b8476ef6bb57113".parse()?,
149                "0x60d1a85f81172a7418206f4b16e1e07e40c91cf58783f63f18a25efc81442dcb".parse()?,
150            ]),
151            ..Default::default()
152        }),
153        after: None,
154        first: None,
155    };
156
157    let operation = Query::build(vars);
158    insta::assert_snapshot!(operation.query, @r###"
159    query Query($filter: ObjectFilter, $after: String, $first: Int) {
160      objects(filter: $filter, first: $first, after: $after) {
161        nodes {
162          address
163          bcs
164        }
165        pageInfo {
166          hasNextPage
167          endCursor
168        }
169      }
170    }
171    "###);
172    Ok(())
173}