1use std::{borrow::Cow, fmt, fmt::Debug, sync::Arc};
2
3use futures_util::{
4 Future, FutureExt, Stream, StreamExt, TryStreamExt, future::BoxFuture, stream::BoxStream,
5};
6use indexmap::IndexMap;
7
8use crate::{
9 ContextSelectionSet, Data, Name, QueryPathNode, QueryPathSegment, Response, Result,
10 ServerResult, Value,
11 dynamic::{
12 FieldValue, InputValue, ObjectAccessor, ResolverContext, Schema, SchemaError, TypeRef,
13 resolve::resolve,
14 },
15 extensions::ResolveInfo,
16 parser::types::Selection,
17 registry::{Deprecation, MetaField, MetaType, Registry},
18 subscription::BoxFieldStream,
19};
20
21type BoxResolveFut<'a> = BoxFuture<'a, Result<BoxStream<'a, Result<FieldValue<'a>>>>>;
22
23pub struct SubscriptionFieldFuture<'a>(pub(crate) BoxResolveFut<'a>);
25
26impl<'a> SubscriptionFieldFuture<'a> {
27 pub fn new<Fut, S, T>(future: Fut) -> Self
29 where
30 Fut: Future<Output = Result<S>> + Send + 'a,
31 S: Stream<Item = Result<T>> + Send + 'a,
32 T: Into<FieldValue<'a>> + Send + 'a,
33 {
34 Self(
35 async move {
36 let res = future.await?.map_ok(Into::into);
37 Ok(res.boxed())
38 }
39 .boxed(),
40 )
41 }
42}
43
44type BoxResolverFn =
45 Arc<(dyn for<'a> Fn(ResolverContext<'a>) -> SubscriptionFieldFuture<'a> + Send + Sync)>;
46
47pub struct SubscriptionField {
49 pub(crate) name: String,
50 pub(crate) description: Option<String>,
51 pub(crate) arguments: IndexMap<String, InputValue>,
52 pub(crate) ty: TypeRef,
53 pub(crate) resolver_fn: BoxResolverFn,
54 pub(crate) deprecation: Deprecation,
55}
56
57impl SubscriptionField {
58 pub fn new<N, T, F>(name: N, ty: T, resolver_fn: F) -> Self
60 where
61 N: Into<String>,
62 T: Into<TypeRef>,
63 F: for<'a> Fn(ResolverContext<'a>) -> SubscriptionFieldFuture<'a> + Send + Sync + 'static,
64 {
65 Self {
66 name: name.into(),
67 description: None,
68 arguments: Default::default(),
69 ty: ty.into(),
70 resolver_fn: Arc::new(resolver_fn),
71 deprecation: Deprecation::NoDeprecated,
72 }
73 }
74
75 impl_set_description!();
76 impl_set_deprecation!();
77
78 #[inline]
80 pub fn argument(mut self, input_value: InputValue) -> Self {
81 self.arguments.insert(input_value.name.clone(), input_value);
82 self
83 }
84}
85
86impl Debug for SubscriptionField {
87 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88 f.debug_struct("Field")
89 .field("name", &self.name)
90 .field("description", &self.description)
91 .field("arguments", &self.arguments)
92 .field("ty", &self.ty)
93 .field("deprecation", &self.deprecation)
94 .finish()
95 }
96}
97
98#[derive(Debug)]
100pub struct Subscription {
101 pub(crate) name: String,
102 pub(crate) description: Option<String>,
103 pub(crate) fields: IndexMap<String, SubscriptionField>,
104}
105
106impl Subscription {
107 #[inline]
109 pub fn new(name: impl Into<String>) -> Self {
110 Self {
111 name: name.into(),
112 description: None,
113 fields: Default::default(),
114 }
115 }
116
117 impl_set_description!();
118
119 #[inline]
121 pub fn field(mut self, field: SubscriptionField) -> Self {
122 assert!(
123 !self.fields.contains_key(&field.name),
124 "Field `{}` already exists",
125 field.name
126 );
127 self.fields.insert(field.name.clone(), field);
128 self
129 }
130
131 #[inline]
133 pub fn type_name(&self) -> &str {
134 &self.name
135 }
136
137 pub(crate) fn register(&self, registry: &mut Registry) -> Result<(), SchemaError> {
138 let mut fields = IndexMap::new();
139
140 for field in self.fields.values() {
141 let mut args = IndexMap::new();
142
143 for argument in field.arguments.values() {
144 args.insert(argument.name.clone(), argument.to_meta_input_value());
145 }
146
147 fields.insert(
148 field.name.clone(),
149 MetaField {
150 name: field.name.clone(),
151 description: field.description.clone(),
152 args,
153 ty: field.ty.to_string(),
154 deprecation: field.deprecation.clone(),
155 cache_control: Default::default(),
156 external: false,
157 requires: None,
158 provides: None,
159 visible: None,
160 shareable: false,
161 inaccessible: false,
162 tags: vec![],
163 override_from: None,
164 compute_complexity: None,
165 directive_invocations: vec![],
166 requires_scopes: vec![],
167 },
168 );
169 }
170
171 registry.types.insert(
172 self.name.clone(),
173 MetaType::Object {
174 name: self.name.clone(),
175 description: self.description.clone(),
176 fields,
177 cache_control: Default::default(),
178 extends: false,
179 shareable: false,
180 resolvable: true,
181 keys: None,
182 visible: None,
183 inaccessible: false,
184 interface_object: false,
185 tags: vec![],
186 is_subscription: true,
187 rust_typename: None,
188 directive_invocations: vec![],
189 requires_scopes: vec![],
190 },
191 );
192
193 Ok(())
194 }
195
196 pub(crate) fn collect_streams<'a>(
197 &self,
198 schema: &Schema,
199 ctx: &ContextSelectionSet<'a>,
200 streams: &mut Vec<BoxFieldStream<'a>>,
201 root_value: &'a FieldValue<'static>,
202 ) {
203 for selection in &ctx.item.node.items {
204 if let Selection::Field(field) = &selection.node {
205 if let Some(field_def) = self.fields.get(field.node.name.node.as_str()) {
206 let schema = schema.clone();
207 let field_type = field_def.ty.clone();
208 let resolver_fn = field_def.resolver_fn.clone();
209 let ctx = ctx.clone();
210
211 streams.push(
212 async_stream::try_stream! {
213 let ctx_field = ctx.with_field(field);
214 let field_name = ctx_field.item.node.response_key().node.clone();
215 let arguments = ObjectAccessor(Cow::Owned(
216 field
217 .node
218 .arguments
219 .iter()
220 .map(|(name, value)| {
221 ctx_field
222 .resolve_input_value(value.clone())
223 .map(|value| (name.node.clone(), value))
224 })
225 .collect::<ServerResult<IndexMap<Name, Value>>>()?,
226 ));
227
228 let mut stream = resolver_fn(ResolverContext {
229 ctx: &ctx_field,
230 args: arguments,
231 parent_value: root_value,
232 })
233 .0
234 .await
235 .map_err(|err| ctx_field.set_error_path(err.into_server_error(ctx_field.item.pos)))?;
236
237 while let Some(value) = stream.next().await.transpose().map_err(|err| ctx_field.set_error_path(err.into_server_error(ctx_field.item.pos)))? {
238 let f = |execute_data: Option<Data>| {
239 let schema = schema.clone();
240 let field_name = field_name.clone();
241 let field_type = field_type.clone();
242 let ctx_field = ctx_field.clone();
243
244 async move {
245 let mut ctx_field = ctx_field.clone();
246 ctx_field.execute_data = execute_data.as_ref();
247 let ri = ResolveInfo {
248 path_node: &QueryPathNode {
249 parent: None,
250 segment: QueryPathSegment::Name(&field_name),
251 },
252 parent_type: schema.0.env.registry.subscription_type.as_ref().unwrap(),
253 return_type: &field_type.to_string(),
254 name: field.node.name.node.as_str(),
255 alias: field.node.alias.as_ref().map(|alias| alias.node.as_str()),
256 is_for_introspection: false,
257 field: &field.node,
258 };
259 let resolve_fut = resolve(&schema, &ctx_field, &field_type, Some(&value));
260 futures_util::pin_mut!(resolve_fut);
261 let value = ctx_field.query_env.extensions.resolve(ri, &mut resolve_fut).await;
262
263 match value {
264 Ok(value) => {
265 let mut map = IndexMap::new();
266 map.insert(field_name.clone(), value.unwrap_or_default());
267 Response::new(Value::Object(map))
268 },
269 Err(err) => Response::from_errors(vec![err]),
270 }
271 }
272 };
273 let resp = ctx_field.query_env.extensions.execute(ctx_field.query_env.operation_name.as_deref(), f).await;
274 let is_err = !resp.errors.is_empty();
275 yield resp;
276 if is_err {
277 break;
278 }
279 }
280 }.map(|res| {
281 res.unwrap_or_else(|err| Response::from_errors(vec![err]))
282 })
283 .boxed(),
284 );
285 }
286 }
287 }
288 }
289}
290
291#[cfg(test)]
292mod tests {
293 use std::time::Duration;
294
295 use futures_util::StreamExt;
296
297 use crate::{Value, dynamic::*, value};
298
299 #[tokio::test]
300 async fn subscription() {
301 struct MyObjData {
302 value: i32,
303 }
304
305 let my_obj = Object::new("MyObject").field(Field::new(
306 "value",
307 TypeRef::named_nn(TypeRef::INT),
308 |ctx| {
309 FieldFuture::new(async {
310 Ok(Some(Value::from(
311 ctx.parent_value.try_downcast_ref::<MyObjData>()?.value,
312 )))
313 })
314 },
315 ));
316
317 let query = Object::new("Query").field(Field::new(
318 "value",
319 TypeRef::named_nn(TypeRef::INT),
320 |_| FieldFuture::new(async { Ok(FieldValue::none()) }),
321 ));
322
323 let subscription = Subscription::new("Subscription").field(SubscriptionField::new(
324 "obj",
325 TypeRef::named_nn(my_obj.type_name()),
326 |_| {
327 SubscriptionFieldFuture::new(async {
328 Ok(async_stream::try_stream! {
329 for i in 0..10 {
330 tokio::time::sleep(Duration::from_millis(100)).await;
331 yield FieldValue::owned_any(MyObjData { value: i });
332 }
333 })
334 })
335 },
336 ));
337
338 let schema = Schema::build(query.type_name(), None, Some(subscription.type_name()))
339 .register(my_obj)
340 .register(query)
341 .register(subscription)
342 .finish()
343 .unwrap();
344
345 let mut stream = schema.execute_stream("subscription { obj { value } }");
346 for i in 0..10 {
347 assert_eq!(
348 stream.next().await.unwrap().into_result().unwrap().data,
349 value!({
350 "obj": { "value": i }
351 })
352 );
353 }
354 }
355
356 #[tokio::test]
357 async fn borrow_context() {
358 struct State {
359 value: i32,
360 }
361
362 let query =
363 Object::new("Query").field(Field::new("value", TypeRef::named(TypeRef::INT), |_| {
364 FieldFuture::new(async { Ok(FieldValue::NONE) })
365 }));
366
367 let subscription = Subscription::new("Subscription").field(SubscriptionField::new(
368 "values",
369 TypeRef::named_nn(TypeRef::INT),
370 |ctx| {
371 SubscriptionFieldFuture::new(async move {
372 Ok(async_stream::try_stream! {
373 for i in 0..10 {
374 tokio::time::sleep(Duration::from_millis(100)).await;
375 yield FieldValue::value(ctx.data_unchecked::<State>().value + i);
376 }
377 })
378 })
379 },
380 ));
381
382 let schema = Schema::build("Query", None, Some(subscription.type_name()))
383 .register(query)
384 .register(subscription)
385 .data(State { value: 123 })
386 .finish()
387 .unwrap();
388
389 let mut stream = schema.execute_stream("subscription { values }");
390 for i in 0..10 {
391 assert_eq!(
392 stream.next().await.unwrap().into_result().unwrap().data,
393 value!({ "values": i + 123 })
394 );
395 }
396 }
397}