refactor: Search pipeline run using Pydantic query parameter

Author: yuechao-qin

cloud_pipelines_backend/api_server_sql.py

@@ -4,7 +4,10 @@
 import json
 import logging
 import typing
-from typing import Any, Optional
+from typing import Annotated, Any, Final, Optional
+
+from fastapi import Query
+from pydantic import BaseModel
 
 if typing.TYPE_CHECKING:
     from cloud_pipelines.orchestration.storage_providers import (
@@ -31,6 +34,9 @@ def _get_current_time() -> datetime.datetime:
 from . import errors
 from .errors import ItemNotFoundError
 
+_PAGE_TOKEN_OFFSET_KEY: Final[str] = "offset"
+_PAGE_TOKEN_FILTER_KEY: Final[str] = "filter"
+
 
 # ==== PipelineJobService
 @dataclasses.dataclass(kw_only=True)
@@ -65,6 +71,13 @@ class ListPipelineJobsResponse:
     next_page_token: str | None = None
 
 
+class ListPipelineRunsParams(BaseModel):
+    filter: str | None = None
+    page_token: str | None = None
+    include_pipeline_names: bool = False
+    include_execution_stats: bool = False
+
+
 import sqlalchemy as sql
 from sqlalchemy import orm
 
@@ -165,45 +178,19 @@ def list(
         self,
         *,
         session: orm.Session,
-        page_token: str | None = None,
-        # page_size: int  = 10,
-        filter: str | None = None,
         current_user: str | None = None,
-        include_pipeline_names: bool = False,
-        include_execution_stats: bool = False,
+        params: Annotated[ListPipelineRunsParams, Query()],
     ) -> ListPipelineJobsResponse:
-        page_token_dict = _decode_page_token(page_token)
-        OFFSET_KEY = "offset"
-        offset = page_token_dict.get(OFFSET_KEY, 0)
-        page_size = 10
+        filter_value, offset = _resolve_filter_value(
+            filter=params.filter,
+            page_token=params.page_token,
+        )
+        where_clauses, next_page_filter_value = _build_filter_where_clauses(
+            filter_value=filter_value,
+            current_user=current_user,
+        )
 
-        FILTER_KEY = "filter"
-        if page_token:
-            filter = page_token_dict.get(FILTER_KEY, None)
-        where_clauses = []
-        parsed_filter = _parse_filter(filter) if filter else {}
-        for key, value in parsed_filter.items():
-            if key == "_text":
-                raise NotImplementedError("Text search is not implemented yet.")
-            elif key == "created_by":
-                if value == "me":
-                    if current_user is None:
-                        # raise ApiServiceError(
-                        #     f"The `created_by:me` filter requires `current_user`."
-                        # )
-                        current_user = ""
-                    value = current_user
-                    # TODO: Maybe make this a bit more robust.
-                    # We need to change the filter since it goes into the next_page_token.
-                    filter = filter.replace(
-                        "created_by:me", f"created_by:{current_user}"
-                    )
-                if value:
-                    where_clauses.append(bts.PipelineRun.created_by == value)
-                else:
-                    where_clauses.append(bts.PipelineRun.created_by == None)
-            else:
-                raise NotImplementedError(f"Unsupported filter {filter}.")
+        page_size = 10
         pipeline_runs = list(
             session.scalars(
                 sql.select(bts.PipelineRun)
@@ -214,50 +201,62 @@ def list(
             ).all()
         )
         next_page_offset = offset + page_size
-        next_page_token_dict = {OFFSET_KEY: next_page_offset, FILTER_KEY: filter}
+        next_page_token_dict = {
+            _PAGE_TOKEN_OFFSET_KEY: next_page_offset,
+            _PAGE_TOKEN_FILTER_KEY: next_page_filter_value,
+        }
         next_page_token = _encode_page_token(next_page_token_dict)
         if len(pipeline_runs) < page_size:
             next_page_token = None
 
-        def create_pipeline_run_response(
-            pipeline_run: bts.PipelineRun,
-        ) -> PipelineRunResponse:
-            response = PipelineRunResponse.from_db(pipeline_run)
-            if include_pipeline_names:
-                pipeline_name = None
-                extra_data = pipeline_run.extra_data or {}
-                if self.PIPELINE_NAME_EXTRA_DATA_KEY in extra_data:
-                    pipeline_name = extra_data[self.PIPELINE_NAME_EXTRA_DATA_KEY]
-                else:
-                    execution_node = session.get(
-                        bts.ExecutionNode, pipeline_run.root_execution_id
-                    )
-                    if execution_node:
-                        task_spec = structures.TaskSpec.from_json_dict(
-                            execution_node.task_spec
-                        )
-                        component_spec = task_spec.component_ref.spec
-                        if component_spec:
-                            pipeline_name = component_spec.name
-                response.pipeline_name = pipeline_name
-            if include_execution_stats:
-                execution_status_stats = self._calculate_execution_status_stats(
-                    session=session, root_execution_id=pipeline_run.root_execution_id
-                )
-                response.execution_status_stats = {
-                    status.value: count
-                    for status, count in execution_status_stats.items()
-                }
-            return response
-
         return ListPipelineJobsResponse(
             pipeline_runs=[
-                create_pipeline_run_response(pipeline_run)
+                self._create_pipeline_run_response(
+                    session=session,
+                    pipeline_run=pipeline_run,
+                    include_pipeline_names=params.include_pipeline_names,
+                    include_execution_stats=params.include_execution_stats,
+                )
                 for pipeline_run in pipeline_runs
             ],
             next_page_token=next_page_token,
         )
 
+    def _create_pipeline_run_response(
+        self,
+        *,
+        session: orm.Session,
+        pipeline_run: bts.PipelineRun,
+        include_pipeline_names: bool,
+        include_execution_stats: bool,
+    ) -> PipelineRunResponse:
+        response = PipelineRunResponse.from_db(pipeline_run)
+        if include_pipeline_names:
+            pipeline_name = None
+            extra_data = pipeline_run.extra_data or {}
+            if self.PIPELINE_NAME_EXTRA_DATA_KEY in extra_data:
+                pipeline_name = extra_data[self.PIPELINE_NAME_EXTRA_DATA_KEY]
+            else:
+                execution_node = session.get(
+                    bts.ExecutionNode, pipeline_run.root_execution_id
+                )
+                if execution_node:
+                    task_spec = structures.TaskSpec.from_json_dict(
+                        execution_node.task_spec
+                    )
+                    component_spec = task_spec.component_ref.spec
+                    if component_spec:
+                        pipeline_name = component_spec.name
+            response.pipeline_name = pipeline_name
+        if include_execution_stats:
+            execution_status_stats = self._calculate_execution_status_stats(
+                session=session, root_execution_id=pipeline_run.root_execution_id
+            )
+            response.execution_status_stats = {
+                status.value: count for status, count in execution_status_stats.items()
+            }
+        return response
+
     def _calculate_execution_status_stats(
         self, session: orm.Session, root_execution_id: bts.IdType
     ) -> dict[bts.ContainerExecutionStatus, int]:
@@ -349,6 +348,58 @@ def delete_annotation(
         session.commit()
 
 
+def _resolve_filter_value(
+    *,
+    filter: str | None,
+    page_token: str | None,
+) -> tuple[str | None, int]:
+    """Decode page_token and return the effective (filter_value, offset).
+
+    If a page_token is present, its stored filter takes precedence over the
+    raw filter parameter (the token carries the resolved filter forward across pages).
+    """
+    page_token_dict = _decode_page_token(page_token)
+    offset = page_token_dict.get(_PAGE_TOKEN_OFFSET_KEY, 0)
+    if page_token:
+        filter = page_token_dict.get(_PAGE_TOKEN_FILTER_KEY, None)
+    return filter, offset
+
+
+def _build_filter_where_clauses(
+    *,
+    filter_value: str | None,
+    current_user: str | None,
+) -> tuple[list[sql.ColumnElement], str | None]:
+    """Parse a filter string into SQLAlchemy WHERE clauses.
+
+    Returns (where_clauses, next_page_filter_value). The second value is the
+    filter string with shorthand values resolved (e.g. "created_by:me" becomes
+    "created_by:alice@example.com") so it can be embedded in the next page token.
+    """
+    where_clauses: list[sql.ColumnElement] = []
+    parsed_filter = _parse_filter(filter_value) if filter_value else {}
+    for key, value in parsed_filter.items():
+        if key == "_text":
+            raise NotImplementedError("Text search is not implemented yet.")
+        elif key == "created_by":
+            if value == "me":
+                if current_user is None:
+                    current_user = ""
+                value = current_user
+                # TODO: Maybe make this a bit more robust.
+                # We need to change the filter since it goes into the next_page_token.
+                filter_value = filter_value.replace(
+                    "created_by:me", f"created_by:{current_user}"
+                )
+            if value:
+                where_clauses.append(bts.PipelineRun.created_by == value)
+            else:
+                where_clauses.append(bts.PipelineRun.created_by == None)
+        else:
+            raise NotImplementedError(f"Unsupported filter {filter_value}.")
+    return where_clauses, filter_value
+
+
 def _decode_page_token(page_token: str) -> dict[str, Any]:
     return json.loads(base64.b64decode(page_token)) if page_token else {}