poc: BigQuery - Add column updated_at to table execution_node

cloud_pipelines_backend/backend_types_sql.py

@@ -58,6 +58,10 @@ def generate_unique_id() -> str:
     return ("%012x" % milliseconds) + random_bytes.hex()
 
 
+def _utcnow() -> datetime.datetime:
+    return datetime.datetime.now(datetime.timezone.utc)
+
+
 id_column = orm.mapped_column(
     sql.String(20), primary_key=True, init=False, insert_default=generate_unique_id
 )
@@ -358,7 +362,11 @@ class ExecutionNode(_TableBase):
         repr=False,
     )
 
-    # updated_at: orm.Mapped[datetime.datetime | None] = orm.mapped_column(default=None)
+    updated_at: orm.Mapped[datetime.datetime | None] = orm.mapped_column(
+        init=False,
+        insert_default=_utcnow,
+        onupdate=_utcnow,
+    )
 
     # execution_kind = orm.Mapped[typing.Literal["CONTAINER", "GRAPH"]]
 

cloud_pipelines_backend/database_migrations.py

@@ -795,3 +795,86 @@ def migrate_secret_value_column(
             f" column={column.name}, current_type={curr_col_type},"
             f" target_type={type_sql}, dialect={dialect}"
         )
+
+
+def migrate_add_execution_node_updated_at_column(
+    *,
+    db_engine: sqlalchemy.Engine,
+) -> bool:
+    """Idempotent: add execution_node.updated_at if the column is missing.
+
+    Inspects the actual DB columns and skips (silently -- this runs on
+    every startup, and logging a no-op every time would be noise) if
+    `updated_at` is already present.
+
+    Uses Alembic's batch_alter_table so that one code path covers every
+    dialect: MySQL, PostgreSQL, and SQLite.
+
+    Returns:
+        True if the column was just added this call, False if it already
+        existed.
+    """
+    table = bts.ExecutionNode.__table__
+    column = table.c.updated_at
+
+    inspector = sqlalchemy.inspect(db_engine)
+    db_columns = {c["name"] for c in inspector.get_columns(table.name)}
+    if column.name in db_columns:
+        return False
+
+    _logger.info(f"Adding column: {table.name}.{column.name}")
+    try:
+        with db_engine.connect() as conn:
+            ctx = alembic_migration.MigrationContext.configure(conn)
+            op = alembic_operations.Operations(ctx)
+            with op.batch_alter_table(table.name) as batch_op:
+                batch_op.add_column(
+                    sqlalchemy.Column(column.name, column.type, nullable=True)
+                )
+            conn.commit()
+        _logger.info(f"Adding column: {table.name}.{column.name} - complete")
+    except Exception:
+        _logger.exception(f"Adding column failed: {table.name}.{column.name}")
+        raise
+    return True
+
+
+def backfill_execution_node_updated_at(
+    *,
+    session: orm.Session,
+    auto_commit: bool,
+) -> int:
+    """Idempotent backfill: stamp execution_node.updated_at for pre-existing rows.
+
+    UPDATE execution_node SET updated_at = :now WHERE updated_at IS NULL
+
+    Every pre-existing row (created before this column existed) gets a
+    single flat timestamp -- the time this backfill runs -- rather than a
+    per-row reconstructed value.
+
+    Idempotent: a second call updates 0 rows (`WHERE updated_at IS NULL`
+    only matches rows no earlier call has already stamped).
+
+    Args:
+        session: SQLAlchemy session. Caller controls the transaction
+            when auto_commit=False.
+        auto_commit: Must be explicitly set. True commits after update.
+            False defers commit to the caller (e.g. migrate_db, which
+            batches this with other migration steps before committing).
+
+    Returns:
+        Number of rows updated.
+    """
+    _logger.info("Starting backfill for `execution_node.updated_at`")
+
+    stmt = (
+        sqlalchemy.update(bts.ExecutionNode)
+        .where(bts.ExecutionNode.updated_at.is_(None))
+        .values(updated_at=bts._utcnow())
+    )
+    result = session.execute(stmt)
+    rowcount = result.rowcount
+    _logger.info(f"Backfill execution_node.updated_at: {rowcount} rows updated")
+    if auto_commit:
+        session.commit()
+    return rowcount

cloud_pipelines_backend/database_ops.py

@@ -105,6 +105,11 @@ def migrate_db(
             break
 
     database_migrations.migrate_secret_value_column(db_engine=db_engine)
+    did_add_execution_node_updated_at_column = (
+        database_migrations.migrate_add_execution_node_updated_at_column(
+            db_engine=db_engine
+        )
+    )
 
     if do_skip_backfill:
         _logger.info("Skipping annotation backfills")
@@ -118,5 +123,10 @@ def migrate_db(
             database_migrations.run_all_annotation_backfills(
                 session=session,
             )
+            if did_add_execution_node_updated_at_column:
+                database_migrations.backfill_execution_node_updated_at(
+                    session=session,
+                    auto_commit=True,
+                )
 
     _logger.info("Exit migrate DB")

tests/test_api_server_sql.py

@@ -1,5 +1,6 @@
 import datetime
 import json
+import time
 from collections.abc import Callable
 
 import pytest
@@ -1957,3 +1958,44 @@ def test_get_with_execution_stats(self):
             assert summary.total_executions == 2
             assert summary.ended_executions == 1
             assert summary.has_ended is False
+
+
+class TestTerminateUpdatesExecutionNodeUpdatedAt:
+    """Regression test for the terminate() / updated_at gap: terminate()
+    only ever sets a running descendant's extra_data["desired_state"], it
+    never touches container_execution_status directly -- so a listener
+    scoped to status transitions would miss this. onupdate (fires on any
+    UPDATE issued for the row) covers it instead."""
+
+    def test_terminate_bumps_running_descendant_execution_node_updated_at(
+        self,
+    ) -> None:
+        session_factory = _initialize_db_and_get_session_factory()
+        service = api_server_sql.PipelineRunsApiService_Sql()
+
+        with session_factory() as session:
+            root = _create_execution_node(session=session)
+            running_child = _create_execution_node(
+                session=session,
+                parent=root,
+                status=bts.ContainerExecutionStatus.RUNNING,
+            )
+            _link_ancestor(
+                session=session, execution_node=running_child, ancestor_node=root
+            )
+            run = _create_pipeline_run(session=session, root_execution=root)
+            run_id = run.id
+            child_id = running_child.id
+            session.commit()
+            first_updated_at = running_child.updated_at
+        assert first_updated_at is not None
+
+        time.sleep(0.001)
+        with session_factory() as session:
+            service.terminate(session=session, id=run_id, skip_user_check=True)
+
+        with session_factory() as session:
+            child = session.get(bts.ExecutionNode, child_id)
+            assert child.extra_data["desired_state"] == "TERMINATED"
+            assert child.updated_at is not None
+            assert child.updated_at > first_updated_at

tests/test_database_migrations.py

@@ -21,6 +21,8 @@
 import logging
 from unittest import mock
 
+from alembic import migration as alembic_migration
+from alembic import operations as alembic_operations
 import pytest
 import sqlalchemy
 from sqlalchemy import orm
@@ -2021,3 +2023,189 @@ def test_migrate_secret_value_column_idempotent(
     our_msgs = [m for m in caplog.messages if m.startswith("migrate column to TEXT:")]
     assert our_msgs[-2] == "migrate column to TEXT: complete"
     assert our_msgs[-1] == "migrate column to TEXT: skipped (already TEXT)"
+
+
+# ---------------------------------------------------------------------------
+# execution_node.updated_at column: migration, backfill, revert
+# ---------------------------------------------------------------------------
+
+
+def _create_execution_node_table_without_updated_at(
+    *,
+    db_engine: sqlalchemy.Engine,
+) -> None:
+    """Create all tables normally, then drop execution_node.updated_at to
+    simulate a pre-migration production DB (column not yet added)."""
+    bts._TableBase.metadata.create_all(db_engine)
+    with db_engine.connect() as conn:
+        ctx = alembic_migration.MigrationContext.configure(conn)
+        op = alembic_operations.Operations(ctx)
+        with op.batch_alter_table("execution_node") as batch_op:
+            batch_op.drop_column("updated_at")
+        conn.commit()
+
+
+def _get_execution_node_columns(
+    *,
+    db_engine: sqlalchemy.Engine,
+) -> set[str]:
+    inspector = sqlalchemy.inspect(db_engine)
+    return {c["name"] for c in inspector.get_columns("execution_node")}
+
+
+class TestMigrateAddExecutionNodeUpdatedAtColumn:
+    def test_adds_column_when_missing(
+        self,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        db_engine = database_ops.create_db_engine(database_uri="sqlite://")
+        _create_execution_node_table_without_updated_at(db_engine=db_engine)
+        assert "updated_at" not in _get_execution_node_columns(db_engine=db_engine)
+
+        with caplog.at_level(logging.INFO):
+            did_add = database_migrations.migrate_add_execution_node_updated_at_column(
+                db_engine=db_engine
+            )
+
+        assert did_add is True
+        assert "updated_at" in _get_execution_node_columns(db_engine=db_engine)
+        assert "Adding column: execution_node.updated_at" in caplog.text
+
+    def test_second_call_is_noop_and_logs_nothing(
+        self,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        db_engine = database_ops.create_db_engine(database_uri="sqlite://")
+        _create_execution_node_table_without_updated_at(db_engine=db_engine)
+
+        database_migrations.migrate_add_execution_node_updated_at_column(
+            db_engine=db_engine
+        )
+        caplog.clear()
+
+        with caplog.at_level(
+            logging.INFO, logger="cloud_pipelines_backend.database_migrations"
+        ):
+            did_add = database_migrations.migrate_add_execution_node_updated_at_column(
+                db_engine=db_engine
+            )
+
+        assert did_add is False
+        assert caplog.messages == []
+
+    def test_column_is_queryable_after_migration(self) -> None:
+        db_engine = database_ops.create_db_engine(database_uri="sqlite://")
+        _create_execution_node_table_without_updated_at(db_engine=db_engine)
+        database_migrations.migrate_add_execution_node_updated_at_column(
+            db_engine=db_engine
+        )
+
+        session_factory = orm.sessionmaker(db_engine)
+        service = api_server_sql.PipelineRunsApiService_Sql()
+        run = _create_run(
+            session_factory=session_factory,
+            service=service,
+            root_task=_make_task_spec(),
+        )
+        with session_factory() as session:
+            db_run = session.get(bts.PipelineRun, run.id)
+            exec_node = session.get(bts.ExecutionNode, db_run.root_execution_id)
+            # insert_default should have stamped this on creation, even though
+            # the column didn't exist until the migration above ran first.
+            assert exec_node.updated_at is not None
+
+
+class TestBackfillExecutionNodeUpdatedAt:
+    def test_backfills_null_rows_with_flat_timestamp(self) -> None:
+        db_engine = database_ops.create_db_engine(database_uri="sqlite://")
+        _create_execution_node_table_without_updated_at(db_engine=db_engine)
+        database_migrations.migrate_add_execution_node_updated_at_column(
+            db_engine=db_engine
+        )
+
+        session_factory = orm.sessionmaker(db_engine)
+        service = api_server_sql.PipelineRunsApiService_Sql()
+        run = _create_run(
+            session_factory=session_factory,
+            service=service,
+            root_task=_make_task_spec(),
+        )
+        # New rows already get updated_at via insert_default -- null it out to
+        # simulate a row created before the column existed.
+        with session_factory() as session:
+            db_run = session.get(bts.PipelineRun, run.id)
+            exec_node = session.get(bts.ExecutionNode, db_run.root_execution_id)
+            exec_node.updated_at = None
+            session.commit()
+
+        with session_factory() as session:
+            count = database_migrations.backfill_execution_node_updated_at(
+                session=session,
+                auto_commit=True,
+            )
+        assert count == 1
+
+        with session_factory() as session:
+            db_run = session.get(bts.PipelineRun, run.id)
+            exec_node = session.get(bts.ExecutionNode, db_run.root_execution_id)
+            assert exec_node.updated_at is not None
+
+    def test_leaves_existing_values_untouched(self) -> None:
+        db_engine = database_ops.create_db_engine(database_uri="sqlite://")
+        bts._TableBase.metadata.create_all(db_engine)
+        session_factory = orm.sessionmaker(db_engine)
+        service = api_server_sql.PipelineRunsApiService_Sql()
+        run = _create_run(
+            session_factory=session_factory,
+            service=service,
+            root_task=_make_task_spec(),
+        )
+        with session_factory() as session:
+            db_run = session.get(bts.PipelineRun, run.id)
+            exec_node = session.get(bts.ExecutionNode, db_run.root_execution_id)
+            original_updated_at = exec_node.updated_at
+        assert original_updated_at is not None
+
+        with session_factory() as session:
+            database_migrations.backfill_execution_node_updated_at(
+                session=session,
+                auto_commit=True,
+            )
+
+        with session_factory() as session:
+            db_run = session.get(bts.PipelineRun, run.id)
+            exec_node = session.get(bts.ExecutionNode, db_run.root_execution_id)
+            assert exec_node.updated_at == original_updated_at
+
+    def test_second_call_updates_zero_rows(self) -> None:
+        db_engine = database_ops.create_db_engine(database_uri="sqlite://")
+        _create_execution_node_table_without_updated_at(db_engine=db_engine)
+        database_migrations.migrate_add_execution_node_updated_at_column(
+            db_engine=db_engine
+        )
+        session_factory = orm.sessionmaker(db_engine)
+        service = api_server_sql.PipelineRunsApiService_Sql()
+        run = _create_run(
+            session_factory=session_factory,
+            service=service,
+            root_task=_make_task_spec(),
+        )
+        with session_factory() as session:
+            db_run = session.get(bts.PipelineRun, run.id)
+            exec_node = session.get(bts.ExecutionNode, db_run.root_execution_id)
+            exec_node.updated_at = None
+            session.commit()
+
+        with session_factory() as session:
+            first = database_migrations.backfill_execution_node_updated_at(
+                session=session,
+                auto_commit=True,
+            )
+        assert first == 1
+
+        with session_factory() as session:
+            second = database_migrations.backfill_execution_node_updated_at(
+                session=session,
+                auto_commit=True,
+            )
+        assert second == 0