feat: Search created by user 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(errors.FilterValidationError)
+    def handle_filter_validation_error(
+        request: fastapi.Request, exc: errors.FilterValidationError
+    ):
+        return fastapi.responses.JSONResponse(
+            status_code=422,
+            content={"detail": str(exc)},
+        )
+
     @app.exception_handler(NotImplementedError)
     def handle_not_implemented_error(
         request: fastapi.Request, exc: NotImplementedError

cloud_pipelines_backend/api_server_sql.py

@@ -30,9 +30,10 @@ def _get_current_time() -> datetime.datetime:
 from . import component_structures as structures
 from . import backend_types_sql as bts
 from . import errors
-from .errors import ItemNotFoundError
+from .errors import InvalidAnnotationKeyError, ItemNotFoundError
 from . import filter_query_sql
-from .filter_query_sql import PageToken
+
+_SYSTEM_KEY_RESERVED_MSG = f"Annotation keys starting with {filter_query_sql.SYSTEM_KEY_PREFIX!r} are reserved for system use."
 
 
 # ==== PipelineJobService
@@ -119,6 +120,19 @@ def create(
                 },
             )
             session.add(pipeline_run)
+            # Mirror created_by into the annotations table so it's searchable
+            # via filter_query like any other annotation.
+            if created_by is not None:
+                # Flush to populate pipeline_run.id (server-generated) before inserting the annotation FK.
+                # TODO: Use ORM relationship instead of explicit flush + manual FK assignment.
+                session.flush()
+                session.add(
+                    bts.PipelineRunAnnotation(
+                        pipeline_run_id=pipeline_run.id,
+                        key=filter_query_sql.SystemKey.CREATED_BY,
+                        value=created_by,
+                    )
+                )
             session.commit()
 
         session.refresh(pipeline_run)
@@ -180,16 +194,13 @@ def list(
         params: Annotated[ListPipelineRunsParams, Query()],
     ) -> ListPipelineJobsResponse:
         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))
+        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,
+        )
 
         pipeline_runs = list(
             session.scalars(
@@ -309,6 +320,8 @@ def set_annotation(
         user_name: str | None = None,
         skip_user_check: bool = False,
     ):
+        if key.startswith(filter_query_sql.SYSTEM_KEY_PREFIX):
+            raise InvalidAnnotationKeyError(_SYSTEM_KEY_RESERVED_MSG)
         pipeline_run = session.get(bts.PipelineRun, id)
         if not pipeline_run:
             raise ItemNotFoundError(f"Pipeline run {id} not found.")
@@ -331,6 +344,8 @@ def delete_annotation(
         user_name: str | None = None,
         skip_user_check: bool = False,
     ):
+        if key.startswith(filter_query_sql.SYSTEM_KEY_PREFIX):
+            raise InvalidAnnotationKeyError(_SYSTEM_KEY_RESERVED_MSG)
         pipeline_run = session.get(bts.PipelineRun, id)
         if not pipeline_run:
             raise ItemNotFoundError(f"Pipeline run {id} not found.")

cloud_pipelines_backend/database_ops.py

@@ -1,6 +1,8 @@
 import sqlalchemy
+from sqlalchemy import orm
 
 from . import backend_types_sql as bts
+from . import filter_query_sql
 
 
 def create_db_engine_and_migrate_db(
@@ -89,6 +91,36 @@ def migrate_db(db_engine: sqlalchemy.Engine):
     )
     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)
+
+    backfill_created_by_annotations(db_engine)
+
+
+def backfill_created_by_annotations(db_engine: sqlalchemy.Engine):
+    """Copy pipeline_run.created_by into pipeline_run_annotation so
+    annotation-based search works for created_by.
+
+    Idempotent -- skips rows that already have the annotation.
+    """
+    with orm.Session(db_engine) as session:
+        stmt = sqlalchemy.insert(bts.PipelineRunAnnotation).from_select(
+            ["pipeline_run_id", "key", "value"],
+            sqlalchemy.select(
+                bts.PipelineRun.id,
+                sqlalchemy.literal(filter_query_sql.SystemKey.CREATED_BY),
+                bts.PipelineRun.created_by,
+            ).where(
+                bts.PipelineRun.created_by.isnot(None),
+                # NOT EXISTS makes the backfill idempotent
+                ~sqlalchemy.exists(
+                    sqlalchemy.select(bts.PipelineRunAnnotation.pipeline_run_id).where(
+                        bts.PipelineRunAnnotation.pipeline_run_id == bts.PipelineRun.id,
+                        bts.PipelineRunAnnotation.key
+                        == filter_query_sql.SystemKey.CREATED_BY,
+                    )
+                ),
+            ),
+        )
+        session.execute(stmt)
+        session.commit()

cloud_pipelines_backend/errors.py

@@ -8,3 +8,17 @@ class ItemAlreadyExistsError(Exception):
 
 class PermissionError(Exception):
     pass
+
+
+class FilterValidationError(Exception):
+    """Base for all filter/annotation validation errors -> 422."""
+
+    pass
+
+
+class MutuallyExclusiveFilterError(FilterValidationError):
+    pass
+
+
+class InvalidAnnotationKeyError(FilterValidationError):
+    pass

cloud_pipelines_backend/filter_query_sql.py

@@ -1,21 +1,44 @@
 import base64
 import dataclasses
 import json
+from enum import StrEnum
+from typing import Final
 
 import sqlalchemy as sql
 
 from . import backend_types_sql as bts
+from .errors import (
+    InvalidAnnotationKeyError,
+    MutuallyExclusiveFilterError,
+)
 from .filter_query_models import (
     AndPredicate,
     FilterQuery,
     KeyExistsPredicate,
     NotPredicate,
     OrPredicate,
+    Predicate,
     ValueContainsPredicate,
+    ValueEquals,
     ValueEqualsPredicate,
     ValueInPredicate,
 )
 
+SYSTEM_KEY_PREFIX: Final[str] = "system/"
+
+
+class SystemKey(StrEnum):
+    CREATED_BY = f"{SYSTEM_KEY_PREFIX}pipeline_run.created_by"
+
+
+SYSTEM_KEY_SUPPORTED_PREDICATES: dict[SystemKey, set[type]] = {
+    SystemKey.CREATED_BY: {
+        KeyExistsPredicate,
+        ValueEqualsPredicate,
+        ValueInPredicate,
+    },
+}
+
 # ---------------------------------------------------------------------------
 # PageToken
 # ---------------------------------------------------------------------------
@@ -39,8 +62,89 @@ def decode(cls, token: str | None) -> "PageToken":
         return cls(**json.loads(base64.b64decode(token)))
 
 
-class MutuallyExclusiveFilterError(ValueError):
-    pass
+# ---------------------------------------------------------------------------
+# SystemKey validation and resolution
+# ---------------------------------------------------------------------------
+
+
+def _get_predicate_key(*, predicate: Predicate) -> str | None:
+    """Extract the annotation key from a leaf predicate, or None for logical operators."""
+    match predicate:
+        case KeyExistsPredicate():
+            return predicate.key_exists.key
+        case ValueEqualsPredicate():
+            return predicate.value_equals.key
+        case ValueContainsPredicate():
+            return predicate.value_contains.key
+        case ValueInPredicate():
+            return predicate.value_in.key
+        case _:
+            return None
+
+
+def _check_predicate_allowed(*, predicate: Predicate) -> None:
+    """Raise if a system key is used with an unsupported predicate type."""
+    key = _get_predicate_key(predicate=predicate)
+    if key is None:
+        return
+
+    try:
+        system_key = SystemKey(key)
+    except ValueError:
+        return
+
+    supported = SYSTEM_KEY_SUPPORTED_PREDICATES.get(system_key, set())
+    if type(predicate) not in supported:
+        raise InvalidAnnotationKeyError(
+            f"Predicate {type(predicate).__name__} is not supported "
+            f"for system key {system_key!r}. "
+            f"Supported: {[t.__name__ for t in supported]}"
+        )
+
+
+def _resolve_system_key_value(
+    *,
+    key: str,
+    value: str,
+    current_user: str | None,
+) -> str:
+    """Resolve special placeholder values for system keys."""
+    if key == SystemKey.CREATED_BY and value == "me":
+        return current_user if current_user is not None else ""
+    return value
+
+
+def _maybe_resolve_system_values(
+    *,
+    predicate: ValueEqualsPredicate,
+    current_user: str | None,
+) -> ValueEqualsPredicate:
+    """Resolve special values in a ValueEqualsPredicate."""
+    key = predicate.value_equals.key
+    value = predicate.value_equals.value
+    resolved = _resolve_system_key_value(
+        key=key,
+        value=value,
+        current_user=current_user,
+    )
+    if resolved != value:
+        return ValueEqualsPredicate(value_equals=ValueEquals(key=key, value=resolved))
+    return predicate
+
+
+def _validate_and_resolve_predicate(
+    *,
+    predicate: Predicate,
+    current_user: str | None,
+) -> Predicate:
+    """Validate system key support, then resolve special values."""
+    _check_predicate_allowed(predicate=predicate)
+    if isinstance(predicate, ValueEqualsPredicate):
+        return _maybe_resolve_system_values(
+            predicate=predicate,
+            current_user=current_user,
+        )
+    return predicate
 
 
 # ---------------------------------------------------------------------------
@@ -79,7 +183,12 @@ def build_list_filters(
 
     if filter_query_value:
         parsed = FilterQuery.model_validate_json(filter_query_value)
-        where_clauses.append(filter_query_to_where_clause(filter_query=parsed))
+        where_clauses.append(
+            filter_query_to_where_clause(
+                filter_query=parsed,
+                current_user=current_user,
+            )
+        )
 
     next_page_token = PageToken(
         offset=offset + page_size,
@@ -93,10 +202,13 @@ def build_list_filters(
 def filter_query_to_where_clause(
     *,
     filter_query: FilterQuery,
+    current_user: str | None = None,
 ) -> sql.ColumnElement:
     predicates = filter_query.and_ or filter_query.or_
     is_and = filter_query.and_ is not None
-    clauses = [_predicate_to_clause(predicate=p) for p in predicates]
+    clauses = [
+        _predicate_to_clause(predicate=p, current_user=current_user) for p in predicates
+    ]
     return sql.and_(*clauses) if is_and else sql.or_(*clauses)
 
 
@@ -161,17 +273,35 @@ def _build_filter_where_clauses(
 
 def _predicate_to_clause(
     *,
-    predicate,
+    predicate: Predicate,
+    current_user: str | None = None,
 ) -> sql.ColumnElement:
+    predicate = _validate_and_resolve_predicate(
+        predicate=predicate,
+        current_user=current_user,
+    )
+
     match predicate:
         case AndPredicate():
             return sql.and_(
-                *[_predicate_to_clause(predicate=p) for p in predicate.and_]
+                *[
+                    _predicate_to_clause(predicate=p, current_user=current_user)
+                    for p in predicate.and_
+                ]
             )
         case OrPredicate():
-            return sql.or_(*[_predicate_to_clause(predicate=p) for p in predicate.or_])
+            return sql.or_(
+                *[
+                    _predicate_to_clause(predicate=p, current_user=current_user)
+                    for p in predicate.or_
+                ]
+            )
         case NotPredicate():
-            return sql.not_(_predicate_to_clause(predicate=predicate.not_))
+            return sql.not_(
+                _predicate_to_clause(
+                    predicate=predicate.not_, current_user=current_user
+                )
+            )
         case KeyExistsPredicate():
             return _key_exists_to_clause(predicate=predicate)
         case ValueEqualsPredicate():