1use std::collections::{BTreeMap, BTreeSet};
2
3use schemars::{schema::RootSchema, JsonSchema};
4use thiserror::Error;
5
6pub use cosmwasm_schema_derive::QueryResponses;
7
8pub 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
75pub 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}