diff --git a/cloud_pipelines_backend/api_server_sql.py b/cloud_pipelines_backend/api_server_sql.py index cf71a16..451f079 100644 --- a/cloud_pipelines_backend/api_server_sql.py +++ b/cloud_pipelines_backend/api_server_sql.py @@ -106,7 +106,7 @@ def _fail_if_changing_system_annotation(self, *, key: str) -> None: if key.startswith(filter_query_sql.SYSTEM_KEY_PREFIX): raise errors.ApiValidationError(self._SYSTEM_KEY_RESERVED_MSG) - def create( + def _create_in_transaction( self, session: orm.Session, root_task: structures.TaskSpec, @@ -115,44 +115,70 @@ def create( # Arbitrary metadata. Can be used to specify user. annotations: Optional[dict[str, Any]] = None, created_by: str | None = None, - ) -> PipelineRunResponse: + ) -> bts.PipelineRun: + """Creates a pipeline run inside a transaction the caller already owns. + + Flushes, so the returned run has its ID populated, but never commits: + the caller decides when the work becomes durable. Use this when a run + must be written atomically with the caller's own rows. Callers that just + want a run created should use `create` instead. + """ # TODO: Validate the pipeline spec # TODO: Load and validate all components # TODO: Fetch missing components and populate component specs pipeline_name = root_task.component_ref.spec.name - with session.begin(): + root_execution_node = _recursively_create_all_executions_and_artifacts_root( + session=session, + root_task_spec=root_task, + ) - root_execution_node = _recursively_create_all_executions_and_artifacts_root( - session=session, - root_task_spec=root_task, - ) + # Store into DB. + current_time = _get_current_time() + pipeline_run = bts.PipelineRun( + root_execution=root_execution_node, + created_at=current_time, + updated_at=current_time, + annotations=annotations, + created_by=created_by, + extra_data={ + self._PIPELINE_NAME_EXTRA_DATA_KEY: pipeline_name, + }, + ) + session.add(pipeline_run) + # Flush to populate pipeline_run.id (server-generated) before inserting annotation FKs. + # TODO: Use ORM relationship instead of explicit flush + manual FK assignment. + session.flush() + _mirror_system_annotations( + session=session, + pipeline_run_id=pipeline_run.id, + created_by=created_by, + pipeline_name=pipeline_name, + annotations=annotations, + ) + return pipeline_run - # Store into DB. - current_time = _get_current_time() - pipeline_run = bts.PipelineRun( - root_execution=root_execution_node, - created_at=current_time, - updated_at=current_time, - annotations=annotations, - created_by=created_by, - extra_data={ - self._PIPELINE_NAME_EXTRA_DATA_KEY: pipeline_name, - }, - ) - session.add(pipeline_run) - # Flush to populate pipeline_run.id (server-generated) before inserting annotation FKs. - # TODO: Use ORM relationship instead of explicit flush + manual FK assignment. - session.flush() - _mirror_system_annotations( + def create( + self, + session: orm.Session, + root_task: structures.TaskSpec, + # Component library to avoid repeating component specs inside task specs + components: Optional[list[structures.ComponentReference]] = None, + # Arbitrary metadata. Can be used to specify user. + annotations: Optional[dict[str, Any]] = None, + created_by: str | None = None, + ) -> PipelineRunResponse: + # `session.begin()` commits when the block exits, so no explicit commit + # is needed here. + with session.begin(): + pipeline_run = self._create_in_transaction( session=session, - pipeline_run_id=pipeline_run.id, - created_by=created_by, - pipeline_name=pipeline_name, + root_task=root_task, + components=components, annotations=annotations, + created_by=created_by, ) - session.commit() session.refresh(pipeline_run) return PipelineRunResponse.from_db(pipeline_run) diff --git a/tests/test_api_server_sql.py b/tests/test_api_server_sql.py index f8f23b9..900f775 100644 --- a/tests/test_api_server_sql.py +++ b/tests/test_api_server_sql.py @@ -433,6 +433,56 @@ def test_create_mirrors_absent_values_as_empty_string( ) +def _count_rows(*, session: orm.Session, table: type) -> int: + return session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(table)) + + +class TestCreateInTransaction: + """Pins the contract of `_create_in_transaction` for callers that own the transaction. + + The method is private, so nothing outside this file is promised it exists or + that it stays free of an internal commit. These tests are what turns that + into a promise: a rename, a removal, or a commit creeping back in fails here + rather than in a consumer that batches the run with its own rows. + """ + + def test_flushes_so_the_caller_can_use_the_run_id(self, session_factory, service): + with session_factory() as session: + session.begin() + pipeline_run = service._create_in_transaction( + session, root_task=_make_task_spec("in-transaction") + ) + assert pipeline_run.id is not None + assert pipeline_run.root_execution_id is not None + session.rollback() + + def test_rollback_leaves_no_rows(self, session_factory, service): + with session_factory() as session: + session.begin() + service._create_in_transaction( + session, root_task=_make_task_spec("rolled-back") + ) + session.rollback() + + with session_factory() as session: + assert _count_rows(session=session, table=bts.PipelineRun) == 0 + assert _count_rows(session=session, table=bts.ExecutionNode) == 0 + + def test_the_callers_commit_makes_the_run_durable(self, session_factory, service): + with session_factory() as session: + # The same shape `create` uses: the block commits on exit, so the + # test never commits by hand. + with session.begin(): + pipeline_run = service._create_in_transaction( + session, root_task=_make_task_spec("committed-by-caller") + ) + run_id = pipeline_run.id + + with session_factory() as session: + assert session.get(bts.PipelineRun, run_id) is not None + assert _count_rows(session=session, table=bts.PipelineRun) == 1 + + class TestCreateMirrorsUserAnnotations: def test_create_mirrors_user_annotations( self,