cosmo/backend/app/services/celestial_body_service.py

"""Persistence operations for celestial bodies."""

from typing import Any, Dict, List, Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import AsyncSessionLocal
from app.models.db import CelestialBody


class CelestialBodyService:
    """Create, read, update, and delete celestial bodies."""

    @staticmethod
    async def get_all_bodies(
        session: Optional[AsyncSession] = None,
        body_type: Optional[str] = None,
        system_id: Optional[int] = None,
    ) -> List[CelestialBody]:
        async def query(active_session: AsyncSession) -> List[CelestialBody]:
            statement = select(CelestialBody)
            if body_type:
                statement = statement.where(CelestialBody.type == body_type)
            if system_id is not None:
                statement = statement.where(CelestialBody.system_id == system_id)
            result = await active_session.execute(statement.order_by(CelestialBody.name))
            return list(result.scalars().all())

        if session is not None:
            return await query(session)
        async with AsyncSessionLocal() as active_session:
            return await query(active_session)

    @staticmethod
    async def get_body_by_id(
        body_id: str,
        session: Optional[AsyncSession] = None,
    ) -> Optional[CelestialBody]:
        async def query(active_session: AsyncSession) -> Optional[CelestialBody]:
            result = await active_session.execute(
                select(CelestialBody).where(CelestialBody.id == body_id)
            )
            return result.scalar_one_or_none()

        if session is not None:
            return await query(session)
        async with AsyncSessionLocal() as active_session:
            return await query(active_session)

    @staticmethod
    async def create_body(
        body_data: Dict[str, Any],
        session: Optional[AsyncSession] = None,
    ) -> CelestialBody:
        async def create(active_session: AsyncSession) -> CelestialBody:
            body = CelestialBody(**body_data)
            active_session.add(body)
            await active_session.commit()
            await active_session.refresh(body)
            return body

        if session is not None:
            return await create(session)
        async with AsyncSessionLocal() as active_session:
            return await create(active_session)

    @staticmethod
    async def update_body(
        body_id: str,
        update_data: Dict[str, Any],
        session: Optional[AsyncSession] = None,
    ) -> Optional[CelestialBody]:
        async def update(active_session: AsyncSession) -> Optional[CelestialBody]:
            result = await active_session.execute(
                select(CelestialBody).where(CelestialBody.id == body_id)
            )
            body = result.scalar_one_or_none()
            if body is None:
                return None

            for key, value in update_data.items():
                if hasattr(body, key):
                    setattr(body, key, value)

            await active_session.commit()
            await active_session.refresh(body)
            return body

        if session is not None:
            return await update(session)
        async with AsyncSessionLocal() as active_session:
            return await update(active_session)

    @staticmethod
    async def delete_body(
        body_id: str,
        session: Optional[AsyncSession] = None,
    ) -> bool:
        async def delete(active_session: AsyncSession) -> bool:
            result = await active_session.execute(
                select(CelestialBody).where(CelestialBody.id == body_id)
            )
            body = result.scalar_one_or_none()
            if body is None:
                return False

            await active_session.delete(body)
            await active_session.commit()
            return True

        if session is not None:
            return await delete(session)
        async with AsyncSessionLocal() as active_session:
            return await delete(active_session)


celestial_body_service = CelestialBodyService()