mcma-backend/app/infrastructure/db/repositories/artist_repository.py

"""Artist repository — adapter over ``AsyncSession``."""

import uuid

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

from app.domain.entities.track import Artist
from app.infrastructure.db.models.album import AlbumModel
from app.infrastructure.db.models.artist import ArtistModel
from app.infrastructure.db.models.track import TrackModel


def _to_entity(row: ArtistModel) -> Artist:
    return Artist(
        id=row.id,
        name=row.name,
        source=row.source,
        source_id=row.source_id,
        created_at=row.created_at,
        updated_at=row.updated_at,
    )


class SqlAlchemyArtistRepository:
    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_or_create(self, name: str) -> Artist:
        row = (
            await self._session.execute(select(ArtistModel).where(ArtistModel.name == name))
        ).scalar_one_or_none()
        if row is None:
            row = ArtistModel(name=name)
            self._session.add(row)
            await self._session.flush()
            await self._session.refresh(row)
        return _to_entity(row)

    async def get_or_create_remote(self, *, name: str, source: str, source_id: str) -> Artist:
        """Resolve an artist by ``(source, source_id)`` first (re-browse/save
        dedup), falling back to ``name`` and gap-filling the remote ids onto an
        existing row, else creating a new remote-bound row."""
        row = (
            await self._session.execute(
                select(ArtistModel).where(
                    ArtistModel.source == source,
                    ArtistModel.source_id == source_id,
                )
            )
        ).scalar_one_or_none()
        if row is None:
            row = (
                await self._session.execute(select(ArtistModel).where(ArtistModel.name == name))
            ).scalar_one_or_none()
        if row is None:
            row = ArtistModel(name=name, source=source, source_id=source_id)
            self._session.add(row)
        elif row.source is None and row.source_id is None:
            row.source = source
            row.source_id = source_id
        await self._session.flush()
        await self._session.refresh(row)
        return _to_entity(row)

    async def get_by_id(self, artist_id: uuid.UUID) -> Artist | None:
        row = await self._session.get(ArtistModel, artist_id)
        return _to_entity(row) if row is not None else None

    async def get_many(self, ids: list[uuid.UUID]) -> list[Artist]:
        if not ids:
            return []
        rows = (
            (await self._session.execute(select(ArtistModel).where(ArtistModel.id.in_(ids))))
            .scalars()
            .all()
        )
        return [_to_entity(r) for r in rows]

    async def list_similar(self, *, artist_id: uuid.UUID, limit: int) -> list[Artist]:
        # Artists whose tracks fall in the seed artist's genres, ranked by how
        # many such tracks they have. Defined before ``list`` so the ``list[Artist]``
        # return annotation isn't shadowed by the method named ``list``.
        seed_genres = (
            select(TrackModel.genre)
            .where(TrackModel.artist_id == artist_id, TrackModel.genre.is_not(None))
            .distinct()
        )
        stmt = (
            select(ArtistModel)
            .join(TrackModel, TrackModel.artist_id == ArtistModel.id)
            .where(TrackModel.genre.in_(seed_genres), ArtistModel.id != artist_id)
            .group_by(ArtistModel.id)
            .order_by(func.count(TrackModel.id).desc())
            .limit(limit)
        )
        rows = (await self._session.execute(stmt)).scalars().all()
        return [_to_entity(r) for r in rows]

    async def list(self, *, q: str | None, limit: int, offset: int) -> list[Artist]:
        stmt = select(ArtistModel)
        if q:
            stmt = stmt.where(ArtistModel.name.ilike(f"%{q}%"))
        stmt = stmt.order_by(ArtistModel.name).limit(limit).offset(offset)
        rows = (await self._session.execute(stmt)).scalars().all()
        return [_to_entity(r) for r in rows]

    async def count(self, *, q: str | None) -> int:
        stmt = select(func.count()).select_from(ArtistModel)
        if q:
            stmt = stmt.where(ArtistModel.name.ilike(f"%{q}%"))
        return (await self._session.execute(stmt)).scalar_one()

    async def album_count(self, artist_id: uuid.UUID) -> int:
        return (
            await self._session.execute(
                select(func.count())
                .select_from(AlbumModel)
                .where(AlbumModel.artist_id == artist_id)
            )
        ).scalar_one()

    async def track_count(self, artist_id: uuid.UUID) -> int:
        return (
            await self._session.execute(
                select(func.count())
                .select_from(TrackModel)
                .where(TrackModel.artist_id == artist_id)
            )
        ).scalar_one()