use async_trait::async_trait;
use serde_json::Value;
use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use tracing::warn;
#[derive(Debug, thiserror::Error)]
pub enum ContextError {
#[error("Required context data not found for key: {0}")]
DataNotFound(String),
#[error("Required data loader not found: {0}")]
LoaderNotFound(String),
}
#[derive(Debug, thiserror::Error)]
pub enum LoaderError {
#[error("Loader error: {0}")]
Load(String),
#[error("Not found: {0}")]
NotFound(String),
#[error("Invalid data: {0}")]
InvalidData(String),
}
#[async_trait]
pub trait DataLoader: Send + Sync + 'static {
type Key: Send;
type Value: Send;
async fn load(&self, key: Self::Key) -> Result<Self::Value, LoaderError>;
async fn load_many(&self, keys: Vec<Self::Key>) -> Result<Vec<Self::Value>, LoaderError>;
}
pub struct GraphQLContext {
custom_data: Arc<RwLock<HashMap<String, Value>>>,
data_loaders: Arc<RwLock<HashMap<TypeId, Box<dyn Any + Send + Sync>>>>,
}
impl GraphQLContext {
pub fn new() -> Self {
Self {
custom_data: Arc::new(RwLock::new(HashMap::new())),
data_loaders: Arc::new(RwLock::new(HashMap::new())),
}
}
fn custom_data_read(&self) -> RwLockReadGuard<'_, HashMap<String, Value>> {
self.custom_data.read().unwrap_or_else(|e| {
warn!("custom_data RwLock was poisoned, recovering read lock");
e.into_inner()
})
}
fn custom_data_write(&self) -> RwLockWriteGuard<'_, HashMap<String, Value>> {
self.custom_data.write().unwrap_or_else(|e| {
warn!("custom_data RwLock was poisoned, recovering write lock");
e.into_inner()
})
}
fn loaders_read(&self) -> RwLockReadGuard<'_, HashMap<TypeId, Box<dyn Any + Send + Sync>>> {
self.data_loaders.read().unwrap_or_else(|e| {
warn!("data_loaders RwLock was poisoned, recovering read lock");
e.into_inner()
})
}
fn loaders_write(&self) -> RwLockWriteGuard<'_, HashMap<TypeId, Box<dyn Any + Send + Sync>>> {
self.data_loaders.write().unwrap_or_else(|e| {
warn!("data_loaders RwLock was poisoned, recovering write lock");
e.into_inner()
})
}
pub fn set_data(&self, key: String, value: Value) {
let mut data = self.custom_data_write();
data.insert(key, value);
}
pub fn get_data(&self, key: &str) -> Option<Value> {
let data = self.custom_data_read();
data.get(key).cloned()
}
pub fn require_data(&self, key: &str) -> async_graphql::Result<Value> {
self.get_data(key)
.ok_or_else(|| ContextError::DataNotFound(key.to_string()).into())
}
pub fn remove_data(&self, key: &str) -> Option<Value> {
let mut data = self.custom_data_write();
data.remove(key)
}
pub fn clear_data(&self) {
let mut data = self.custom_data_write();
data.clear();
}
pub fn add_data_loader<T: DataLoader>(&self, loader: Arc<T>) {
let mut loaders = self.loaders_write();
loaders.insert(TypeId::of::<T>(), Box::new(loader));
}
pub fn get_data_loader<T: DataLoader>(&self) -> Option<Arc<T>> {
let loaders = self.loaders_read();
loaders
.get(&TypeId::of::<T>())
.and_then(|loader| loader.downcast_ref::<Arc<T>>().cloned())
}
pub fn require_data_loader<T: DataLoader>(&self) -> async_graphql::Result<Arc<T>> {
self.get_data_loader::<T>().ok_or_else(|| {
ContextError::LoaderNotFound(std::any::type_name::<T>().to_string()).into()
})
}
pub fn remove_data_loader<T: DataLoader>(&self) {
let mut loaders = self.loaders_write();
loaders.remove(&TypeId::of::<T>());
}
pub fn clear_loaders(&self) {
let mut loaders = self.loaders_write();
loaders.clear();
}
}
impl Default for GraphQLContext {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[derive(Debug)]
struct TestLoader;
#[async_trait]
impl DataLoader for TestLoader {
type Key = String;
type Value = i32;
async fn load(&self, key: Self::Key) -> Result<Self::Value, LoaderError> {
key.parse::<i32>()
.map_err(|e| LoaderError::InvalidData(e.to_string()))
}
async fn load_many(&self, keys: Vec<Self::Key>) -> Result<Vec<Self::Value>, LoaderError> {
keys.into_iter()
.map(|k| {
k.parse::<i32>()
.map_err(|e| LoaderError::InvalidData(e.to_string()))
})
.collect()
}
}
#[rstest]
fn test_context_new() {
let context = GraphQLContext::new();
assert!(context.get_data("any_key").is_none());
}
#[rstest]
fn test_set_and_get_data() {
let context = GraphQLContext::new();
let value = serde_json::json!({"name": "test", "value": 42});
context.set_data("test_key".to_string(), value.clone());
let retrieved = context.get_data("test_key");
assert_eq!(retrieved, Some(value));
}
#[rstest]
fn test_get_nonexistent_data() {
let context = GraphQLContext::new();
let result = context.get_data("nonexistent");
assert_eq!(result, None);
}
#[rstest]
fn test_remove_data() {
let context = GraphQLContext::new();
let value = serde_json::json!("test_value");
context.set_data("key".to_string(), value.clone());
let removed = context.remove_data("key");
assert_eq!(removed, Some(value));
assert_eq!(context.get_data("key"), None);
}
#[rstest]
fn test_clear_data() {
let context = GraphQLContext::new();
context.set_data("key1".to_string(), serde_json::json!(1));
context.set_data("key2".to_string(), serde_json::json!(2));
context.set_data("key3".to_string(), serde_json::json!(3));
context.clear_data();
assert_eq!(context.get_data("key1"), None);
assert_eq!(context.get_data("key2"), None);
assert_eq!(context.get_data("key3"), None);
}
#[rstest]
fn test_add_and_get_data_loader() {
let context = GraphQLContext::new();
let loader = Arc::new(TestLoader);
context.add_data_loader(loader);
let retrieved = context.get_data_loader::<TestLoader>();
assert!(retrieved.is_some());
}
#[rstest]
fn test_get_nonexistent_loader() {
let context = GraphQLContext::new();
let result = context.get_data_loader::<TestLoader>();
assert!(result.is_none());
}
#[rstest]
fn test_remove_data_loader() {
let context = GraphQLContext::new();
let loader = Arc::new(TestLoader);
context.add_data_loader(loader);
context.remove_data_loader::<TestLoader>();
let result = context.get_data_loader::<TestLoader>();
assert!(result.is_none());
}
#[rstest]
fn test_clear_loaders() {
struct Loader1;
struct Loader2;
#[async_trait]
impl DataLoader for Loader1 {
type Key = i32;
type Value = String;
async fn load(&self, key: Self::Key) -> Result<Self::Value, LoaderError> {
Ok(key.to_string())
}
async fn load_many(
&self,
keys: Vec<Self::Key>,
) -> Result<Vec<Self::Value>, LoaderError> {
Ok(keys.iter().map(|k| k.to_string()).collect())
}
}
#[async_trait]
impl DataLoader for Loader2 {
type Key = String;
type Value = i32;
async fn load(&self, _key: Self::Key) -> Result<Self::Value, LoaderError> {
Ok(0)
}
async fn load_many(
&self,
keys: Vec<Self::Key>,
) -> Result<Vec<Self::Value>, LoaderError> {
Ok(vec![0; keys.len()])
}
}
let context = GraphQLContext::new();
context.add_data_loader(Arc::new(Loader1));
context.add_data_loader(Arc::new(Loader2));
context.clear_loaders();
assert!(context.get_data_loader::<Loader1>().is_none());
assert!(context.get_data_loader::<Loader2>().is_none());
}
#[rstest]
fn test_multiple_data_values() {
let context = GraphQLContext::new();
context.set_data("int".to_string(), serde_json::json!(123));
context.set_data("string".to_string(), serde_json::json!("hello"));
context.set_data("array".to_string(), serde_json::json!([1, 2, 3]));
context.set_data("object".to_string(), serde_json::json!({"key": "value"}));
assert_eq!(context.get_data("int"), Some(serde_json::json!(123)));
assert_eq!(context.get_data("string"), Some(serde_json::json!("hello")));
assert_eq!(
context.get_data("array"),
Some(serde_json::json!([1, 2, 3]))
);
assert_eq!(
context.get_data("object"),
Some(serde_json::json!({"key": "value"}))
);
}
#[rstest]
#[tokio::test]
async fn test_data_loader_load() {
let loader = TestLoader;
let result = loader.load("42".to_string()).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), 42);
}
#[rstest]
#[tokio::test]
async fn test_data_loader_load_many() {
let loader = TestLoader;
let keys = vec!["1".to_string(), "2".to_string(), "3".to_string()];
let result = loader.load_many(keys).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), vec![1, 2, 3]);
}
#[rstest]
#[tokio::test]
async fn test_data_loader_error() {
let loader = TestLoader;
let result = loader.load("invalid".to_string()).await;
assert!(result.is_err());
match result {
Err(LoaderError::InvalidData(_)) => {}
_ => panic!("Expected InvalidData error"),
}
}
#[rstest]
fn test_context_default() {
let context = GraphQLContext::default();
assert!(context.get_data("any_key").is_none());
}
#[rstest]
fn test_overwrite_data() {
let context = GraphQLContext::new();
context.set_data("key".to_string(), serde_json::json!(1));
context.set_data("key".to_string(), serde_json::json!(2));
assert_eq!(context.get_data("key"), Some(serde_json::json!(2)));
}
#[rstest]
fn test_require_data_returns_value_when_present() {
let context = GraphQLContext::new();
context.set_data("user_id".to_string(), serde_json::json!("user-42"));
let result = context.require_data("user_id");
assert!(result.is_ok());
assert_eq!(result.unwrap(), serde_json::json!("user-42"));
}
#[rstest]
fn test_require_data_returns_error_when_missing() {
let context = GraphQLContext::new();
let result = context.require_data("nonexistent_key");
assert!(result.is_err());
let err = result.unwrap_err();
assert!(
err.message.contains("nonexistent_key"),
"Error should mention the missing key, got: {}",
err.message
);
assert!(
err.message.contains("Required context data not found"),
"Error should describe the issue, got: {}",
err.message
);
}
#[rstest]
fn test_require_data_does_not_panic_on_missing_key() {
let context = GraphQLContext::new();
let result = context.require_data("missing");
assert!(result.is_err());
}
#[rstest]
fn test_require_data_loader_returns_loader_when_present() {
let context = GraphQLContext::new();
context.add_data_loader(Arc::new(TestLoader));
let result = context.require_data_loader::<TestLoader>();
assert!(result.is_ok());
}
#[rstest]
fn test_require_data_loader_returns_error_when_missing() {
let context = GraphQLContext::new();
let result = context.require_data_loader::<TestLoader>();
assert!(result.is_err());
let err = result.unwrap_err();
assert!(
err.message.contains("Required data loader not found"),
"Error should describe the issue, got: {}",
err.message
);
}
#[rstest]
fn test_require_data_loader_does_not_panic_on_missing_loader() {
let context = GraphQLContext::new();
let result = context.require_data_loader::<TestLoader>();
assert!(result.is_err());
}
#[rstest]
fn test_require_data_loader_after_removal_returns_error() {
let context = GraphQLContext::new();
context.add_data_loader(Arc::new(TestLoader));
context.remove_data_loader::<TestLoader>();
let result = context.require_data_loader::<TestLoader>();
assert!(result.is_err());
}
#[rstest]
fn test_context_error_data_not_found_display() {
let err = ContextError::DataNotFound("my_key".to_string());
let message = err.to_string();
assert_eq!(message, "Required context data not found for key: my_key");
}
#[rstest]
fn test_context_error_loader_not_found_display() {
let err = ContextError::LoaderNotFound("MyLoader".to_string());
let message = err.to_string();
assert_eq!(message, "Required data loader not found: MyLoader");
}
#[rstest]
fn test_context_error_converts_to_graphql_error() {
let err = ContextError::DataNotFound("test_key".to_string());
let gql_err: async_graphql::Error = err.into();
assert!(gql_err.message.contains("test_key"));
assert!(gql_err.message.contains("Required context data not found"));
}
}