async_graphql/dynamic/
subscription.rs

1use std::{borrow::Cow, fmt, fmt::Debug, sync::Arc};
2
3use futures_util::{
4    future::BoxFuture, stream::BoxStream, Future, FutureExt, Stream, StreamExt, TryStreamExt,
5};
6use indexmap::IndexMap;
7
8use crate::{
9    dynamic::{
10        resolve::resolve, FieldValue, InputValue, ObjectAccessor, ResolverContext, Schema,
11        SchemaError, TypeRef,
12    },
13    extensions::ResolveInfo,
14    parser::types::Selection,
15    registry::{Deprecation, MetaField, MetaType, Registry},
16    subscription::BoxFieldStream,
17    ContextSelectionSet, Data, Name, QueryPathNode, QueryPathSegment, Response, Result,
18    ServerResult, Value,
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                },
167            );
168        }
169
170        registry.types.insert(
171            self.name.clone(),
172            MetaType::Object {
173                name: self.name.clone(),
174                description: self.description.clone(),
175                fields,
176                cache_control: Default::default(),
177                extends: false,
178                shareable: false,
179                resolvable: true,
180                keys: None,
181                visible: None,
182                inaccessible: false,
183                interface_object: false,
184                tags: vec![],
185                is_subscription: true,
186                rust_typename: None,
187                directive_invocations: vec![],
188            },
189        );
190
191        Ok(())
192    }
193
194    pub(crate) fn collect_streams<'a>(
195        &self,
196        schema: &Schema,
197        ctx: &ContextSelectionSet<'a>,
198        streams: &mut Vec<BoxFieldStream<'a>>,
199        root_value: &'a FieldValue<'static>,
200    ) {
201        for selection in &ctx.item.node.items {
202            if let Selection::Field(field) = &selection.node {
203                if let Some(field_def) = self.fields.get(field.node.name.node.as_str()) {
204                    let schema = schema.clone();
205                    let field_type = field_def.ty.clone();
206                    let resolver_fn = field_def.resolver_fn.clone();
207                    let ctx = ctx.clone();
208
209                    streams.push(
210                        async_stream::try_stream! {
211                            let ctx_field = ctx.with_field(field);
212                            let field_name = ctx_field.item.node.response_key().node.clone();
213                            let arguments = ObjectAccessor(Cow::Owned(
214                                field
215                                    .node
216                                    .arguments
217                                    .iter()
218                                    .map(|(name, value)| {
219                                        ctx_field
220                                            .resolve_input_value(value.clone())
221                                            .map(|value| (name.node.clone(), value))
222                                    })
223                                    .collect::<ServerResult<IndexMap<Name, Value>>>()?,
224                            ));
225
226                            let mut stream = resolver_fn(ResolverContext {
227                                ctx: &ctx_field,
228                                args: arguments,
229                                parent_value: root_value,
230                            })
231                            .0
232                            .await
233                            .map_err(|err| ctx_field.set_error_path(err.into_server_error(ctx_field.item.pos)))?;
234
235                            while let Some(value) = stream.next().await.transpose().map_err(|err| ctx_field.set_error_path(err.into_server_error(ctx_field.item.pos)))? {
236                                let f = |execute_data: Option<Data>| {
237                                    let schema = schema.clone();
238                                    let field_name = field_name.clone();
239                                    let field_type = field_type.clone();
240                                    let ctx_field = ctx_field.clone();
241
242                                    async move {
243                                        let mut ctx_field = ctx_field.clone();
244                                        ctx_field.execute_data = execute_data.as_ref();
245                                        let ri = ResolveInfo {
246                                            path_node: &QueryPathNode {
247                                                parent: None,
248                                                segment: QueryPathSegment::Name(&field_name),
249                                            },
250                                            parent_type: schema.0.env.registry.subscription_type.as_ref().unwrap(),
251                                            return_type: &field_type.to_string(),
252                                            name: field.node.name.node.as_str(),
253                                            alias: field.node.alias.as_ref().map(|alias| alias.node.as_str()),
254                                            is_for_introspection: false,
255                                            field: &field.node,
256                                        };
257                                        let resolve_fut = resolve(&schema, &ctx_field, &field_type, Some(&value));
258                                        futures_util::pin_mut!(resolve_fut);
259                                        let value = ctx_field.query_env.extensions.resolve(ri, &mut resolve_fut).await;
260
261                                        match value {
262                                            Ok(value) => {
263                                                let mut map = IndexMap::new();
264                                                map.insert(field_name.clone(), value.unwrap_or_default());
265                                                Response::new(Value::Object(map))
266                                            },
267                                            Err(err) => Response::from_errors(vec![err]),
268                                        }
269                                    }
270                                };
271                                let resp = ctx_field.query_env.extensions.execute(ctx_field.query_env.operation_name.as_deref(), f).await;
272                                let is_err = !resp.errors.is_empty();
273                                yield resp;
274                                if is_err {
275                                    break;
276                                }
277                            }
278                        }.map(|res| {
279                            res.unwrap_or_else(|err| Response::from_errors(vec![err]))
280                        })
281                        .boxed(),
282                    );
283                }
284            }
285        }
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use std::time::Duration;
292
293    use futures_util::StreamExt;
294
295    use crate::{dynamic::*, value, Value};
296
297    #[tokio::test]
298    async fn subscription() {
299        struct MyObjData {
300            value: i32,
301        }
302
303        let my_obj = Object::new("MyObject").field(Field::new(
304            "value",
305            TypeRef::named_nn(TypeRef::INT),
306            |ctx| {
307                FieldFuture::new(async {
308                    Ok(Some(Value::from(
309                        ctx.parent_value.try_downcast_ref::<MyObjData>()?.value,
310                    )))
311                })
312            },
313        ));
314
315        let query = Object::new("Query").field(Field::new(
316            "value",
317            TypeRef::named_nn(TypeRef::INT),
318            |_| FieldFuture::new(async { Ok(FieldValue::none()) }),
319        ));
320
321        let subscription = Subscription::new("Subscription").field(SubscriptionField::new(
322            "obj",
323            TypeRef::named_nn(my_obj.type_name()),
324            |_| {
325                SubscriptionFieldFuture::new(async {
326                    Ok(async_stream::try_stream! {
327                        for i in 0..10 {
328                            tokio::time::sleep(Duration::from_millis(100)).await;
329                            yield FieldValue::owned_any(MyObjData { value: i });
330                        }
331                    })
332                })
333            },
334        ));
335
336        let schema = Schema::build(query.type_name(), None, Some(subscription.type_name()))
337            .register(my_obj)
338            .register(query)
339            .register(subscription)
340            .finish()
341            .unwrap();
342
343        let mut stream = schema.execute_stream("subscription { obj { value } }");
344        for i in 0..10 {
345            assert_eq!(
346                stream.next().await.unwrap().into_result().unwrap().data,
347                value!({
348                    "obj": { "value": i }
349                })
350            );
351        }
352    }
353
354    #[tokio::test]
355    async fn borrow_context() {
356        struct State {
357            value: i32,
358        }
359
360        let query =
361            Object::new("Query").field(Field::new("value", TypeRef::named(TypeRef::INT), |_| {
362                FieldFuture::new(async { Ok(FieldValue::NONE) })
363            }));
364
365        let subscription = Subscription::new("Subscription").field(SubscriptionField::new(
366            "values",
367            TypeRef::named_nn(TypeRef::INT),
368            |ctx| {
369                SubscriptionFieldFuture::new(async move {
370                    Ok(async_stream::try_stream! {
371                        for i in 0..10 {
372                            tokio::time::sleep(Duration::from_millis(100)).await;
373                            yield FieldValue::value(ctx.data_unchecked::<State>().value + i);
374                        }
375                    })
376                })
377            },
378        ));
379
380        let schema = Schema::build("Query", None, Some(subscription.type_name()))
381            .register(query)
382            .register(subscription)
383            .data(State { value: 123 })
384            .finish()
385            .unwrap();
386
387        let mut stream = schema.execute_stream("subscription { values }");
388        for i in 0..10 {
389            assert_eq!(
390                stream.next().await.unwrap().into_result().unwrap().data,
391                value!({ "values": i + 123 })
392            );
393        }
394    }
395}
async_graphql/dynamic/subscription.rs

async_graphql/dynamic/
subscription.rs
