use crate::QueryError;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use super::context::QueryContext;
pub type MiddlewareResult<T> = Result<T, QueryError>;
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub struct Next<'a> {
pub(crate) inner:
Box<dyn FnOnce(QueryContext) -> BoxFuture<'a, MiddlewareResult<QueryResponse>> + Send + 'a>,
}
impl<'a> Next<'a> {
pub fn run(self, ctx: QueryContext) -> BoxFuture<'a, MiddlewareResult<QueryResponse>> {
(self.inner)(ctx)
}
}
#[derive(Debug, Clone)]
pub struct QueryResponse {
pub data: serde_json::Value,
pub rows_affected: Option<u64>,
pub execution_time_us: u64,
pub from_cache: bool,
pub metadata: serde_json::Map<String, serde_json::Value>,
}
impl QueryResponse {
pub fn new(data: serde_json::Value) -> Self {
Self {
data,
rows_affected: None,
execution_time_us: 0,
from_cache: false,
metadata: serde_json::Map::new(),
}
}
pub fn empty() -> Self {
Self::new(serde_json::Value::Null)
}
pub fn with_affected(count: u64) -> Self {
Self {
data: serde_json::Value::Null,
rows_affected: Some(count),
execution_time_us: 0,
from_cache: false,
metadata: serde_json::Map::new(),
}
}
pub fn with_execution_time(mut self, us: u64) -> Self {
self.execution_time_us = us;
self
}
pub fn from_cache(mut self) -> Self {
self.from_cache = true;
self
}
pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
self.metadata.insert(key.into(), value);
self
}
}
pub trait Middleware: Send + Sync {
fn handle<'a>(
&'a self,
ctx: QueryContext,
next: Next<'a>,
) -> BoxFuture<'a, MiddlewareResult<QueryResponse>>;
fn name(&self) -> &'static str {
std::any::type_name::<Self>()
}
fn enabled(&self) -> bool {
true
}
}
pub type SharedMiddleware = Arc<dyn Middleware>;
pub trait IntoSharedMiddleware {
fn into_shared(self) -> SharedMiddleware;
}
impl<T: Middleware + 'static> IntoSharedMiddleware for T {
fn into_shared(self) -> SharedMiddleware {
Arc::new(self)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_query_response_builder() {
let response = QueryResponse::new(serde_json::json!({"id": 1}))
.with_execution_time(1000)
.with_metadata("cache_hit", serde_json::Value::Bool(false));
assert_eq!(response.execution_time_us, 1000);
assert!(!response.from_cache);
assert!(response.metadata.contains_key("cache_hit"));
}
#[test]
fn test_query_response_affected() {
let response = QueryResponse::with_affected(5);
assert_eq!(response.rows_affected, Some(5));
}
#[test]
fn test_query_response_from_cache() {
let response = QueryResponse::empty().from_cache();
assert!(response.from_cache);
}
}