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