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