FabledCurator/backend/app/services/migrators/cleanup.py

"""Targeted cleanup migrator: delete every image attributed to one Artist.

Built for the IR-migration rescue case where the filesystem scan derived
a bogus 'imagerepo' artist from a mismatched bind-mount layout. Every
image attributed to that artist (40k+ rows) needs to be removed — DB
rows, original files under `/images/<bucket>/...`, and thumbnails under
`/images/thumbs/...` — before the operator remounts and re-scans.

CASCADE handles image_tag, image_provenance, series_page, and
tag_suggestion_rejection child rows; import_task.result_image_id is
SET NULL by FK. We also delete ImportTask rows whose source_path starts
with the (still-existing) IR scan prefix so the next scan isn't fooled
by them.
"""
from __future__ import annotations

import logging
from pathlib import Path

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

from ...models import Artist, ImageRecord, ImportBatch, ImportTask

log = logging.getLogger(__name__)

_BATCH_SIZE = 500


def _zero_counts() -> dict:
    return {
        "rows_processed": 0, "rows_inserted": 0, "rows_skipped": 0,
        "files_copied": 0, "bytes_copied": 0, "conflicts": 0,
    }


def _thumb_path(images_root: Path, sha256_hex: str) -> tuple[Path, Path]:
    """Return both possible thumbnail paths (.jpg and .png). We try both
    because the extension is chosen at generate-time based on the source
    image's mode (alpha → .png, otherwise → .jpg)."""
    bucket = sha256_hex[:3]
    base = images_root / "thumbs" / bucket / sha256_hex
    return base.with_suffix(".jpg"), base.with_suffix(".png")


def _delete_file(path: Path) -> bool:
    """Best-effort unlink; True if the file was actually removed."""
    try:
        path.unlink(missing_ok=True)
        return True
    except OSError as exc:
        log.warning("cleanup: failed to unlink %s: %s", path, exc)
        return False


async def cleanup_artist_async(
    db: AsyncSession,
    *,
    slug: str,
    images_root: Path | None = None,
    dry_run: bool = False,
    source_path_prefix: str | None = None,
) -> dict:
    """Delete every image attributed to the Artist with this slug,
    along with the artist row itself and any associated import tasks.

    Args:
        slug: artist.slug to target (e.g. 'imagerepo').
        images_root: defaults to /images.
        dry_run: skip filesystem + DB writes; still walk rows for counts.
        source_path_prefix: if set, ImportTask rows whose source_path
            starts with this string are deleted too (use the IR scan
            mount prefix, e.g. '/import/imagerepo').
    """
    root = images_root if images_root is not None else Path("/images")

    artist = (await db.execute(
        select(Artist).where(Artist.slug == slug)
    )).scalar_one_or_none()
    if artist is None:
        raise ValueError(f"no Artist with slug={slug!r}")

    artist_id = artist.id
    artist_name = artist.name

    total_images = (await db.execute(
        select(func.count(ImageRecord.id)).where(ImageRecord.artist_id == artist_id)
    )).scalar_one()

    counts = _zero_counts()
    files_deleted = 0
    thumbs_deleted = 0
    images_deleted = 0

    # Batched delete loop. CASCADE handles image_tag, image_provenance,
    # series_page, tag_suggestion_rejection. import_task.result_image_id
    # is SET NULL by FK.
    while True:
        rows = (await db.execute(
            select(ImageRecord.id, ImageRecord.path, ImageRecord.sha256)
            .where(ImageRecord.artist_id == artist_id)
            .limit(_BATCH_SIZE)
        )).all()
        if not rows:
            break

        ids = [r.id for r in rows]
        counts["rows_processed"] += len(ids)

        if not dry_run:
            for r in rows:
                if r.path:
                    if _delete_file(Path(r.path)):
                        files_deleted += 1
                if r.sha256:
                    jpg, png = _thumb_path(root, r.sha256)
                    if _delete_file(jpg):
                        thumbs_deleted += 1
                    if _delete_file(png):
                        thumbs_deleted += 1

            await db.execute(
                delete(ImageRecord).where(ImageRecord.id.in_(ids))
            )
            await db.commit()

        images_deleted += len(ids)

        if dry_run:
            # Nothing was actually deleted from the DB; bail after one
            # pass so we don't loop forever.
            break

    import_tasks_deleted = 0
    if source_path_prefix and not dry_run:
        # Delete ImportTask rows whose source_path is under the bad mount
        # prefix. These are mostly orphaned now (result_image_id was set
        # NULL by CASCADE) but their presence still blocks the
        # idempotency check in scan_directory if the operator remounts
        # the same prefix.
        like_pattern = source_path_prefix.rstrip("/") + "/%"
        result = await db.execute(
            delete(ImportTask).where(ImportTask.source_path.like(like_pattern))
        )
        import_tasks_deleted = result.rowcount or 0
        await db.commit()

    # Sweep ImportBatch rows that are now empty.
    empty_batches_deleted = 0
    if not dry_run:
        empty_batch_ids = (await db.execute(
            select(ImportBatch.id).where(
                ~select(ImportTask.id)
                .where(ImportTask.batch_id == ImportBatch.id)
                .exists()
            )
        )).scalars().all()
        if empty_batch_ids:
            result = await db.execute(
                delete(ImportBatch).where(ImportBatch.id.in_(empty_batch_ids))
            )
            empty_batches_deleted = result.rowcount or 0
            await db.commit()

    # Finally, the artist row.
    if not dry_run:
        await db.execute(delete(Artist).where(Artist.id == artist_id))
        await db.commit()

    return {
        "counts": counts,
        "artist": {"id": artist_id, "name": artist_name, "slug": slug},
        "summary": {
            "images_targeted": total_images,
            "images_deleted": images_deleted,
            "files_deleted": files_deleted,
            "thumbs_deleted": thumbs_deleted,
            "import_tasks_deleted": import_tasks_deleted,
            "empty_batches_deleted": empty_batches_deleted,
            "dry_run": dry_run,
        },
    }