diff --git a/cloud_pipelines_backend/backend_types_sql.py b/cloud_pipelines_backend/backend_types_sql.py index a2f8f88..c47e943 100644 --- a/cloud_pipelines_backend/backend_types_sql.py +++ b/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"]] diff --git a/cloud_pipelines_backend/database_migrations.py b/cloud_pipelines_backend/database_migrations.py index 8a75462..bf0be03 100644 --- a/cloud_pipelines_backend/database_migrations.py +++ b/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 diff --git a/cloud_pipelines_backend/database_ops.py b/cloud_pipelines_backend/database_ops.py index 75ed714..991800c 100644 --- a/cloud_pipelines_backend/database_ops.py +++ b/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") diff --git a/tests/test_api_server_sql.py b/tests/test_api_server_sql.py index f8f23b9..0e2b787 100644 --- a/tests/test_api_server_sql.py +++ b/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 diff --git a/tests/test_database_migrations.py b/tests/test_database_migrations.py index 08b7387..343384e 100644 --- a/tests/test_database_migrations.py +++ b/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 diff --git a/tests/test_sql_event_listeners.py b/tests/test_sql_event_listeners.py index 832f1eb..572ccbc 100644 --- a/tests/test_sql_event_listeners.py +++ b/tests/test_sql_event_listeners.py @@ -1,5 +1,6 @@ """Tests for SQLAlchemy event listeners in orchestrator_sql and instrumentation.metrics.""" +import time import unittest.mock import pytest @@ -67,3 +68,53 @@ def test_second_status_change_records_duration_metric( "execution.status.from": bts.ContainerExecutionStatus.QUEUED, "execution.status.to": bts.ContainerExecutionStatus.RUNNING, } + + +class TestExecutionNodeUpdatedAt: + """Verify execution_node.updated_at is self-maintaining via + insert_default/onupdate on the mapped_column -- no listener needed.""" + + def test_creation_sets_updated_at_without_explicit_value( + self, session: orm.Session + ) -> None: + node = bts.ExecutionNode(task_spec={}) + session.add(node) + session.commit() + + assert node.updated_at is not None + + def test_status_change_bumps_updated_at(self, session: orm.Session) -> None: + node = bts.ExecutionNode(task_spec={}) + session.add(node) + session.commit() + first_updated_at = node.updated_at + assert first_updated_at is not None + + time.sleep(0.001) + node.container_execution_status = bts.ContainerExecutionStatus.QUEUED + session.commit() + + assert node.updated_at is not None + assert node.updated_at > first_updated_at + + def test_extra_data_only_change_bumps_updated_at( + self, session: orm.Session + ) -> None: + """A plain extra_data mutation with no status touch also bumps + updated_at -- this is the case a status-scoped listener would miss. + `PipelineRunsApiService_Sql.terminate()` only ever sets + extra_data["desired_state"], never container_execution_status + directly, so onupdate (fires on any UPDATE to the row) is what + makes that gap get covered too.""" + node = bts.ExecutionNode(task_spec={}, extra_data={}) + session.add(node) + session.commit() + first_updated_at = node.updated_at + assert first_updated_at is not None + + time.sleep(0.001) + node.extra_data["desired_state"] = "TERMINATED" + session.commit() + + assert node.updated_at is not None + assert node.updated_at > first_updated_at