1use super::entities::invitation::Model;
2use super::entities::team;
3use super::{SqlxStore, lock_exclusive};
4use crate::error::record_not_updated;
5use crate::model::{self, ActiveRow, SqlxModel};
6use crate::pool::Exec;
7use crate::schema::{AuthSchema, SqlxSessionModel, SqlxUserModel};
8use crate::sql::Sql;
9use alibi_core::entity::AuthUser;
10use alibi_core::error::{AuthError, AuthResult};
11use alibi_core::store::InvitationStore;
12use alibi_core::{CreateInvitation, Invitation, InvitationStatus};
13use async_trait::async_trait;
14use chrono::Utc;
15
16impl<S: AuthSchema> SqlxStore<S> {
17 async fn find_invitation(
18 &self,
19 id: &str,
20 ) -> AuthResult<Option<crate::organization_models::Row<Model>>> {
21 let mut sql = self.organization_models.invitation.by_id(self.exec(), id)?;
22 model::limit_one(&mut sql);
23 self.organization_models
24 .invitation
25 .fetch_optional(self.exec(), sql)
26 .await
27 }
28
29 async fn count_invitation_rows(&self, sql: Sql) -> AuthResult<i64> {
30 Ok(self
31 .exec()
32 .fetch_scalar::<i64>(sql)
33 .await?
34 .unwrap_or_default())
35 }
36}
37
38#[async_trait]
39impl<S> InvitationStore for SqlxStore<S>
40where
41 S: AuthSchema + Send + Sync,
42 S::User: SqlxUserModel,
43 S::Session: SqlxSessionModel,
44{
45 async fn create_invitation(&self, invitation: CreateInvitation) -> AuthResult<Invitation> {
46 self.create_invitation_with_options(
47 invitation,
48 alibi_core::store::InvitationCreateOptions::default(),
49 )
50 .await
51 }
52 async fn create_invitation_with_options(
53 &self,
54 invitation: CreateInvitation,
55 options: alibi_core::store::InvitationCreateOptions,
56 ) -> AuthResult<Invitation> {
57 let mut active = ActiveRow::new();
58 active.set(
59 "id",
60 options
61 .id
62 .unwrap_or_else(|| self.organization_models.invitation.new_id()),
63 );
64 active.set("organization_id", invitation.organization_id);
65 active.set("email", invitation.email);
66 active.set("role", invitation.role);
67 active.set("team_id", invitation.team_id);
68 active.set("status", options.status.unwrap_or_default().to_string());
69 active.set("inviter_id", invitation.inviter_id);
70 active.set("expires_at", invitation.expires_at);
71 active.set("created_at", options.created_at.unwrap_or_else(Utc::now));
72 self.organization_models
73 .invitation
74 .insert(self.exec(), &active)
75 .await
76 .map(|model| Invitation::from(&model))
77 }
78
79 async fn get_invitation_by_id(&self, id: &str) -> AuthResult<Option<Invitation>> {
80 Ok(self
81 .find_invitation(id)
82 .await?
83 .map(|model| Invitation::from(&model)))
84 }
85
86 async fn update_invitation_team_ids(
87 &self,
88 id: &str,
89 team_ids: Option<String>,
90 ) -> AuthResult<Invitation> {
91 let row = self
92 .find_invitation(id)
93 .await?
94 .ok_or_else(|| AuthError::not_found("Invitation not found"))?;
95 let mut active = row.into_active();
96 active.set("team_id", team_ids);
97 self.organization_models
98 .invitation
99 .update(self.exec(), &active)
100 .await?
101 .map(|row_2| Invitation::from(&row_2))
102 .ok_or_else(record_not_updated)
103 }
104
105 #[expect(
106 clippy::too_many_lines,
107 reason = "Keep invitation quota checks and insertion together inside their single transaction"
108 )]
109 async fn accept_invitation_with_teams(
110 &self,
111 invitation_id: &str,
112 user_id: &str,
113 session_token: &str,
114 team_limits: &[(String, Option<f64>)],
115 membership_limit: Option<usize>,
116 ) -> AuthResult<Option<(Invitation, alibi_core::types::Member)>> {
117 let transaction = self.pool().begin(true).await?;
118 let outcome = async {
119 let tx = &transaction;
120 let exec = Exec::Tx(tx);
121 let mut select = self
122 .organization_models
123 .invitation
124 .by_id(exec, invitation_id)?;
125 model::limit_one(&mut select);
126 lock_exclusive(&mut select);
127 let Some(invitation) = self
128 .organization_models
129 .invitation
130 .fetch_optional(exec, select)
131 .await?
132 else {
133 return Ok(None);
134 };
135 if invitation.status != "pending" || invitation.expires_at < Utc::now() {
136 return Ok(None);
137 }
138 let user =
139 super::users::find_user_by_id::<S::User>(exec, user_id, super::users::Lock::Shared)
140 .await?
141 .ok_or(AuthError::UserNotFound)?;
142 if user
143 .email()
144 .is_none_or(|email| email.to_lowercase() != invitation.email.to_lowercase())
145 {
146 return Err(AuthError::forbidden("This invitation is not for you"));
147 }
148 let session_table = <S::Session as SqlxModel>::TABLE;
149 let mut session = model::select_model::<S::Session>(exec);
150 session.push(" WHERE ");
151 session.compare(
152 session_table,
153 S::Session::token_column(),
154 " = ",
155 session_token,
156 );
157 session.push(" AND ");
158 session.compare(
159 session_table,
160 S::Session::user_id_column(),
161 " = ",
162 S::Session::parse_user_id(user_id)?,
163 );
164 session.push(" AND ");
165 session.compare(session_table, S::Session::active_column(), " = ", true);
166 model::limit_one(&mut session);
167 lock_exclusive(&mut session);
168 let session = exec
169 .fetch_optional::<S::Session>(session)
170 .await?
171 .ok_or(AuthError::SessionNotFound)?;
172 if alibi_core::AuthSession::expires_at(&session) < Utc::now() {
173 return Err(AuthError::SessionNotFound);
174 }
175 let mut owner = self
176 .organization_models
177 .organization
178 .by_id(exec, invitation.organization_id.as_str())?;
179 model::limit_one(&mut owner);
180 lock_exclusive(&mut owner);
181 drop(
182 self.organization_models
183 .organization
184 .fetch_optional(exec, owner)
185 .await?
186 .ok_or_else(|| AuthError::bad_request("Organization not found"))?,
187 );
188 if let Some(limit) = membership_limit {
189 let mut count = Sql::with(exec.engine(), "SELECT COUNT(*) FROM ");
190 count.ident(self.organization_models.member.table());
191 count.push(" WHERE ");
192 self.organization_models.member.compare(
193 &mut count,
194 "organization_id",
195 " = ",
196 invitation.organization_id.as_str(),
197 )?;
198 let members = exec.fetch_scalar::<i64>(count).await?.unwrap_or_default();
199 if u64::try_from(members).unwrap_or_default()
200 >= u64::try_from(limit)
201 .map_err(|_error| AuthError::internal("Membership limit exceeds u64"))?
202 {
203 return Err(AuthError::bad_request(
204 "Organization membership limit reached",
205 ));
206 }
207 }
208 let requested = invitation
209 .team_id
210 .as_deref()
211 .filter(|ids| !ids.is_empty())
212 .map(|ids| ids.split(',').collect::<Vec<_>>())
213 .unwrap_or_default();
214 for team_id in &requested {
215 let mut room = model::by_id::<team::Model>(exec, *team_id);
216 room.push(" AND ");
217 room.compare(
218 team::Model::TABLE,
219 "organization_id",
220 " = ",
221 invitation.organization_id.as_str(),
222 );
223 model::limit_one(&mut room);
224 if exec.fetch_optional::<team::Model>(room).await?.is_none() {
225 return Err(AuthError::bad_request("Team not found"));
226 }
227 let maximum = team_limits
228 .iter()
229 .find(|(id, _)| id == team_id)
230 .and_then(|(_, limit)| *limit);
231 if matches!(
232 self.add_team_member_in_tx(tx, team_id, user_id, maximum)
233 .await?,
234 alibi_core::types::AddTeamMemberResult::LimitReached
235 ) {
236 return Err(AuthError::Upstream {
237 status: 403,
238 code: "TEAM_MEMBER_LIMIT_REACHED",
239 message: "Team member limit reached",
240 });
241 }
242 }
243 let mut created = ActiveRow::new();
244 created.set("id", self.organization_models.member.new_id());
245 created.set("organization_id", invitation.organization_id.clone());
246 created.set("user_id", user_id);
247 created.set("role", invitation.active_role());
248 created.set("created_at", Utc::now());
249 let created = self
250 .organization_models
251 .member
252 .insert(exec, &created)
253 .await?;
254 let mut active = session.into_active();
255 S::Session::set_active_organization_id(
256 &mut active,
257 Some(invitation.organization_id.clone()),
258 );
259 if requested.len() == 1 {
260 S::Session::set_active_team_id(
261 &mut active,
262 requested.first().map(|id| (*id).to_owned()),
263 )?;
264 }
265 S::Session::set_updated_at(&mut active, Utc::now());
266 drop(
267 model::update::<S::Session>(exec, &active)
268 .await?
269 .ok_or_else(record_not_updated)?,
270 );
271 let mut changed = Sql::with(exec.engine(), "UPDATE ");
272 changed.ident(self.organization_models.invitation.table());
273 changed.push(" SET ");
274 self.organization_models
275 .invitation
276 .assign(&mut changed, "status", "accepted")?;
277 changed.push(" WHERE ");
278 self.organization_models.invitation.compare(
279 &mut changed,
280 "id",
281 " = ",
282 invitation_id,
283 )?;
284 changed.push(" AND ");
285 self.organization_models.invitation.compare(
286 &mut changed,
287 "status",
288 " = ",
289 "pending",
290 )?;
291 if exec.execute(changed).await? != 1 {
292 return Err(AuthError::bad_request("Invitation not found"));
293 }
294 let mut invitation = Invitation::from(&invitation);
295 invitation.status = InvitationStatus::Accepted;
296 Ok(Some((
297 invitation,
298 alibi_core::types::Member::from(&created),
299 )))
300 }
301 .await;
302 match outcome {
303 Ok(value) => {
304 transaction.commit().await?;
305 Ok(value)
306 }
307 Err(error) => {
308 transaction.rollback().await?;
309 Err(error)
310 }
311 }
312 }
313
314 async fn update_invitation_status_if_status(
315 &self,
316 id: &str,
317 expected: InvitationStatus,
318 status: InvitationStatus,
319 ) -> AuthResult<Option<Invitation>> {
320 let mut sql = Sql::with(self.exec().engine(), "UPDATE ");
322 sql.ident(self.organization_models.invitation.table());
323 sql.push(" SET ");
324 self.organization_models
325 .invitation
326 .assign(&mut sql, "status", status.to_string())?;
327 sql.push(" WHERE ");
328 self.organization_models
329 .invitation
330 .compare(&mut sql, "id", " = ", id)?;
331 sql.push(" AND ");
332 self.organization_models.invitation.compare(
333 &mut sql,
334 "status",
335 " = ",
336 expected.to_string(),
337 )?;
338 self.organization_models.invitation.returning(&mut sql);
339 Ok(self
340 .organization_models
341 .invitation
342 .fetch_all(self.exec(), sql)
343 .await?
344 .first()
345 .map(Invitation::from))
346 }
347
348 async fn pending_invitation_page(
349 &self,
350 org_id: &str,
351 email: Option<&str>,
352 ) -> AuthResult<Vec<Invitation>> {
353 let mut sql = self.organization_models.invitation.select(self.exec());
354 sql.push(" WHERE ");
355 self.organization_models
356 .invitation
357 .compare(&mut sql, "organization_id", " = ", org_id)?;
358 sql.push(" AND ");
359 self.organization_models.invitation.compare(
360 &mut sql,
361 "status",
362 " = ",
363 InvitationStatus::Pending.to_string(),
364 )?;
365 if let Some(email) = email {
366 sql.push(" AND ");
367 self.organization_models.invitation.compare(
368 &mut sql,
369 "email",
370 " = ",
371 email.to_lowercase(),
372 )?;
373 }
374 sql.push(" LIMIT ");
375 sql.bind(self.find_many_limit());
376 Ok(self
377 .organization_models
378 .invitation
379 .fetch_all(self.exec(), sql)
380 .await?
381 .iter()
382 .map(Invitation::from)
383 .collect())
384 }
385 async fn update_invitation_expiry(
386 &self,
387 id: &str,
388 expires_at: chrono::DateTime<Utc>,
389 ) -> AuthResult<Invitation> {
390 let model = self
391 .find_invitation(id)
392 .await?
393 .ok_or_else(|| AuthError::not_found("Invitation not found"))?;
394 let mut active = model.into_active();
395 active.set("expires_at", expires_at);
396 self.organization_models
397 .invitation
398 .update(self.exec(), &active)
399 .await?
400 .map(|model| Invitation::from(&model))
401 .ok_or_else(record_not_updated)
402 }
403 async fn get_pending_invitation(
404 &self,
405 org_id: &str,
406 email: &str,
407 ) -> AuthResult<Option<Invitation>> {
408 let mut sql = self.organization_models.invitation.select(self.exec());
409 sql.push(" WHERE ");
410 self.organization_models
411 .invitation
412 .compare(&mut sql, "organization_id", " = ", org_id)?;
413 sql.push(" AND ");
414 self.organization_models.invitation.compare(
415 &mut sql,
416 "email",
417 " = ",
418 email.to_lowercase(),
419 )?;
420 sql.push(" AND ");
421 self.organization_models.invitation.compare(
422 &mut sql,
423 "status",
424 " = ",
425 InvitationStatus::Pending.to_string(),
426 )?;
427 sql.push(" AND ");
428 self.organization_models
429 .invitation
430 .compare(&mut sql, "expires_at", " > ", Utc::now())?;
431 model::limit_one(&mut sql);
432 Ok(self
433 .organization_models
434 .invitation
435 .fetch_optional(self.exec(), sql)
436 .await?
437 .map(|model| Invitation::from(&model)))
438 }
439
440 async fn update_invitation_status(
441 &self,
442 id: &str,
443 status: InvitationStatus,
444 ) -> AuthResult<Invitation> {
445 let Some(model) = self.find_invitation(id).await? else {
446 return Err(AuthError::not_found("Invitation not found"));
447 };
448
449 let mut active = model.into_active();
450 active.set("status", status.to_string());
451 self.organization_models
452 .invitation
453 .update(self.exec(), &active)
454 .await?
455 .map(|model_2| Invitation::from(&model_2))
456 .ok_or_else(record_not_updated)
457 }
458
459 async fn list_organization_invitations(&self, org_id: &str) -> AuthResult<Vec<Invitation>> {
460 let mut sql = self.organization_models.invitation.select(self.exec());
461 sql.push(" WHERE ");
462 self.organization_models
463 .invitation
464 .compare(&mut sql, "organization_id", " = ", org_id)?;
465 sql.push(" LIMIT ");
466 sql.bind(self.find_many_limit());
467 Ok(self
468 .organization_models
469 .invitation
470 .fetch_all(self.exec(), sql)
471 .await?
472 .iter()
473 .map(Invitation::from)
474 .collect())
475 }
476
477 async fn count_pending_organization_invitations(&self, org_id: &str) -> AuthResult<i64> {
478 let mut sql = Sql::with(self.exec().engine(), "SELECT COUNT(*) FROM ");
479 sql.ident(self.organization_models.invitation.table());
480 sql.push(" WHERE ");
481 self.organization_models
482 .invitation
483 .compare(&mut sql, "organization_id", " = ", org_id)?;
484 sql.push(" AND ");
485 self.organization_models.invitation.compare(
486 &mut sql,
487 "status",
488 " = ",
489 InvitationStatus::Pending.to_string(),
490 )?;
491 sql.push(" AND ");
492 self.organization_models
493 .invitation
494 .compare(&mut sql, "expires_at", " > ", Utc::now())?;
495 self.count_invitation_rows(sql).await
496 }
497
498 async fn list_user_invitations(&self, email: &str) -> AuthResult<Vec<Invitation>> {
499 let mut sql = self.organization_models.invitation.select(self.exec());
500 sql.push(" WHERE ");
501 self.organization_models.invitation.compare(
502 &mut sql,
503 "email",
504 " = ",
505 email.to_lowercase(),
506 )?;
507 sql.push(" LIMIT ");
508 sql.bind(self.find_many_limit());
509 Ok(self
510 .organization_models
511 .invitation
512 .fetch_all(self.exec(), sql)
513 .await?
514 .iter()
515 .map(Invitation::from)
516 .collect())
517 }
518}