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

alibi_sqlx/store/
teams.rs
