ss-tools/backend/tests/integration/test_superset_sqllab_e2e.py

# #region SupersetSqllabE2E [C:5] [TYPE Module] [SEMANTICS test,superset,sqllab,e2e,execution]
# @BRIEF End-to-end: register real PostgreSQL in Superset → execute SQL via SupersetClient.
# @RELATION BINDS_TO -> [SupersetSqlLabExecutor]
# @RELATION BINDS_TO -> [SupersetClient]
#
# @TEST_CONTRACT SupersetSqlLabE2E ->
# {
#   invariants: [
#     "SupersetClient.create_database() registers a real Postgres DB via REST API",
#     "SupersetClient handles JWT + CSRF internally for mutation endpoints",
#     "SQL Lab execute returns query_id for valid SQL on registered DB",
#     "SupersetSqlLabExecutor can resolve DB and execute SQL end-to-end"
#   ]
# }
#
# @REQUIRES Docker daemon running (testcontainers starts superset + postgres)
import sys
from pathlib import Path

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

import pytest
from sqlalchemy import text
from urllib.parse import urlparse


# #region TestSupersetSqllabEndToEnd [C:4] [TYPE Class]
class TestSupersetSqllabEndToEnd:
    """Full E2E: PostgreSQL → Superset registration → SQL Lab via SupersetClient."""

    @staticmethod
    def _setup_test_table(pg_engine):
        with pg_engine.begin() as conn:
            conn.execute(text("DROP TABLE IF EXISTS ss_e2e_test"))
            conn.execute(text(
                "CREATE TABLE ss_e2e_test ("
                "  id SERIAL PRIMARY KEY,"
                "  name VARCHAR(100) NOT NULL,"
                "  value INTEGER NOT NULL"
                ")"
            ))
            conn.execute(text(
                "INSERT INTO ss_e2e_test (name, value) VALUES "
                "('alpha', 10), ('beta', 20), ('gamma', 30)"
            ))

    @staticmethod
    def _cleanup_test_table(pg_engine):
        with pg_engine.begin() as conn:
            conn.execute(text("DROP TABLE IF EXISTS ss_e2e_test"))

    # #region test_register_db_and_execute_sql [C:2] [TYPE Function]
    @pytest.mark.asyncio
    async def test_register_db_and_execute_sql(
        self, superset_client, superset_db_url, pg_engine,
    ):
        """Register Postgres DB → execute SQL via SupersetClient (handles CSRF)."""
        self._setup_test_table(pg_engine)
        db_id = None
        try:
            # Build Postgres URI with Docker bridge IP
            parsed = urlparse(superset_db_url)
            register_uri = f"postgresql://test:test@{parsed.hostname}:{parsed.port}/test_translate"

            # Register DB via SupersetClient.create_database()
            reg_resp = await superset_client.create_database(
                database_name="E2E Test DB",
                sqlalchemy_uri=register_uri,
                expose_in_sqllab=True,
                allow_dml=False,
            )
            db_id = reg_resp.get("id")
            if not db_id:
                # Fallback: list and find
                _, dbs = await superset_client.get_databases()
                target = next(
                    (d for d in dbs if "E2E" in d.get("database_name", "")),
                    None,
                )
                if target:
                    db_id = target["id"]
                else:
                    pytest.skip(f"DB reg failed: {str(reg_resp)[:300]}")

            # Execute SQL via client.request()
            exec_resp = await superset_client.client.request(
                "POST", "/sqllab/execute/",
                data={
                    "database_id": db_id,
                    "sql": "SELECT * FROM ss_e2e_test ORDER BY id",
                    "runAsync": False,
                },
            )

            # Verify inline results
            result_part = exec_resp.get("result", exec_resp)
            if isinstance(result_part, dict) and "columns" in result_part:
                data_rows = result_part.get("data", [])
                assert len(data_rows) == 3
                col_names = [
                    c.get("column_name", c.get("name", ""))
                    for c in result_part["columns"]
                ]
                assert "name" in col_names
                assert "value" in col_names
                return  # success

            # Poll for async results
            query_id = exec_resp.get("query", {}).get("id") or exec_resp.get("id")
            if query_id:
                import asyncio
                for attempt in range(15):
                    await asyncio.sleep(1)
                    result_data = await superset_client.client.request(
                        "GET", f"/sqllab/results/{query_id}/",
                    )
                    for key in ("result", "data", "query"):
                        part = result_data.get(key, {})
                        if isinstance(part, dict) and "columns" in part:
                            data_rows = part.get("data", [])
                            if len(data_rows) == 3:
                                return  # success
                    if result_data.get("status") in ("pending", "running"):
                        continue
                pytest.fail(f"No valid results after 15s for query {query_id}")
            else:
                pytest.skip(f"No query_id: {list(exec_resp.keys())}")
        finally:
            self._cleanup_test_table(pg_engine)
            if db_id:
                try:
                    await superset_client.delete_database(db_id)
                except Exception:
                    pass
    # #endregion test_register_db_and_execute_sql
# #endregion TestSupersetSqllabEndToEnd


# #region TestSupersetSqlLabExecutorE2E [C:4] [TYPE Class]
class TestSupersetSqlLabExecutorE2E:
    """SupersetSqlLabExecutor against a real registered database."""

    # #region test_executor_against_real_db [C:2] [TYPE Function]
    @pytest.mark.asyncio
    async def test_executor_against_real_db(
        self, superset_client,
        superset_admin_password, superset_db_url, pg_engine,
    ):
        """SupersetSqlLabExecutor resolves DB and executes SQL end-to-end."""
        from unittest.mock import MagicMock
        from src.plugins.translate.superset_executor import SupersetSqlLabExecutor

        TestSupersetSqllabEndToEnd._setup_test_table(pg_engine)

        config_manager = MagicMock()
        config_manager.get_environments.return_value = [superset_client.env]
        config_manager.get_environment.return_value = superset_client.env
        config_manager.get_config.return_value = MagicMock()

        executor = SupersetSqlLabExecutor(config_manager, "test_env")
        db_id = None

        try:
            # Register DB via SupersetClient.create_database()
            parsed = urlparse(superset_db_url)
            register_uri = f"postgresql://test:test@{parsed.hostname}:{parsed.port}/test_translate"
            reg_resp = await superset_client.create_database(
                database_name="Executor E2E DB",
                sqlalchemy_uri=register_uri,
                expose_in_sqllab=True,
                allow_dml=False,
            )
            db_id = reg_resp.get("id")
            if not db_id:
                pytest.skip(f"DB reg: {str(reg_resp)[:200]}")

            # Resolve database_id via executor
            resolved_id = await executor.resolve_database_id(
                database_name="Executor E2E DB",
            )
            assert isinstance(resolved_id, int) and resolved_id > 0
            assert resolved_id == db_id

            # Execute SQL with runAsync=False (inline results in raw_response)
            result = await executor.execute_sql(
                "SELECT * FROM ss_e2e_test ORDER BY id",
                database_id=db_id, run_async=False,
            )

            # Inline results from execute_sql's raw_response
            raw = result.get("raw_response", {})
            result_part = raw.get("result", raw)
            if isinstance(result_part, dict) and "columns" in result_part:
                data_rows = result_part.get("data", [])
                assert len(data_rows) == 3, \
                    f"Expected 3 rows, got {len(data_rows)}"
                col_names = [
                    c.get("column_name", c.get("name", ""))
                    for c in result_part["columns"]
                ]
                assert "name" in col_names
                assert "value" in col_names
                return  # success

            # Fallback: poll with runAsync=True
            query_id = result.get("query_id")
            if query_id:
                import asyncio
                for attempt in range(15):
                    await asyncio.sleep(1)
                    qr = await executor.get_query_results(query_id)
                    if qr and qr.get("status") == "success" and qr.get("results"):
                        assert len(qr["results"]) == 3
                        return
                pytest.fail(f"No results for query {query_id}")
            else:
                pytest.skip(f"No query_id or inline results: {list(result.keys())}")
        finally:
            TestSupersetSqllabEndToEnd._cleanup_test_table(pg_engine)
            if db_id:
                try:
                    await superset_client.delete_database(db_id)
                except Exception:
                    pass
    # #endregion test_executor_against_real_db
# #endregion TestSupersetSqlLabExecutorE2E
# #endregion SupersetSqllabE2E