ss-tools/backend/tests/plugins/translate/test_service_target_schema.py

# #region Test.TargetSchemaValidation [C:3] [TYPE Module] [SEMANTICS test, translate, schema, validation]
# @BRIEF Tests for service_target_schema.py — _build_expected_columns, _extract_columns_from_rows,
#         _parse_sqllab_result, validate_target_table_schema.
# @RELATION BINDS_TO -> [TargetSchemaValidation]
# @TEST_CONTRACT: _build_expected_columns -> list[TargetSchemaColumnInfo] | deduplicated column list
# @TEST_CONTRACT: _extract_columns_from_rows -> list[dict] | normalised column info
# @TEST_CONTRACT: _parse_sqllab_result -> tuple | data_rows + table_exists
# @TEST_CONTRACT: validate_target_table_schema -> TargetSchemaValidationResponse | full schema diff

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "src"))

import os
os.environ.setdefault("SECRET_KEY", "test-secret-key-for-testing")
os.environ.setdefault("DATABASE_URL", "sqlite:///:memory:")
os.environ.setdefault("AUTH_DATABASE_URL", "sqlite:///:memory:")

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from src.schemas.translate import (
    TargetSchemaColumnInfo,
    TargetSchemaValidationRequest,
    TargetSchemaValidationResponse,
)
from src.plugins.translate.service_target_schema import (
    _build_expected_columns,
    _extract_columns_from_rows,
    _parse_sqllab_result,
    validate_target_table_schema,
)


class TestBuildExpectedColumns:
    """_build_expected_columns — build deduplicated column list from request."""

    def _make_req(self, overrides: dict | None = None) -> TargetSchemaValidationRequest:
        defaults = {
            "environment_id": "env-1",
            "target_database_id": "db-1",
            "target_table": "my_table",
            "target_schema": "public",
            "target_key_cols": ["id"],
            "target_column": None,
            "translation_column": "name",
            "target_language_column": "lang",
            "target_source_column": "src",
            "target_source_language_column": "src_lang",
        }
        if overrides:
            defaults.update(overrides)
        return TargetSchemaValidationRequest(**defaults)

    def test_with_all_columns(self):
        """All column types included in order."""
        req = self._make_req()
        cols = _build_expected_columns(req)
        names = [c.name for c in cols]
        assert names == ["id", "name", "lang", "src", "src_lang", "context", "is_original"]

    def test_basic_minimal(self):
        """Minimal config: only key cols and translation column."""
        req = self._make_req({
            "target_key_cols": ["id"],
            "translation_column": "name",
            "target_column": None,
            "target_language_column": None,
            "target_source_column": None,
            "target_source_language_column": None,
        })
        cols = _build_expected_columns(req)
        names = [c.name for c in cols]
        assert names == ["id", "name", "context", "is_original"]

    def test_empty_key_cols(self):
        """Empty key cols list handled."""
        req = self._make_req({
            "target_key_cols": [],
            "translation_column": "name",
            "target_column": None,
            "target_language_column": None,
            "target_source_column": None,
            "target_source_language_column": None,
        })
        cols = _build_expected_columns(req)
        names = [c.name for c in cols]
        assert names == ["name", "context", "is_original"]

    def test_target_column_overrides_translation(self):
        """target_column used instead of translation_column when set."""
        req = self._make_req({
            "target_key_cols": [],
            "translation_column": "source_col",
            "target_column": "target_col",
            "target_language_column": None,
            "target_source_column": None,
            "target_source_language_column": None,
        })
        cols = _build_expected_columns(req)
        names = [c.name for c in cols]
        assert "target_col" in names
        assert "source_col" not in names

    def test_deduplication(self):
        """Duplicate column names are removed."""
        req = self._make_req({
            "target_key_cols": ["id", "id"],
            "target_column": "name",
            "translation_column": "name",
            "target_language_column": None,
            "target_source_column": None,
            "target_source_language_column": None,
        })
        cols = _build_expected_columns(req)
        names = [c.name for c in cols]
        assert names == ["id", "name", "context", "is_original"]

    def test_none_key_cols(self):
        """None key cols handled."""
        req = self._make_req({
            "target_key_cols": None,
            "translation_column": "name",
            "target_column": None,
            "target_language_column": None,
            "target_source_column": None,
            "target_source_language_column": None,
        })
        cols = _build_expected_columns(req)
        names = [c.name for c in cols]
        assert names == ["name", "context", "is_original"]

    def test_all_nullable_fields(self):
        """All optional column fields are None."""
        req = self._make_req({
            "target_key_cols": None,
            "translation_column": None,
            "target_column": None,
            "target_language_column": None,
            "target_source_column": None,
            "target_source_language_column": None,
        })
        cols = _build_expected_columns(req)
        names = [c.name for c in cols]
        # Only the always-included columns
        assert names == ["context", "is_original"]

    def test_return_type(self):
        """Returns list of TargetSchemaColumnInfo."""
        req = self._make_req()
        cols = _build_expected_columns(req)
        for c in cols:
            assert isinstance(c, TargetSchemaColumnInfo)
            assert c.name
            assert c.data_type is None
            assert c.is_nullable is True


class TestExtractColumnsFromRows:
    """_extract_columns_from_rows — parse column info from data rows."""

    def test_postgres_format(self):
        """PostgreSQL information_schema.columns format."""
        rows = [
            {"column_name": "id", "data_type": "integer", "is_nullable": "YES"},
            {"column_name": "name", "data_type": "varchar", "is_nullable": "NO"},
        ]
        cols = _extract_columns_from_rows(rows)
        assert len(cols) == 2
        assert cols[0]["name"] == "id"
        assert cols[0]["type"] == "integer"
        assert cols[0]["is_nullable"] is True
        assert cols[1]["name"] == "name"
        assert cols[1]["type"] == "varchar"
        assert cols[1]["is_nullable"] is False

    def test_clickhouse_format(self):
        """ClickHouse DESCRIBE TABLE format."""
        rows = [
            {"name": "id", "type": "Int32"},
            {"name": "name", "type": "String"},
        ]
        cols = _extract_columns_from_rows(rows)
        assert len(cols) == 2
        assert cols[0]["name"] == "id"
        assert cols[0]["type"] == "Int32"
        assert cols[0]["is_nullable"] is True  # default for CH

    def test_mysql_format(self):
        """MySQL SHOW COLUMNS format."""
        rows = [
            {"Field": "id", "Type": "int(11)"},
            {"Field": "name", "Type": "varchar(255)"},
        ]
        cols = _extract_columns_from_rows(rows)
        assert len(cols) == 2
        assert cols[0]["name"] == "id"
        assert cols[0]["type"] == "int(11)"

    def test_non_dict_row_skipped(self):
        """Non-dict rows are skipped."""
        rows = [{"column_name": "id"}, "not a dict", 42]
        cols = _extract_columns_from_rows(rows)
        assert len(cols) == 1

    def test_missing_column_name_skipped(self):
        """Rows without any known column key are skipped."""
        rows = [{"unknown_key": "value"}]
        cols = _extract_columns_from_rows(rows)
        assert len(cols) == 0

    def test_empty_rows(self):
        """Empty list returns empty list."""
        assert _extract_columns_from_rows([]) == []

    def test_is_nullable_string_handling(self):
        """is_nullable string 'YES'/'NO' converted to bool."""
        rows = [
            {"column_name": "a", "is_nullable": "YES"},
            {"column_name": "b", "is_nullable": "NO"},
            {"column_name": "c", "is_nullable": "yes"},
            {"column_name": "d", "is_nullable": None},
        ]
        cols = _extract_columns_from_rows(rows)
        assert cols[0]["is_nullable"] is True
        assert cols[1]["is_nullable"] is False
        assert cols[2]["is_nullable"] is True
        assert cols[3]["is_nullable"] is True  # None -> default True

    def test_type_none_handled(self):
        """type is None when not found in row."""
        rows = [{"column_name": "id"}]
        cols = _extract_columns_from_rows(rows)
        assert cols[0]["type"] is None


class TestParseSqllabResult:
    """_parse_sqllab_result — extract data rows from SQL Lab response."""

    def test_format_direct_data(self):
        """Format 1: direct 'data' key."""
        result = {"data": [{"col": 1}, {"col": 2}]}
        rows, exists = _parse_sqllab_result(result)
        assert len(rows) == 2
        assert exists is True

    def test_format_async_result(self):
        """Format 2: nested 'result' dict."""
        result = {"result": {"data": [{"col": 1}]}}
        rows, exists = _parse_sqllab_result(result)
        assert len(rows) == 1

    def test_format_get_query_results(self):
        """Format 3: 'results' dict with 'data'."""
        result = {"results": {"data": [{"col": 1}]}}
        rows, exists = _parse_sqllab_result(result)
        assert len(rows) == 1

    def test_format_get_query_results_nested(self):
        """Format 3: 'results' -> 'result' -> 'data'."""
        result = {"results": {"result": {"data": [{"col": 1}]}}}
        rows, exists = _parse_sqllab_result(result)
        assert len(rows) == 1

    def test_format_results_list(self):
        """Format 4: 'results' as list."""
        result = {"results": [{"col": 1}]}
        rows, exists = _parse_sqllab_result(result)
        assert len(rows) == 1

    def test_empty_result(self):
        """None/empty result returns empty."""
        rows, exists = _parse_sqllab_result({})
        assert rows == []
        assert exists is False

    def test_none_result(self):
        """None result returns empty."""
        rows, exists = _parse_sqllab_result(None)
        assert rows == []
        assert exists is False

    def test_no_data_key(self):
        """Result without data/result/results returns empty."""
        result = {"status": "success"}
        rows, exists = _parse_sqllab_result(result)
        assert rows == []
        assert exists is False

    def test_raw_response_is_falsy(self):
        """raw_response is a falsy value -> returns empty."""
        result = {"raw_response": None}
        rows, exists = _parse_sqllab_result(result)
        assert rows == []
        assert exists is False

    def test_raw_response_is_empty_string(self):
        """raw_response is empty string -> returns empty."""
        result = {"raw_response": ""}
        rows, exists = _parse_sqllab_result(result)
        assert rows == []
        assert exists is False

    def test_with_raw_response_wrapper(self):
        """Result with raw_response wrapper unwrapped."""
        result = {"raw_response": {"data": [{"col": 1}]}}
        rows, exists = _parse_sqllab_result(result)
        assert len(rows) == 1

    def test_normalize_list_rows_to_dicts(self):
        """List-of-lists with columns metadata normalized to dicts."""
        result = {
            "data": [["id1", "name1"], ["id2", "name2"]],
            "columns": [{"name": "id"}, {"name": "name"}],
        }
        rows, exists = _parse_sqllab_result(result)
        assert len(rows) == 2
        assert rows[0] == {"id": "id1", "name": "name1"}
        assert rows[1] == {"id": "id2", "name": "name2"}

    def test_table_not_exists_empty_data(self):
        """Empty data means table does not exist."""
        result = {"data": []}
        rows, exists = _parse_sqllab_result(result)
        assert rows == []
        assert exists is False

    def test_data_not_list(self):
        """data key exists but is not a list."""
        result = {"data": "not a list"}
        rows, exists = _parse_sqllab_result(result)
        assert rows == []
        assert exists is False


class TestValidateTargetTableSchema:
    """validate_target_table_schema — full schema validation flow."""

    @pytest.mark.asyncio
    async def test_invalid_table_name_sql_injection(self):
        """Invalid table name (SQL injection) returns error response."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="users; DROP TABLE",
            target_schema="public",
        )
        config_manager = MagicMock()
        result = await validate_target_table_schema(req, config_manager)
        assert result.table_exists is False
        assert "Invalid target table name" in result.error
        assert len(result.missing_columns) > 0
        assert result.all_present is False

    @pytest.mark.asyncio
    async def test_invalid_schema_name(self):
        """Invalid schema name returns error response."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="schema; DROP",
        )
        config_manager = MagicMock()
        result = await validate_target_table_schema(req, config_manager)
        assert result.table_exists is False
        assert "Invalid target schema name" in result.error
        assert result.all_present is False

    @pytest.mark.asyncio
    async def test_direct_db_path(self):
        """Direct DB path uses DbExecutor instead of SQL Lab."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="public",
            insert_method="direct_db",
            connection_id="conn-1",
        )
        config_manager = MagicMock()

        mock_column = MagicMock()
        mock_column.name = "id"
        mock_column.data_type = "integer"
        mock_executor = AsyncMock()
        mock_executor.fetch_schema.return_value = [mock_column]

        with patch("src.plugins.translate.service_target_schema.ConnectionService") as MockCS, \
             patch("src.plugins.translate.service_target_schema.DbExecutor", return_value=mock_executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.table_exists is True
        assert len(result.actual_columns) == 1
        assert result.actual_columns[0].name == "id"
        assert result.actual_columns[0].data_type == "integer"

    @pytest.mark.asyncio
    async def test_direct_db_no_columns(self):
        """Direct DB path with None fetch returns table_exists=False."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="public",
            insert_method="direct_db",
            connection_id="conn-1",
        )
        config_manager = MagicMock()

        mock_executor = AsyncMock()
        mock_executor.fetch_schema.return_value = None

        with patch("src.plugins.translate.service_target_schema.ConnectionService"), \
             patch("src.plugins.translate.service_target_schema.DbExecutor", return_value=mock_executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.table_exists is False
        assert "Failed to fetch schema" in (result.error or "")

    def _make_sqllab_executor(self, backend="postgresql", db_name="test_db"):
        """Create a mock executor with proper sync/async methods."""
        executor = MagicMock()
        executor.resolve_database_id = AsyncMock()
        executor.get_database_backend.return_value = backend
        executor.get_database_name.return_value = db_name
        executor.execute_and_poll = AsyncMock()
        return executor

    @pytest.mark.asyncio
    async def test_sqllab_path_success(self):
        """SQL Lab path: query succeeds, columns match."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="public",
            target_key_cols=["id"],
            translation_column="name",
            target_column=None,
            target_language_column=None,
            target_source_column=None,
            target_source_language_column=None,
        )
        config_manager = MagicMock()

        executor = self._make_sqllab_executor()
        executor.execute_and_poll.return_value = {
            "status": "success",
            "data": [
                {"column_name": "id", "data_type": "integer", "is_nullable": "NO"},
                {"column_name": "name", "data_type": "varchar", "is_nullable": "YES"},
                {"column_name": "context", "data_type": "text", "is_nullable": "YES"},
                {"column_name": "is_original", "data_type": "boolean", "is_nullable": "YES"},
            ],
        }

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.table_exists is True
        assert result.all_present is True
        assert len(result.expected_columns) == 4  # id, name, context, is_original
        assert len(result.missing_columns) == 0

    @pytest.mark.asyncio
    async def test_sqllab_path_missing_columns(self):
        """SQL Lab path: some expected columns missing."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="public",
            target_key_cols=["id"],
            translation_column="name",
        )
        config_manager = MagicMock()

        executor = self._make_sqllab_executor()
        executor.execute_and_poll.return_value = {
            "status": "success",
            "data": [
                {"column_name": "id", "data_type": "integer", "is_nullable": "NO"},
            ],
        }

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.table_exists is True
        assert result.all_present is False
        assert len(result.missing_columns) > 0
        missing_names = [c.name for c in result.missing_columns]
        assert "name" in missing_names

    @pytest.mark.asyncio
    async def test_sqllab_path_clickhouse(self):
        """ClickHouse backend uses system.columns."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="public",
            target_key_cols=["id"],
            translation_column="name",
        )
        config_manager = MagicMock()

        executor = self._make_sqllab_executor(backend="clickhouse", db_name="test_ch")
        executor.execute_and_poll.return_value = {
            "status": "success",
            "data": [
                {"name": "id", "type": "Int32"},
                {"name": "name", "type": "String"},
                {"name": "context", "type": "String"},
                {"name": "is_original", "type": "UInt8"},
            ],
        }

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.table_exists is True
        assert result.all_present is True
        # Verify clickhouse SQL was used
        call_sql = executor.execute_and_poll.call_args[1]["sql"]
        assert "system.columns" in call_sql
        assert "SELECT name, type" in call_sql

    @pytest.mark.asyncio
    async def test_sqllab_path_failed_status(self):
        """SQL Lab query fails status -> error response."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
        )
        config_manager = MagicMock()

        executor = self._make_sqllab_executor()
        executor.execute_and_poll.return_value = {
            "status": "failed",
            "error_message": "Table does not exist",
        }

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.table_exists is False
        assert "Superset SQL Lab error" in (result.error or "")
        assert result.database_name == "test_db"
        assert result.database_backend == "postgresql"

    @pytest.mark.asyncio
    async def test_sqllab_path_exception(self):
        """Exception during SQL Lab execution -> error response."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
        )
        config_manager = MagicMock()

        executor = MagicMock()
        executor.resolve_database_id = AsyncMock(side_effect=Exception("Connection refused"))

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.table_exists is False
        assert "Failed to validate" in (result.error or "")

    @pytest.mark.asyncio
    async def test_sqllab_with_extra_columns(self):
        """Extra columns in target table are reported."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_key_cols=["id"],
            translation_column="name",
        )
        config_manager = MagicMock()

        executor = self._make_sqllab_executor()
        executor.execute_and_poll.return_value = {
            "status": "success",
            "data": [
                {"column_name": "id", "data_type": "integer"},
                {"column_name": "name", "data_type": "varchar"},
                {"column_name": "context", "data_type": "text"},
                {"column_name": "is_original", "data_type": "boolean"},
                {"column_name": "extra_col", "data_type": "text"},
            ],
        }

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        assert result.all_present is True
        extra_names = [c.name for c in result.extra_columns]
        assert "extra_col" in extra_names

    @pytest.mark.asyncio
    async def test_clickhousedb_normalized_to_clickhouse(self):
        """Backend 'clickhousedb' normalized to 'clickhouse' for SQL generation."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="default",
        )
        config_manager = MagicMock()

        executor = self._make_sqllab_executor(backend="clickhousedb", db_name="test_ch")
        executor.execute_and_poll.return_value = {
            "status": "success",
            "data": [],
        }

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        call_sql = executor.execute_and_poll.call_args[1]["sql"]
        assert "system.columns" in call_sql

    @pytest.mark.asyncio
    async def test_greenplum_normalized_to_postgresql(self):
        """Backend 'greenplum' normalized to 'postgresql' for SQL generation."""
        req = TargetSchemaValidationRequest(
            environment_id="env-1",
            target_database_id="db-1",
            target_table="my_table",
            target_schema="public",
        )
        config_manager = MagicMock()

        executor = self._make_sqllab_executor(backend="greenplum", db_name="test_gp")
        executor.execute_and_poll.return_value = {
            "status": "success",
            "data": [],
        }

        with patch("src.plugins.translate.service_target_schema.SupersetSqlLabExecutor", return_value=executor):
            result = await validate_target_table_schema(req, config_manager)

        call_sql = executor.execute_and_poll.call_args[1]["sql"]
        assert "information_schema.columns" in call_sql
# #endregion Test.TargetSchemaValidation