1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

use crate::context::DeferList;
use crate::registry::Registry;
use crate::{ContextSelectionSet, OutputValueType, Positioned, QueryResponse, Result, Type};
use async_graphql_parser::query::Field;
use itertools::Itertools;
use std::borrow::Cow;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;

/// Streamed type
///
/// Similar to Deferred, but you can defer every item of the list type, only takes effect when the @stream directive exists on the field.
pub struct Streamed<T: Type + Send + Sync + Clone + 'static>(Vec<T>);

impl<T: Type + Send + Sync + Clone + 'static> From<Vec<T>> for Streamed<T> {
    fn from(value: Vec<T>) -> Self {
        Self(value)
    }
}

impl<T: Type + Send + Sync + Clone + 'static> Type for Streamed<T> {
    fn type_name() -> Cow<'static, str> {
        Vec::<T>::type_name()
    }

    fn create_type_info(registry: &mut Registry) -> String {
        Vec::<T>::create_type_info(registry)
    }
}

#[async_trait::async_trait]
impl<T: OutputValueType + Send + Sync + Clone + 'static> OutputValueType for Streamed<T> {
    async fn resolve(
        &self,
        ctx: &ContextSelectionSet<'_>,
        field: &Positioned<Field>,
    ) -> Result<serde_json::Value> {
        if let Some(defer_list) = ctx.defer_list {
            if ctx.is_stream(&field.directives) {
                let list = self.0.clone();
                let mut field = field.clone();

                // remove @stream directive
                if let Some((idx, _)) = field
                    .node
                    .directives
                    .iter()
                    .find_position(|d| d.name.as_str() == "stream")
                {
                    field.node.directives.remove(idx);
                }

                let field = Arc::new(field);

                let path_prefix = ctx
                    .path_node
                    .as_ref()
                    .map(|path| path.to_json())
                    .unwrap_or_default();

                for (idx, item) in list.into_iter().enumerate() {
                    let path_prefix = {
                        let mut path_prefix = path_prefix.clone();
                        path_prefix.push(serde_json::Value::Number(idx.into()));
                        path_prefix
                    };
                    let field = field.clone();
                    let schema_env = ctx.schema_env.clone();
                    let query_env = ctx.query_env.clone();

                    defer_list.append(async move {
                        let inc_resolve_id = AtomicUsize::default();
                        let defer_list = DeferList {
                            path_prefix: path_prefix.clone(),
                            futures: Default::default(),
                        };
                        let ctx = query_env.create_context(
                            &schema_env,
                            None,
                            &field.selection_set,
                            &inc_resolve_id,
                            Some(&defer_list),
                        );
                        let data = item.resolve(&ctx, &field).await?;

                        Ok((
                            QueryResponse {
                                path: Some(path_prefix),
                                data,
                                extensions: None,
                                cache_control: Default::default(),
                            },
                            defer_list,
                        ))
                    });
                }
                return Ok(serde_json::Value::Array(Vec::new()));
            }
        }
        OutputValueType::resolve(&self.0, ctx, field).await
    }
}