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

# [DEF:TranslateJobService:Module]
# @COMPLEXITY: 4
# @SEMANTICS: translate, service, crud, validation
# @PURPOSE: Service layer for translation job CRUD with datasource column validation and database dialect detection.
# @LAYER: Domain
# @RELATION: DEPENDS_ON -> [TranslationJob]
# @RELATION: DEPENDS_ON -> [SupersetClient]
# @RELATION: DEPENDS_ON -> [ConfigManager]
# @PRE: Database session and config manager are available.
# @POST: Translation jobs are created/updated/deleted with column validation and dialect caching.
# @SIDE_EFFECT: Queries Superset for column metadata and database dialect at save time.
# @RATIONALE: Snapshot isolation — in-progress runs use config snapshot; config edits affect future runs only.
# @REJECTED: Invalidating in-progress runs on config edit would break scheduled run continuity.

from typing import Any, Dict, List, Optional, Tuple
from sqlalchemy.orm import Session
from datetime import datetime, timezone
import uuid

from ...core.logger import logger
from ...core.config_manager import ConfigManager
from ...core.superset_client import SupersetClient
from ...models.translate import TranslationJob, TranslationJobDictionary
from ...schemas.translate import (
    TranslateJobCreate,
    TranslateJobUpdate,
    TranslateJobResponse,
    DatasourceColumnsResponse,
    DatasourceColumnResponse,
)

# Supported database dialects for translation
SUPPORTED_DIALECTS = {
    "postgresql", "mysql", "clickhouse", "sqlite", "mssql",
    "oracle", "snowflake", "bigquery", "redshift", "presto",
    "trino", "druid", "hive", "spark", "databricks",
}


# [DEF:get_dialect_from_database:Function]
# @PURPOSE: Extract normalized dialect string from a Superset database record.
# @PRE: database_record is a dict from Superset API.
# @POST: Returns normalized dialect string or raises ValueError if unsupported.
def get_dialect_from_database(database_record: Dict[str, Any]) -> str:
    """Extract and validate dialect from Superset database record."""
    backend = (
        database_record.get("backend")
        or database_record.get("engine")
        or ""
    ).lower().strip()

    if not backend:
        raise ValueError("Could not determine database dialect from connection")

    # Map Superset backend names to normalized dialect
    dialect_map = {
        "postgresql": "postgresql",
        "greenplum": "postgresql",
        "mysql": "mysql",
        "clickhouse": "clickhouse",
        "clickhousedb": "clickhouse",
        "sqlite": "sqlite",
        "mssql": "mssql",
        "oracle": "oracle",
        "snowflake": "snowflake",
        "bigquery": "bigquery",
        "redshift": "redshift",
        "presto": "presto",
        "trino": "trino",
        "druid": "druid",
        "hive": "hive",
        "spark": "spark",
        "databricks": "databricks",
    }

    normalized = dialect_map.get(backend, backend)
    if normalized not in SUPPORTED_DIALECTS:
        raise ValueError(
            f"Unsupported database dialect: '{backend}'. "
            f"Supported dialects: {', '.join(sorted(SUPPORTED_DIALECTS))}"
        )
    return normalized
# [/DEF:get_dialect_from_database:Function]


# [DEF:fetch_datasource_metadata:Function]
# @PURPOSE: Fetch datasource columns and database dialect from Superset.
# @PRE: dataset_id is a valid Superset dataset ID and environment has valid credentials.
# @POST: Returns (columns_list, dialect_string) or raises on failure.
def fetch_datasource_metadata(
    dataset_id: int,
    env_id: str,
    config_manager: ConfigManager,
) -> Tuple[List[Dict[str, Any]], str]:
    """Fetch column metadata and database dialect for a datasource from Superset."""
    # Find environment config
    environments = config_manager.get_environments()
    env_config = next((e for e in environments if e.id == env_id), None)
    if not env_config:
        raise ValueError(f"Superset environment '{env_id}' not found in configuration")

    # Create Superset client and fetch dataset detail
    client = SupersetClient(env_config)
    dataset_detail = client.get_dataset_detail(dataset_id)

    # Extract columns
    raw_columns = dataset_detail.get("columns", [])
    columns = []
    for col in raw_columns:
        col_name = col.get("name") or col.get("column_name")
        if not col_name:
            continue
        columns.append({
            "name": str(col_name),
            "type": col.get("type"),
            "is_physical": col.get("is_physical", True),
            "is_dttm": col.get("is_dttm", False),
            "description": col.get("description", ""),
        })

    # Extract database dialect from datasource
    database_info = dataset_detail.get("database", {})
    if isinstance(database_info, dict):
        dialect = get_dialect_from_database(database_info)
    else:
        # Fallback: try to fetch database directly
        try:
            db_id = dataset_detail.get("database_id")
            if db_id:
                db_record = client.get_database(int(db_id))
                db_result = db_record.get("result", db_record) if isinstance(db_record, dict) else db_record
                dialect = get_dialect_from_database(db_result)
            else:
                raise ValueError("No database information available for this datasource")
        except Exception as e:
            logger.warning(f"[translate_service] Could not fetch database dialect: {e}")
            raise ValueError(f"Could not determine database dialect: {e}")

    return columns, dialect
# [/DEF:fetch_datasource_metadata:Function]


# [DEF:detect_virtual_columns:Function]
# @PURPOSE: Identify virtual (calculated) columns from column metadata.
# @PRE: columns is a list of dicts with 'is_physical' key.
# @POST: Returns list of virtual column names.
def detect_virtual_columns(columns: List[Dict[str, Any]]) -> List[str]:
    """Return names of columns that are virtual (not physical)."""
    return [col["name"] for col in columns if not col.get("is_physical", True)]
# [/DEF:detect_virtual_columns:Function]


# [DEF:TranslateJobService:Class]
# @PURPOSE: Service for translation job CRUD with validation and Superset integration.
class TranslateJobService:

    def __init__(self, db: Session, config_manager: ConfigManager, current_user: Optional[str] = None):
        self.db = db
        self.config_manager = config_manager
        self.current_user = current_user

    # [DEF:list_jobs:Function]
    # @PURPOSE: List translation jobs with optional status filter and pagination.
    # @POST: Returns tuple of (total_count, list_of_jobs).
    def list_jobs(
        self,
        page: int = 1,
        page_size: int = 20,
        status_filter: Optional[str] = None,
    ) -> Tuple[int, List[TranslationJob]]:
        query = self.db.query(TranslationJob)

        if status_filter:
            query = query.filter(TranslationJob.status == status_filter)

        total = query.count()
        offset = (page - 1) * page_size
        jobs = query.order_by(TranslationJob.created_at.desc()).offset(offset).limit(page_size).all()

        return total, jobs
    # [/DEF:list_jobs:Function]

    # [DEF:get_job:Function]
    # @PURPOSE: Get a single translation job by ID.
    # @POST: Returns TranslationJob or raises ValueError.
    def get_job(self, job_id: str) -> TranslationJob:
        job = self.db.query(TranslationJob).filter(TranslationJob.id == job_id).first()
        if not job:
            raise ValueError(f"Translation job '{job_id}' not found")
        return job
    # [/DEF:get_job:Function]

    # [DEF:create_job:Function]
    # @PURPOSE: Create a new translation job with column validation.
    # @PRE: payload contains valid job configuration.
    # @POST: Returns the created TranslationJob with database_dialect cached.
    # @SIDE_EFFECT: Validates columns via SupersetClient if source_datasource_id is provided.
    # @SIDE_EFFECT: Caches database_dialect from Superset connection.
    def create_job(self, payload: TranslateJobCreate) -> TranslationJob:
        logger.info(f"[TranslateJobService] Creating job '{payload.name}'")

        # Validate: must have a translation column if datasource is configured
        if payload.source_datasource_id and not payload.translation_column:
            raise ValueError("A translation column is required when a datasource is selected")

        # Validate upsert strategy
        valid_strategies = {"MERGE", "INSERT", "UPDATE"}
        if payload.upsert_strategy not in valid_strategies:
            raise ValueError(
                f"Invalid upsert_strategy '{payload.upsert_strategy}'. "
                f"Must be one of: {', '.join(sorted(valid_strategies))}"
            )

        # Detect database dialect and validate columns if datasource is specified
        dialect = payload.database_dialect
        if payload.source_datasource_id and (payload.environment_id or payload.source_dialect):
            # If no explicit dialect, try to detect it
            if not dialect:
                try:
                    env_id = payload.environment_id or payload.source_dialect
                    _, detected_dialect = fetch_datasource_metadata(
                        int(payload.source_datasource_id),
                        env_id,
                        self.config_manager,
                    )
                    dialect = detected_dialect
                except Exception as e:
                    logger.warning(f"[TranslateJobService] Dialect detection failed: {e}")
                    dialect = payload.source_dialect

        # Build job instance
        job = TranslationJob(
            id=str(uuid.uuid4()),
            name=payload.name,
            description=payload.description,
            source_dialect=payload.source_dialect,
            target_dialect=payload.target_dialect,
            database_dialect=dialect,
            source_datasource_id=payload.source_datasource_id,
            source_table=payload.source_table,
            target_schema=payload.target_schema,
            target_table=payload.target_table,
            source_key_cols=payload.source_key_cols or [],
            target_key_cols=payload.target_key_cols or [],
            translation_column=payload.translation_column,
            context_columns=payload.context_columns or [],
            target_language=payload.target_language,
            provider_id=payload.provider_id,
            batch_size=payload.batch_size,
            upsert_strategy=payload.upsert_strategy,
            environment_id=payload.environment_id,
            status="DRAFT",
            created_by=self.current_user,
        )
        self.db.add(job)
        self.db.flush()

        # Attach dictionaries
        if payload.dictionary_ids:
            for dict_id in payload.dictionary_ids:
                assoc = TranslationJobDictionary(
                    id=str(uuid.uuid4()),
                    job_id=job.id,
                    dictionary_id=dict_id,
                )
                self.db.add(assoc)

        self.db.commit()
        self.db.refresh(job)
        logger.info(f"[TranslateJobService] Created job '{job.id}'")
        return job
    # [/DEF:create_job:Function]

    # [DEF:update_job:Function]
    # @PURPOSE: Update an existing translation job.
    # @PRE: payload contains fields to update.
    # @POST: Returns the updated TranslationJob.
    # @SIDE_EFFECT: Re-detects database_dialect if source_datasource_id changed.
    def update_job(self, job_id: str, payload: TranslateJobUpdate) -> TranslationJob:
        logger.info(f"[TranslateJobService] Updating job '{job_id}'")
        job = self.get_job(job_id)

        update_data = payload.model_dump(exclude_unset=True)
        dict_ids = update_data.pop("dictionary_ids", None)

        for field, value in update_data.items():
            if hasattr(job, field):
                setattr(job, field, value)

        # Re-detect dialect if datasource changed
        if payload.source_datasource_id and not payload.database_dialect:
            try:
                env_id = (payload.environment_id or payload.source_dialect or job.environment_id or job.source_dialect)
                _, detected_dialect = fetch_datasource_metadata(
                    int(payload.source_datasource_id),
                    env_id,
                    self.config_manager,
                )
                job.database_dialect = detected_dialect
            except Exception as e:
                logger.warning(f"[TranslateJobService] Dialect re-detection failed: {e}")

        job.updated_at = datetime.now(timezone.utc)

        # Update dictionary associations if provided
        if dict_ids is not None:
            self.db.query(TranslationJobDictionary).filter(
                TranslationJobDictionary.job_id == job_id
            ).delete()
            for dict_id in dict_ids:
                assoc = TranslationJobDictionary(
                    id=str(uuid.uuid4()),
                    job_id=job_id,
                    dictionary_id=dict_id,
                )
                self.db.add(assoc)

        self.db.commit()
        self.db.refresh(job)
        logger.info(f"[TranslateJobService] Updated job '{job_id}'")
        return job
    # [/DEF:update_job:Function]

    # [DEF:delete_job:Function]
    # @PURPOSE: Delete a translation job and its associations.
    # @PRE: job_id must exist.
    # @POST: Job and all related records are deleted.
    def delete_job(self, job_id: str) -> None:
        logger.info(f"[TranslateJobService] Deleting job '{job_id}'")
        job = self.get_job(job_id)

        # Delete dictionary associations
        self.db.query(TranslationJobDictionary).filter(
            TranslationJobDictionary.job_id == job_id
        ).delete()

        self.db.delete(job)
        self.db.commit()
        logger.info(f"[TranslateJobService] Deleted job '{job_id}'")
    # [/DEF:delete_job:Function]

    # [DEF:duplicate_job:Function]
    # @PURPOSE: Duplicate a translation job with a new name.
    # @PRE: job_id must exist.
    # @POST: Returns the duplicated TranslationJob with status DRAFT.
    def duplicate_job(self, job_id: str, new_name: Optional[str] = None) -> TranslationJob:
        logger.info(f"[TranslateJobService] Duplicating job '{job_id}'")
        source = self.get_job(job_id)

        # Copy all fields except id, created_at, updated_at, status
        new_job = TranslationJob(
            id=str(uuid.uuid4()),
            name=new_name or f"{source.name} (Copy)",
            description=source.description,
            source_dialect=source.source_dialect,
            target_dialect=source.target_dialect,
            database_dialect=source.database_dialect,
            source_datasource_id=source.source_datasource_id,
            source_table=source.source_table,
            target_schema=source.target_schema,
            target_table=source.target_table,
            source_key_cols=source.source_key_cols,
            target_key_cols=source.target_key_cols,
            translation_column=source.translation_column,
            context_columns=source.context_columns,
            target_language=source.target_language,
            provider_id=source.provider_id,
            batch_size=source.batch_size,
            upsert_strategy=source.upsert_strategy,
            status="DRAFT",
            created_by=self.current_user,
        )
        self.db.add(new_job)
        self.db.flush()

        # Copy dictionary associations
        old_dicts = self.db.query(TranslationJobDictionary).filter(
            TranslationJobDictionary.job_id == job_id
        ).all()
        for assoc in old_dicts:
            new_assoc = TranslationJobDictionary(
                id=str(uuid.uuid4()),
                job_id=new_job.id,
                dictionary_id=assoc.dictionary_id,
            )
            self.db.add(new_assoc)

        self.db.commit()
        self.db.refresh(new_job)
        logger.info(f"[TranslateJobService] Duplicated job '{job_id}' -> '{new_job.id}'")
        return new_job
    # [/DEF:duplicate_job:Function]

    # [DEF:get_job_dictionary_ids:Function]
    # @PURPOSE: Get dictionary IDs attached to a job.
    # @POST: Returns list of dictionary IDs.
    def get_job_dictionary_ids(self, job_id: str) -> List[str]:
        assocs = self.db.query(TranslationJobDictionary).filter(
            TranslationJobDictionary.job_id == job_id
        ).all()
        return [a.dictionary_id for a in assocs]
    # [/DEF:get_job_dictionary_ids:Function]

    # [DEF:fetch_available_datasources:Function]
    # @PURPOSE: List available Superset datasets for translation job creation.
    def fetch_available_datasources(self, env_id: str, search: Optional[str] = None) -> list:
        """List Superset datasets available for translation."""
        from ...core.superset_client import SupersetClient
        env = self.config_manager.get_environment(env_id)
        if not env:
            raise ValueError(f"Environment '{env_id}' not found")
        client = SupersetClient(env)
        _, datasets = client.get_datasets()
        result = []
        for ds in datasets:
            name = ds.get("table_name", "")
            if search and search.lower() not in name.lower():
                continue
            db_info = ds.get("database", {})
            backend = db_info.get("backend", "")
            dialect = _extract_dialect(backend)
            result.append({
                "id": ds.get("id"),
                "table_name": name,
                "schema": ds.get("schema"),
                "database_name": db_info.get("database_name", "Unknown"),
                "database_dialect": dialect,
                "description": ds.get("description", ""),
            })
        return result
    # [/DEF:fetch_available_datasources:Function]
# [/DEF:TranslateJobService:Class]


# [DEF:_extract_dialect:Function]
# @PURPOSE: Extract database dialect from Superset backend URI.
def _extract_dialect(backend: str) -> str:
    """Extract dialect name from a Superset database backend URI (e.g. 'postgresql+psycopg2://...')."""
    if not backend:
        return "unknown"
    try:
        scheme = backend.split("://")[0]
        dialect = scheme.split("+")[0]
        return dialect.lower()
    except Exception:
        return "unknown"


# [DEF:job_to_response:Function]
# @PURPOSE: Convert a TranslationJob ORM model to a TranslateJobResponse schema with dictionary_ids.
def job_to_response(job: TranslationJob, dict_ids: Optional[List[str]] = None) -> TranslateJobResponse:
    return TranslateJobResponse(
        id=job.id,
        name=job.name,
        description=job.description,
        source_dialect=job.source_dialect,
        target_dialect=job.target_dialect,
        database_dialect=job.database_dialect,
        source_datasource_id=job.source_datasource_id,
        source_table=job.source_table,
        target_schema=job.target_schema,
        target_table=job.target_table,
        source_key_cols=job.source_key_cols or [],
        target_key_cols=job.target_key_cols or [],
        translation_column=job.translation_column,
        context_columns=job.context_columns or [],
        target_language=job.target_language,
        provider_id=job.provider_id,
        batch_size=job.batch_size or 50,
        upsert_strategy=job.upsert_strategy or "MERGE",
        status=job.status,
        created_by=job.created_by,
        created_at=job.created_at,
        updated_at=job.updated_at,
        dictionary_ids=dict_ids or [],
        environment_id=job.environment_id,
    )
# [/DEF:job_to_response:Function]


# [DEF:DatasourceColumnsService:Function]
# @PURPOSE: Fetch datasource column metadata from Superset and return structured response.
# @PRE: datasource_id is a valid Superset dataset ID.
# @POST: Returns DatasourceColumnsResponse with column metadata and database dialect.
# @SIDE_EFFECT: Queries Superset API for dataset detail and database info.
def get_datasource_columns(
    datasource_id: int,
    env_id: str,
    config_manager: ConfigManager,
) -> DatasourceColumnsResponse:
    """Fetch and return column metadata for a given Superset datasource."""
    logger.info(f"[get_datasource_columns] Fetching columns for datasource {datasource_id}")

    # Find environment config
    environments = config_manager.get_environments()
    env_config = next((e for e in environments if e.id == env_id), None)
    if not env_config:
        raise ValueError(f"Superset environment '{env_id}' not found")

    # Create Superset client
    client = SupersetClient(env_config)
    dataset_detail = client.get_dataset_detail(datasource_id)

    # Extract database dialect
    database_info = dataset_detail.get("database", {})
    dialect = None
    if isinstance(database_info, dict):
        try:
            dialect = get_dialect_from_database(database_info)
        except (ValueError, KeyError):
            dialect = None
    if dialect is None:
        database_id = dataset_detail.get("database_id")
        if database_id:
            db_record = client.get_database(int(database_id))
            db_result = db_record.get("result", db_record) if isinstance(db_record, dict) else db_record
            dialect = get_dialect_from_database(db_result)
        else:
            raise ValueError("Could not determine database dialect for this datasource")

    # Extract columns
    raw_columns = dataset_detail.get("columns", [])
    columns = []
    for col in raw_columns:
        col_name = col.get("name") or col.get("column_name")
        if not col_name:
            continue
        columns.append(DatasourceColumnResponse(
            name=str(col_name),
            type=col.get("type"),
            is_physical=col.get("is_physical", True),
            is_dttm=col.get("is_dttm", False),
            description=col.get("description", ""),
        ))

    return DatasourceColumnsResponse(
        datasource_id=datasource_id,
        datasource_name=dataset_detail.get("table_name"),
        schema_name=dataset_detail.get("schema"),
        database_dialect=dialect,
        columns=columns,
    )
# [/DEF:DatasourceColumnsService:Function]

# [/DEF:TranslateJobService:Module]