feat: Search created by user in pipeline run API

Author: yuechao-qin

cloud_pipelines_backend/api_server_sql.py

@@ -33,6 +33,7 @@ def _get_current_time() -> datetime.datetime:
 
 
 _DEFAULT_PAGE_SIZE: Final[int] = 10
+_SYSTEM_KEY_RESERVED_MSG = f"Annotation keys starting with {filter_query_sql.SYSTEM_KEY_PREFIX!r} are reserved for system use."
 
 
 # ==== PipelineJobService
@@ -107,6 +108,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)
@@ -297,6 +311,8 @@ def set_annotation(
         user_name: str | None = None,
         skip_user_check: bool = False,
     ):
+        if key.startswith(filter_query_sql.SYSTEM_KEY_PREFIX):
+            raise errors.InvalidAnnotationKeyError(_SYSTEM_KEY_RESERVED_MSG)
         pipeline_run = session.get(bts.PipelineRun, id)
         if not pipeline_run:
             raise errors.ItemNotFoundError(f"Pipeline run {id} not found.")
@@ -319,6 +335,8 @@ def delete_annotation(
         user_name: str | None = None,
         skip_user_check: bool = False,
     ):
+        if key.startswith(filter_query_sql.SYSTEM_KEY_PREFIX):
+            raise errors.InvalidAnnotationKeyError(_SYSTEM_KEY_RESERVED_MSG)
         pipeline_run = session.get(bts.PipelineRun, id)
         if not pipeline_run:
             raise errors.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(
@@ -83,3 +85,59 @@ def migrate_db(db_engine: sqlalchemy.Engine):
         if index.name == bts.IX_ANNOTATION_RUN_ID_KEY_VALUE:
             index.create(db_engine, checkfirst=True)
             break
+
+    backfill_created_by_annotations(db_engine=db_engine)
+
+
+def is_annotation_key_already_backfilled(
+    *,
+    db_engine: sqlalchemy.Engine,
+    key: str,
+) -> bool:
+    """Return True if at least one annotation with the given key exists."""
+    with orm.Session(db_engine) as session:
+        return session.query(
+            sqlalchemy.exists(
+                sqlalchemy.select(sqlalchemy.literal(1))
+                .select_from(bts.PipelineRunAnnotation)
+                .where(
+                    bts.PipelineRunAnnotation.key == key,
+                )
+            )
+        ).scalar()
+
+
+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.
+
+    Skips entirely if any created_by annotation key already exists (i.e. the
+    write-path is populating them, so the backfill has already run or is
+    no longer needed).
+    """
+    if is_annotation_key_already_backfilled(
+        db_engine=db_engine, key=filter_query_sql.SystemKey.CREATED_BY
+    ):
+        return
+
+    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

@@ -18,3 +18,7 @@ class ApiValidationError(Exception):
 
 class MutuallyExclusiveFilterError(ApiValidationError):
     pass
+
+
+class InvalidAnnotationKeyError(ApiValidationError):
+    pass

cloud_pipelines_backend/filter_query_sql.py

@@ -1,13 +1,30 @@
 import base64
 import dataclasses
 import json
+import enum
+from typing import Final
 
 import sqlalchemy as sql
 
 from . import backend_types_sql as bts
 from . import errors
 from . import filter_query_models
 
+SYSTEM_KEY_PREFIX: Final[str] = "system/"
+
+
+class SystemKey(enum.StrEnum):
+    CREATED_BY = f"{SYSTEM_KEY_PREFIX}pipeline_run.created_by"
+
+
+SYSTEM_KEY_SUPPORTED_PREDICATES: dict[SystemKey, set[type]] = {
+    SystemKey.CREATED_BY: {
+        filter_query_models.KeyExistsPredicate,
+        filter_query_models.ValueEqualsPredicate,
+        filter_query_models.ValueInPredicate,
+    },
+}
+
 # ---------------------------------------------------------------------------
 # PageToken
 # ---------------------------------------------------------------------------
@@ -31,6 +48,93 @@ def decode(cls, token: str | None) -> "PageToken":
         return cls(**json.loads(base64.b64decode(token)))
 
 
+# ---------------------------------------------------------------------------
+# SystemKey validation and resolution
+# ---------------------------------------------------------------------------
+
+
+def _get_predicate_key(*, predicate: filter_query_models.Predicate) -> str | None:
+    """Extract the annotation key from a leaf predicate, or None for logical operators."""
+    match predicate:
+        case filter_query_models.KeyExistsPredicate():
+            return predicate.key_exists.key
+        case filter_query_models.ValueEqualsPredicate():
+            return predicate.value_equals.key
+        case filter_query_models.ValueContainsPredicate():
+            return predicate.value_contains.key
+        case filter_query_models.ValueInPredicate():
+            return predicate.value_in.key
+        case _:
+            return None
+
+
+def _check_predicate_allowed(*, predicate: filter_query_models.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 errors.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: filter_query_models.ValueEqualsPredicate,
+    current_user: str | None,
+) -> filter_query_models.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 filter_query_models.ValueEqualsPredicate(
+            value_equals=filter_query_models.ValueEquals(key=key, value=resolved)
+        )
+    return predicate
+
+
+def _validate_and_resolve_predicate(
+    *,
+    predicate: filter_query_models.Predicate,
+    current_user: str | None,
+) -> filter_query_models.Predicate:
+    """Validate system key support, then resolve special values."""
+    _check_predicate_allowed(predicate=predicate)
+    if isinstance(predicate, filter_query_models.ValueEqualsPredicate):
+        return _maybe_resolve_system_values(
+            predicate=predicate,
+            current_user=current_user,
+        )
+    return predicate
+
+
 # ---------------------------------------------------------------------------
 # Public API
 # ---------------------------------------------------------------------------
@@ -67,7 +171,12 @@ def build_list_filters(
 
     if filter_query_value:
         parsed = filter_query_models.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,
@@ -81,10 +190,13 @@ def build_list_filters(
 def filter_query_to_where_clause(
     *,
     filter_query: filter_query_models.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)
 
 
@@ -149,17 +261,35 @@ def _build_filter_where_clauses(
 
 def _predicate_to_clause(
     *,
-    predicate,
+    predicate: filter_query_models.Predicate,
+    current_user: str | None = None,
 ) -> sql.ColumnElement:
+    predicate = _validate_and_resolve_predicate(
+        predicate=predicate,
+        current_user=current_user,
+    )
+
     match predicate:
         case filter_query_models.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 filter_query_models.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 filter_query_models.NotPredicate():
-            return sql.not_(_predicate_to_clause(predicate=predicate.not_))
+            return sql.not_(
+                _predicate_to_clause(
+                    predicate=predicate.not_, current_user=current_user
+                )
+            )
         case filter_query_models.KeyExistsPredicate():
             return _key_exists_to_clause(predicate=predicate)
         case filter_query_models.ValueEqualsPredicate():