Skip to main content

reifydb_sub_server/
actor.rs

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