Repository: aipotheosis-labs/aci
File: backend/aci/common/db/crud/linked_accounts.py (Python)

Code before:

def update_linked_account_last_used_at(
    db_session: Session,
    last_used_at: datetime,
    linked_account: LinkedAccount,
) -> LinkedAccount:
    linked_account.last_used_at = last_used_at
    db_session.flush()
    db_session.refresh(linked_account)
    return linked_account


def delete_linked_accounts(db_session: Session, project_id: UUID, app_name: str) -> int:
    statement = (
        select(LinkedAccount)
        .join(App, LinkedAccount.app_id == App.id)
        .filter(LinkedAccount.project_id == project_id, App.name == app_name)
    )
    linked_accounts_to_delete = db_session.execute(statement).scalars().all()
    for linked_account in linked_accounts_to_delete:
        db_session.delete(linked_account)
    db_session.flush()
    return len(linked_accounts_to_delete)


def get_total_number_of_unique_linked_account_owner_ids(db_session: Session, org_id: UUID) -> int:
    statement = select(func.count(distinct(LinkedAccount.linked_account_owner_id))).where(
        LinkedAccount.project_id.in_(select(Project.id).filter(Project.org_id == org_id))
    )
    return db_session.execute(statement).scalar_one()


def linked_account_owner_id_exists_in_org(
    db_session: Session, org_id: UUID, linked_account_owner_id: str
) -> bool:
    statement = select(LinkedAccount).filter(
        LinkedAccount.linked_account_owner_id == linked_account_owner_id,
        LinkedAccount.project_id.in_(select(Project.id).filter(Project.org_id == org_id)),
    )
    return db_session.execute(statement).scalar_one_or_none() is not None


Review comment (bug):
technically, both sql operations here need locking to avoid a race condition. (if multiple requests hit at the same time, they could all see a count below the quota, and proceed to insert, exceeding the quota.)
Either a mutex on application level or row-level locking here.
However, this might not be super necessary as it's pretty low stake. (so maybe just leave a comment at where this is called, specifying the potential race condition?)

Code after:

def update_linked_account_last_used_at(
    db_session: Session,
    last_used_at: datetime,
    linked_account: LinkedAccount,
) -> LinkedAccount:
    linked_account.last_used_at = last_used_at
    db_session.flush()
    db_session.refresh(linked_account)
    return linked_account


def delete_linked_accounts(db_session: Session, project_id: UUID, app_name: str) -> int:
    statement = (
        select(LinkedAccount)
        .join(App, LinkedAccount.app_id == App.id)
        .filter(LinkedAccount.project_id == project_id, App.name == app_name)
    )
    linked_accounts_to_delete = db_session.execute(statement).scalars().all()
    for linked_account in linked_accounts_to_delete:
        db_session.delete(linked_account)
    db_session.flush()
    return len(linked_accounts_to_delete)


def get_total_number_of_unique_linked_account_owner_ids(db_session: Session, org_id: UUID) -> int:
    """
    TODO: Add a lock to prevent the race condition.
    Get the total number of unique linked account owner IDs for an organization.

    WARNING: Race condition potential! This function is vulnerable to race conditions in
    concurrent environments. If this function is called concurrently with operations that
    add or remove linked accounts:

    1. Thread A starts counting unique linked_account_owner_ids
    2. Thread B adds a new linked account with a new owner_id
    3. Thread A completes its count, unaware of the newly added account
    """
    statement = select(func.count(distinct(LinkedAccount.linked_account_owner_id))).where(
        LinkedAccount.project_id.in_(select(Project.id).filter(Project.org_id == org_id))
    )
    return db_session.execute(statement).scalar_one()


def linked_account_owner_id_exists_in_org(
    db_session: Session, org_id: UUID, linked_account_owner_id: str
) -> bool:
    statement = select(LinkedAccount).filter(
        LinkedAccount.linked_account_owner_id == linked_account_owner_id,
        LinkedAccount.project_id.in_(select(Project.id).filter(Project.org_id == org_id)),
    )