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

Code before:
    return db_session.execute(select(Agent).filter_by(id=agent_id)).scalar_one_or_none()


def get_agent_by_api_key_id(db_session: Session, api_key_id: UUID) -> Agent | None:
    return db_session.execute(
        select(Agent).join(APIKey, Agent.id == APIKey.agent_id).filter(APIKey.id == str(api_key_id))
    ).scalar_one_or_none()


def get_agents_whose_allowed_apps_contains(db_session: Session, app_name: str) -> list[Agent]:
    statement = select(Agent).where(Agent.allowed_apps.contains([app_name]))
    return list(db_session.execute(statement).scalars().all())


def get_api_key_by_agent_id(db_session: Session, agent_id: UUID) -> APIKey | None:
    return db_session.execute(select(APIKey).filter_by(agent_id=agent_id)).scalar_one_or_none()


def get_api_key(db_session: Session, key: str) -> APIKey | None:
    key_hmac = encryption.hmac_sha256(key)
    return db_session.execute(select(APIKey).filter_by(key_hmac=key_hmac)).scalar_one_or_none()


def get_project_agent_org_by_api_key(
    db_session: Session, key: str
) -> tuple[UUID | None, UUID | None, UUID | None, UUID | None]:
    """
    Get project_id, agent_id, api_key_id, and org_id in a single query given an API key string.
    Returns a tuple of (api_key_id, agent_id, project_id, org_id) - any of which can be None if not found.
    """
    key_hmac = encryption.hmac_sha256(key)

    result = db_session.execute(
        select(APIKey.id, Agent.id, Project.id, Project.org_id)
        .join(Agent, APIKey.agent_id == Agent.id, isouter=True)
        .join(Project, Agent.project_id == Project.id, isouter=True)
        .filter(APIKey.key_hmac == key_hmac)
    ).first()

    if not result:
        return None, None, None, None

    api_key_id, agent_id, project_id, org_id = result
    return api_key_id, agent_id, project_id, org_id


def get_all_api_key_ids_for_project(db_session: Session, project_id: UUID) -> list[UUID]:
    agents = get_agents_by_project(db_session, project_id)
    project_api_key_ids = []
    for agent in agents:
        api_key = get_api_key_by_agent_id(db_session, agent.id)

Review comment (suggestion):
instead of `tuple[UUID | None, UUID | None, UUID | None, UUID | None]`,
it should just be `tuple[UUID,UUID,UUID,UUID] | None`

Code after:
    return db_session.execute(select(Agent).filter_by(id=agent_id)).scalar_one_or_none()


def get_agent_by_api_key_id(db_session: Session, api_key_id: UUID) -> Agent | None:
    return db_session.execute(
        select(Agent).join(APIKey, Agent.id == APIKey.agent_id).filter(APIKey.id == str(api_key_id))
    ).scalar_one_or_none()


def get_agents_whose_allowed_apps_contains(db_session: Session, app_name: str) -> list[Agent]:
    statement = select(Agent).where(Agent.allowed_apps.contains([app_name]))
    return list(db_session.execute(statement).scalars().all())


def get_api_key_by_agent_id(db_session: Session, agent_id: UUID) -> APIKey | None:
    return db_session.execute(select(APIKey).filter_by(agent_id=agent_id)).scalar_one_or_none()


def get_api_key(db_session: Session, key: str) -> APIKey | None:
    key_hmac = encryption.hmac_sha256(key)
    return db_session.execute(select(APIKey).filter_by(key_hmac=key_hmac)).scalar_one_or_none()


def get_request_context_by_api_key(
    db_session: Session, key: str
) -> tuple[UUID | None, UUID | None, UUID | None, UUID | None]:
    """
    Get project_id, agent_id, api_key_id, and org_id in a single query given an API key string.
    Returns a tuple of (api_key_id, agent_id, project_id, org_id) - any of which can be None if not found.
    """
    key_hmac = encryption.hmac_sha256(key)

    result = db_session.execute(
        select(APIKey.id, Agent.id, Project.id, Project.org_id)
        .join(Agent, APIKey.agent_id == Agent.id, isouter=True)
        .join(Project, Agent.project_id == Project.id, isouter=True)
        .filter(APIKey.key_hmac == key_hmac)
    ).first()

    if not result:
        return None, None, None, None

    api_key_id, agent_id, project_id, org_id = result
    return api_key_id, agent_id, project_id, org_id


def get_all_api_key_ids_for_project(db_session: Session, project_id: UUID) -> list[UUID]:
    agents = get_agents_by_project(db_session, project_id)
    project_api_key_ids = []
    for agent in agents:
        api_key = get_api_key_by_agent_id(db_session, agent.id)