From e152073d750059b3f21ace93b3e919cf7b06a99c Mon Sep 17 00:00:00 2001 From: Yue Chao Qin Date: Thu, 26 Feb 2026 00:51:03 -0800 Subject: [PATCH] feat: Search within date range in pipeline run API --- .../filter_query_models.py | 6 +- cloud_pipelines_backend/filter_query_sql.py | 62 ++- tests/test_api_server_sql.py | 362 +++++++++++++++- tests/test_filter_query_sql.py | 396 +++++++++++++++++- 4 files changed, 814 insertions(+), 12 deletions(-) diff --git a/cloud_pipelines_backend/filter_query_models.py b/cloud_pipelines_backend/filter_query_models.py index 6daed83..537874f 100644 --- a/cloud_pipelines_backend/filter_query_models.py +++ b/cloud_pipelines_backend/filter_query_models.py @@ -99,9 +99,13 @@ def key(self) -> str: return self.value_equals.key -class TimeRangePredicate(_BaseModel): +class TimeRangePredicate(KeyPredicateBase): time_range: TimeRange + @property + def key(self) -> str: + return self.time_range.key + LeafPredicate = ( KeyExistsPredicate diff --git a/cloud_pipelines_backend/filter_query_sql.py b/cloud_pipelines_backend/filter_query_sql.py index ef887e1..7f2b6d5 100644 --- a/cloud_pipelines_backend/filter_query_sql.py +++ b/cloud_pipelines_backend/filter_query_sql.py @@ -1,4 +1,5 @@ import base64 +import datetime import json import enum from typing import Any, Final @@ -16,6 +17,7 @@ class PipelineRunAnnotationSystemKey(enum.StrEnum): CREATED_BY = f"{_PIPELINE_RUN_KEY_PREFIX}created_by" NAME = f"{_PIPELINE_RUN_KEY_PREFIX}name" + CREATED_AT = f"{_PIPELINE_RUN_KEY_PREFIX}date.created_at" SYSTEM_KEY_SUPPORTED_PREDICATES: dict[PipelineRunAnnotationSystemKey, set[type]] = { @@ -30,6 +32,9 @@ class PipelineRunAnnotationSystemKey(enum.StrEnum): filter_query_models.ValueContainsPredicate, filter_query_models.ValueInPredicate, }, + PipelineRunAnnotationSystemKey.CREATED_AT: { + filter_query_models.TimeRangePredicate, + }, } # --------------------------------------------------------------------------- @@ -301,8 +306,9 @@ def _predicate_to_clause( return _value_contains_to_clause(predicate=predicate) case filter_query_models.ValueInPredicate(): return _value_in_to_clause(predicate=predicate) + case filter_query_models.TimeRangePredicate(): + return _time_range_to_clause(predicate=predicate) case _: - # TODO: TimeRangePredicate -- not supported currently, will be supported in the future. raise NotImplementedError( f"Predicate type {type(predicate).__name__} is not yet implemented." ) @@ -363,3 +369,57 @@ def _value_in_to_clause( bts.PipelineRunAnnotation.value.in_(predicate.value_in.values), ], ) + + +# --------------------------------------------------------------------------- +# Column-based predicates (bypass annotation table) +# --------------------------------------------------------------------------- + + +def _time_range_to_clause( + *, predicate: filter_query_models.TimeRangePredicate +) -> sql.ColumnElement: + """Build a WHERE clause for pipeline_run.created_at from a time range. + + Pydantic's AwareDatetime preserves the original timezone offset, so we + must normalize to naive UTC before comparing against the DB column. + + The DB stores "naive UTC" datetimes -- the values represent UTC but carry + no timezone label. For example, the DB stores '2024-01-01 02:30:00', not + '2024-01-01 02:30:00+00:00'. The UtcDateTime type decorator (in + backend_types_sql.py) strips tzinfo on write and re-attaches UTC on read. + + Conversion pipeline for input '2024-01-01T08:00:00+05:30': + + API request (JSON string) + '2024-01-01T08:00:00+05:30' + | + v + Pydantic AwareDatetime (preserves offset) + datetime(2024, 1, 1, 8, 0, 0, tzinfo=+05:30) + | + v .astimezone(utc) -- converts 08:00 - 05:30 = 02:30 + UTC-aware datetime + datetime(2024, 1, 1, 2, 30, 0, tzinfo=UTC) + | + v .replace(tzinfo=None) -- strips timezone label + Naive datetime + datetime(2024, 1, 1, 2, 30, 0) + | + v SQLAlchemy literal_binds -- adds microsecond precision + SQL string + '2024-01-01 02:30:00.000000' <-- matches DB storage format + """ + tr = predicate.time_range + if tr.key != PipelineRunAnnotationSystemKey.CREATED_AT: + raise errors.ApiValidationError( + "time_range only supports key " + f"{PipelineRunAnnotationSystemKey.CREATED_AT!r}, got {tr.key!r}" + ) + # Convert aware datetimes to naive UTC to match DB storage format. + start_utc = tr.start_time.astimezone(datetime.timezone.utc).replace(tzinfo=None) + clauses: list[sql.ColumnElement] = [bts.PipelineRun.created_at >= start_utc] + if tr.end_time is not None: + end_utc = tr.end_time.astimezone(datetime.timezone.utc).replace(tzinfo=None) + clauses.append(bts.PipelineRun.created_at < end_utc) + return sql.and_(*clauses) diff --git a/tests/test_api_server_sql.py b/tests/test_api_server_sql.py index a685b1a..8d7c330 100644 --- a/tests/test_api_server_sql.py +++ b/tests/test_api_server_sql.py @@ -1,6 +1,8 @@ +import datetime import json import pytest +import sqlalchemy from sqlalchemy import orm from cloud_pipelines_backend import api_server_sql @@ -908,7 +910,74 @@ def test_no_match(self, session_factory, service): ) assert len(result.pipeline_runs) == 0 - def test_time_range_raises_not_implemented(self, session_factory, service): + def _create_run_at(self, *, session_factory, service, created_at, **kwargs): + """Create a run and override its created_at timestamp.""" + run = _create_run( + session_factory, service, root_task=_make_task_spec(), **kwargs + ) + with session_factory() as session: + session.execute( + sqlalchemy.update(bts.PipelineRun) + .where(bts.PipelineRun.id == run.id) + .values(created_at=created_at) + ) + session.commit() + return run + + def _utc(self, *, year, month, day, hour=0, minute=0, second=0): + return datetime.datetime( + year, + month, + day, + hour, + minute, + second, + tzinfo=datetime.timezone.utc, + ) + + def test_list_filter_query_time_range(self, session_factory, service): + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1), + ) + feb = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=2, day=1), + ) + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=3, day=1), + ) + + fq = json.dumps( + { + "and": [ + { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-01-15T00:00:00Z", + "end_time": "2024-02-15T00:00:00Z", + } + } + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 1 + assert result.pipeline_runs[0].id == feb.id + + def test_list_filter_query_time_range_start_boundary( + self, session_factory, service + ): + run = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1), + ) fq = json.dumps( { "and": [ @@ -922,11 +991,292 @@ def test_time_range_raises_not_implemented(self, session_factory, service): } ) with session_factory() as session: - with pytest.raises(NotImplementedError, match="TimeRangePredicate"): - service.list( - session=session, - filter_query=fq, - ) + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 1 + assert result.pipeline_runs[0].id == run.id + + def test_list_filter_query_time_range_end_boundary(self, session_factory, service): + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=2, day=1), + ) + + fq = json.dumps( + { + "and": [ + { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-02-01T00:00:00Z", + } + } + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 0 + + def test_list_filter_query_time_range_start_only(self, session_factory, service): + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1), + ) + mar = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=3, day=1), + ) + + fq = json.dumps( + { + "and": [ + { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-02-01T00:00:00Z", + } + } + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 1 + assert result.pipeline_runs[0].id == mar.id + + def test_list_filter_query_time_range_not(self, session_factory, service): + jan = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1), + ) + feb = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=2, day=1), + ) + mar = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=3, day=1), + ) + + fq = json.dumps( + { + "and": [ + { + "not": { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-01-15T00:00:00Z", + "end_time": "2024-02-15T00:00:00Z", + } + } + } + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + result_ids = {r.id for r in result.pipeline_runs} + assert feb.id not in result_ids + assert jan.id in result_ids + assert mar.id in result_ids + + def test_list_filter_query_time_range_after_annotation( + self, session_factory, service + ): + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1), + created_by="alice", + ) + feb = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=2, day=1), + created_by="alice", + ) + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=3, day=1), + created_by="bob", + ) + + fq = json.dumps( + { + "and": [ + { + "value_equals": { + "key": "system/pipeline_run.created_by", + "value": "alice", + } + }, + { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-01-15T00:00:00Z", + "end_time": "2024-03-01T00:00:00Z", + } + }, + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 1 + assert result.pipeline_runs[0].id == feb.id + + def test_list_filter_query_time_range_with_annotation( + self, session_factory, service + ): + jan = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1), + created_by="test-user", + ) + self._set_annotation( + session_factory=session_factory, + service=service, + run_id=jan.id, + key="team", + value="ml-ops", + ) + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=2, day=1), + created_by="test-user", + ) + mar = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=3, day=1), + created_by="test-user", + ) + self._set_annotation( + session_factory=session_factory, + service=service, + run_id=mar.id, + key="team", + value="ml-ops", + ) + + fq = json.dumps( + { + "and": [ + {"value_equals": {"key": "team", "value": "ml-ops"}}, + { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-02-01T00:00:00Z", + } + }, + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 1 + assert result.pipeline_runs[0].id == mar.id + + def test_list_filter_query_time_range_before_annotation( + self, session_factory, service + ): + jan = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1), + created_by="test-user", + ) + self._set_annotation( + session_factory=session_factory, + service=service, + run_id=jan.id, + key="team", + value="ml-ops", + ) + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=2, day=1), + created_by="test-user", + ) + mar = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=3, day=1), + created_by="test-user", + ) + self._set_annotation( + session_factory=session_factory, + service=service, + run_id=mar.id, + key="team", + value="ml-ops", + ) + + fq = json.dumps( + { + "and": [ + { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-02-01T00:00:00Z", + } + }, + {"value_equals": {"key": "team", "value": "ml-ops"}}, + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 1 + assert result.pipeline_runs[0].id == mar.id + + def test_list_filter_query_time_range_offset_timezone( + self, session_factory, service + ): + self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1, hour=2, minute=0), + ) + run_b = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1, hour=2, minute=30), + ) + run_c = self._create_run_at( + session_factory=session_factory, + service=service, + created_at=self._utc(year=2024, month=1, day=1, hour=6, minute=0), + ) + + fq = json.dumps( + { + "and": [ + { + "time_range": { + "key": "system/pipeline_run.date.created_at", + "start_time": "2024-01-01T08:00:00+05:30", + } + } + ] + } + ) + with session_factory() as session: + result = service.list(session=session, filter_query=fq) + assert len(result.pipeline_runs) == 2 + returned_ids = {r.id for r in result.pipeline_runs} + assert returned_ids == {run_b.id, run_c.id} def test_pagination_preserves_filter_query(self, session_factory, service): for _ in range(12): diff --git a/tests/test_filter_query_sql.py b/tests/test_filter_query_sql.py index d893210..ba90a27 100644 --- a/tests/test_filter_query_sql.py +++ b/tests/test_filter_query_sql.py @@ -1,5 +1,6 @@ import json +import pydantic import pytest import sqlalchemy as sql from sqlalchemy.dialects import sqlite as sqlite_dialect @@ -166,17 +167,365 @@ def test_nested_and_or(self): ) -class TestUnsupportedPredicate: - def test_time_range_raises_not_implemented(self): +class TestTimeRangePredicate: + _KEY = filter_query_sql.PipelineRunAnnotationSystemKey.CREATED_AT + + _EXISTS_VALUE_EQUALS_TEAM = ( + "(EXISTS (SELECT pipeline_run_annotation.pipeline_run_id" + " FROM pipeline_run_annotation, pipeline_run" + " WHERE pipeline_run_annotation.pipeline_run_id = pipeline_run.id" + " AND pipeline_run_annotation.\"key\" = 'team'" + " AND pipeline_run_annotation.value = 'ml-ops'))" + ) + _EXISTS_KEY_EXISTS_ENV = ( + "(EXISTS (SELECT pipeline_run_annotation.pipeline_run_id" + " FROM pipeline_run_annotation, pipeline_run" + " WHERE pipeline_run_annotation.pipeline_run_id = pipeline_run.id" + " AND pipeline_run_annotation.\"key\" = 'env'))" + ) + _EXISTS_KEY_EXISTS_DEPRECATED = ( + "(EXISTS (SELECT pipeline_run_annotation.pipeline_run_id" + " FROM pipeline_run_annotation, pipeline_run" + " WHERE pipeline_run_annotation.pipeline_run_id = pipeline_run.id" + " AND pipeline_run_annotation.\"key\" = 'deprecated'))" + ) + + def test_time_range_with_start_and_end(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-15T10:30:00Z", + "end_time": "2024-02-15T23:59:59.999Z", + } + } + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "pipeline_run.created_at >= '2024-01-15 10:30:00.000000'" + " AND pipeline_run.created_at < '2024-02-15 23:59:59.999000'" + ) + + def test_time_range_start_only(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-06-01T00:00:00+00:00", + } + } + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == "pipeline_run.created_at >= '2024-06-01 00:00:00.000000'" + + def test_time_range_with_milliseconds(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T12:00:00.500Z", + "end_time": "2024-01-02T12:00:00.999Z", + } + } + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "pipeline_run.created_at >= '2024-01-01 12:00:00.500000'" + " AND pipeline_run.created_at < '2024-01-02 12:00:00.999000'" + ) + + def test_time_range_with_offset_timezone(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T08:00:00+05:30", + } + } + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ("pipeline_run.created_at >= '2024-01-01 02:30:00.000000'") + + def test_time_range_not(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "not": { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-15T00:00:00Z", + "end_time": "2024-02-15T00:00:00Z", + } + } + } + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "NOT (pipeline_run.created_at >= '2024-01-15 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-02-15 00:00:00.000000')" + ) + + def test_time_range_combined_with_annotation(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-02-01T00:00:00Z", + } + }, + {"value_equals": {"key": "team", "value": "ml-ops"}}, + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "pipeline_run.created_at >= '2024-01-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-02-01 00:00:00.000000'" + f" AND {self._EXISTS_VALUE_EQUALS_TEAM}" + ) + + def test_time_range_not_first_in_list(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + {"value_equals": {"key": "team", "value": "ml-ops"}}, + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-02-01T00:00:00Z", + } + }, + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + f"{self._EXISTS_VALUE_EQUALS_TEAM}" + " AND pipeline_run.created_at >= '2024-01-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-02-01 00:00:00.000000'" + ) + + def test_time_range_last_in_list(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + {"key_exists": {"key": "env"}}, + {"value_equals": {"key": "team", "value": "ml-ops"}}, + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-06-01T00:00:00Z", + } + }, + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + f"{self._EXISTS_KEY_EXISTS_ENV}" + f" AND {self._EXISTS_VALUE_EQUALS_TEAM}" + " AND pipeline_run.created_at >= '2024-01-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-06-01 00:00:00.000000'" + ) + + def test_time_range_nested_in_or(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "or": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-02-01T00:00:00Z", + } + }, + {"value_equals": {"key": "team", "value": "ml-ops"}}, + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "pipeline_run.created_at >= '2024-01-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-02-01 00:00:00.000000'" + f" OR {self._EXISTS_VALUE_EQUALS_TEAM}" + ) + + def test_time_range_deeply_nested(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "or": [ + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-02-01T00:00:00Z", + } + }, + { + "value_equals": { + "key": "team", + "value": "ml-ops", + } + }, + ] + }, + {"key_exists": {"key": "deprecated"}}, + ] + } + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "pipeline_run.created_at >= '2024-01-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-02-01 00:00:00.000000'" + f" AND {self._EXISTS_VALUE_EQUALS_TEAM}" + f" OR {self._EXISTS_KEY_EXISTS_DEPRECATED}" + ) + + def test_multiple_time_ranges_in_and(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-06-01T00:00:00Z", + } + }, + { + "time_range": { + "key": self._KEY, + "start_time": "2024-03-01T00:00:00Z", + "end_time": "2024-04-01T00:00:00Z", + } + }, + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "pipeline_run.created_at >= '2024-01-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-06-01 00:00:00.000000'" + " AND pipeline_run.created_at >= '2024-03-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-04-01 00:00:00.000000'" + ) + + def test_multiple_time_ranges_in_or(self): + fq = filter_query_models.FilterQuery.model_validate( + { + "or": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00Z", + "end_time": "2024-02-01T00:00:00Z", + } + }, + { + "time_range": { + "key": self._KEY, + "start_time": "2024-03-01T00:00:00Z", + "end_time": "2024-04-01T00:00:00Z", + } + }, + ] + }, + ) + compiled = _compile( + filter_query_sql.filter_query_to_where_clause(filter_query=fq) + ) + assert compiled == ( + "pipeline_run.created_at >= '2024-01-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-02-01 00:00:00.000000'" + " OR pipeline_run.created_at >= '2024-03-01 00:00:00.000000'" + " AND pipeline_run.created_at < '2024-04-01 00:00:00.000000'" + ) + + def test_time_range_invalid_key_rejected(self): predicate = filter_query_models.TimeRangePredicate( time_range=filter_query_models.TimeRange( - key="system/pipeline_run.date.created_at", + key="custom/my_date", start_time="2024-01-01T00:00:00Z", ) ) - with pytest.raises(NotImplementedError, match="TimeRangePredicate"): + with pytest.raises(errors.ApiValidationError, match="only supports key"): filter_query_sql._predicate_to_clause(predicate=predicate) + def test_time_range_system_key_registered(self): + supported = filter_query_sql.SYSTEM_KEY_SUPPORTED_PREDICATES + assert filter_query_sql.PipelineRunAnnotationSystemKey.CREATED_AT in supported + assert ( + filter_query_models.TimeRangePredicate + in supported[filter_query_sql.PipelineRunAnnotationSystemKey.CREATED_AT] + ) + + def test_time_range_naive_datetime_rejected(self): + with pytest.raises(pydantic.ValidationError): + filter_query_models.FilterQuery.model_validate( + { + "and": [ + { + "time_range": { + "key": self._KEY, + "start_time": "2024-01-01T00:00:00", + } + } + ] + }, + ) + class TestPageToken: def test_decode_none(self): @@ -570,3 +919,42 @@ def test_check_predicate_allowed_name_value_in(self): ) ) filter_query_sql._check_predicate_allowed(predicate=pred) + + def test_check_predicate_allowed_key_exists_on_created_at_rejected(self): + pred = filter_query_models.KeyExistsPredicate( + key_exists=filter_query_models.KeyExists( + key=filter_query_sql.PipelineRunAnnotationSystemKey.CREATED_AT + ) + ) + with pytest.raises(errors.ApiValidationError, match="not supported"): + filter_query_sql._check_predicate_allowed(predicate=pred) + + def test_check_predicate_allowed_value_equals_on_created_at_rejected(self): + pred = filter_query_models.ValueEqualsPredicate( + value_equals=filter_query_models.ValueEquals( + key=filter_query_sql.PipelineRunAnnotationSystemKey.CREATED_AT, + value="2024-01-01", + ) + ) + with pytest.raises(errors.ApiValidationError, match="not supported"): + filter_query_sql._check_predicate_allowed(predicate=pred) + + def test_check_predicate_allowed_time_range_on_created_by_rejected(self): + pred = filter_query_models.TimeRangePredicate( + time_range=filter_query_models.TimeRange( + key=filter_query_sql.PipelineRunAnnotationSystemKey.CREATED_BY, + start_time="2024-01-01T00:00:00Z", + ) + ) + with pytest.raises(errors.ApiValidationError, match="not supported"): + filter_query_sql._check_predicate_allowed(predicate=pred) + + def test_check_predicate_allowed_time_range_on_name_rejected(self): + pred = filter_query_models.TimeRangePredicate( + time_range=filter_query_models.TimeRange( + key=filter_query_sql.PipelineRunAnnotationSystemKey.NAME, + start_time="2024-01-01T00:00:00Z", + ) + ) + with pytest.raises(errors.ApiValidationError, match="not supported"): + filter_query_sql._check_predicate_allowed(predicate=pred)