Skip to content

Add full custom query for observations #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dmitrkozlovsk merged 1 commit into from
Oct 18, 2025
Merged

Add full custom query for observations #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/domain/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
)
Comment on lines 176 to +184
Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Oct 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Add migration for Observation.full_custom_query column

The ORM model now defines a full_custom_query column, but no schema migration accompanies the change. Existing installations will have an observations table without this column, so attempts to create or update an observation that includes full_custom_query will fail with database errors (e.g. “no such column: observations.full_custom_query”). A migration or other schema update is required to keep the database consistent with the model.

Useful? React with 👍 / 👎.


metric_tags = Column(JSON, nullable=True, comment="Filter metrics by tags")
metric_groups = Column(JSON, nullable=True, comment="Filter metrics by groups")
Expand Down
3 changes: 3 additions & 0 deletions src/services/runners/renderer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand Down
5 changes: 5 additions & 0 deletions src/ui/observations/inputs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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"
)
Comment on lines 346 to +351
Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Oct 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve full custom query when editing observations

The new form field for full_custom_query is always initialised with an empty string (st.text_area(..., value="", ...)) while the submission code later passes that value to the update handler. When an observation that already has a custom query is opened for editing, the text area shows up blank and submitting the form (even without touching that field) overwrites the existing SQL with an empty string. After any edit the observation silently reverts to template-based rendering. Pre-populate the field with the existing value, similar to custom_test_ids_query, so that updates don’t erase user-provided custom queries.

Useful? React with 👍 / 👎.

Comment on lines +349 to +351
Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Oct 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve full_custom_query when updating observations

The new text area for full_custom_query is always initialized with an empty string. When a user edits an existing observation the saved query is not shown, and asdict later sends the empty value back to update_observation, wiping the stored SQL even if the user never touched the field. This causes data loss whenever any other attribute is changed. Prefilling the widget with predefined.full_custom_query (defaulting to None) would keep existing queries intact.

Useful? React with 👍 / 👎.

Comment on lines +349 to +351
Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Oct 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P0 Badge Full custom query lacks read-only validation before execution

full_custom_query is taken straight from the Streamlit form and executed by the runner, but unlike custom_test_ids_query it is never passed through ValidationUtils.validate_sql_query. A user can therefore submit mutating statements (INSERT, DROP, etc.) that will execute under application credentials. Because other SQL fields are explicitly validated to be read-only, this omission is a security regression. The field should be validated or otherwise constrained before rendering/execution.

Useful? React with 👍 / 👎.


metric_tags: list[str] | None = st.multiselect(
"Metric Tags",
Expand Down Expand Up @@ -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,
),
Expand Down
43 changes: 25 additions & 18 deletions tests/services/runners/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
35 changes: 15 additions & 20 deletions tests/services/runners/test_executors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
14 changes: 12 additions & 2 deletions tests/services/runners/test_renderer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"],
)
Expand All @@ -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