Files
FabledCurator/backend/app/tasks/import_file.py
T
2026-05-19 11:14:49 -04:00

147 lines
5.2 KiB
Python

"""import_media_file task: imports one file via the Importer service and
updates the ImportTask state machine + ImportBatch counters atomically.
"""
from datetime import UTC, datetime
from pathlib import Path
from sqlalchemy import create_engine, select, update
from sqlalchemy.orm import sessionmaker
from ..celery_app import celery
from ..config import get_config
from ..models import ImportBatch, ImportSettings, ImportTask
from ..services.importer import Importer
from ..services.thumbnailer import Thumbnailer
def _sync_session_factory():
cfg = get_config()
engine = create_engine(cfg.database_url_sync, future=True, pool_pre_ping=True)
return sessionmaker(engine, expire_on_commit=False)
IMAGES_ROOT = Path("/images")
def _map_result_to_status(result):
"""(ImportTask.status, should_requeue_ml_and_thumb) for an ImportResult.
'superseded' = the kept row's file/ML changed → complete + re-derive.
'attached' = a non-art file preserved → complete, no ML/thumb."""
if result.status in ("imported", "superseded"):
return ("complete", True)
if result.status == "attached":
return ("complete", False)
if result.status == "skipped":
return ("skipped", False)
return ("failed", False)
@celery.task(name="backend.app.tasks.import_file.import_media_file", bind=True)
def import_media_file(self, import_task_id: int) -> dict:
"""Returns a dict so the eager-mode tests can assert without DB."""
SessionLocal = _sync_session_factory()
with SessionLocal() as session:
task = session.get(ImportTask, import_task_id)
if task is None:
return {"status": "missing", "task_id": import_task_id}
task.status = "processing"
task.started_at = datetime.now(UTC)
session.add(task)
session.commit()
settings = session.execute(
select(ImportSettings).where(ImportSettings.id == 1)
).scalar_one()
import_root = Path(settings.import_scan_path)
importer = Importer(
session=session,
images_root=IMAGES_ROOT,
import_root=import_root,
thumbnailer=Thumbnailer(images_root=IMAGES_ROOT),
settings=settings,
)
try:
result = importer.import_one(Path(task.source_path))
except Exception as exc: # pragma: no cover — pipeline crash
task.status = "failed"
task.error = f"{type(exc).__name__}: {exc}"
task.finished_at = datetime.now(UTC)
session.execute(
update(ImportBatch)
.where(ImportBatch.id == task.batch_id)
.values(failed=ImportBatch.failed + 1)
)
session.add(task)
session.commit()
raise
if result.status in ("imported", "superseded"):
task.status = "complete"
task.result_image_id = result.image_id
counter_col_name = "imported"
counter_col = ImportBatch.imported
elif result.status == "attached":
task.status = "complete"
counter_col_name = "attachments"
counter_col = ImportBatch.attachments
elif result.status == "skipped":
task.status = "skipped"
task.error = (
f"{result.skip_reason.value}: {result.error}"
if result.skip_reason
else result.error
)
task.result_image_id = result.image_id
counter_col_name = "skipped"
counter_col = ImportBatch.skipped
else:
task.status = "failed"
task.error = result.error
counter_col_name = "failed"
counter_col = ImportBatch.failed
task.finished_at = datetime.now(UTC)
session.execute(
update(ImportBatch)
.where(ImportBatch.id == task.batch_id)
.values({counter_col_name: counter_col + 1})
)
session.add(task)
session.commit()
# Enqueue thumbnail + ML for newly imported AND superseded images
# (a superseded row has cleared ML + no thumbnail).
if result.status in ("imported", "superseded"):
from .ml import tag_and_embed
from .thumbnail import generate_thumbnail
ids = list(result.member_image_ids)
if result.image_id is not None and result.image_id not in ids:
ids.append(result.image_id)
for img_id in ids:
generate_thumbnail.delay(img_id)
tag_and_embed.delay(img_id)
# If this was the last task in the batch, mark the batch complete.
remaining = session.execute(
select(ImportTask.id)
.where(ImportTask.batch_id == task.batch_id)
.where(ImportTask.status.in_(["pending", "queued", "processing"]))
.limit(1)
).scalar_one_or_none()
if remaining is None:
session.execute(
update(ImportBatch)
.where(ImportBatch.id == task.batch_id)
.values(
status="complete",
finished_at=datetime.now(UTC),
)
)
session.commit()
return {"task_id": import_task_id, "status": task.status}