diff --git a/cloud_pipelines_backend/backend_types_sql.py b/cloud_pipelines_backend/backend_types_sql.py index a4ba08d..572b4f6 100644 --- a/cloud_pipelines_backend/backend_types_sql.py +++ b/cloud_pipelines_backend/backend_types_sql.py @@ -432,6 +432,9 @@ class ExecutionNode(_TableBase): ) _status_changed: bool = dataclasses.field(default=False, init=False, repr=False) + _run_end_check_pending: bool = dataclasses.field( + default=False, init=False, repr=False + ) # The name of a special input that links to the the `TaskSpec.is_enabled` artifact. @@ -515,6 +518,7 @@ class PipelineRunAnnotation(_TableBase): _IX_ANNOTATION_RUN_ID_KEY_VALUE: Final[str] = ( "ix_pipeline_run_annotation_pipeline_run_id_key_value" ) + _IX_ANNOTATION_KEY_VALUE: Final[str] = "ix_pipeline_run_annotation_key_value" pipeline_run_id: orm.Mapped[IdType] = orm.mapped_column( sql.ForeignKey(PipelineRun.id), primary_key=True, @@ -531,6 +535,7 @@ class PipelineRunAnnotation(_TableBase): "key", "value", ), + sql.Index(_IX_ANNOTATION_KEY_VALUE, "key", "value"), ) diff --git a/cloud_pipelines_backend/database_ops.py b/cloud_pipelines_backend/database_ops.py index 75ed714..6aa645b 100644 --- a/cloud_pipelines_backend/database_ops.py +++ b/cloud_pipelines_backend/database_ops.py @@ -100,9 +100,11 @@ def migrate_db( index.create(db_engine, checkfirst=True) for index in bts.PipelineRunAnnotation.__table__.indexes: - if index.name == bts.PipelineRunAnnotation._IX_ANNOTATION_RUN_ID_KEY_VALUE: + if index.name in ( + bts.PipelineRunAnnotation._IX_ANNOTATION_RUN_ID_KEY_VALUE, + bts.PipelineRunAnnotation._IX_ANNOTATION_KEY_VALUE, + ): index.create(db_engine, checkfirst=True) - break database_migrations.migrate_secret_value_column(db_engine=db_engine) diff --git a/cloud_pipelines_backend/filter_query_sql.py b/cloud_pipelines_backend/filter_query_sql.py index b306a57..095c51a 100644 --- a/cloud_pipelines_backend/filter_query_sql.py +++ b/cloud_pipelines_backend/filter_query_sql.py @@ -26,6 +26,7 @@ class PipelineRunAnnotationSystemKey(str, enum.Enum): CREATED_BY = f"{_PIPELINE_RUN_KEY_PREFIX}created_by" PIPELINE_NAME = f"{_PIPELINE_RUN_KEY_PREFIX}name" CREATED_AT = f"{_PIPELINE_RUN_KEY_PREFIX}date.created_at" + ENDED_AT = f"{_PIPELINE_RUN_KEY_PREFIX}ended_at" SYSTEM_KEY_SUPPORTED_PREDICATES: dict[PipelineRunAnnotationSystemKey, set[type]] = { @@ -43,6 +44,9 @@ class PipelineRunAnnotationSystemKey(str, enum.Enum): PipelineRunAnnotationSystemKey.CREATED_AT: { filter_query_models.TimeRangePredicate, }, + PipelineRunAnnotationSystemKey.ENDED_AT: { + filter_query_models.KeyExistsPredicate, + }, } # --------------------------------------------------------------------------- diff --git a/cloud_pipelines_backend/orchestrator_sql.py b/cloud_pipelines_backend/orchestrator_sql.py index 4d05178..c12ad9c 100644 --- a/cloud_pipelines_backend/orchestrator_sql.py +++ b/cloud_pipelines_backend/orchestrator_sql.py @@ -20,6 +20,7 @@ from . import backend_types_sql as bts from . import component_structures as structures +from . import filter_query_sql from .launchers import common_annotations from .launchers import interfaces as launcher_interfaces from .instrumentation import bugsnag_instrumentation @@ -1374,11 +1375,15 @@ def _maybe_get_small_artifact_value( def _handle_container_execution_status_set( execution: bts.ExecutionNode, value: bts.ContainerExecutionStatus | None, - _old_value: object, + old_value: object, _initiator: object, ) -> None: if value is None: return + if (value in bts.CONTAINER_STATUSES_ENDED) != ( + old_value in bts.CONTAINER_STATUSES_ENDED + ): + execution._run_end_check_pending = True if execution.extra_data is None: execution.extra_data = {} history: list = execution.extra_data.get(_HISTORY_KEY, []) @@ -1393,3 +1398,104 @@ def _handle_container_execution_status_set( _HISTORY_KEY: history + [entry], } execution._status_changed = True + + +@sql_event.listens_for(orm.Session, "after_flush") +def _handle_after_flush(session: orm.Session, _flush_context: object) -> None: + """Sync the ended_at annotation of runs whose executions entered or left ended.""" + execution_ids = set() + for obj in list(session.new) + list(session.dirty): + if isinstance(obj, bts.ExecutionNode) and obj._run_end_check_pending: + execution_ids.add(obj.id) + obj._run_end_check_pending = False + if not execution_ids: + return + connection = session.connection() + pipeline_runs = _get_pipeline_runs_containing_executions( + connection=connection, execution_ids=execution_ids + ) + for pipeline_run_id, root_execution_id in pipeline_runs.items(): + _reconcile_pipeline_run_ended_at_annotation( + connection=connection, + pipeline_run_id=pipeline_run_id, + root_execution_id=root_execution_id, + ) + + +def _get_pipeline_runs_containing_executions( + *, + connection: sql.Connection, + execution_ids: set[bts.IdType], +) -> dict[bts.IdType, bts.IdType]: + """Map each run containing any of the executions to its root execution id.""" + link = bts.ExecutionToAncestorExecutionLink + as_descendant = ( + sql.select(bts.PipelineRun.id, bts.PipelineRun.root_execution_id) + .join(link, link.ancestor_execution_id == bts.PipelineRun.root_execution_id) + .where(link.execution_id.in_(execution_ids)) + ) + as_root = sql.select(bts.PipelineRun.id, bts.PipelineRun.root_execution_id).where( + bts.PipelineRun.root_execution_id.in_(execution_ids) + ) + rows = connection.execute(sql.union(as_descendant, as_root)).all() + return {pipeline_run_id: root_id for pipeline_run_id, root_id in rows} + + +def _pipeline_run_has_ended( + *, + connection: sql.Connection, + root_execution_id: bts.IdType, +) -> bool: + """Whether the run has container executions and all of them have ended.""" + status = bts.ExecutionNode.container_execution_status + link = bts.ExecutionToAncestorExecutionLink + root_status = connection.execute( + sql.select(status).where(bts.ExecutionNode.id == root_execution_id) + ).scalar_one_or_none() + if root_status is not None and root_status not in bts.CONTAINER_STATUSES_ENDED: + return False + descendants = ( + sql.select(bts.ExecutionNode.id) + .join(link, link.execution_id == bts.ExecutionNode.id) + .where(link.ancestor_execution_id == root_execution_id) + .where(status.is_not(None)) + ) + unfinished_descendant = connection.execute( + descendants.where(status.not_in(bts.CONTAINER_STATUSES_ENDED)).limit(1) + ).first() + if unfinished_descendant is not None: + return False + if root_status is not None: + return True + return connection.execute(descendants.limit(1)).first() is not None + + +def _reconcile_pipeline_run_ended_at_annotation( + *, + connection: sql.Connection, + pipeline_run_id: bts.IdType, + root_execution_id: bts.IdType, +) -> None: + """Write the ended_at annotation when the run has ended, else remove it.""" + annotation = bts.PipelineRunAnnotation + key = filter_query_sql.PipelineRunAnnotationSystemKey.ENDED_AT + this_annotation = sql.and_( + annotation.pipeline_run_id == pipeline_run_id, annotation.key == key + ) + is_annotated = ( + connection.execute(sql.select(annotation.key).where(this_annotation)).first() + is not None + ) + has_ended = _pipeline_run_has_ended( + connection=connection, root_execution_id=root_execution_id + ) + if has_ended and not is_annotated: + connection.execute( + sql.insert(annotation).values( + pipeline_run_id=pipeline_run_id, + key=key, + value=_get_current_time().isoformat(), + ) + ) + elif is_annotated and not has_ended: + connection.execute(sql.delete(annotation).where(this_annotation)) diff --git a/tests/test_pipeline_run_ended_at.py b/tests/test_pipeline_run_ended_at.py new file mode 100644 index 0000000..7b50057 --- /dev/null +++ b/tests/test_pipeline_run_ended_at.py @@ -0,0 +1,273 @@ +"""Tests for the system/pipeline_run.ended_at annotation kept by orchestrator_sql.""" + +import datetime +from typing import Callable + +import pytest +import sqlalchemy +from sqlalchemy import orm + +from cloud_pipelines_backend import api_server_sql +from cloud_pipelines_backend import backend_types_sql as bts +from cloud_pipelines_backend import component_structures as structures +from cloud_pipelines_backend import database_ops +from cloud_pipelines_backend import filter_query_sql +from cloud_pipelines_backend import orchestrator_sql # noqa: F401 + +_ENDED_AT = filter_query_sql.PipelineRunAnnotationSystemKey.ENDED_AT +_Status = bts.ContainerExecutionStatus + + +def _container_component() -> structures.ComponentSpec: + return structures.ComponentSpec( + implementation=structures.ContainerImplementation( + container=structures.ContainerSpec(image="python") + ), + ) + + +def _graph_task(*, task_ids: list[str]) -> structures.TaskSpec: + tasks = { + task_id: structures.TaskSpec( + component_ref=structures.ComponentReference(spec=_container_component()) + ) + for task_id in task_ids + } + return structures.TaskSpec( + component_ref=structures.ComponentReference( + spec=structures.ComponentSpec( + implementation=structures.GraphImplementation( + graph=structures.GraphSpec(tasks=tasks) + ), + ) + ) + ) + + +@pytest.fixture() +def session_factory() -> Callable[[], orm.Session]: + db_engine = database_ops.create_db_engine_and_migrate_db(database_uri="sqlite://") + return lambda: orm.Session(bind=db_engine) + + +def _create_run( + *, + session_factory: Callable[[], orm.Session], + root_task: structures.TaskSpec, +) -> str: + return ( + api_server_sql.PipelineRunsApiService_Sql() + .create(session=session_factory(), root_task=root_task, created_by="ada") + .id + ) + + +def _task_node(*, session: orm.Session, task_id: str) -> bts.ExecutionNode: + return session.scalars( + sqlalchemy.select(bts.ExecutionNode).where( + bts.ExecutionNode.task_id_in_parent_execution == task_id + ) + ).one() + + +def _ended_at(*, session: orm.Session, pipeline_run_id: str) -> str | None: + return session.scalar( + sqlalchemy.select(bts.PipelineRunAnnotation.value).where( + bts.PipelineRunAnnotation.pipeline_run_id == pipeline_run_id, + bts.PipelineRunAnnotation.key == _ENDED_AT, + ) + ) + + +def _set_statuses( + *, + session_factory: Callable[[], orm.Session], + statuses: dict[str, bts.ContainerExecutionStatus], +) -> None: + with session_factory() as session: + for task_id, status in statuses.items(): + _task_node(session=session, task_id=task_id).container_execution_status = ( + status + ) + session.commit() + + +class TestPipelineRunEndedAt: + def test_no_annotation_while_any_execution_is_unfinished( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a", "b"]) + ) + _set_statuses( + session_factory=session_factory, statuses={"a": _Status.SUCCEEDED} + ) + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=run_id) is None + + def test_annotation_written_when_the_last_execution_ends( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a", "b"]) + ) + _set_statuses( + session_factory=session_factory, statuses={"a": _Status.SUCCEEDED} + ) + _set_statuses(session_factory=session_factory, statuses={"b": _Status.FAILED}) + + with session_factory() as session: + value = _ended_at(session=session, pipeline_run_id=run_id) + assert value is not None + assert datetime.datetime.fromisoformat(value).tzinfo is not None + + def test_every_ended_status_counts( + self, session_factory: Callable[[], orm.Session] + ) -> None: + ended = sorted(bts.CONTAINER_STATUSES_ENDED) + task_ids = [status.value.lower() for status in ended] + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=task_ids) + ) + _set_statuses( + session_factory=session_factory, statuses=dict(zip(task_ids, ended)) + ) + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=run_id) is not None + + def test_annotation_removed_when_an_execution_leaves_ended( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a"]) + ) + _set_statuses( + session_factory=session_factory, statuses={"a": _Status.SUCCEEDED} + ) + _set_statuses(session_factory=session_factory, statuses={"a": _Status.QUEUED}) + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=run_id) is None + + def test_moving_between_ended_statuses_keeps_the_first_timestamp( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a"]) + ) + _set_statuses( + session_factory=session_factory, statuses={"a": _Status.SUCCEEDED} + ) + with session_factory() as session: + first = _ended_at(session=session, pipeline_run_id=run_id) + _set_statuses(session_factory=session_factory, statuses={"a": _Status.FAILED}) + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=run_id) == first + + def test_transition_flushed_before_commit_is_still_recorded( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a"]) + ) + with session_factory() as session: + _task_node(session=session, task_id="a").container_execution_status = ( + _Status.SUCCEEDED + ) + session.flush() + session.scalar(sqlalchemy.select(bts.PipelineRun.id)) + session.commit() + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=run_id) is not None + + def test_rolled_back_transition_leaves_no_annotation( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a"]) + ) + with session_factory() as session: + _task_node(session=session, task_id="a").container_execution_status = ( + _Status.SUCCEEDED + ) + session.flush() + session.rollback() + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=run_id) is None + + def test_container_root_run( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, + root_task=structures.TaskSpec( + component_ref=structures.ComponentReference(spec=_container_component()) + ), + ) + with session_factory() as session: + root_id = session.get(bts.PipelineRun, run_id).root_execution_id + session.get(bts.ExecutionNode, root_id).container_execution_status = ( + _Status.SUCCEEDED + ) + session.commit() + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=run_id) is not None + + def test_other_runs_are_untouched( + self, session_factory: Callable[[], orm.Session] + ) -> None: + ended_run = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a"]) + ) + running_run = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["b"]) + ) + _set_statuses( + session_factory=session_factory, statuses={"a": _Status.SUCCEEDED} + ) + + with session_factory() as session: + assert _ended_at(session=session, pipeline_run_id=ended_run) is not None + assert _ended_at(session=session, pipeline_run_id=running_run) is None + + def test_users_cannot_set_the_annotation( + self, session_factory: Callable[[], orm.Session] + ) -> None: + run_id = _create_run( + session_factory=session_factory, root_task=_graph_task(task_ids=["a"]) + ) + with pytest.raises(Exception): + api_server_sql.PipelineRunsApiService_Sql().set_annotation( + session=session_factory(), + id=run_id, + key=_ENDED_AT, + value="2026-01-01T00:00:00+00:00", + user_name="ada", + ) + + +def test_key_value_index_is_created_on_an_existing_database() -> None: + db_engine = database_ops.create_db_engine(database_uri="sqlite://") + bts._TableBase.metadata.create_all(db_engine) + with db_engine.begin() as connection: + connection.execute( + sqlalchemy.text( + f"DROP INDEX {bts.PipelineRunAnnotation._IX_ANNOTATION_KEY_VALUE}" + ) + ) + + database_ops.migrate_db(db_engine=db_engine, do_skip_backfill=True) + + index_names = { + index["name"] + for index in sqlalchemy.inspect(db_engine).get_indexes( + bts.PipelineRunAnnotation.__tablename__ + ) + } + assert bts.PipelineRunAnnotation._IX_ANNOTATION_KEY_VALUE in index_names