feat: Search annotations in pipeline run API

Author: yuechao-qin

cloud_pipelines_backend/api_router.py

@@ -125,6 +125,15 @@ def handle_item_already_exists_error(
             content={"message": str(exc)},
         )
 
+    @app.exception_handler(NotImplementedError)
+    def handle_not_implemented_error(
+        request: fastapi.Request, exc: NotImplementedError
+    ):
+        return fastapi.responses.JSONResponse(
+            status_code=501,
+            content={"detail": str(exc)},
+        )
+
     get_user_details_dependency = fastapi.Depends(user_details_getter)
 
     def get_user_name(

cloud_pipelines_backend/api_server_sql.py

@@ -1,10 +1,8 @@
-import base64
 import dataclasses
 import datetime
-import json
 import logging
 import typing
-from typing import Annotated, Any, Final, Optional
+from typing import Annotated, Any, Optional
 
 from fastapi import HTTPException, Query
 from pydantic import BaseModel
@@ -33,10 +31,8 @@ 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"
+from . import filter_query_sql
+from .filter_query_sql import PageToken
 
 
 # ==== PipelineJobService
@@ -183,29 +179,18 @@ 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.",
+        page_size = 10
+        try:
+            where_clauses, offset, next_token = filter_query_sql.build_list_filters(
+                filter_value=params.filter,
+                filter_query_value=params.filter_query,
+                page_token_value=params.page_token,
+                current_user=current_user,
+                page_size=page_size,
             )
+        except filter_query_sql.MutuallyExclusiveFilterError as e:
+            raise HTTPException(status_code=422, detail=str(e))
 
-        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,
-        )
-
-        page_size = 10
         pipeline_runs = list(
             session.scalars(
                 sql.select(bts.PipelineRun)
@@ -215,14 +200,10 @@ def list(
                 .limit(page_size)
             ).all()
         )
-        next_page_offset = offset + page_size
-        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
+
+        next_page_token = (
+            next_token.encode() if len(pipeline_runs) >= page_size else None
+        )
 
         return ListPipelineJobsResponse(
             pipeline_runs=[
@@ -363,82 +344,6 @@ 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 {}
-
-
-def _encode_page_token(page_token_dict: dict[str, Any]) -> str:
-    return (base64.b64encode(json.dumps(page_token_dict).encode("utf8"))).decode(
-        "utf-8"
-    )
-
-
-def _parse_filter(filter: str) -> dict[str, str]:
-    # TODO: Improve
-    parts = filter.strip().split()
-    parsed_filter = {}
-    for part in parts:
-        key, sep, value = part.partition(":")
-        if sep:
-            parsed_filter[key] = value
-        else:
-            parsed_filter.setdefault("_text", "")
-            parsed_filter["_text"] += part
-    return parsed_filter
-
-
 # ========== ExecutionNodeApiService_Sql
 
 

cloud_pipelines_backend/database_ops.py

@@ -77,3 +77,18 @@ def migrate_db(db_engine: sqlalchemy.Engine):
     for index in bts.ExecutionNode.__table__.indexes:
         if index.name == "ix_execution_node_container_execution_cache_key":
             index.create(db_engine, checkfirst=True)
+
+    # Covering index for annotation-based filter_query predicates.
+    # Column order: (key, pipeline_run_id, value) enables B-tree seek on key,
+    # ordered pipeline_run_id scan, and avoids table lookups for all predicates.
+    annotation_index = sqlalchemy.Index(
+        "ix_pipeline_run_annotation_key_run_id_value",
+        bts.PipelineRunAnnotation.key,
+        bts.PipelineRunAnnotation.pipeline_run_id,
+        bts.PipelineRunAnnotation.value,
+    )
+    annotation_index.create(db_engine, checkfirst=True)
+
+    # TODO: I believe this is needed with the comments above?
+    # Workaround for https://github.com/sqlalchemy/sqlalchemy/issues/12965
+    bts.PipelineRunAnnotation.__table__.indexes.discard(annotation_index)