reifydb_sub_server/
interceptor.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::{collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc};
5
6use futures_util::FutureExt;
7use reifydb_core::{actors::server::Operation, metric::ExecutionMetrics};
8use reifydb_type::{params::Params, value::identity::IdentityId};
9use tracing::error;
10
11use crate::execute::ExecuteError;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
14pub enum Protocol {
15	#[default]
16	Http,
17	WebSocket,
18	Grpc,
19}
20
21#[derive(Debug, Clone, Default)]
22pub struct RequestMetadata {
23	headers: HashMap<String, String>,
24	protocol: Protocol,
25}
26
27impl RequestMetadata {
28	pub fn new(protocol: Protocol) -> Self {
29		Self {
30			headers: HashMap::new(),
31			protocol,
32		}
33	}
34
35	pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) {
36		self.headers.insert(key.into().to_ascii_lowercase(), value.into());
37	}
38
39	pub fn get(&self, key: &str) -> Option<&str> {
40		self.headers.get(&key.to_ascii_lowercase()).map(|s| s.as_str())
41	}
42
43	pub fn protocol(&self) -> Protocol {
44		self.protocol
45	}
46
47	pub fn headers(&self) -> &HashMap<String, String> {
48		&self.headers
49	}
50}
51
52pub struct RequestContext {
53	pub identity: IdentityId,
54
55	pub operation: Operation,
56
57	pub rql: String,
58
59	pub params: Params,
60
61	pub metadata: RequestMetadata,
62}
63
64pub struct ResponseContext {
65	pub identity: IdentityId,
66
67	pub operation: Operation,
68
69	pub rql: String,
70
71	pub metrics: ExecutionMetrics,
72
73	pub params: Params,
74
75	pub metadata: RequestMetadata,
76
77	pub result: Result<usize, String>,
78}
79
80pub trait RequestInterceptor: Send + Sync + 'static {
81	fn pre_execute<'a>(
82		&'a self,
83		ctx: &'a mut RequestContext,
84	) -> Pin<Box<dyn Future<Output = Result<(), ExecuteError>> + Send + 'a>>;
85
86	fn post_execute<'a>(&'a self, ctx: &'a ResponseContext) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
87}
88
89#[derive(Clone)]
90pub struct RequestInterceptorChain {
91	interceptors: Arc<Vec<Arc<dyn RequestInterceptor>>>,
92}
93
94impl RequestInterceptorChain {
95	pub fn new(interceptors: Vec<Arc<dyn RequestInterceptor>>) -> Self {
96		Self {
97			interceptors: Arc::new(interceptors),
98		}
99	}
100
101	pub fn empty() -> Self {
102		Self {
103			interceptors: Arc::new(Vec::new()),
104		}
105	}
106
107	pub fn is_empty(&self) -> bool {
108		self.interceptors.is_empty()
109	}
110
111	pub async fn pre_execute(&self, ctx: &mut RequestContext) -> Result<(), ExecuteError> {
112		for interceptor in self.interceptors.iter() {
113			interceptor.pre_execute(ctx).await?;
114		}
115		Ok(())
116	}
117
118	pub async fn post_execute(&self, ctx: &ResponseContext) {
119		for interceptor in self.interceptors.iter().rev() {
120			if let Err(panic) = AssertUnwindSafe(interceptor.post_execute(ctx)).catch_unwind().await {
121				let msg = panic
122					.downcast_ref::<&str>()
123					.copied()
124					.or_else(|| panic.downcast_ref::<String>().map(|s| s.as_str()))
125					.unwrap_or("unknown panic");
126				error!("post_execute interceptor panicked: {}", msg);
127			}
128		}
129	}
130}
131
132impl Default for RequestInterceptorChain {
133	fn default() -> Self {
134		Self::empty()
135	}
136}
reifydb_sub_server/interceptor.rs

reifydb_sub_server/
interceptor.rs
