Skip to main content

reifydb_sub_server/
actor.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::collections::HashMap;
5
6use reifydb_auth::service::{AuthResponse, AuthService};
7use reifydb_core::{
8	actors::server::{
9		ServerAuthResponse, ServerLogoutResponse, ServerMessage, ServerResponse, ServerSubscribeResponse,
10	},
11	execution::ExecutionResult,
12};
13use reifydb_engine::{engine::StandardEngine, session::RetryStrategy};
14use reifydb_runtime::{
15	actor::{
16		context::Context,
17		reply::Reply,
18		traits::{Actor, Directive},
19	},
20	context::clock::Clock,
21};
22use reifydb_value::{
23	params::Params,
24	value::{duration::Duration, identity::IdentityId},
25};
26
27pub struct ServerActor {
28	engine: StandardEngine,
29	auth_service: AuthService,
30	clock: Clock,
31	retry: RetryStrategy,
32}
33
34impl ServerActor {
35	pub fn new(engine: StandardEngine, auth_service: AuthService, clock: Clock) -> Self {
36		Self {
37			engine,
38			auth_service,
39			clock,
40			retry: RetryStrategy::default(),
41		}
42	}
43
44	fn dispatch_execute(
45		&self,
46		identity: IdentityId,
47		rql: String,
48		params: Params,
49		reply: Reply<ServerResponse>,
50		execute: impl Fn(&StandardEngine, IdentityId, &str, Params) -> ExecutionResult,
51	) {
52		let t = self.clock.instant();
53		let result = self
54			.retry
55			.execute(self.engine.rng(), &rql, || execute(&self.engine, identity, &rql, params.clone()));
56		if let Some(err) = result.error {
57			reply.send(ServerResponse::EngineError {
58				diagnostic: Box::new(err.diagnostic()),
59				rql,
60			});
61		} else {
62			reply.send(ServerResponse::Success {
63				frames: result.frames,
64				duration: Duration::from_std(t.elapsed()),
65				metrics: result.metrics,
66			});
67		}
68	}
69
70	#[inline]
71	fn handle_subscribe(&self, identity: IdentityId, rql: String, reply: Reply<ServerSubscribeResponse>) {
72		let t = self.clock.instant();
73		let result = self.engine.subscribe_as(identity, &rql, Params::None);
74		if let Some(err) = result.error {
75			reply.send(ServerSubscribeResponse::EngineError {
76				diagnostic: Box::new(err.diagnostic()),
77				rql,
78			});
79		} else {
80			reply.send(ServerSubscribeResponse::Subscribed {
81				frames: result.frames,
82				duration: Duration::from_std(t.elapsed()),
83				metrics: result.metrics,
84			});
85		}
86	}
87
88	#[inline]
89	fn handle_authenticate(
90		&self,
91		method: String,
92		credentials: HashMap<String, String>,
93		reply: Reply<ServerAuthResponse>,
94	) {
95		match self.auth_service.authenticate(&method, credentials) {
96			Ok(AuthResponse::Authenticated {
97				identity,
98				token,
99			}) => reply.send(ServerAuthResponse::Authenticated {
100				identity,
101				token,
102			}),
103			Ok(AuthResponse::Challenge {
104				challenge_id,
105				payload,
106			}) => reply.send(ServerAuthResponse::Challenge {
107				challenge_id,
108				payload,
109			}),
110			Ok(AuthResponse::Failed {
111				reason,
112			}) => reply.send(ServerAuthResponse::Failed {
113				reason,
114			}),
115			Err(e) => reply.send(ServerAuthResponse::Error(e.to_string())),
116		}
117	}
118
119	#[inline]
120	fn handle_logout(&self, token: String, reply: Reply<ServerLogoutResponse>) {
121		if self.auth_service.revoke_token(&token) {
122			reply.send(ServerLogoutResponse::Ok);
123		} else {
124			reply.send(ServerLogoutResponse::InvalidToken);
125		}
126	}
127}
128
129impl Actor for ServerActor {
130	type State = ();
131	type Message = ServerMessage;
132
133	fn init(&self, _ctx: &Context<Self::Message>) -> Self::State {}
134
135	fn handle(&self, _state: &mut (), msg: ServerMessage, _ctx: &Context<ServerMessage>) -> Directive {
136		match msg {
137			ServerMessage::Query {
138				identity,
139				rql,
140				params,
141				reply,
142			} => {
143				self.dispatch_execute(identity, rql, params, reply, |e, id, s, p| e.query_as(id, s, p));
144			}
145			ServerMessage::Command {
146				identity,
147				rql,
148				params,
149				reply,
150			} => {
151				self.dispatch_execute(identity, rql, params, reply, |e, id, s, p| {
152					e.command_as(id, s, p)
153				});
154			}
155			ServerMessage::Admin {
156				identity,
157				rql,
158				params,
159				reply,
160			} => {
161				self.dispatch_execute(identity, rql, params, reply, |e, id, s, p| e.admin_as(id, s, p));
162			}
163			ServerMessage::Subscribe {
164				identity,
165				rql,
166				reply,
167			} => {
168				self.handle_subscribe(identity, rql, reply);
169			}
170			ServerMessage::Authenticate {
171				method,
172				credentials,
173				reply,
174			} => {
175				self.handle_authenticate(method, credentials, reply);
176			}
177			ServerMessage::Logout {
178				token,
179				reply,
180			} => {
181				self.handle_logout(token, reply);
182			}
183		}
184		Directive::Continue
185	}
186}