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                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}