ss-tools/backend/src/plugins/translate/orchestrator_retry.py

# #region TranslationRunRetryManager [C:4] [TYPE Class] [SEMANTICS translate, retry]
# @defgroup Translate Module group.
# @BRIEF Manage retry of translation run batches.
# @PRE Database session and config manager are available.
# @POST Retried batches are re-executed.
# @SIDE_EFFECT DB writes, event log entries.
# @RELATION DEPENDS_ON -> [TranslationRun]
# @RELATION DEPENDS_ON -> [TranslationBatch]
# @RELATION DEPENDS_ON -> [TranslationRecord]
# @RELATION DEPENDS_ON -> [TranslationExecutor]
# @RELATION DEPENDS_ON -> [TranslationEventLog]
# @RATIONALE Cancel and retry-insert extracted to orchestrator_cancel module for INV_7 compliance.
# @REJECTED Keeping cancel and retry-insert logic in the same class was rejected — the module exceeded INV_7 (400 lines); extracting cancel to orchestrator_cancel keeps each module independently maintainable and testable.

from datetime import UTC, datetime

from sqlalchemy.orm import Session

from ...core.config_manager import ConfigManager
from ...core.logger import belief_scope, logger
from ...models.translate import (
    TranslationBatch,
    TranslationJob,
    TranslationRecord,
    TranslationRun,
)
from .events import TranslationEventLog
from .executor import TranslationExecutor
from .orchestrator_cancel import cancel_run as _cancel_run, retry_insert as _retry_insert


class TranslationRunRetryManager:
    """Manage retry and cancellation of translation runs."""

    def __init__(
        self,
        db: Session,
        config_manager: ConfigManager,
        event_log: TranslationEventLog,
        current_user: str | None = None,
    ):
        self.db = db
        self.config_manager = config_manager
        self.event_log = event_log
        self.current_user = current_user

    # region retry_failed_batches [TYPE Function]
    # @PURPOSE: Retry failed batches in a run.
    # @SIDE_EFFECT  Re-executes batch translations; DB writes.
    async def retry_failed_batches(self, run_id: str) -> TranslationRun:
        with belief_scope("TranslationRunRetryManager.retry_failed_batches"):
            run = self.db.query(TranslationRun).filter(TranslationRun.id == run_id).first()
            if not run:
                raise ValueError(f"Run '{run_id}' not found")
            job = self.db.query(TranslationJob).filter(TranslationJob.id == run.job_id).first()
            if not job:
                raise ValueError(f"Job '{run.job_id}' not found")

            failed_batches = (
                self.db.query(TranslationBatch)
                .filter(
                    TranslationBatch.run_id == run_id,
                    TranslationBatch.status.in_(["FAILED", "COMPLETED_WITH_ERRORS"]),
                )
                .all()
            )
            if not failed_batches:
                raise ValueError(f"No failed batches found for run '{run_id}'")

            logger.reason("Retrying failed batches", {"run_id": run_id, "batch_count": len(failed_batches)})
            self.event_log.log_event(
                job_id=job.id, run_id=run.id, event_type="RUN_RETRYING",
                payload={"batch_count": len(failed_batches)}, created_by=self.current_user,
            )

            executor = TranslationExecutor(self.db, self.config_manager, self.current_user)
            run.status = "RUNNING"
            self.db.flush()

            for batch in failed_batches:
                failed_records = (
                    self.db.query(TranslationRecord)
                    .filter(
                        TranslationRecord.batch_id == batch.id,
                        TranslationRecord.status == "FAILED",
                    )
                    .all()
                )
                if not failed_records:
                    continue
                retry_rows = [
                    {
                        "row_index": rec.source_object_id or "0",
                        "source_text": rec.source_sql or "",
                        "approved_translation": None,
                        "source_object_name": rec.source_object_name or "",
                    }
                    for rec in failed_records
                ]
                result = await executor._process_batch(
                    job=job, run_id=run_id, batch_index=batch.batch_index,
                    batch_rows=retry_rows,
                )
                run.successful_records = (run.successful_records or 0) + result["successful"]
                run.failed_records = (run.failed_records or 0) + result["failed"]
                run.skipped_records = (run.skipped_records or 0) + result["skipped"]
                self.db.flush()

            run.status = (
                "COMPLETED" if run.failed_records == 0
                else ("FAILED" if run.successful_records == 0 else "COMPLETED")
            )
            run.completed_at = datetime.now(UTC)
            self.db.flush()

            self.event_log.log_event(
                job_id=job.id, run_id=run.id, event_type="RUN_COMPLETED",
                payload={"retry": True}, created_by=self.current_user,
            )
            self.db.commit()
            return run
    # endregion retry_failed_batches

    # region retry_insert [TYPE Function]
    # @PURPOSE: Retry the SQL insert phase for a completed run.
    # @SIDE_EFFECT  Superset API call; DB writes.
    async def retry_insert(self, run_id: str) -> TranslationRun:
        return await _retry_insert(self.db, self.config_manager, self.event_log, self.current_user, run_id)
    # endregion retry_insert

    # region cancel_run [TYPE Function]
    # @PURPOSE: Cancel a running translation.
    # @SIDE_EFFECT  DB writes; event log.
    def cancel_run(self, run_id: str) -> TranslationRun:
        return _cancel_run(self.db, self.event_log, self.current_user, run_id)
    # endregion cancel_run

# #endregion TranslationRunRetryManager