async_graphql/dynamic/
subscription.rs

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
23/// A future that returned from field resolver
24pub struct SubscriptionFieldFuture<'a>(pub(crate) BoxResolveFut<'a>);
25
26impl<'a> SubscriptionFieldFuture<'a> {
27    /// Create a ResolverFuture
28    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
47/// A GraphQL subscription field
48pub 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    /// Create a GraphQL subscription field
59    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    /// Add an argument to the subscription field
79    #[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/// A GraphQL subscription type
99#[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    /// Create a GraphQL object type
108    #[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    /// Add an field to the object
120    #[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    /// Returns the type name
132    #[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                && let Some(field_def) = self.fields.get(field.node.name.node.as_str())
206            {
207                let schema = schema.clone();
208                let field_type = field_def.ty.clone();
209                let resolver_fn = field_def.resolver_fn.clone();
210                let ctx = ctx.clone();
211
212                streams.push(
213                    asynk_strim::try_stream_fn(move |mut yielder| async move {
214                        let ctx_field = ctx.with_field(field);
215                        let field_name = ctx_field.item.node.response_key().node.clone();
216                        let arguments = ObjectAccessor(Cow::Owned(
217                            field
218                                .node
219                                .arguments
220                                .iter()
221                                .map(|(name, value)| {
222                                    ctx_field
223                                        .resolve_input_value(value.clone())
224                                        .map(|value| (name.node.clone(), value))
225                                })
226                                .collect::<ServerResult<IndexMap<Name, Value>>>()?,
227                        ));
228
229                        let mut stream = resolver_fn(ResolverContext {
230                            ctx: &ctx_field,
231                            args: arguments,
232                            parent_value: root_value,
233                        })
234                        .0
235                        .await
236                        .map_err(|err| {
237                            ctx_field.set_error_path(err.into_server_error(ctx_field.item.pos))
238                        })?;
239
240                        while let Some(value) = stream.next().await.transpose().map_err(|err| {
241                            ctx_field.set_error_path(err.into_server_error(ctx_field.item.pos))
242                        })? {
243                            let f = |execute_data: Option<Data>| {
244                                let schema = schema.clone();
245                                let field_name = field_name.clone();
246                                let field_type = field_type.clone();
247                                let ctx_field = ctx_field.clone();
248
249                                async move {
250                                    let mut ctx_field = ctx_field.clone();
251                                    ctx_field.execute_data = execute_data.as_ref();
252                                    let ri = ResolveInfo {
253                                        path_node: &QueryPathNode {
254                                            parent: None,
255                                            segment: QueryPathSegment::Name(&field_name),
256                                        },
257                                        parent_type: schema
258                                            .0
259                                            .env
260                                            .registry
261                                            .subscription_type
262                                            .as_ref()
263                                            .unwrap(),
264                                        return_type: &field_type.to_string(),
265                                        name: field.node.name.node.as_str(),
266                                        alias: field
267                                            .node
268                                            .alias
269                                            .as_ref()
270                                            .map(|alias| alias.node.as_str()),
271                                        is_for_introspection: false,
272                                        field: &field.node,
273                                    };
274                                    let resolve_fut =
275                                        resolve(&schema, &ctx_field, &field_type, Some(&value));
276                                    futures_util::pin_mut!(resolve_fut);
277                                    let value = ctx_field
278                                        .query_env
279                                        .extensions
280                                        .resolve(ri, &mut resolve_fut)
281                                        .await;
282
283                                    match value {
284                                        Ok(value) => {
285                                            let mut map = IndexMap::new();
286                                            map.insert(
287                                                field_name.clone(),
288                                                value.unwrap_or_default(),
289                                            );
290                                            Response::new(Value::Object(map))
291                                        }
292                                        Err(err) => Response::from_errors(vec![err]),
293                                    }
294                                }
295                            };
296                            let resp = ctx_field
297                                .query_env
298                                .extensions
299                                .execute(ctx_field.query_env.operation_name.as_deref(), f)
300                                .await;
301                            let is_err = !resp.errors.is_empty();
302                            yielder.yield_ok(resp).await;
303                            if is_err {
304                                break;
305                            }
306                        }
307
308                        Ok(())
309                    })
310                    .map(|res| res.unwrap_or_else(|err| Response::from_errors(vec![err])))
311                    .boxed(),
312                );
313            }
314        }
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use std::time::Duration;
321
322    use futures_util::StreamExt;
323
324    use crate::{Value, dynamic::*, value};
325
326    #[tokio::test]
327    async fn subscription() {
328        struct MyObjData {
329            value: i32,
330        }
331
332        let my_obj = Object::new("MyObject").field(Field::new(
333            "value",
334            TypeRef::named_nn(TypeRef::INT),
335            |ctx| {
336                FieldFuture::new(async {
337                    Ok(Some(Value::from(
338                        ctx.parent_value.try_downcast_ref::<MyObjData>()?.value,
339                    )))
340                })
341            },
342        ));
343
344        let query = Object::new("Query").field(Field::new(
345            "value",
346            TypeRef::named_nn(TypeRef::INT),
347            |_| FieldFuture::new(async { Ok(FieldValue::none()) }),
348        ));
349
350        let subscription = Subscription::new("Subscription").field(SubscriptionField::new(
351            "obj",
352            TypeRef::named_nn(my_obj.type_name()),
353            |_| {
354                SubscriptionFieldFuture::new(async {
355                    Ok(asynk_strim::try_stream_fn(|mut yielder| async move {
356                        for i in 0..10 {
357                            tokio::time::sleep(Duration::from_millis(100)).await;
358                            yielder
359                                .yield_ok(FieldValue::owned_any(MyObjData { value: i }))
360                                .await;
361                        }
362
363                        Ok(())
364                    }))
365                })
366            },
367        ));
368
369        let schema = Schema::build(query.type_name(), None, Some(subscription.type_name()))
370            .register(my_obj)
371            .register(query)
372            .register(subscription)
373            .finish()
374            .unwrap();
375
376        let mut stream = schema.execute_stream("subscription { obj { value } }");
377        for i in 0..10 {
378            assert_eq!(
379                stream.next().await.unwrap().into_result().unwrap().data,
380                value!({
381                    "obj": { "value": i }
382                })
383            );
384        }
385    }
386
387    #[tokio::test]
388    async fn borrow_context() {
389        struct State {
390            value: i32,
391        }
392
393        let query =
394            Object::new("Query").field(Field::new("value", TypeRef::named(TypeRef::INT), |_| {
395                FieldFuture::new(async { Ok(FieldValue::NONE) })
396            }));
397
398        let subscription = Subscription::new("Subscription").field(SubscriptionField::new(
399            "values",
400            TypeRef::named_nn(TypeRef::INT),
401            |ctx| {
402                SubscriptionFieldFuture::new(async move {
403                    Ok(asynk_strim::try_stream_fn(|mut yielder| async move {
404                        for i in 0..10 {
405                            tokio::time::sleep(Duration::from_millis(100)).await;
406                            yielder
407                                .yield_ok(FieldValue::value(
408                                    ctx.data_unchecked::<State>().value + i,
409                                ))
410                                .await;
411                        }
412
413                        Ok(())
414                    }))
415                })
416            },
417        ));
418
419        let schema = Schema::build("Query", None, Some(subscription.type_name()))
420            .register(query)
421            .register(subscription)
422            .data(State { value: 123 })
423            .finish()
424            .unwrap();
425
426        let mut stream = schema.execute_stream("subscription { values }");
427        for i in 0..10 {
428            assert_eq!(
429                stream.next().await.unwrap().into_result().unwrap().data,
430                value!({ "values": i + 123 })
431            );
432        }
433    }
434}
async_graphql/dynamic/subscription.rs

async_graphql/dynamic/
subscription.rs
