Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yuechao-qin wants to merge 1 commit into
base: master
Choose a base branch
from
+104 −28
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 54 additions & 28 deletions cloud_pipelines_backend/api_server_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
"""
Comment on lines +119 to +125

@Volv-G Volv-G Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

This precondition isn't enforced, and violating it fails silently.

Thanks to autobegin, calling this on a session with no open transaction appears to work: it autobegins, hands back a fully populated pipeline_run.id, and then discards everything when the session closes without a commit. Verified at this head:

in_transaction before: False
in_transaction after:  True   (autobegin)    id handed out: True
-> session closed without commit
persisted PipelineRun rows: 0

A caller gets a real-looking run ID for a run that will never exist. That's the same silent-drop failure mode this PR exists to remove, just relocated from the service into the caller — and it's the sort of mistake a new caller makes exactly once, expensively, in production.

Three lines make it loud:

if not session.in_transaction():
    raise ApiServiceError(
        "_create_in_transaction requires the caller to have an open transaction; "
        "use `create` if you just want a run created."
    )

I checked this is purely additive against every current call path: create() opens with session.begin(): before calling, all three new tests call session.begin() first, and any caller that has already staged a row has autobegun. in_transaction() is True in every legitimate case.

# 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,
)
Comment on lines +132 to +135

@Volv-G Volv-G Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

If this raises, the caller is left holding partially-built rows.

_recursively_create_all_executions_and_artifacts does its first session.add(...) at line 1655, before all 10 of its subsequent raise ApiServiceError sites (1677, 1709, 1779, 1791, 1831, 1863, 1883, 1888, 1904, 1918). create() was protected by with session.begin():, which rolls back on exception; _create_in_transaction has no such boundary.

Reproduced at this head with a root component that has a required input and no argument supplied:

create()                 -> ApiServiceError, ExecutionNode rows = 0
_create_in_transaction() -> ApiServiceError; caller commits its own rows,
                            ExecutionNode rows = 1, PipelineRun rows = 0

An orphan execution_node with no pipeline_run referencing it. A caller that catches the error and commits its own work anyway — which is a reasonable thing for the trigger sink to want, since it may well wish to keep the fence row that records the attempt — silently persists junk.

Two ways out: (1) document that any exception from this method obliges the caller to roll back the entire transaction, or (2) wrap the body in with session.begin_nested(): so a failure unwinds only this method's writes and leaves the caller's rows intact. Option 2 looks worth the one SAVEPOINT, since batching with caller-owned rows is the whole reason this method exists — and it makes the two goals compatible instead of mutually exclusive.


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()
Comment on lines +149 to +152

@Volv-G Volv-G Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

This flush isn't scoped to this method — it flushes whatever the caller staged, too.

session.flush() is a whole-session flush. Verified at this head: a caller that stages a row and doesn't flush finds it written to the DB by this call, and if that row violates a constraint the IntegrityError is raised from inside _create_in_transaction:

caller stages its own row, does NOT flush   -> session.new = 1
_create_in_transaction(...)                 -> caller's row is now in the DB
and when the caller's staged row collides:
                                            -> IntegrityError raised from INSIDE this method

That matters for the trigger sink in oasis-backend#562. The fence there is a UNIQUE (subscription_id, cycle) constraint and losing the race is supposed to raise IntegrityError — it's the normal, expected no-op path. If the sink stages the fence row and then calls this, the race-loss exception no longer comes out of the sink's own insert:

  • a sink that wraps only its own session.add/flush in try/except IntegrityError won't catch the race loss at all;
  • a sink that wraps this call can't distinguish "I lost the fence race" (expected, ignore) from "creating the pipeline run failed" (real, alert) without inspecting the constraint name.

I notice the TODO on the line above already wants this flush gone. Until then, the docstring seems like the right place to state its blast radius, since the method explicitly invites the stage-then-call pattern ("must be written atomically with the caller's own rows"):

Flushes the whole session, so rows the caller staged but has not flushed
are written too; a constraint violation on the caller's own rows will
therefore surface from this call. Flush before calling if you need to tell
your own conflicts apart from run-creation failures.

The sink-side fix is a one-liner — session.flush() right after staging the fence row — but only if this behaviour is documented, otherwise there's nothing to tell the sink author they need it.

[nit] While here: pipeline_run.id isn't server-generated. backend_types_sql.py:134 declares it insert_default=generate_unique_id, a Python-side callable, and IdType is a str — no autoincrement, no server_default, no sequence. The flush is still genuinely required (verified: id is None before, '01a04b0f…' after), so only the parenthetical is wrong. Worth correcting to (Python-side insert_default) precisely because this comment is the stated justification for a flush that turns out to have wider reach than advertised.

_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)
Expand Down
50 changes: 50 additions & 0 deletions tests/test_api_server_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +459 to +469

@Volv-G Volv-G Aug 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(AI-assisted)

These pin "the run rolls back" but not "the run and my rows roll back together" — which is the guarantee the method exists to provide.

The three new tests are good: real engine, and test_the_callers_commit_makes_the_run_durable proves the StaticPool in-memory DB is genuinely shared across sessions, so this rollback test isn't vacuous. But every test here writes only the run, so the atomicity claim in the docstring — a run written atomically with the caller's own rows — has no coverage.

def test_the_run_and_the_callers_own_rows_share_a_fate(self, session_factory, service):
    with session_factory() as session:
        session.begin()
        session.add(bts.PipelineRunAnnotation(pipeline_run_id="caller-row", key="k", value="v"))
        service._create_in_transaction(session, root_task=_make_task_spec("atomic"))
        session.rollback()
    with session_factory() as session:
        assert _count_rows(session=session, table=bts.PipelineRun) == 0
        assert _count_rows(session=session, table=bts.PipelineRunAnnotation) == 0

Two smaller gaps:

  • This test counts PipelineRun and ExecutionNode but not PipelineRunAnnotation, and those rows are added by _mirror_system_annotations after the flush — exactly the writes a partial-commit regression would leak. One more assert _count_rows(..., table=bts.PipelineRunAnnotation) == 0 closes it.
  • Nothing covers two calls in one transaction. I verified it works today (two distinct IDs, 2 runs and 4 annotation rows committed together), but nothing holds that behaviour still, and a trigger firing could plausibly start more than one run per cycle.


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,
Expand Down
Loading