cosmwasm_schema/
query_response.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use schemars::{schema::RootSchema, JsonSchema};
4use thiserror::Error;
5
6pub use cosmwasm_schema_derive::QueryResponses;
7
8/// A trait for tying QueryMsg variants (different contract queries) to their response types.
9/// This is mostly useful for the generated contracted API description when using `cargo schema`.
10///
11/// Using the derive macro is the preferred way of implementing this trait.
12///
13/// # Examples
14/// ```
15/// use cosmwasm_schema::QueryResponses;
16/// use schemars::JsonSchema;
17///
18/// #[derive(JsonSchema)]
19/// struct AccountInfo {
20///     IcqHandle: String,
21/// }
22///
23/// #[derive(JsonSchema, QueryResponses)]
24/// enum QueryMsg {
25///     #[returns(Vec<String>)]
26///     Denoms {},
27///     #[returns(AccountInfo)]
28///     AccountInfo { account: String },
29/// }
30/// ```
31///
32/// You can compose multiple queries using `#[query_responses(nested)]`. This might be useful
33/// together with `#[serde(untagged)]`. If the `nested` flag is set, no `returns` attributes
34/// are necessary on the enum variants. Instead, the response types are collected from the
35/// nested enums.
36///
37/// ```
38/// # use cosmwasm_schema::QueryResponses;
39/// # use schemars::JsonSchema;
40/// #[derive(JsonSchema, QueryResponses)]
41/// #[query_responses(nested)]
42/// #[serde(untagged)]
43/// enum QueryMsg {
44///     MsgA(QueryA),
45///     MsgB(QueryB),
46/// }
47///
48/// #[derive(JsonSchema, QueryResponses)]
49/// enum QueryA {
50///     #[returns(Vec<String>)]
51///     Denoms {},
52/// }
53///
54/// #[derive(JsonSchema, QueryResponses)]
55/// enum QueryB {
56///     #[returns(AccountInfo)]
57///     AccountInfo { account: String },
58/// }
59///
60/// # #[derive(JsonSchema)]
61/// # struct AccountInfo {
62/// #     IcqHandle: String,
63/// # }
64/// ```
65pub trait QueryResponses: JsonSchema {
66    fn response_schemas() -> Result<BTreeMap<String, RootSchema>, IntegrityError> {
67        let response_schemas = Self::response_schemas_impl();
68
69        Ok(response_schemas)
70    }
71
72    fn response_schemas_impl() -> BTreeMap<String, RootSchema>;
73}
74
75/// Combines multiple response schemas into one. Panics if there are name collisions.
76/// Used internally in the implementation of [`QueryResponses`] when using `#[query_responses(nested)]`
77pub fn combine_subqueries<const N: usize, T>(
78    subqueries: [BTreeMap<String, RootSchema>; N],
79) -> BTreeMap<String, RootSchema> {
80    let sub_count = subqueries.iter().flatten().count();
81    let map: BTreeMap<_, _> = subqueries.into_iter().flatten().collect();
82    if map.len() != sub_count {
83        panic!(
84            "name collision in subqueries for {}",
85            std::any::type_name::<T>()
86        )
87    }
88    map
89}
90
91#[derive(Debug, Error, PartialEq, Eq)]
92pub enum IntegrityError {
93    #[error("the structure of the QueryMsg schema was unexpected")]
94    InvalidQueryMsgSchema,
95    #[error("external reference in schema found, but they are not supported")]
96    ExternalReference { reference: String },
97    #[error(
98        "inconsistent queries - QueryMsg schema has {query_msg:?}, but query responses have {responses:?}"
99    )]
100    InconsistentQueries {
101        query_msg: BTreeSet<String>,
102        responses: BTreeSet<String>,
103    },
104}
105
106#[cfg(test)]
107mod tests {
108    use schemars::schema_for;
109
110    use super::*;
111
112    #[derive(Debug, JsonSchema)]
113    #[serde(rename_all = "snake_case")]
114    #[allow(dead_code)]
115    pub enum GoodMsg {
116        BalanceFor { account: String },
117        AccountIdFor(String),
118        Supply {},
119        Liquidity,
120        AccountCount(),
121    }
122
123    impl QueryResponses for GoodMsg {
124        fn response_schemas_impl() -> BTreeMap<String, RootSchema> {
125            BTreeMap::from([
126                ("balance_for".to_string(), schema_for!(u128)),
127                ("account_id_for".to_string(), schema_for!(u128)),
128                ("supply".to_string(), schema_for!(u128)),
129                ("liquidity".to_string(), schema_for!(u128)),
130                ("account_count".to_string(), schema_for!(u128)),
131            ])
132        }
133    }
134
135    #[test]
136    fn good_msg_works() {
137        let response_schemas = GoodMsg::response_schemas().unwrap();
138        assert_eq!(
139            response_schemas,
140            BTreeMap::from([
141                ("balance_for".to_string(), schema_for!(u128)),
142                ("account_id_for".to_string(), schema_for!(u128)),
143                ("supply".to_string(), schema_for!(u128)),
144                ("liquidity".to_string(), schema_for!(u128)),
145                ("account_count".to_string(), schema_for!(u128))
146            ])
147        );
148    }
149
150    #[derive(Debug, JsonSchema)]
151    #[serde(rename_all = "snake_case")]
152    #[allow(dead_code)]
153    pub enum EmptyMsg {}
154
155    impl QueryResponses for EmptyMsg {
156        fn response_schemas_impl() -> BTreeMap<String, RootSchema> {
157            BTreeMap::from([])
158        }
159    }
160
161    #[test]
162    fn empty_msg_works() {
163        let response_schemas = EmptyMsg::response_schemas().unwrap();
164        assert_eq!(response_schemas, BTreeMap::from([]));
165    }
166
167    #[derive(Debug, JsonSchema)]
168    #[serde(rename_all = "kebab-case")]
169    #[allow(dead_code)]
170    pub enum BadMsg {
171        BalanceFor { account: String },
172    }
173
174    impl QueryResponses for BadMsg {
175        fn response_schemas_impl() -> BTreeMap<String, RootSchema> {
176            BTreeMap::from([("balance_for".to_string(), schema_for!(u128))])
177        }
178    }
179
180    #[derive(Debug, JsonSchema)]
181    #[serde(rename_all = "snake_case")]
182    #[allow(dead_code)]
183    pub enum ExtMsg {
184        Extension {},
185    }
186
187    #[derive(Debug, JsonSchema)]
188    #[serde(untagged, rename_all = "snake_case")]
189    #[allow(dead_code)]
190    pub enum UntaggedMsg {
191        Good(GoodMsg),
192        Ext(ExtMsg),
193        Empty(EmptyMsg),
194    }
195
196    impl QueryResponses for UntaggedMsg {
197        fn response_schemas_impl() -> BTreeMap<String, RootSchema> {
198            BTreeMap::from([
199                ("balance_for".to_string(), schema_for!(u128)),
200                ("account_id_for".to_string(), schema_for!(u128)),
201                ("supply".to_string(), schema_for!(u128)),
202                ("liquidity".to_string(), schema_for!(u128)),
203                ("account_count".to_string(), schema_for!(u128)),
204                ("extension".to_string(), schema_for!(())),
205            ])
206        }
207    }
208
209    #[test]
210    fn untagged_msg_works() {
211        let response_schemas = UntaggedMsg::response_schemas().unwrap();
212        assert_eq!(
213            response_schemas,
214            BTreeMap::from([
215                ("balance_for".to_string(), schema_for!(u128)),
216                ("account_id_for".to_string(), schema_for!(u128)),
217                ("supply".to_string(), schema_for!(u128)),
218                ("liquidity".to_string(), schema_for!(u128)),
219                ("account_count".to_string(), schema_for!(u128)),
220                ("extension".to_string(), schema_for!(())),
221            ])
222        );
223    }
224}