|
47 | 47 | from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter |
48 | 48 | from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator |
49 | 49 | from pydantic import TypeAdapter |
50 | | -from sqlalchemy import delete, update |
| 50 | +from sqlalchemy import create_engine, delete, event, func, select, text, update |
| 51 | +from sqlalchemy.exc import OperationalError |
| 52 | +from sqlalchemy.pool import NullPool |
51 | 53 | from structlog.typing import FilteringBoundLogger |
52 | 54 |
|
53 | 55 | from airflow._shared.timezones import timezone |
@@ -685,6 +687,71 @@ def test_heartbeat_watchdog_disabled_when_threshold_is_zero(supervisor_builder, |
685 | 687 | perform_heartbeat_mock.assert_called_once() |
686 | 688 |
|
687 | 689 |
|
| 690 | +def _is_trigger_delete(statement: str) -> bool: |
| 691 | + return "delete from trigger" in " ".join(statement.lower().replace('"', "").replace("`", "").split()) |
| 692 | + |
| 693 | + |
| 694 | +def test_clean_unused_retries_failed_batches_in_new_transactions(jobless_supervisor, session): |
| 695 | + session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})) |
| 696 | + session.commit() |
| 697 | + |
| 698 | + delete_attempt_count = 0 |
| 699 | + begin_count = 0 |
| 700 | + |
| 701 | + def fail_first_two_deletes(connection, cursor, statement, parameters, context, executemany): |
| 702 | + nonlocal delete_attempt_count |
| 703 | + if _is_trigger_delete(statement): |
| 704 | + delete_attempt_count += 1 |
| 705 | + if delete_attempt_count <= 2: |
| 706 | + raise OperationalError(statement, parameters, RuntimeError("injected cleanup failure")) |
| 707 | + |
| 708 | + def count_begin(cleanup_session, transaction, connection): |
| 709 | + nonlocal begin_count |
| 710 | + begin_count += 1 |
| 711 | + |
| 712 | + event.listen(session.bind, "before_cursor_execute", fail_first_two_deletes) |
| 713 | + event.listen(session.__class__, "after_begin", count_begin) |
| 714 | + try: |
| 715 | + jobless_supervisor.clean_unused() |
| 716 | + finally: |
| 717 | + event.remove(session.bind, "before_cursor_execute", fail_first_two_deletes) |
| 718 | + event.remove(session.__class__, "after_begin", count_begin) |
| 719 | + |
| 720 | + assert (delete_attempt_count, begin_count) == (3, 3) |
| 721 | + session.expire_all() |
| 722 | + assert session.scalar(select(func.count()).select_from(Trigger)) == 0 |
| 723 | + |
| 724 | + |
| 725 | +@pytest.mark.backend("postgres") |
| 726 | +def test_clean_unused_recovers_after_the_database_connection_drops(jobless_supervisor, session): |
| 727 | + session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})) |
| 728 | + session.commit() |
| 729 | + |
| 730 | + killer_engine = create_engine(session.bind.url, poolclass=NullPool) |
| 731 | + terminated = [] |
| 732 | + |
| 733 | + def terminate_connection_before_delete(connection, cursor, statement, parameters, context, executemany): |
| 734 | + if not terminated and _is_trigger_delete(statement): |
| 735 | + backend_pid = connection.connection.dbapi_connection.info.backend_pid |
| 736 | + with killer_engine.connect() as killer: |
| 737 | + terminated.append( |
| 738 | + killer.execute( |
| 739 | + text("SELECT pg_terminate_backend(:pid, 5000)"), {"pid": backend_pid} |
| 740 | + ).scalar() |
| 741 | + ) |
| 742 | + |
| 743 | + event.listen(session.bind, "before_cursor_execute", terminate_connection_before_delete) |
| 744 | + try: |
| 745 | + jobless_supervisor.clean_unused() |
| 746 | + finally: |
| 747 | + event.remove(session.bind, "before_cursor_execute", terminate_connection_before_delete) |
| 748 | + killer_engine.dispose() |
| 749 | + |
| 750 | + assert terminated == [True] |
| 751 | + session.expire_all() |
| 752 | + assert session.scalar(select(func.count()).select_from(Trigger)) == 0 |
| 753 | + |
| 754 | + |
688 | 755 | def test_metric_tags_default_uses_job_hostname(supervisor_builder): |
689 | 756 | """metric_tags() defaults to {"hostname": job.hostname} when a job is present.""" |
690 | 757 | supervisor = supervisor_builder() |
|
0 commit comments