alibi-sqlx 0.10.0

use super::entities::{team, team_member};
use super::{SqlxStore, lock_exclusive};
use crate::error::record_not_updated;
use crate::model::{self, ActiveRow, SqlxModel};
use crate::pool::{Engine, Exec, SqlxTransaction};
use crate::schema::{AuthSchema, SqlxUserModel};
use crate::sql::Sql;
use alibi_core::error::{AuthError, AuthResult};
use alibi_core::store::{TeamStore, team_membership_key};
use alibi_core::types::{AddTeamMemberResult, CreateTeam, Team, TeamMember, UpdateTeam};
use async_trait::async_trait;
use chrono::Utc;
use uuid::Uuid;

fn team_member_lookup(exec: Exec<'_>, team_id: &str, user_id: &str) -> Sql {
    let mut sql = model::select_model::<team_member::Model>(exec);
    sql.push(" WHERE ");
    sql.compare(team_member::Model::TABLE, "teamId", " = ", team_id);
    sql.push(" AND ");
    sql.compare(team_member::Model::TABLE, "userId", " = ", user_id);
    model::limit_one(&mut sql);
    sql
}

impl<S> SqlxStore<S>
where
    S: AuthSchema,
    S::User: SqlxUserModel,
{
    pub(super) async fn get_team_with_connection(
        &self,
        exec: Exec<'_>,
        organization_id: Option<&str>,
        team_id: &str,
    ) -> AuthResult<Option<Team>> {
        let mut sql = model::by_id::<team::Model>(exec, team_id);
        if let Some(org) = organization_id {
            sql.push(" AND ");
            sql.compare(team::Model::TABLE, "organizationId", " = ", org);
        }
        model::limit_one(&mut sql);
        Ok(exec
            .fetch_optional::<team::Model>(sql)
            .await?
            .map(Into::into))
    }

    pub(super) async fn add_team_member_in_tx(
        &self,
        tx: &SqlxTransaction,
        team_id: &str,
        user_id: &str,
        maximum: Option<f64>,
    ) -> AuthResult<AddTeamMemberResult> {
        let exec = Exec::tx(tx);
        if super::users::find_user_by_id::<S::User>(exec, user_id, super::users::Lock::Shared)
            .await?
            .is_none()
        {
            return Err(AuthError::UserNotFound);
        }
        let mut room = model::by_id::<team::Model>(exec, team_id);
        model::limit_one(&mut room);
        lock_exclusive(&mut room);
        let room = exec
            .fetch_optional::<team::Model>(room)
            .await?
            .ok_or_else(|| AuthError::bad_request("Team not found"))?;
        if let Some(member) = exec
            .fetch_optional::<team_member::Model>(team_member_lookup(exec, team_id, user_id))
            .await?
        {
            return Ok(AddTeamMemberResult::Existing(member.into()));
        }
        let mut count = Sql::with(exec.engine(), "SELECT COUNT(*) FROM ");
        count.ident(team_member::Model::TABLE);
        count.push(" WHERE ");
        count.compare(team_member::Model::TABLE, "teamId", " = ", team_id);
        let count = u64::try_from(exec.fetch_scalar::<i64>(count).await?.unwrap_or_default())
            .unwrap_or_default();
        // Source only repairs a stale low durable counter before reserving a seat.
        // Preserve high counters, and bind the raw policy to the database's `<`
        // predicate: notably SQLite binds NaN as NULL, which cannot reserve a seat.
        let reserved_count = i64::try_from(count)
            .map_err(|_error| AuthError::internal("Team membership count overflow"))?
            .max(room.member_count);
        let mut active = room.into_active();
        active.set("memberCount", reserved_count);
        _ = model::update::<team::Model>(exec, &active)
            .await?
            .ok_or_else(record_not_updated)?;
        if let Some(maximum) = maximum {
            let mut seat = Sql::with(exec.engine(), "SELECT COUNT(*) FROM ");
            seat.ident(team::Model::TABLE);
            seat.push(" WHERE ");
            seat.compare(team::Model::TABLE, "id", " = ", team_id);
            seat.push(" AND ");
            if exec.engine() == Engine::Postgres {
                // Source's text Number parameter is parsed as the physical
                // bigint counter, rather than promoting that counter to float8.
                seat.column(team::Model::TABLE, "memberCount");
                seat.push(" < ");
                seat.bind(ryu_js::Buffer::new().format(maximum).to_owned());
                seat.push("::bigint");
            } else {
                seat.compare(team::Model::TABLE, "memberCount", " < ", maximum);
            }
            if exec.fetch_scalar::<i64>(seat).await?.unwrap_or_default() == 0 {
                return Ok(AddTeamMemberResult::LimitReached);
            }
        }
        let key = team_membership_key(team_id, user_id)?;
        let mut member = ActiveRow::new();
        member.set(
            "id",
            self.generated_id(exec, "teamMember", "teamMember", "id")
                .await?
                .unwrap_or_else(|| Uuid::new_v4().to_string()),
        );
        member.set("teamId", team_id);
        member.set("userId", user_id);
        member.set("membershipKey", Some(key));
        member.set("createdAt", Utc::now());
        let member = model::insert::<team_member::Model>(exec, &member).await?;
        active.set(
            "memberCount",
            reserved_count
                .checked_add(1)
                .ok_or_else(|| AuthError::internal("Team membership count overflow"))?,
        );
        _ = model::update::<team::Model>(exec, &active)
            .await?
            .ok_or_else(record_not_updated)?;
        Ok(AddTeamMemberResult::Added(member.into()))
    }
}

#[async_trait]
impl<S> TeamStore for SqlxStore<S>
where
    S: AuthSchema,
    S::User: SqlxUserModel,
{
    async fn create_team(&self, data: CreateTeam) -> AuthResult<Team> {
        let now = data.updated_at.unwrap_or_else(Utc::now);
        let mut active = ActiveRow::new();
        active.set(
            "id",
            self.generated_id(self.exec(), "team", "team", "id")
                .await?
                .unwrap_or_else(|| Uuid::new_v4().to_string()),
        );
        active.set("name", data.name);
        active.set("organizationId", data.organization_id);
        active.set("memberCount", 0_i64);
        active.set("createdAt", now);
        active.set("updatedAt", data.updated_at);
        model::insert::<team::Model>(self.exec(), &active)
            .await
            .map(Into::into)
    }
    async fn get_team(
        &self,
        organization_id: Option<&str>,
        team_id: &str,
    ) -> AuthResult<Option<Team>> {
        self.get_team_with_connection(self.exec(), organization_id, team_id)
            .await
    }
    async fn list_teams(&self, organization_id: &str) -> AuthResult<Vec<Team>> {
        let mut sql = model::select_model::<team::Model>(self.exec());
        sql.push(" WHERE ");
        sql.compare(team::Model::TABLE, "organizationId", " = ", organization_id);
        sql.push(" LIMIT ");
        sql.bind(self.find_many_limit());
        Ok(self
            .exec()
            .fetch_all::<team::Model>(sql)
            .await?
            .into_iter()
            .map(Into::into)
            .collect())
    }
    async fn update_team(
        &self,
        organization_id: &str,
        team_id: &str,
        update: UpdateTeam,
    ) -> AuthResult<Team> {
        let mut sql = model::by_id::<team::Model>(self.exec(), team_id);
        sql.push(" AND ");
        sql.compare(team::Model::TABLE, "organizationId", " = ", organization_id);
        model::limit_one(&mut sql);
        let model = self
            .exec()
            .fetch_optional::<team::Model>(sql)
            .await?
            .ok_or_else(|| AuthError::bad_request("Team not found"))?;
        let mut active = model.into_active();
        if let Some(name) = update.name {
            active.set("name", name);
        }
        active.set("updatedAt", Some(Utc::now()));
        model::update::<team::Model>(self.exec(), &active)
            .await?
            .map(Into::into)
            .ok_or_else(record_not_updated)
    }
    async fn delete_team(&self, organization_id: &str, team_id: &str) -> AuthResult<bool> {
        let (organization_id, team_id) = (organization_id.to_owned(), team_id.to_owned());
        self.in_transaction(true, async move |tx| {
            let exec = Exec::tx(tx);
            let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
            delete.ident(team::Model::TABLE);
            delete.push(" WHERE ");
            delete.compare(team::Model::TABLE, "id", " = ", team_id.as_str());
            delete.push(" AND ");
            delete.compare(
                team::Model::TABLE,
                "organizationId",
                " = ",
                organization_id.as_str(),
            );
            if exec.execute(delete).await? == 0 {
                return Ok(false);
            }
            let mut members = Sql::with(exec.engine(), "DELETE FROM ");
            members.ident(team_member::Model::TABLE);
            members.push(" WHERE ");
            members.compare(team_member::Model::TABLE, "teamId", " = ", team_id.as_str());
            _ = exec.execute(members).await?;
            let mut pending = self.organization_models.invitation.select(exec);
            pending.push(" WHERE ");
            self.organization_models.invitation.compare(
                &mut pending,
                "organizationId",
                " = ",
                organization_id.as_str(),
            )?;
            pending.push(" AND ");
            self.organization_models.invitation.compare(
                &mut pending,
                "status",
                " = ",
                "pending",
            )?;
            pending.push(" AND ");
            self.organization_models.invitation.compare(
                &mut pending,
                "expiresAt",
                " > ",
                Utc::now(),
            )?;
            for invite in self
                .organization_models
                .invitation
                .fetch_all(exec, pending)
                .await?
            {
                let Some(ids) = invite.team_id.as_deref() else {
                    continue;
                };
                if !ids.split(',').any(|id| id == team_id) {
                    continue;
                }
                let remaining = ids
                    .split(',')
                    .filter(|id| *id != team_id)
                    .collect::<Vec<_>>()
                    .join(",");
                let mut active = invite.into_active();
                active.set("teamId", (!remaining.is_empty()).then_some(remaining));
                _ = self
                    .organization_models
                    .invitation
                    .update(exec, &active)
                    .await?
                    .ok_or_else(record_not_updated)?;
            }
            Ok(true)
        })
        .await
    }
    async fn get_team_member(
        &self,
        team_id: &str,
        user_id: &str,
    ) -> AuthResult<Option<TeamMember>> {
        Ok(self
            .exec()
            .fetch_optional::<team_member::Model>(team_member_lookup(self.exec(), team_id, user_id))
            .await?
            .map(Into::into))
    }
    async fn add_team_member(
        &self,
        team_id: &str,
        user_id: &str,
        maximum: Option<f64>,
    ) -> AuthResult<AddTeamMemberResult> {
        let tx = self.begin(true).await?;
        let result = self
            .add_team_member_in_tx(&tx, team_id, user_id, maximum)
            .await?;
        tx.commit().await?;
        Ok(result)
    }
    async fn remove_team_member(&self, team_id: &str, user_id: &str) -> AuthResult<usize> {
        let (team_id, user_id) = (team_id.to_owned(), user_id.to_owned());
        let removed = self
            .in_transaction(true, async move |tx| {
                let exec = Exec::tx(tx);
                let mut room = model::by_id::<team::Model>(exec, team_id.as_str());
                model::limit_one(&mut room);
                lock_exclusive(&mut room);
                _ = exec.fetch_optional::<team::Model>(room).await?;
                let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
                delete.ident(team_member::Model::TABLE);
                delete.push(" WHERE ");
                delete.compare(team_member::Model::TABLE, "teamId", " = ", team_id.as_str());
                delete.push(" AND ");
                delete.compare(team_member::Model::TABLE, "userId", " = ", user_id.as_str());
                let removed = exec.execute(delete).await?;
                let count = i64::try_from(removed)
                    .map_err(|_error| AuthError::internal("Team membership count overflow"))?;
                if count > 0 {
                    release_seats(exec, &team_id, count).await?;
                }
                Ok(removed)
            })
            .await?;
        usize::try_from(removed)
            .map_err(|_error| AuthError::internal("Team membership count overflow"))
    }
    async fn list_team_members(&self, team_id: &str) -> AuthResult<Vec<TeamMember>> {
        let mut sql = model::select_model::<team_member::Model>(self.exec());
        sql.push(" WHERE ");
        sql.compare(team_member::Model::TABLE, "teamId", " = ", team_id);
        sql.push(" LIMIT ");
        sql.bind(self.find_many_limit());
        Ok(self
            .exec()
            .fetch_all::<team_member::Model>(sql)
            .await?
            .into_iter()
            .map(Into::into)
            .collect())
    }
    async fn list_user_teams(&self, user_id: &str) -> AuthResult<Vec<Team>> {
        let mut sql = model::select_model::<team_member::Model>(self.exec());
        sql.push(" WHERE ");
        sql.compare(team_member::Model::TABLE, "userId", " = ", user_id);
        sql.push(" LIMIT ");
        sql.bind(self.find_many_limit());
        let memberships: Vec<team_member::Model> = self.exec().fetch_all(sql).await?;
        if memberships.is_empty() {
            return Ok(Vec::new());
        }
        let mut rooms = model::select_model::<team::Model>(self.exec());
        rooms.push(" WHERE ");
        rooms.column(team::Model::TABLE, "id");
        rooms.push(" IN ");
        rooms.bind_list(memberships.iter().map(|row| row.team_id.clone()));
        let rooms = self
            .exec()
            .fetch_all::<team::Model>(rooms)
            .await?
            .into_iter()
            .map(|row| (row.id.clone(), row))
            .collect::<std::collections::HashMap<_, _>>();
        // The upstream join maps membership rows directly, preserving their
        // adapter order. An IN query's team order must not replace that order.
        Ok(memberships
            .into_iter()
            .filter_map(|membership| rooms.get(&membership.team_id).cloned().map(Into::into))
            .collect())
    }
}

/// `UPDATE team SET member_count = member_count - n WHERE id = ? AND member_count >= n`.
async fn release_seats(exec: Exec<'_>, team_id: &str, count: i64) -> AuthResult<()> {
    let mut sql = Sql::with(exec.engine(), "UPDATE ");
    sql.ident(team::Model::TABLE);
    sql.push(" SET ");
    sql.ident("memberCount");
    sql.push(" = ");
    sql.ident("memberCount");
    sql.push(" - ");
    sql.bind(count);
    sql.push(" WHERE ");
    sql.compare(team::Model::TABLE, "id", " = ", team_id);
    sql.push(" AND ");
    sql.compare(team::Model::TABLE, "memberCount", " >= ", count);
    _ = exec.execute(sql).await?;
    Ok(())
}

/// Delete owned memberships and release exactly the seats those rows occupied.
pub(super) async fn remove_owned_team_members(
    tx: &SqlxTransaction,
    user_id: &str,
    organization_id: Option<&str>,
) -> AuthResult<()> {
    let exec = Exec::tx(tx);
    if !super::migrator::has_table(exec, "teamMember").await? {
        return Ok(());
    }
    let mut teams = model::select_model::<team::Model>(exec);
    if let Some(org) = organization_id {
        teams.push(" WHERE ");
        teams.compare(team::Model::TABLE, "organizationId", " = ", org);
    }
    teams.push(" ORDER BY ");
    teams.column(team::Model::TABLE, "id");
    teams.push(" ASC");
    lock_exclusive(&mut teams);
    let rooms = exec.fetch_all::<team::Model>(teams).await?;
    release_owned_team_members(tx, user_id, rooms).await
}

/// Remove memberships from an already selected adapter page and release seats.
pub(super) async fn release_owned_team_members(
    tx: &SqlxTransaction,
    user_id: &str,
    rooms: Vec<team::Model>,
) -> AuthResult<()> {
    let exec = Exec::tx(tx);
    for room in rooms {
        let mut delete = Sql::with(exec.engine(), "DELETE FROM ");
        delete.ident(team_member::Model::TABLE);
        delete.push(" WHERE ");
        delete.compare(team_member::Model::TABLE, "teamId", " = ", room.id.as_str());
        delete.push(" AND ");
        delete.compare(team_member::Model::TABLE, "userId", " = ", user_id);
        let deleted = exec.execute(delete).await?;
        if deleted > 0 {
            let count = i64::try_from(deleted)
                .map_err(|_error| AuthError::internal("Team membership count overflow"))?;
            release_seats(exec, &room.id, count).await?;
        }
    }
    Ok(())
}