cosmwasm_schema/
query_response.rs

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