sui_gql_client/queries/
full_objects.rs

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