Extract _create_in_transaction so a caller can own the pipeline-run transaction

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)

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,