partners.py

"""A module that returns the partners available in :term:`GSO`."""

from typing import Any

from orchestrator.db import db
from sqlalchemy.exc import NoResultFound, SQLAlchemyError

from gso.db.models import PartnerTable
from gso.schema.partner import PartnerCreate


class PartnerNotFoundError(Exception):
    """Exception raised when a partner is not found."""


def get_all_partners() -> list[dict]:
    """Fetch all partners from the database and serialize them to JSON."""
    partners = PartnerTable.query.all()
    return [partner.__json__() for partner in partners]


def get_partner_by_name(name: str) -> dict[str, Any]:
    """Try to get a partner by their name."""
    try:
        partner = PartnerTable.query.filter(PartnerTable.name == name).one()
        return partner.__json__()
    except NoResultFound as e:
        msg = f"partner {name} not found"
        raise PartnerNotFoundError(msg) from e


def create_partner(
    partner_data: PartnerCreate,
) -> dict:
    """Create a new partner and add it to the database using Pydantic schema for validation.

    :param partner_data: Partner data validated by Pydantic schema.
    :return: JSON representation of the created partner.
    """
    try:
        new_partner = PartnerTable(**partner_data.model_dump())

        db.session.add(new_partner)
        db.session.commit()

        return new_partner.__json__()
    except SQLAlchemyError:
        db.session.rollback()
        raise
    finally:
        db.session.close()


def delete_partner_by_name(name: str) -> None:
    """Delete a partner by their name."""
    try:
        partner = PartnerTable.query.filter(PartnerTable.name == name).one()
        db.session.delete(partner)
        db.session.commit()
    except NoResultFound as e:
        msg = f"partner {name} not found"
        raise PartnerNotFoundError(msg) from e
    except SQLAlchemyError:
        db.session.rollback()
        raise
    finally:
        db.session.close()