Add full custom query for observations

Author: dmitrkozlovsk

src/domain/models.py

@@ -176,6 +176,12 @@ class Observation(Base):
         comment="Custom query to calculate the test_ids for the experiment. "
         "The query will be used in jinja template instead of the default query",
     )
+    full_custom_query = Column(
+        Text,
+        nullable=True,
+        comment="Full custom query to calculate metrics. "
+        "If provided, this query will be used instead of the template-based query.",
+    )
 
     metric_tags = Column(JSON, nullable=True, comment="Filter metrics by tags")
     metric_groups = Column(JSON, nullable=True, comment="Filter metrics by groups")

src/services/runners/renderer.py

@@ -46,6 +46,9 @@ def render(
         Returns:
             str: The rendered SQL query string.
         """
+        if obs.full_custom_query:
+            return str(obs.full_custom_query)
+
         calc_scenario_path = self.templates_config.scenarios.get(str(obs.calculation_scenario))
         if not calc_scenario_path:
             error_message = (

src/ui/observations/inputs.py

@@ -138,6 +138,7 @@ class ObservationFormData:
     audience_tables: list[str] | None
     filters: list[str] | None
     custom_test_ids_query: str | None
+    full_custom_query: str | None
     metric_tags: list[str] | None
     metric_groups: list[str] | None
 
@@ -345,6 +346,9 @@ def render(
                 "Supports JOINs, CTEs, subqueries, and complex aggregations. "
                 "Modifying operations (INSERT, UPDATE, DELETE) are not allowed.",
             )
+            full_custom_query = st.text_area(
+                "Full Custom Query", value="", key="full_custom_query_input_key"
+            )
 
             metric_tags: list[str] | None = st.multiselect(
                 "Metric Tags",
@@ -425,6 +429,7 @@ def render(
                     audience_tables=[t.strip() for t in audience_tables.split("\n") if t.strip()],
                     filters=[f.strip() for f in filters.split("\n") if f.strip()],
                     custom_test_ids_query=custom_test_ids_query,
+                    full_custom_query=full_custom_query,
                     metric_tags=metric_tags,
                     metric_groups=metric_groups,
                 ),

tests/services/runners/conftest.py

@@ -82,24 +82,31 @@ def fake_metric_results():
 @pytest.fixture
 def mock_observation():
     """Fixture providing mock observation object for testing."""
-    return Observation(
-        id=1,
-        experiment_id=1,
-        name="Test Observation",
-        db_experiment_name="test_experiment_db_name",
-        split_id="user_id",
-        calculation_scenario="base",
-        exposure_start_datetime=datetime(2024, 1, 1),
-        exposure_end_datetime=datetime(2024, 1, 31),
-        calc_start_datetime=datetime(2024, 1, 1),
-        calc_end_datetime=datetime(2024, 1, 31),
-        exposure_event="view",
-        audience_tables=["active_users"],
-        filters=["platform='web'", "(country='US' or country='CA')"],
-        custom_test_ids_query=None,
-        metric_tags=["main"],
-        metric_groups=["core"],
-    )
+
+    def _mock_observation(**kwargs):
+        default_attrs = {
+            "id": 1,
+            "experiment_id": 1,
+            "name": "Test Observation",
+            "db_experiment_name": "test_experiment_db_name",
+            "split_id": "user_id",
+            "calculation_scenario": "base",
+            "exposure_start_datetime": datetime(2024, 1, 1),
+            "exposure_end_datetime": datetime(2024, 1, 31),
+            "calc_start_datetime": datetime(2024, 1, 1),
+            "calc_end_datetime": datetime(2024, 1, 31),
+            "exposure_event": "view",
+            "audience_tables": ["active_users"],
+            "filters": ["platform='web'", "(country='US' or country='CA')"],
+            "custom_test_ids_query": None,
+            "full_custom_query": None,
+            "metric_tags": ["main"],
+            "metric_groups": ["core"],
+        }
+        default_attrs.update(kwargs)
+        return Observation(**default_attrs)
+
+    return _mock_observation
 
 
 @pytest.fixture

tests/services/runners/test_executors.py

@@ -16,20 +16,17 @@ def test_mock_calculation_runner_run_calculations_regular(
 ):
     """Test calculation runner executes regular calculations successfully."""
     mock_calculation_runner.exec._connector.fetch_results.return_value = mock_metric_results
+    obs = mock_observation()
     # Execute
-    result = mock_calculation_runner.run_calculation(
-        obs=mock_observation, purpose=CalculationPurpose.REGULAR
-    )
+    result = mock_calculation_runner.run_calculation(obs=obs, purpose=CalculationPurpose.REGULAR)
 
     assert result.job_id == 1
     assert result.success is True
     assert len(result.metric_results) == 2
     assert result.metric_results[0].metric_name == "conversion_rate"
     # get result of job
     executed_job = JobHandler(engine).get(1)
-    rendered_query = mock_calculation_runner.renderer.render(
-        obs=mock_observation, purpose=CalculationPurpose.REGULAR
-    )
+    rendered_query = mock_calculation_runner.renderer.render(obs=obs, purpose=CalculationPurpose.REGULAR)
     mock_calculation_runner.exec._connector.fetch_results.assert_called_once_with(rendered_query)
     assert executed_job.status == JobStatus.COMPLETED
 
@@ -45,9 +42,10 @@ def test_mock_calculation_runner_run_calculations_planning(
 ):
     """Test calculation runner executes planning calculations without storing."""
     mock_calculation_runner.exec._connector.fetch_results.return_value = mock_metric_results
+    obs = mock_observation()
     # Execute
     result = mock_calculation_runner.run_calculation(
-        obs=mock_observation,
+        obs=obs,
         purpose=CalculationPurpose.PLANNING,
     )
     assert isinstance(result.metric_results[0], MetricResult)
@@ -66,7 +64,8 @@ def test_mock_calculation_runner_run_calculations_planning_failed(
 ):
     """Test calculation runner handles database errors properly."""
     mock_calculation_runner.exec._connector.fetch_results.side_effect = Exception("Database error")
-    mock_calculation_runner.run_calculation(obs=mock_observation, purpose=CalculationPurpose.REGULAR)
+    obs = mock_observation()
+    mock_calculation_runner.run_calculation(obs=obs, purpose=CalculationPurpose.REGULAR)
     executed_job = JobHandler(engine).get(1)
     assert executed_job.status == JobStatus.FAILED
     assert "Database error" in executed_job.error_message
@@ -76,10 +75,9 @@ def test_mock_calculation_runner_run_calculations_planning_failed(
 
 def test_calculation_runner_job_creation_failure(mock_calculation_runner, mock_observation):
     """Test that CalculationRunner handles job creation failure."""
+    obs = mock_observation()
     with patch("src.services.runners.executors.JobHandler.create", return_value=None):
-        result = mock_calculation_runner.run_calculation(
-            obs=mock_observation, purpose=CalculationPurpose.REGULAR
-        )
+        result = mock_calculation_runner.run_calculation(obs=obs, purpose=CalculationPurpose.REGULAR)
 
     assert result.success is False
     assert result.job_id is None
@@ -88,10 +86,9 @@ def test_calculation_runner_job_creation_failure(mock_calculation_runner, mock_o
 
 def test_calculation_runner_job_creation_db_exception(mock_calculation_runner, mock_observation):
     """Test that CalculationRunner handles job creation failure due to DB exception."""
+    obs = mock_observation()
     with patch("src.services.runners.executors.JobHandler.create", side_effect=Exception("DB is down")):
-        result = mock_calculation_runner.run_calculation(
-            obs=mock_observation, purpose=CalculationPurpose.REGULAR
-        )
+        result = mock_calculation_runner.run_calculation(obs=obs, purpose=CalculationPurpose.REGULAR)
 
     assert result.success is False
     assert result.job_id is None
@@ -100,10 +97,9 @@ def test_calculation_runner_job_creation_db_exception(mock_calculation_runner, m
 
 def test_calculation_runner_render_failure(mock_calculation_runner, mock_observation, engine, tables):
     """Test that CalculationRunner handles query rendering failure."""
+    obs = mock_observation()
     with patch.object(mock_calculation_runner.renderer, "render", side_effect=Exception("Template error")):
-        result = mock_calculation_runner.run_calculation(
-            obs=mock_observation, purpose=CalculationPurpose.REGULAR
-        )
+        result = mock_calculation_runner.run_calculation(obs=obs, purpose=CalculationPurpose.REGULAR)
 
     assert result.success is False
     assert result.job_id == 1
@@ -119,14 +115,13 @@ def test_calculation_runner_store_metrics_failure(
 ):
     """Test that CalculationRunner handles storing metrics failure."""
     mock_calculation_runner.exec._connector.fetch_results.return_value = mock_metric_results
+    obs = mock_observation()
     with patch.object(
         mock_calculation_runner.jobs,
         "store_metrics",
         side_effect=Exception("DB connection failed"),
     ):
-        result = mock_calculation_runner.run_calculation(
-            obs=mock_observation, purpose=CalculationPurpose.REGULAR
-        )
+        result = mock_calculation_runner.run_calculation(obs=obs, purpose=CalculationPurpose.REGULAR)
 
     assert result.success is False
     assert result.job_id == 1

tests/services/runners/test_renderer.py

@@ -3,7 +3,8 @@
 
 def test_renderer(query_renderer, mock_observation):
     """Test rendering base calculation query for planning purpose"""
-    query = query_renderer.render(obs=mock_observation, purpose=CalculationPurpose.PLANNING)
+    obs = mock_observation()
+    query = query_renderer.render(obs=obs, purpose=CalculationPurpose.PLANNING)
     # TODO: fix mock observation and fix fake template according to custom_test_ids_query
     assert isinstance(query, str)
     assert "user_id as split_id" in query
@@ -15,8 +16,9 @@ def test_renderer(query_renderer, mock_observation):
 
 def test_renderer_with_metric_names(query_renderer, mock_observation):
     """Test rendering base calculation query with specific metric names"""
+    obs = mock_observation()
     query = query_renderer.render(
-        obs=mock_observation,
+        obs=obs,
         purpose=CalculationPurpose.PLANNING,
         experiment_metric_names=["click_through_rate"],
     )
@@ -27,3 +29,11 @@ def test_renderer_with_metric_names(query_renderer, mock_observation):
         assert exp_alias not in query
     for user_alias in ["product_purchase_cnt", "session_duration"]:
         assert user_alias not in query
+
+
+def test_renderer_with_full_custom_query(query_renderer, mock_observation):
+    """Test that full_custom_query bypasses template rendering"""
+    custom_query = "SELECT 1"
+    obs = mock_observation(full_custom_query=custom_query)
+    query = query_renderer.render(obs=obs, purpose=CalculationPurpose.REGULAR)
+    assert query == custom_query