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

# #region RunSourceFetcher [C:3] [TYPE Module] [SEMANTICS translate, source, fetch, superset, preview]
# @defgroup Translate Module group.
# @BRIEF Fetch source rows for translation runs from Superset datasource or preview session.
#         Extracted from RunExecutionService to comply with INV_7 (< 400 lines/module).
# @LAYER Domain
# @RELATION DEPENDS_ON -> [TranslationJob], [TranslationPreviewSession]
# @RELATION DEPENDS_ON -> [SupersetClient]
# @RELATION DEPENDS_ON -> [ConfigManager]
# @PRE job_id exists in DB.
# @POST Returns list of dicts with source data (text, row_index, approved_translation, source_data).
# @SIDE_EFFECT Makes HTTP call to Superset chart data API when datasource is configured.
# @RATIONALE Extracted from RunExecutionService (620 lines) — _fetch_source_rows alone was ~100 lines.
# @REJECTED Inline fetch logic in execute_run — made execute_run exceed 150 lines.

import json
from typing import Any

from sqlalchemy.orm import Session

from ...core.config_manager import ConfigManager
from ...core.logger import belief_scope, logger
from ...models.translate import TranslationJob, TranslationPreviewRecord, TranslationPreviewSession

MAX_ROWS_PER_RUN = 10000


# #region fetch_source_rows [C:3] [TYPE Function] [SEMANTICS translate, source, fetch]
# @ingroup Translate
# @BRIEF Fetch source rows from Superset datasource or preview session fallback.
async def fetch_source_rows(db: Session, config_manager: ConfigManager, job_id: str, run_id: str) -> list[dict[str, Any]]:
    """Fetch source rows from Superset datasource or preview session fallback."""
    with belief_scope("RunSourceFetcher.fetch_source_rows"):
        job = db.query(TranslationJob).filter(TranslationJob.id == job_id).first()

        if job and job.source_datasource_id:
            try:
                logger.reason("Fetching full dataset from Superset datasource", {
                    "run_id": run_id, "datasource_id": job.source_datasource_id,
                    "environment_id": job.environment_id,
                })
                environments = config_manager.get_environments()
                target_env_id = job.environment_id or job.source_dialect or ""
                env_config = next((e for e in environments if e.id == target_env_id), None)
                if not env_config and environments:
                    env_config = environments[0]

                if env_config:
                    from ...core.utils.client_registry import get_superset_client
                    client = await get_superset_client(env_config)
                    dataset_detail = await client.get_dataset_detail(int(job.source_datasource_id))
                    query_context = client.build_dataset_preview_query_context(
                        dataset_id=int(job.source_datasource_id), dataset_record=dataset_detail,
                        template_params={}, effective_filters=[],
                    )
                    queries = query_context.get("queries", [])
                    if queries:
                        queries[0]["row_limit"] = MAX_ROWS_PER_RUN
                        queries[0].pop("result_type", None)
                        queries[0]["metrics"] = []
                    query_context["result_type"] = "samples"
                    form_data = query_context.get("form_data", {})
                    form_data.pop("query_mode", None)

                    response = await client.network.request(
                        method="POST", endpoint="/api/v1/chart/data",
                        data=json.dumps(query_context), headers={"Content-Type": "application/json"},
                    )
                    rows = _extract_chart_data_rows(response)
                    if rows:
                        logger.reason(f"Fetched {len(rows)} rows from Superset datasource", {"run_id": run_id})
                        source_rows = []
                        for idx, row in enumerate(rows):
                            sd = dict(row) if row else None
                            st = str(row.get(job.translation_column, "")) if job.translation_column else json.dumps(row)
                            source_rows.append({
                                "row_index": str(idx), "source_text": st,
                                "approved_translation": None, "source_object_name": f"Row {idx}",
                                "source_data": sd,
                            })
                        return source_rows
                    logger.explore("Superset datasource returned no rows", {"run_id": run_id, "datasource_id": job.source_datasource_id})
                else:
                    logger.explore("No environment config found", {"env_id": target_env_id})
            except Exception as e:
                logger.explore("Failed to fetch full dataset from Superset, falling back to preview",
                               {"run_id": run_id, "error": str(e)})

        session = (
            db.query(TranslationPreviewSession)
            .filter(TranslationPreviewSession.job_id == job_id, TranslationPreviewSession.status == "APPLIED")
            .order_by(TranslationPreviewSession.created_at.desc())
            .first()
        )
        if not session:
            logger.explore("No accepted preview session found", {"job_id": job_id})
            return []

        records = (
            db.query(TranslationPreviewRecord)
            .filter(TranslationPreviewRecord.session_id == session.id,
                    TranslationPreviewRecord.status.in_(["APPROVED", "PENDING"]))
            .all()
        )
        source_rows = []
        for rec in records:
            sd = None
            if hasattr(rec, "source_data") and rec.source_data:
                sd = dict(rec.source_data)
            source_rows.append({
                "row_index": rec.source_object_id or "0",
                "source_text": rec.source_sql or "",
                "approved_translation": rec.target_sql if rec.status == "APPROVED" else None,
                "source_object_name": rec.source_object_name or "",
                "source_data": sd,
            })
        logger.reason(f"Fetched {len(source_rows)} source rows from preview fallback",
                      {"run_id": run_id, "session_id": session.id})
        return source_rows
# #endregion fetch_source_rows


# #region _extract_chart_data_rows [C:1] [TYPE Function]
def _extract_chart_data_rows(response: dict[str, Any]) -> list[dict[str, Any]]:
    """Extract data rows from Superset chart data API response."""
    result = response.get("result")
    if isinstance(result, list):
        for item in result:
            if isinstance(item, dict):
                data = item.get("data")
                if isinstance(data, list) and data:
                    return data
    if isinstance(result, dict):
        data = result.get("data")
        if isinstance(data, list) and data:
            return data
    data = response.get("data")
    if isinstance(data, list) and data:
        return data
    if isinstance(result, list):
        return result
    return []
# #endregion _extract_chart_data_rows
# #endregion RunSourceFetcher