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, "teamId", " = ", team_id);
19    sql.push(" AND ");
20    sql.compare(team_member::Model::TABLE, "userId", " = ", 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, "organizationId", " = ", 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, "teamId", " = ", 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("memberCount", 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, "memberCount");
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, "memberCount", " < ", 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", "teamMember", "id")
117                .await?
118                .unwrap_or_else(|| Uuid::new_v4().to_string()),
119        );
120        member.set("teamId", team_id);
121        member.set("userId", user_id);
122        member.set("membershipKey", Some(key));
123        member.set("createdAt", Utc::now());
124        let member = model::insert::<team_member::Model>(exec, &member).await?;
125        active.set(
126            "memberCount",
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("organizationId", data.organization_id);
155        active.set("memberCount", 0_i64);
156        active.set("createdAt", now);
157        active.set("updatedAt", 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(team::Model::TABLE, "organizationId", " = ", organization_id);
174        sql.push(" LIMIT ");
175        sql.bind(self.find_many_limit());
176        Ok(self
177            .exec()
178            .fetch_all::<team::Model>(sql)
179            .await?
180            .into_iter()
181            .map(Into::into)
182            .collect())
183    }
184    async fn update_team(
185        &self,
186        organization_id: &str,
187        team_id: &str,
188        update: UpdateTeam,
189    ) -> AuthResult<Team> {
190        let mut sql = model::by_id::<team::Model>(self.exec(), team_id);
191        sql.push(" AND ");
192        sql.compare(team::Model::TABLE, "organizationId", " = ", organization_id);
193        model::limit_one(&mut sql);
194        let model = self
195            .exec()
196            .fetch_optional::<team::Model>(sql)
197            .await?
198            .ok_or_else(|| AuthError::bad_request("Team not found"))?;
199        let mut active = model.into_active();
200        if let Some(name) = update.name {
201            active.set("name", name);
202        }
203        active.set("updatedAt", Some(Utc::now()));
204        model::update::<team::Model>(self.exec(), &active)
205            .await?
206            .map(Into::into)
207            .ok_or_else(record_not_updated)
208    }
209    async fn delete_team(&self, organization_id: &str, team_id: &str) -> AuthResult<bool> {
210        let (organization_id, team_id) = (organization_id.to_owned(), team_id.to_owned());
211        self.in_transaction(true, async move |tx| {
212            let exec = Exec::tx(tx);
213            let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
214            delete.ident(team::Model::TABLE);
215            delete.push(" WHERE ");
216            delete.compare(team::Model::TABLE, "id", " = ", team_id.as_str());
217            delete.push(" AND ");
218            delete.compare(
219                team::Model::TABLE,
220                "organizationId",
221                " = ",
222                organization_id.as_str(),
223            );
224            if exec.execute(delete).await? == 0 {
225                return Ok(false);
226            }
227            let mut members = Sql::with(exec.engine(), "DELETE FROM ");
228            members.ident(team_member::Model::TABLE);
229            members.push(" WHERE ");
230            members.compare(team_member::Model::TABLE, "teamId", " = ", team_id.as_str());
231            _ = exec.execute(members).await?;
232            let mut pending = self.organization_models.invitation.select(exec);
233            pending.push(" WHERE ");
234            self.organization_models.invitation.compare(
235                &mut pending,
236                "organizationId",
237                " = ",
238                organization_id.as_str(),
239            )?;
240            pending.push(" AND ");
241            self.organization_models.invitation.compare(
242                &mut pending,
243                "status",
244                " = ",
245                "pending",
246            )?;
247            pending.push(" AND ");
248            self.organization_models.invitation.compare(
249                &mut pending,
250                "expiresAt",
251                " > ",
252                Utc::now(),
253            )?;
254            for invite in self
255                .organization_models
256                .invitation
257                .fetch_all(exec, pending)
258                .await?
259            {
260                let Some(ids) = invite.team_id.as_deref() else {
261                    continue;
262                };
263                if !ids.split(',').any(|id| id == team_id) {
264                    continue;
265                }
266                let remaining = ids
267                    .split(',')
268                    .filter(|id| *id != team_id)
269                    .collect::<Vec<_>>()
270                    .join(",");
271                let mut active = invite.into_active();
272                active.set("teamId", (!remaining.is_empty()).then_some(remaining));
273                _ = self
274                    .organization_models
275                    .invitation
276                    .update(exec, &active)
277                    .await?
278                    .ok_or_else(record_not_updated)?;
279            }
280            Ok(true)
281        })
282        .await
283    }
284    async fn get_team_member(
285        &self,
286        team_id: &str,
287        user_id: &str,
288    ) -> AuthResult<Option<TeamMember>> {
289        Ok(self
290            .exec()
291            .fetch_optional::<team_member::Model>(team_member_lookup(self.exec(), team_id, user_id))
292            .await?
293            .map(Into::into))
294    }
295    async fn add_team_member(
296        &self,
297        team_id: &str,
298        user_id: &str,
299        maximum: Option<f64>,
300    ) -> AuthResult<AddTeamMemberResult> {
301        let tx = self.begin(true).await?;
302        let result = self
303            .add_team_member_in_tx(&tx, team_id, user_id, maximum)
304            .await?;
305        tx.commit().await?;
306        Ok(result)
307    }
308    async fn remove_team_member(&self, team_id: &str, user_id: &str) -> AuthResult<usize> {
309        let (team_id, user_id) = (team_id.to_owned(), user_id.to_owned());
310        let removed = self
311            .in_transaction(true, async move |tx| {
312                let exec = Exec::tx(tx);
313                let mut room = model::by_id::<team::Model>(exec, team_id.as_str());
314                model::limit_one(&mut room);
315                lock_exclusive(&mut room);
316                _ = exec.fetch_optional::<team::Model>(room).await?;
317                let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
318                delete.ident(team_member::Model::TABLE);
319                delete.push(" WHERE ");
320                delete.compare(team_member::Model::TABLE, "teamId", " = ", team_id.as_str());
321                delete.push(" AND ");
322                delete.compare(team_member::Model::TABLE, "userId", " = ", user_id.as_str());
323                let removed = exec.execute(delete).await?;
324                let count = i64::try_from(removed)
325                    .map_err(|_error| AuthError::internal("Team membership count overflow"))?;
326                if count > 0 {
327                    release_seats(exec, &team_id, count).await?;
328                }
329                Ok(removed)
330            })
331            .await?;
332        usize::try_from(removed)
333            .map_err(|_error| AuthError::internal("Team membership count overflow"))
334    }
335    async fn list_team_members(&self, team_id: &str) -> AuthResult<Vec<TeamMember>> {
336        let mut sql = model::select_model::<team_member::Model>(self.exec());
337        sql.push(" WHERE ");
338        sql.compare(team_member::Model::TABLE, "teamId", " = ", team_id);
339        sql.push(" LIMIT ");
340        sql.bind(self.find_many_limit());
341        Ok(self
342            .exec()
343            .fetch_all::<team_member::Model>(sql)
344            .await?
345            .into_iter()
346            .map(Into::into)
347            .collect())
348    }
349    async fn list_user_teams(&self, user_id: &str) -> AuthResult<Vec<Team>> {
350        let mut sql = model::select_model::<team_member::Model>(self.exec());
351        sql.push(" WHERE ");
352        sql.compare(team_member::Model::TABLE, "userId", " = ", user_id);
353        sql.push(" LIMIT ");
354        sql.bind(self.find_many_limit());
355        let memberships: Vec<team_member::Model> = self.exec().fetch_all(sql).await?;
356        if memberships.is_empty() {
357            return Ok(Vec::new());
358        }
359        let mut rooms = model::select_model::<team::Model>(self.exec());
360        rooms.push(" WHERE ");
361        rooms.column(team::Model::TABLE, "id");
362        rooms.push(" IN ");
363        rooms.bind_list(memberships.iter().map(|row| row.team_id.clone()));
364        let rooms = self
365            .exec()
366            .fetch_all::<team::Model>(rooms)
367            .await?
368            .into_iter()
369            .map(|row| (row.id.clone(), row))
370            .collect::<std::collections::HashMap<_, _>>();
371        // The upstream join maps membership rows directly, preserving their
372        // adapter order. An IN query's team order must not replace that order.
373        Ok(memberships
374            .into_iter()
375            .filter_map(|membership| rooms.get(&membership.team_id).cloned().map(Into::into))
376            .collect())
377    }
378}
379
380/// `UPDATE team SET member_count = member_count - n WHERE id = ? AND member_count >= n`.
381async fn release_seats(exec: Exec<'_>, team_id: &str, count: i64) -> AuthResult<()> {
382    let mut sql = Sql::with(exec.engine(), "UPDATE ");
383    sql.ident(team::Model::TABLE);
384    sql.push(" SET ");
385    sql.ident("memberCount");
386    sql.push(" = ");
387    sql.ident("memberCount");
388    sql.push(" - ");
389    sql.bind(count);
390    sql.push(" WHERE ");
391    sql.compare(team::Model::TABLE, "id", " = ", team_id);
392    sql.push(" AND ");
393    sql.compare(team::Model::TABLE, "memberCount", " >= ", count);
394    _ = exec.execute(sql).await?;
395    Ok(())
396}
397
398/// Delete owned memberships and release exactly the seats those rows occupied.
399pub(super) async fn remove_owned_team_members(
400    tx: &SqlxTransaction,
401    user_id: &str,
402    organization_id: Option<&str>,
403) -> AuthResult<()> {
404    let exec = Exec::tx(tx);
405    if !super::migrator::has_table(exec, "teamMember").await? {
406        return Ok(());
407    }
408    let mut teams = model::select_model::<team::Model>(exec);
409    if let Some(org) = organization_id {
410        teams.push(" WHERE ");
411        teams.compare(team::Model::TABLE, "organizationId", " = ", org);
412    }
413    teams.push(" ORDER BY ");
414    teams.column(team::Model::TABLE, "id");
415    teams.push(" ASC");
416    lock_exclusive(&mut teams);
417    let rooms = exec.fetch_all::<team::Model>(teams).await?;
418    release_owned_team_members(tx, user_id, rooms).await
419}
420
421/// Remove memberships from an already selected adapter page and release seats.
422pub(super) async fn release_owned_team_members(
423    tx: &SqlxTransaction,
424    user_id: &str,
425    rooms: Vec<team::Model>,
426) -> AuthResult<()> {
427    let exec = Exec::tx(tx);
428    for room in rooms {
429        let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
430        delete.ident(team_member::Model::TABLE);
431        delete.push(" WHERE ");
432        delete.compare(team_member::Model::TABLE, "teamId", " = ", room.id.as_str());
433        delete.push(" AND ");
434        delete.compare(team_member::Model::TABLE, "userId", " = ", user_id);
435        let deleted = exec.execute(delete).await?;
436        if deleted > 0 {
437            let count = i64::try_from(deleted)
438                .map_err(|_error| AuthError::internal("Team membership count overflow"))?;
439            release_seats(exec, &room.id, count).await?;
440        }
441    }
442    Ok(())
443}
alibi_sqlx/store/teams.rs

alibi_sqlx/store/
teams.rs
