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

# #region DatasourceMetadataService [C:4] [TYPE Module] [SEMANTICS superset, datasource, columns, dialect]
# @BRIEF Fetch datasource column metadata and database dialect from Superset.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [SupersetClient]
# @RELATION DEPENDS_ON -> [ConfigManager]
# @PRE Database session and config manager are available.
# @POST Datasource metadata is fetched with column details and normalized dialect.

from typing import Any

from ...core.config_manager import ConfigManager
from ...core.logger import logger
from ...core.superset_client import SupersetClient
from ...schemas.translate import DatasourceColumnResponse, DatasourceColumnsResponse

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


# #region get_dialect_from_database [C:2] [TYPE Function]
# @BRIEF Extract normalized dialect string from a Superset database record.
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")

    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
# #endregion get_dialect_from_database


# #region fetch_datasource_metadata [C:3] [TYPE Function]
# @BRIEF Fetch datasource columns and database dialect from Superset.
async 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."""
    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")

    client = SupersetClient(env_config)
    dataset_detail = await client.get_dataset_detail(dataset_id)

    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", ""),
        })

    database_info = dataset_detail.get("database", {})
    if isinstance(database_info, dict):
        dialect = get_dialect_from_database(database_info)
    else:
        try:
            db_id = dataset_detail.get("database_id")
            if db_id:
                db_record = await 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:
            raise ValueError(f"Could not determine database dialect: {e}")

    return columns, dialect
# #endregion fetch_datasource_metadata


# #region detect_virtual_columns [C:2] [TYPE Function]
# @BRIEF Identify virtual (calculated) columns from column metadata.
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)]
# #endregion detect_virtual_columns


# #region get_datasource_columns [C:3] [TYPE Function]
# @BRIEF Fetch datasource column metadata from Superset and return structured response.
async 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}")

    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")

    client = SupersetClient(env_config)
    dataset_detail = await client.get_dataset_detail(datasource_id)

    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 = await 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")

    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,
    )
# #endregion get_datasource_columns
# #endregion DatasourceMetadataService