1use 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}