"""Alias resolution + CRUD. A tag_alias maps (model_name, model_category) -> canonical Tag. Resolution happens at suggestion-read time so the raw image_prediction rows stay unmolested. """ from collections.abc import Sequence from dataclasses import dataclass from sqlalchemy import delete, select from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession from ...models import Tag, TagAlias @dataclass(frozen=True) class AliasRow: alias_string: str alias_category: str canonical_tag_id: int canonical_tag_name: str class AliasService: def __init__(self, session: AsyncSession): self.session = session async def resolve(self, name: str, category: str) -> Tag | None: """Return the canonical Tag for (name, category), or None if no alias.""" stmt = ( select(Tag) .join(TagAlias, TagAlias.canonical_tag_id == Tag.id) .where(TagAlias.alias_string == name) .where(TagAlias.alias_category == category) ) return (await self.session.execute(stmt)).scalar_one_or_none() async def resolve_many( self, pairs: list[tuple[str, str]] ) -> dict[tuple[str, str], Tag]: """Batch-resolve. Returns only the pairs that have an alias. Used by SuggestionService so it does one query instead of N. """ if not pairs: return {} strings = {p[0] for p in pairs} stmt = ( select(TagAlias, Tag) .join(Tag, Tag.id == TagAlias.canonical_tag_id) .where(TagAlias.alias_string.in_(strings)) ) rows = (await self.session.execute(stmt)).all() wanted = set(pairs) out: dict[tuple[str, str], Tag] = {} for alias, tag in rows: key = (alias.alias_string, alias.alias_category) if key in wanted: out[key] = tag return out async def create( self, alias_string: str, alias_category: str, canonical_tag_id: int ) -> None: """Idempotent create (ON CONFLICT DO NOTHING).""" stmt = insert(TagAlias).values( alias_string=alias_string, alias_category=alias_category, canonical_tag_id=canonical_tag_id, ) stmt = stmt.on_conflict_do_nothing( index_elements=["alias_string", "alias_category"] ) await self.session.execute(stmt) async def remove(self, alias_string: str, alias_category: str) -> None: await self.session.execute( delete(TagAlias) .where(TagAlias.alias_string == alias_string) .where(TagAlias.alias_category == alias_category) ) async def list_for_tag(self, canonical_tag_id: int) -> Sequence[AliasRow]: """Aliases that resolve TO this tag — drives the tag-side 'Aliases' view (see/remove the model keys that fold into a tag).""" stmt = ( select( TagAlias.alias_string, TagAlias.alias_category, TagAlias.canonical_tag_id, Tag.name, ) .join(Tag, Tag.id == TagAlias.canonical_tag_id) .where(TagAlias.canonical_tag_id == canonical_tag_id) .order_by(TagAlias.alias_string.asc()) ) rows = (await self.session.execute(stmt)).all() return [ AliasRow( alias_string=r[0], alias_category=r[1], canonical_tag_id=r[2], canonical_tag_name=r[3], ) for r in rows ] async def list_all(self) -> Sequence[AliasRow]: stmt = ( select( TagAlias.alias_string, TagAlias.alias_category, TagAlias.canonical_tag_id, Tag.name, ) .join(Tag, Tag.id == TagAlias.canonical_tag_id) .order_by(TagAlias.alias_string.asc()) ) rows = (await self.session.execute(stmt)).all() return [ AliasRow( alias_string=r[0], alias_category=r[1], canonical_tag_id=r[2], canonical_tag_name=r[3], ) for r in rows ]