feat: Create Pydantic models for Search Pipeline Run API

Author: yuechao-qin

cloud_pipelines_backend/api_server_sql.py

@@ -6,7 +6,7 @@
 import typing
 from typing import Annotated, Any, Final, Optional
 
-from fastapi import Query
+from fastapi import HTTPException, Query
 from pydantic import BaseModel
 
 if typing.TYPE_CHECKING:
@@ -33,6 +33,7 @@ def _get_current_time() -> datetime.datetime:
 from . import backend_types_sql as bts
 from . import errors
 from .errors import ItemNotFoundError
+from .filter_query_models import FilterQuery
 
 _PAGE_TOKEN_OFFSET_KEY: Final[str] = "offset"
 _PAGE_TOKEN_FILTER_KEY: Final[str] = "filter"
@@ -73,6 +74,7 @@ class ListPipelineJobsResponse:
 
 class ListPipelineRunsParams(BaseModel):
     filter: str | None = None
+    filter_query: str | None = None
     page_token: str | None = None
     include_pipeline_names: bool = False
     include_execution_stats: bool = False
@@ -181,6 +183,19 @@ def list(
         current_user: str | None = None,
         params: Annotated[ListPipelineRunsParams, Query()],
     ) -> ListPipelineJobsResponse:
+        if params.filter and params.filter_query:
+            raise HTTPException(
+                status_code=422,
+                detail="Cannot use both 'filter' and 'filter_query'. Use one or the other.",
+            )
+
+        if params.filter_query:
+            FilterQuery.model_validate_json(params.filter_query)
+            raise HTTPException(
+                status_code=501,
+                detail="filter_query is not yet implemented.",
+            )
+
         filter_value, offset = _resolve_filter_value(
             filter=params.filter,
             page_token=params.page_token,

cloud_pipelines_backend/filter_query_models.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+from pydantic import AwareDatetime, BaseModel, Field, StringConstraints, model_validator
+
+NonEmptyStr = Annotated[str, StringConstraints(min_length=1)]
+
+
+class _StrictBase(BaseModel):
+    model_config = {"extra": "forbid"}
+
+
+# --- Leaf argument models ---
+
+
+class KeyExists(BaseModel):
+    key: NonEmptyStr
+
+
+class ValueContains(BaseModel):
+    key: NonEmptyStr
+    value_substring: NonEmptyStr
+
+
+class ValueIn(BaseModel):
+    key: NonEmptyStr
+    values: list[NonEmptyStr] = Field(min_length=1)
+
+
+class ValueEquals(BaseModel):
+    key: NonEmptyStr
+    value: NonEmptyStr
+
+
+class TimeRange(BaseModel):
+    """AwareDatetime requires timezone info (e.g. "2024-01-01T00:00:00Z").
+    Naive datetimes like "2024-01-01T00:00:00" are rejected, preventing
+    ambiguous timestamps that could silently resolve to the wrong timezone."""
+
+    key: NonEmptyStr
+    start_time: AwareDatetime
+    end_time: AwareDatetime | None = None
+
+
+# --- Predicate wrapper models (one field each) ---
+
+
+class KeyExistsPredicate(_StrictBase):
+    key_exists: KeyExists
+
+
+class ValueContainsPredicate(_StrictBase):
+    value_contains: ValueContains
+
+
+class ValueInPredicate(_StrictBase):
+    value_in: ValueIn
+
+
+class ValueEqualsPredicate(_StrictBase):
+    value_equals: ValueEquals
+
+
+class TimeRangePredicate(_StrictBase):
+    time_range: TimeRange
+
+
+LeafPredicate = (
+    KeyExistsPredicate
+    | ValueContainsPredicate
+    | ValueInPredicate
+    | ValueEqualsPredicate
+    | TimeRangePredicate
+)
+
+
+class NotPredicate(_StrictBase):
+    not_: LeafPredicate = Field(alias="not")
+
+
+class AndPredicate(_StrictBase):
+    and_: list["Predicate"] = Field(alias="and", min_length=1)
+
+
+class OrPredicate(_StrictBase):
+    or_: list["Predicate"] = Field(alias="or", min_length=1)
+
+
+Predicate = (
+    KeyExistsPredicate
+    | ValueContainsPredicate
+    | ValueInPredicate
+    | ValueEqualsPredicate
+    | TimeRangePredicate
+    | NotPredicate
+    | AndPredicate
+    | OrPredicate
+)
+
+# Resolve forward reference to "Predicate" in recursive and/or models
+AndPredicate.model_rebuild()
+OrPredicate.model_rebuild()
+
+
+class FilterQuery(_StrictBase):
+    """Root: must be exactly one of {"and": [...]} or {"or": [...]}."""
+
+    and_: list[Predicate] | None = Field(None, alias="and", min_length=1)
+    or_: list[Predicate] | None = Field(None, alias="or", min_length=1)
+
+    @model_validator(mode="after")
+    def _exactly_one_root_operator(self) -> FilterQuery:
+        has_and = self.and_ is not None
+        has_or = self.or_ is not None
+        if has_and == has_or:
+            raise ValueError("FilterQuery root must have exactly one of 'and' or 'or'.")
+        return self

tests/test_api_server_sql.py

@@ -551,3 +551,44 @@ def test_text_search_raises(self):
                 filter_value="some_text_without_colon",
                 current_user=None,
             )
+
+
+class TestFilterQueryApiWiring:
+    def test_filter_query_returns_501(self, session_factory, service):
+        from fastapi import HTTPException
+
+        valid_json = '{"and": [{"key_exists": {"key": "team"}}]}'
+        with session_factory() as session:
+            with pytest.raises(HTTPException) as exc_info:
+                service.list(
+                    session=session,
+                    params=_default_params(filter_query=valid_json),
+                )
+            assert exc_info.value.status_code == 501
+            assert "not yet implemented" in exc_info.value.detail.lower()
+
+    def test_filter_query_validates_before_501(self, session_factory, service):
+        from pydantic import ValidationError
+
+        invalid_json = '{"bad_key": "not_valid"}'
+        with session_factory() as session:
+            with pytest.raises(ValidationError):
+                service.list(
+                    session=session,
+                    params=_default_params(filter_query=invalid_json),
+                )
+
+    def test_mutual_exclusivity_rejected(self, session_factory, service):
+        from fastapi import HTTPException
+
+        with session_factory() as session:
+            with pytest.raises(HTTPException) as exc_info:
+                service.list(
+                    session=session,
+                    params=_default_params(
+                        filter="created_by:alice",
+                        filter_query='{"and": [{"key_exists": {"key": "team"}}]}',
+                    ),
+                )
+            assert exc_info.value.status_code == 422
+            assert "Cannot use both" in exc_info.value.detail