alibi_sqlx/store/
teams.rs

1use super::entities::{team, team_member};
2use super::{SqlxStore, lock_exclusive};
3use crate::error::record_not_updated;
4use crate::model::{self, ActiveRow, SqlxModel};
5use crate::pool::{Engine, Exec, SqlxTransaction};
6use crate::schema::{AuthSchema, SqlxUserModel};
7use crate::sql::Sql;
8use alibi_core::error::{AuthError, AuthResult};
9use alibi_core::store::{TeamStore, team_membership_key};
10use alibi_core::types::{AddTeamMemberResult, CreateTeam, Team, TeamMember, UpdateTeam};
11use async_trait::async_trait;
12use chrono::Utc;
13use uuid::Uuid;
14
15fn team_member_lookup(exec: Exec<'_>, team_id: &str, user_id: &str) -> Sql {
16    let mut sql = model::select_model::<team_member::Model>(exec);
17    sql.push(" WHERE ");
18    sql.compare(team_member::Model::TABLE, "team_id", " = ", team_id);
19    sql.push(" AND ");
20    sql.compare(team_member::Model::TABLE, "user_id", " = ", user_id);
21    model::limit_one(&mut sql);
22    sql
23}
24
25impl<S> SqlxStore<S>
26where
27    S: AuthSchema,
28    S::User: SqlxUserModel,
29{
30    pub(super) async fn get_team_with_connection(
31        &self,
32        exec: Exec<'_>,
33        organization_id: Option<&str>,
34        team_id: &str,
35    ) -> AuthResult<Option<Team>> {
36        let mut sql = model::by_id::<team::Model>(exec, team_id);
37        if let Some(org) = organization_id {
38            sql.push(" AND ");
39            sql.compare(team::Model::TABLE, "organization_id", " = ", org);
40        }
41        model::limit_one(&mut sql);
42        Ok(exec
43            .fetch_optional::<team::Model>(sql)
44            .await?
45            .map(Into::into))
46    }
47
48    pub(super) async fn add_team_member_in_tx(
49        &self,
50        tx: &SqlxTransaction,
51        team_id: &str,
52        user_id: &str,
53        maximum: Option<f64>,
54    ) -> AuthResult<AddTeamMemberResult> {
55        let exec = Exec::tx(tx);
56        if super::users::find_user_by_id::<S::User>(exec, user_id, super::users::Lock::Shared)
57            .await?
58            .is_none()
59        {
60            return Err(AuthError::UserNotFound);
61        }
62        let mut room = model::by_id::<team::Model>(exec, team_id);
63        model::limit_one(&mut room);
64        lock_exclusive(&mut room);
65        let room = exec
66            .fetch_optional::<team::Model>(room)
67            .await?
68            .ok_or_else(|| AuthError::bad_request("Team not found"))?;
69        if let Some(member) = exec
70            .fetch_optional::<team_member::Model>(team_member_lookup(exec, team_id, user_id))
71            .await?
72        {
73            return Ok(AddTeamMemberResult::Existing(member.into()));
74        }
75        let mut count = Sql::with(exec.engine(), "SELECT COUNT(*) FROM ");
76        count.ident(team_member::Model::TABLE);
77        count.push(" WHERE ");
78        count.compare(team_member::Model::TABLE, "team_id", " = ", team_id);
79        let count = u64::try_from(exec.fetch_scalar::<i64>(count).await?.unwrap_or_default())
80            .unwrap_or_default();
81        // Source only repairs a stale low durable counter before reserving a seat.
82        // Preserve high counters, and bind the raw policy to the database's `<`
83        // predicate: notably SQLite binds NaN as NULL, which cannot reserve a seat.
84        let reserved_count = i64::try_from(count)
85            .map_err(|_error| AuthError::internal("Team membership count overflow"))?
86            .max(room.member_count);
87        let mut active = room.into_active();
88        active.set("member_count", reserved_count);
89        _ = model::update::<team::Model>(exec, &active)
90            .await?
91            .ok_or_else(record_not_updated)?;
92        if let Some(maximum) = maximum {
93            let mut seat = Sql::with(exec.engine(), "SELECT COUNT(*) FROM ");
94            seat.ident(team::Model::TABLE);
95            seat.push(" WHERE ");
96            seat.compare(team::Model::TABLE, "id", " = ", team_id);
97            seat.push(" AND ");
98            if exec.engine() == Engine::Postgres {
99                // Source's text Number parameter is parsed as the physical
100                // bigint counter, rather than promoting that counter to float8.
101                seat.column(team::Model::TABLE, "member_count");
102                seat.push(" < ");
103                seat.bind(ryu_js::Buffer::new().format(maximum).to_owned());
104                seat.push("::bigint");
105            } else {
106                seat.compare(team::Model::TABLE, "member_count", " < ", maximum);
107            }
108            if exec.fetch_scalar::<i64>(seat).await?.unwrap_or_default() == 0 {
109                return Ok(AddTeamMemberResult::LimitReached);
110            }
111        }
112        let key = team_membership_key(team_id, user_id)?;
113        let mut member = ActiveRow::new();
114        member.set(
115            "id",
116            self.generated_id(exec, "teamMember", "team_member", "id")
117                .await?
118                .unwrap_or_else(|| Uuid::new_v4().to_string()),
119        );
120        member.set("team_id", team_id);
121        member.set("user_id", user_id);
122        member.set("membership_key", Some(key));
123        member.set("created_at", Utc::now());
124        let member = model::insert::<team_member::Model>(exec, &member).await?;
125        active.set(
126            "member_count",
127            reserved_count
128                .checked_add(1)
129                .ok_or_else(|| AuthError::internal("Team membership count overflow"))?,
130        );
131        _ = model::update::<team::Model>(exec, &active)
132            .await?
133            .ok_or_else(record_not_updated)?;
134        Ok(AddTeamMemberResult::Added(member.into()))
135    }
136}
137
138#[async_trait]
139impl<S> TeamStore for SqlxStore<S>
140where
141    S: AuthSchema,
142    S::User: SqlxUserModel,
143{
144    async fn create_team(&self, data: CreateTeam) -> AuthResult<Team> {
145        let now = data.updated_at.unwrap_or_else(Utc::now);
146        let mut active = ActiveRow::new();
147        active.set(
148            "id",
149            self.generated_id(self.exec(), "team", "team", "id")
150                .await?
151                .unwrap_or_else(|| Uuid::new_v4().to_string()),
152        );
153        active.set("name", data.name);
154        active.set("organization_id", data.organization_id);
155        active.set("member_count", 0_i64);
156        active.set("created_at", now);
157        active.set("updated_at", data.updated_at);
158        model::insert::<team::Model>(self.exec(), &active)
159            .await
160            .map(Into::into)
161    }
162    async fn get_team(
163        &self,
164        organization_id: Option<&str>,
165        team_id: &str,
166    ) -> AuthResult<Option<Team>> {
167        self.get_team_with_connection(self.exec(), organization_id, team_id)
168            .await
169    }
170    async fn list_teams(&self, organization_id: &str) -> AuthResult<Vec<Team>> {
171        let mut sql = model::select_model::<team::Model>(self.exec());
172        sql.push(" WHERE ");
173        sql.compare(
174            team::Model::TABLE,
175            "organization_id",
176            " = ",
177            organization_id,
178        );
179        sql.push(" LIMIT ");
180        sql.bind(self.find_many_limit());
181        Ok(self
182            .exec()
183            .fetch_all::<team::Model>(sql)
184            .await?
185            .into_iter()
186            .map(Into::into)
187            .collect())
188    }
189    async fn update_team(
190        &self,
191        organization_id: &str,
192        team_id: &str,
193        update: UpdateTeam,
194    ) -> AuthResult<Team> {
195        let mut sql = model::by_id::<team::Model>(self.exec(), team_id);
196        sql.push(" AND ");
197        sql.compare(
198            team::Model::TABLE,
199            "organization_id",
200            " = ",
201            organization_id,
202        );
203        model::limit_one(&mut sql);
204        let model = self
205            .exec()
206            .fetch_optional::<team::Model>(sql)
207            .await?
208            .ok_or_else(|| AuthError::bad_request("Team not found"))?;
209        let mut active = model.into_active();
210        if let Some(name) = update.name {
211            active.set("name", name);
212        }
213        active.set("updated_at", Some(Utc::now()));
214        model::update::<team::Model>(self.exec(), &active)
215            .await?
216            .map(Into::into)
217            .ok_or_else(record_not_updated)
218    }
219    async fn delete_team(&self, organization_id: &str, team_id: &str) -> AuthResult<bool> {
220        let (organization_id, team_id) = (organization_id.to_owned(), team_id.to_owned());
221        self.in_transaction(true, async move |tx| {
222            let exec = Exec::tx(tx);
223            let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
224            delete.ident(team::Model::TABLE);
225            delete.push(" WHERE ");
226            delete.compare(team::Model::TABLE, "id", " = ", team_id.as_str());
227            delete.push(" AND ");
228            delete.compare(
229                team::Model::TABLE,
230                "organization_id",
231                " = ",
232                organization_id.as_str(),
233            );
234            if exec.execute(delete).await? == 0 {
235                return Ok(false);
236            }
237            let mut members = Sql::with(exec.engine(), "DELETE FROM ");
238            members.ident(team_member::Model::TABLE);
239            members.push(" WHERE ");
240            members.compare(
241                team_member::Model::TABLE,
242                "team_id",
243                " = ",
244                team_id.as_str(),
245            );
246            _ = exec.execute(members).await?;
247            let mut pending = self.organization_models.invitation.select(exec);
248            pending.push(" WHERE ");
249            self.organization_models.invitation.compare(
250                &mut pending,
251                "organization_id",
252                " = ",
253                organization_id.as_str(),
254            )?;
255            pending.push(" AND ");
256            self.organization_models.invitation.compare(
257                &mut pending,
258                "status",
259                " = ",
260                "pending",
261            )?;
262            pending.push(" AND ");
263            self.organization_models.invitation.compare(
264                &mut pending,
265                "expires_at",
266                " > ",
267                Utc::now(),
268            )?;
269            for invite in self
270                .organization_models
271                .invitation
272                .fetch_all(exec, pending)
273                .await?
274            {
275                let Some(ids) = invite.team_id.as_deref() else {
276                    continue;
277                };
278                if !ids.split(',').any(|id| id == team_id) {
279                    continue;
280                }
281                let remaining = ids
282                    .split(',')
283                    .filter(|id| *id != team_id)
284                    .collect::<Vec<_>>()
285                    .join(",");
286                let mut active = invite.into_active();
287                active.set("team_id", (!remaining.is_empty()).then_some(remaining));
288                _ = self
289                    .organization_models
290                    .invitation
291                    .update(exec, &active)
292                    .await?
293                    .ok_or_else(record_not_updated)?;
294            }
295            Ok(true)
296        })
297        .await
298    }
299    async fn get_team_member(
300        &self,
301        team_id: &str,
302        user_id: &str,
303    ) -> AuthResult<Option<TeamMember>> {
304        Ok(self
305            .exec()
306            .fetch_optional::<team_member::Model>(team_member_lookup(self.exec(), team_id, user_id))
307            .await?
308            .map(Into::into))
309    }
310    async fn add_team_member(
311        &self,
312        team_id: &str,
313        user_id: &str,
314        maximum: Option<f64>,
315    ) -> AuthResult<AddTeamMemberResult> {
316        let tx = self.begin(true).await?;
317        let result = self
318            .add_team_member_in_tx(&tx, team_id, user_id, maximum)
319            .await?;
320        tx.commit().await?;
321        Ok(result)
322    }
323    async fn remove_team_member(&self, team_id: &str, user_id: &str) -> AuthResult<usize> {
324        let (team_id, user_id) = (team_id.to_owned(), user_id.to_owned());
325        let removed = self
326            .in_transaction(true, async move |tx| {
327                let exec = Exec::tx(tx);
328                let mut room = model::by_id::<team::Model>(exec, team_id.as_str());
329                model::limit_one(&mut room);
330                lock_exclusive(&mut room);
331                _ = exec.fetch_optional::<team::Model>(room).await?;
332                let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
333                delete.ident(team_member::Model::TABLE);
334                delete.push(" WHERE ");
335                delete.compare(
336                    team_member::Model::TABLE,
337                    "team_id",
338                    " = ",
339                    team_id.as_str(),
340                );
341                delete.push(" AND ");
342                delete.compare(
343                    team_member::Model::TABLE,
344                    "user_id",
345                    " = ",
346                    user_id.as_str(),
347                );
348                let removed = exec.execute(delete).await?;
349                let count = i64::try_from(removed)
350                    .map_err(|_error| AuthError::internal("Team membership count overflow"))?;
351                if count > 0 {
352                    release_seats(exec, &team_id, count).await?;
353                }
354                Ok(removed)
355            })
356            .await?;
357        usize::try_from(removed)
358            .map_err(|_error| AuthError::internal("Team membership count overflow"))
359    }
360    async fn list_team_members(&self, team_id: &str) -> AuthResult<Vec<TeamMember>> {
361        let mut sql = model::select_model::<team_member::Model>(self.exec());
362        sql.push(" WHERE ");
363        sql.compare(team_member::Model::TABLE, "team_id", " = ", team_id);
364        sql.push(" LIMIT ");
365        sql.bind(self.find_many_limit());
366        Ok(self
367            .exec()
368            .fetch_all::<team_member::Model>(sql)
369            .await?
370            .into_iter()
371            .map(Into::into)
372            .collect())
373    }
374    async fn list_user_teams(&self, user_id: &str) -> AuthResult<Vec<Team>> {
375        let mut sql = model::select_model::<team_member::Model>(self.exec());
376        sql.push(" WHERE ");
377        sql.compare(team_member::Model::TABLE, "user_id", " = ", user_id);
378        sql.push(" LIMIT ");
379        sql.bind(self.find_many_limit());
380        let memberships: Vec<team_member::Model> = self.exec().fetch_all(sql).await?;
381        if memberships.is_empty() {
382            return Ok(Vec::new());
383        }
384        let mut rooms = model::select_model::<team::Model>(self.exec());
385        rooms.push(" WHERE ");
386        rooms.column(team::Model::TABLE, "id");
387        rooms.push(" IN ");
388        rooms.bind_list(memberships.iter().map(|row| row.team_id.clone()));
389        let rooms = self
390            .exec()
391            .fetch_all::<team::Model>(rooms)
392            .await?
393            .into_iter()
394            .map(|row| (row.id.clone(), row))
395            .collect::<std::collections::HashMap<_, _>>();
396        // The upstream join maps membership rows directly, preserving their
397        // adapter order. An IN query's team order must not replace that order.
398        Ok(memberships
399            .into_iter()
400            .filter_map(|membership| rooms.get(&membership.team_id).cloned().map(Into::into))
401            .collect())
402    }
403}
404
405/// `UPDATE team SET member_count = member_count - n WHERE id = ? AND member_count >= n`.
406async fn release_seats(exec: Exec<'_>, team_id: &str, count: i64) -> AuthResult<()> {
407    let mut sql = Sql::with(exec.engine(), "UPDATE ");
408    sql.ident(team::Model::TABLE);
409    sql.push(" SET ");
410    sql.ident("member_count");
411    sql.push(" = ");
412    sql.ident("member_count");
413    sql.push(" - ");
414    sql.bind(count);
415    sql.push(" WHERE ");
416    sql.compare(team::Model::TABLE, "id", " = ", team_id);
417    sql.push(" AND ");
418    sql.compare(team::Model::TABLE, "member_count", " >= ", count);
419    _ = exec.execute(sql).await?;
420    Ok(())
421}
422
423/// Delete owned memberships and release exactly the seats those rows occupied.
424pub(super) async fn remove_owned_team_members(
425    tx: &SqlxTransaction,
426    user_id: &str,
427    organization_id: Option<&str>,
428) -> AuthResult<()> {
429    let exec = Exec::tx(tx);
430    if !super::migrator::has_table(exec, "team_member").await? {
431        return Ok(());
432    }
433    let mut teams = model::select_model::<team::Model>(exec);
434    if let Some(org) = organization_id {
435        teams.push(" WHERE ");
436        teams.compare(team::Model::TABLE, "organization_id", " = ", org);
437    }
438    teams.push(" ORDER BY ");
439    teams.column(team::Model::TABLE, "id");
440    teams.push(" ASC");
441    lock_exclusive(&mut teams);
442    let rooms = exec.fetch_all::<team::Model>(teams).await?;
443    release_owned_team_members(tx, user_id, rooms).await
444}
445
446/// Remove memberships from an already selected adapter page and release seats.
447pub(super) async fn release_owned_team_members(
448    tx: &SqlxTransaction,
449    user_id: &str,
450    rooms: Vec<team::Model>,
451) -> AuthResult<()> {
452    let exec = Exec::tx(tx);
453    for room in rooms {
454        let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
455        delete.ident(team_member::Model::TABLE);
456        delete.push(" WHERE ");
457        delete.compare(
458            team_member::Model::TABLE,
459            "team_id",
460            " = ",
461            room.id.as_str(),
462        );
463        delete.push(" AND ");
464        delete.compare(team_member::Model::TABLE, "user_id", " = ", user_id);
465        let deleted = exec.execute(delete).await?;
466        if deleted > 0 {
467            let count = i64::try_from(deleted)
468                .map_err(|_error| AuthError::internal("Team membership count overflow"))?;
469            release_seats(exec, &room.id, count).await?;
470        }
471    }
472    Ok(())
473}
alibi_sqlx/store/teams.rs

alibi_sqlx/store/
teams.rs
