Add DataContext trait

src/context.rs

@@ -22,6 +22,28 @@ use crate::{
     ServerResult, UploadValue, Value,
 };
 
+/// Data related functions of the context.
+pub trait DataContext<'a> {
+    /// Gets the global data defined in the `Context` or `Schema`.
+    ///
+    /// If both `Schema` and `Query` have the same data type, the data in the `Query` is obtained.
+    ///
+    /// # Errors
+    ///
+    /// Returns a `Error` if the specified type data does not exist.
+    fn data<D: Any + Send + Sync>(&self) -> Result<&'a D>;
+
+    /// Gets the global data defined in the `Context` or `Schema`.
+    ///
+    /// # Panics
+    ///
+    /// It will panic if the specified data type does not exist.
+    fn data_unchecked<D: Any + Send + Sync>(&self) -> &'a D;
+
+    /// Gets the global data defined in the `Context` or `Schema` or `None` if the specified type data does not exist.
+    fn data_opt<D: Any + Send + Sync>(&self) -> Option<&'a D>;
+}
+
 /// Schema/Context data.
 ///
 /// This is a type map, allowing you to store anything inside it.
@@ -259,6 +281,20 @@ impl QueryEnv {
     }
 }
 
+impl<'a, T> DataContext<'a> for ContextBase<'a, T> {
+    fn data<D: Any + Send + Sync>(&self) -> Result<&'a D> {
+        ContextBase::data::<D>(self)
+    }
+
+    fn data_unchecked<D: Any + Send + Sync>(&self) -> &'a D {
+        ContextBase::data_unchecked::<D>(self)
+    }
+
+    fn data_opt<D: Any + Send + Sync>(&self) -> Option<&'a D> {
+        ContextBase::data_opt::<D>(self)
+    }
+}
+
 impl<'a, T> ContextBase<'a, T> {
     #[doc(hidden)]
     pub fn with_field(

src/extensions/mod.rs

@@ -30,8 +30,8 @@ use futures_util::stream::BoxStream;
 
 use crate::parser::types::ExecutableDocument;
 use crate::{
-    Data, Error, QueryPathNode, Request, Response, Result, SchemaEnv, ServerError, ServerResult,
-    ValidationResult, Value, Variables,
+    Data, DataContext, Error, QueryPathNode, Request, Response, Result, SchemaEnv, ServerError,
+    ServerResult, ValidationResult, Value, Variables,
 };
 
 /// Context for extension
@@ -46,6 +46,20 @@ pub struct ExtensionContext<'a> {
     pub query_data: Option<&'a Data>,
 }
 
+impl<'a> DataContext<'a> for ExtensionContext<'a> {
+    fn data<D: Any + Send + Sync>(&self) -> Result<&'a D> {
+        ExtensionContext::data::<D>(self)
+    }
+
+    fn data_unchecked<D: Any + Send + Sync>(&self) -> &'a D {
+        ExtensionContext::data_unchecked::<D>(self)
+    }
+
+    fn data_opt<D: Any + Send + Sync>(&self) -> Option<&'a D> {
+        ExtensionContext::data_opt::<D>(self)
+    }
+}
+
 impl<'a> ExtensionContext<'a> {
     /// Convert the specified [ExecutableDocument] into a query string.
     ///