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 model::limit_one(&mut session);
178 lock_exclusive(&mut session);
179 let session = exec
180 .fetch_optional::<S::Session>(session)
181 .await?
182 .ok_or(AuthError::SessionNotFound)?;
183 if alibi_core::AuthSession::expires_at(&session) < Utc::now() {
184 return Err(AuthError::SessionNotFound);
185 }
186 let mut owner = self
187 .organization_models
188 .organization
189 .by_id(exec, invitation.organization_id.as_str())?;
190 model::limit_one(&mut owner);
191 lock_exclusive(&mut owner);
192 _ = self
193 .organization_models
194 .organization
195 .fetch_optional(exec, owner)
196 .await?
197 .ok_or_else(|| AuthError::bad_request("Organization not found"))?;
198 if let Some(limit) = membership_limit {
199 let mut count = Sql::with(exec.engine(), "SELECT COUNT(*) FROM ");
200 count.ident(self.organization_models.member.table());
201 count.push(" WHERE ");
202 self.organization_models.member.compare(
203 &mut count,
204 "organization_id",
205 " = ",
206 invitation.organization_id.as_str(),
207 )?;
208 let members = exec.fetch_scalar::<i64>(count).await?.unwrap_or_default();
209 if u64::try_from(members).unwrap_or_default()
210 >= u64::try_from(limit)
211 .map_err(|_error| AuthError::internal("Membership limit exceeds u64"))?
212 {
213 return Err(AuthError::bad_request(
214 "Organization membership limit reached",
215 ));
216 }
217 }
218 let requested = invitation
219 .team_id
220 .as_deref()
221 .filter(|ids| !ids.is_empty())
222 .map(|ids| ids.split(',').collect::<Vec<_>>())
223 .unwrap_or_default();
224 for team_id in &requested {
225 let mut room = model::by_id::<team::Model>(exec, *team_id);
226 room.push(" AND ");
227 room.compare(
228 team::Model::TABLE,
229 "organizationId",
230 " = ",
231 invitation.organization_id.as_str(),
232 );
233 model::limit_one(&mut room);
234 if exec.fetch_optional::<team::Model>(room).await?.is_none() {
235 return Err(AuthError::bad_request("Team not found"));
236 }
237 let maximum = team_limits
238 .iter()
239 .find(|(id, _)| id == team_id)
240 .and_then(|(_, limit)| *limit);
241 if matches!(
242 self.add_team_member_in_tx(tx, team_id, user_id, maximum)
243 .await?,
244 AddTeamMemberResult::LimitReached
245 ) {
246 return Err(AuthError::Upstream {
247 status: 403,
248 code: "TEAM_MEMBER_LIMIT_REACHED",
249 message: "Team member limit reached",
250 });
251 }
252 }
253 let mut created = ActiveRow::new();
254 created.set("id", self.organization_models.member.new_id());
255 created.set("organization_id", invitation.organization_id.clone());
256 created.set("user_id", user_id);
257 created.set("role", invitation.active_role());
258 created.set("created_at", Utc::now());
259 let created = self
260 .organization_models
261 .member
262 .insert(exec, &created)
263 .await?;
264 let mut active = session.into_active();
265 S::Session::set_active_organization_id(
266 &mut active,
267 Some(invitation.organization_id.clone()),
268 );
269 if requested.len() == 1 {
270 S::Session::set_active_team_id(
271 &mut active,
272 requested.first().map(|id| (*id).to_owned()),
273 )?;
274 }
275 S::Session::set_updated_at(&mut active, Utc::now());
276 _ = model::update::<S::Session>(exec, &active)
277 .await?
278 .ok_or_else(record_not_updated)?;
279 let mut changed = Sql::with(exec.engine(), "UPDATE ");
280 changed.ident(self.organization_models.invitation.table());
281 changed.push(" SET ");
282 self.organization_models
283 .invitation
284 .assign(&mut changed, "status", "accepted")?;
285 changed.push(" WHERE ");
286 self.organization_models.invitation.compare(
287 &mut changed,
288 "id",
289 " = ",
290 invitation_id,
291 )?;
292 changed.push(" AND ");
293 self.organization_models.invitation.compare(
294 &mut changed,
295 "status",
296 " = ",
297 "pending",
298 )?;
299 if exec.execute(changed).await? != 1 {
300 return Err(AuthError::bad_request("Invitation not found"));
301 }
302 let mut invitation = Invitation::from(&invitation);
303 invitation.status = InvitationStatus::Accepted;
304 Ok(Some((
305 invitation,
306 alibi_core::types::Member::from(&created),
307 )))
308 }
309 .await;
310 match outcome {
311 Ok(value) => {
312 transaction.commit().await?;
313 Ok(value)
314 }
315 Err(error) => {
316 transaction.rollback().await?;
317 Err(error)
318 }
319 }
320 }
321
322 async fn update_invitation_status_if_status(
323 &self,
324 id: &str,
325 expected: InvitationStatus,
326 status: InvitationStatus,
327 ) -> AuthResult<Option<Invitation>> {
328 let mut sql = Sql::with(self.exec().engine(), "UPDATE ");
330 sql.ident(self.organization_models.invitation.table());
331 sql.push(" SET ");
332 self.organization_models
333 .invitation
334 .assign(&mut sql, "status", status.to_string())?;
335 sql.push(" WHERE ");
336 self.organization_models
337 .invitation
338 .compare(&mut sql, "id", " = ", id)?;
339 sql.push(" AND ");
340 self.organization_models.invitation.compare(
341 &mut sql,
342 "status",
343 " = ",
344 expected.to_string(),
345 )?;
346 self.organization_models.invitation.returning(&mut sql);
347 Ok(self
348 .organization_models
349 .invitation
350 .fetch_all(self.exec(), sql)
351 .await?
352 .first()
353 .map(Invitation::from))
354 }
355
356 async fn pending_invitation_page(
357 &self,
358 org_id: &str,
359 email: Option<&str>,
360 ) -> AuthResult<Vec<Invitation>> {
361 let mut sql = self.organization_models.invitation.select(self.exec());
362 sql.push(" WHERE ");
363 self.organization_models
364 .invitation
365 .compare(&mut sql, "organization_id", " = ", org_id)?;
366 sql.push(" AND ");
367 self.organization_models.invitation.compare(
368 &mut sql,
369 "status",
370 " = ",
371 InvitationStatus::Pending.to_string(),
372 )?;
373 if let Some(email) = email {
374 sql.push(" AND ");
375 self.organization_models.invitation.compare(
376 &mut sql,
377 "email",
378 " = ",
379 email.to_lowercase(),
380 )?;
381 }
382 sql.push(" LIMIT ");
383 sql.bind(self.find_many_limit());
384 Ok(self
385 .organization_models
386 .invitation
387 .fetch_all(self.exec(), sql)
388 .await?
389 .iter()
390 .map(Invitation::from)
391 .collect())
392 }
393 async fn update_invitation_expiry(
394 &self,
395 id: &str,
396 expires_at: chrono::DateTime<Utc>,
397 ) -> AuthResult<Invitation> {
398 let model = self
399 .find_invitation(id)
400 .await?
401 .ok_or_else(|| AuthError::not_found("Invitation not found"))?;
402 let mut active = model.into_active();
403 active.set("expires_at", expires_at);
404 self.organization_models
405 .invitation
406 .update(self.exec(), &active)
407 .await?
408 .map(|model| Invitation::from(&model))
409 .ok_or_else(record_not_updated)
410 }
411 async fn get_pending_invitation(
412 &self,
413 org_id: &str,
414 email: &str,
415 ) -> AuthResult<Option<Invitation>> {
416 let mut sql = self.organization_models.invitation.select(self.exec());
417 sql.push(" WHERE ");
418 self.organization_models
419 .invitation
420 .compare(&mut sql, "organization_id", " = ", org_id)?;
421 sql.push(" AND ");
422 self.organization_models.invitation.compare(
423 &mut sql,
424 "email",
425 " = ",
426 email.to_lowercase(),
427 )?;
428 sql.push(" AND ");
429 self.organization_models.invitation.compare(
430 &mut sql,
431 "status",
432 " = ",
433 InvitationStatus::Pending.to_string(),
434 )?;
435 sql.push(" AND ");
436 self.organization_models
437 .invitation
438 .compare(&mut sql, "expires_at", " > ", Utc::now())?;
439 model::limit_one(&mut sql);
440 Ok(self
441 .organization_models
442 .invitation
443 .fetch_optional(self.exec(), sql)
444 .await?
445 .map(|model| Invitation::from(&model)))
446 }
447
448 async fn update_invitation_status(
449 &self,
450 id: &str,
451 status: InvitationStatus,
452 ) -> AuthResult<Invitation> {
453 let Some(model) = self.find_invitation(id).await? else {
454 return Err(AuthError::not_found("Invitation not found"));
455 };
456
457 let mut active = model.into_active();
458 active.set("status", status.to_string());
459 self.organization_models
460 .invitation
461 .update(self.exec(), &active)
462 .await?
463 .map(|model| Invitation::from(&model))
464 .ok_or_else(record_not_updated)
465 }
466
467 async fn list_organization_invitations(&self, org_id: &str) -> AuthResult<Vec<Invitation>> {
468 let mut sql = self.organization_models.invitation.select(self.exec());
469 sql.push(" WHERE ");
470 self.organization_models
471 .invitation
472 .compare(&mut sql, "organization_id", " = ", org_id)?;
473 sql.push(" LIMIT ");
474 sql.bind(self.find_many_limit());
475 Ok(self
476 .organization_models
477 .invitation
478 .fetch_all(self.exec(), sql)
479 .await?
480 .iter()
481 .map(Invitation::from)
482 .collect())
483 }
484
485 async fn count_pending_organization_invitations(&self, org_id: &str) -> AuthResult<i64> {
486 let mut sql = Sql::with(self.exec().engine(), "SELECT COUNT(*) FROM ");
487 sql.ident(self.organization_models.invitation.table());
488 sql.push(" WHERE ");
489 self.organization_models
490 .invitation
491 .compare(&mut sql, "organization_id", " = ", org_id)?;
492 sql.push(" AND ");
493 self.organization_models.invitation.compare(
494 &mut sql,
495 "status",
496 " = ",
497 InvitationStatus::Pending.to_string(),
498 )?;
499 sql.push(" AND ");
500 self.organization_models
501 .invitation
502 .compare(&mut sql, "expires_at", " > ", Utc::now())?;
503 self.count_invitation_rows(sql).await
504 }
505
506 async fn list_user_invitations(&self, email: &str) -> AuthResult<Vec<Invitation>> {
507 let mut sql = self.organization_models.invitation.select(self.exec());
508 sql.push(" WHERE ");
509 self.organization_models.invitation.compare(
510 &mut sql,
511 "email",
512 " = ",
513 email.to_lowercase(),
514 )?;
515 sql.push(" LIMIT ");
516 sql.bind(self.find_many_limit());
517 Ok(self
518 .organization_models
519 .invitation
520 .fetch_all(self.exec(), sql)
521 .await?
522 .iter()
523 .map(Invitation::from)
524 .collect())
525 }
526}