Skip to content

Commit 9d794c9

Browse files
authored
Merge branch 'main' into feature/dynamic-task-iteration
2 parents 42259a9 + e3110a3 commit 9d794c9

15 files changed

Lines changed: 583 additions & 104 deletions

File tree

‎.pre-commit-config.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1161,6 +1161,12 @@ repos:
11611161
language: python
11621162
files: ^.*/example_dags/.*\.py$
11631163
pass_filenames: true
1164+
- id: check-example-dag-shell-commands
1165+
name: Check example Dags do not build shell commands from runtime values
1166+
entry: ./scripts/ci/prek/check_example_dag_shell_commands.py
1167+
language: python
1168+
files: ^.*/example_dags/.*\.py$
1169+
pass_filenames: true
11641170
- id: check-apache-license-rat
11651171
name: Check if licenses are OK for Apache
11661172
entry: ./scripts/ci/prek/check_license.py

‎airflow-core/src/airflow/example_dags/example_dag_decorator.py‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -59,20 +59,24 @@ def example_dag_decorator(url: str = "https://httpbingo.org/get"):
5959
"""
6060
get_ip = GetRequestOperator(task_id="get_ip", url=url)
6161

62-
@task(multiple_outputs=True)
63-
def prepare_command(raw_json: dict[str, Any]) -> dict[str, str]:
62+
@task
63+
def get_external_ip(raw_json: dict[str, Any]) -> str:
6464
external_ip = raw_json["origin"]
6565
try:
66-
ipaddress.ip_address(external_ip)
67-
return {
68-
"command": f"echo 'Seems like today your server executing Airflow is connected from IP {external_ip}'",
69-
}
66+
return str(ipaddress.ip_address(external_ip))
7067
except ValueError:
7168
raise ValueError(f"Invalid IP address: '{external_ip}'.")
7269

73-
command_info = prepare_command(get_ip.output)
70+
external_ip = get_external_ip(get_ip.output)
7471

75-
BashOperator(task_id="echo_ip_info", bash_command=command_info["command"])
72+
# Pass the value to the command through the environment rather than formatting it into the
73+
# command text, so the shell never interprets it as part of the command.
74+
BashOperator(
75+
task_id="echo_ip_info",
76+
bash_command='echo "Seems like today your server executing Airflow is connected from IP $EXTERNAL_IP"',
77+
env={"EXTERNAL_IP": external_ip},
78+
append_env=True,
79+
)
7680

7781

7882
example_dag = example_dag_decorator()

‎airflow-core/src/airflow/jobs/triggerer_job_runner.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@
139139
from airflow.utils.helpers import prune_dict
140140
from airflow.utils.log.logging_mixin import LoggingMixin
141141
from airflow.utils.log.task_log_address import prepare_task_log_contexts, render_task_log_filename
142+
from airflow.utils.retries import run_with_db_retries
142143
from airflow.utils.session import create_session, provide_session
143144

144145
if TYPE_CHECKING:
@@ -818,8 +819,10 @@ def on_trigger_event(self, trigger_id: int, event: TriggerEvent) -> None:
818819
Trigger.submit_event(trigger_id=trigger_id, event=event)
819820

820821
def clean_unused(self) -> None:
821-
"""Remove triggers that are no longer needed."""
822-
Trigger.clean_unused()
822+
"""Remove triggers that are no longer needed, retrying database errors in a new transaction."""
823+
for attempt in run_with_db_retries():
824+
with attempt, create_session() as session:
825+
Trigger.clean_unused(session=session)
823826

824827
def handle_failed_triggers(self):
825828
"""

‎airflow-core/src/airflow/models/trigger.py‎

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,7 @@
3737
from airflow.models.taskinstance import TaskInstance
3838
from airflow.serialization.enums import stringify_encoding_keys
3939
from airflow.triggers.base import BaseTaskEndEvent
40-
from airflow.utils.retries import run_with_db_retries
41-
from airflow.utils.session import NEW_SESSION, create_session, provide_session
40+
from airflow.utils.session import NEW_SESSION, provide_session
4241
from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks
4342
from airflow.utils.state import TaskInstanceState
4443

@@ -249,30 +248,30 @@ def fetch_trigger_ids_with_non_task_associations(cls, *, session: Session = NEW_
249248
return set(session.scalars(query))
250249

251250
@classmethod
252-
def clean_unused(cls) -> None:
251+
@provide_session
252+
def clean_unused(cls, *, session: Session = NEW_SESSION) -> None:
253253
"""
254254
Delete up to one batch of triggers that have no dependent tasks, assets, or callbacks.
255255
256256
Triggers have a one-to-many relationship to task instances, so we need to clean those up first.
257257
Afterward we can drop unreferenced triggers; later triggerer loops drain any remaining batches.
258+
259+
Database errors are not retried; to retry, run the call again in a new transaction.
258260
"""
259261
batch_size = conf.getint("triggerer", "unreferenced_triggers_cleanup_batch_size", fallback=500)
260262
if batch_size <= 0:
261263
raise ValueError("[triggerer] unreferenced_triggers_cleanup_batch_size must be at least 1")
262264

263-
for attempt in run_with_db_retries():
264-
with attempt:
265-
with create_session(scoped=False) as session:
266-
# Update all task instances with trigger IDs that are not DEFERRED to remove them
267-
session.execute(
268-
update(TaskInstance)
269-
.where(
270-
TaskInstance.state != TaskInstanceState.DEFERRED,
271-
TaskInstance.trigger_id.is_not(None),
272-
)
273-
.values(trigger_id=None)
274-
)
275-
cls._delete_unused_batch(batch_size, session=session)
265+
# Update all task instances with trigger IDs that are not DEFERRED to remove them
266+
session.execute(
267+
update(TaskInstance)
268+
.where(
269+
TaskInstance.state != TaskInstanceState.DEFERRED,
270+
TaskInstance.trigger_id.is_not(None),
271+
)
272+
.values(trigger_id=None)
273+
)
274+
cls._delete_unused_batch(batch_size, session=session)
276275

277276
@classmethod
278277
def _delete_unused_batch(cls, batch_size: int, *, session: Session) -> None:

‎airflow-core/tests/unit/jobs/test_triggerer_job.py‎

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,9 @@
4747
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
4848
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
4949
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
5153
from structlog.typing import FilteringBoundLogger
5254

5355
from airflow._shared.timezones import timezone
@@ -685,6 +687,71 @@ def test_heartbeat_watchdog_disabled_when_threshold_is_zero(supervisor_builder,
685687
perform_heartbeat_mock.assert_called_once()
686688

687689

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+
688755
def test_metric_tags_default_uses_job_hostname(supervisor_builder):
689756
"""metric_tags() defaults to {"hostname": job.hostname} when a job is present."""
690757
supervisor = supervisor_builder()

‎airflow-core/tests/unit/models/test_trigger.py‎

Lines changed: 19 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -343,24 +343,26 @@ def count_commit(committed_session):
343343
assert commit_count == 0
344344

345345

346-
@conf_vars({("triggerer", "unreferenced_triggers_cleanup_batch_size"): "2"})
347-
def test_clean_unused_retries_failed_batch_in_fresh_transactions(session):
348-
session.add_all(
349-
[
350-
Trigger(classpath=f"airflow.triggers.testing.SuccessTrigger{index}", kwargs={})
351-
for index in range(5)
352-
]
353-
)
346+
def test_clean_unused_runs_in_the_caller_transaction(session):
347+
session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}))
354348
session.commit()
355349

356-
Trigger.clean_unused()
357-
session.expire_all()
358-
assert session.scalar(select(func.count()).select_from(Trigger)) == 3
350+
Trigger.clean_unused(session=session)
351+
assert session.scalar(select(func.count()).select_from(Trigger)) == 0
352+
session.rollback()
353+
354+
assert session.scalar(select(func.count()).select_from(Trigger)) == 1
355+
356+
357+
def test_clean_unused_raises_database_errors_to_the_caller_without_retrying(session):
358+
session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={}))
359+
session.commit()
360+
# The caller's own pending write must survive the cleanup failing.
361+
job = Job()
362+
session.add(job)
363+
session.flush()
359364

360365
delete_attempt_count = 0
361-
cleanup_sessions = []
362-
commit_count = 0
363-
rollback_count = 0
364366

365367
def fail_delete(connection, cursor, statement, parameters, context, executemany):
366368
nonlocal delete_attempt_count
@@ -369,42 +371,15 @@ def fail_delete(connection, cursor, statement, parameters, context, executemany)
369371
delete_attempt_count += 1
370372
raise OperationalError(statement, parameters, RuntimeError("injected cleanup failure"))
371373

372-
def capture_session(cleanup_session, transaction, connection):
373-
cleanup_sessions.append(cleanup_session)
374-
375-
def count_commit(committed_session):
376-
nonlocal commit_count
377-
commit_count += 1
378-
379-
def count_rollback(rolled_back_session):
380-
nonlocal rollback_count
381-
rollback_count += 1
382-
383374
event.listen(session.bind, "before_cursor_execute", fail_delete)
384-
event.listen(session.__class__, "after_begin", capture_session)
385-
event.listen(session.__class__, "after_commit", count_commit)
386-
event.listen(session.__class__, "after_rollback", count_rollback)
387375
try:
388376
with pytest.raises(OperationalError, match="injected cleanup failure"):
389-
Trigger.clean_unused()
377+
Trigger.clean_unused(session=session)
390378
finally:
391379
event.remove(session.bind, "before_cursor_execute", fail_delete)
392-
event.remove(session.__class__, "after_begin", capture_session)
393-
event.remove(session.__class__, "after_commit", count_commit)
394-
event.remove(session.__class__, "after_rollback", count_rollback)
395380

396-
assert len({id(cleanup_session) for cleanup_session in cleanup_sessions}) == 3
397-
assert (delete_attempt_count, commit_count, rollback_count) == (3, 0, 3)
398-
session.expire_all()
399-
assert session.scalar(select(func.count()).select_from(Trigger)) == 3
400-
401-
Trigger.clean_unused()
402-
session.expire_all()
403-
assert session.scalar(select(func.count()).select_from(Trigger)) == 1
404-
405-
Trigger.clean_unused()
406-
session.expire_all()
407-
assert session.scalar(select(func.count()).select_from(Trigger)) == 0
381+
assert delete_attempt_count == 1
382+
assert job in session
408383

409384

410385
@patch.object(TriggererCallback, "handle_event")

‎providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -853,13 +853,9 @@ def test_get_conn_with_password_none_auth(self, mock_connect):
853853
],
854854
)
855855
def test_get_conn_with_wrong_connection_parameters(self, host, port, schema, message):
856-
connection = Connection(
857-
conn_id="test",
858-
conn_type="hive",
859-
host=host,
860-
port=port,
861-
schema=schema,
862-
)
856+
# Newer core Connection rejects invalid ports on assignment; a spec'd mock lets the hook's own
857+
# check, which still guards older Airflow versions, see them.
858+
connection = mock.Mock(spec=Connection, conn_id="test", host=host, port=port, schema=schema)
863859
hook = HiveCliHook()
864860
if message:
865861
with pytest.raises(Exception, match=message):

‎providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,10 @@ def connection():
135135
@task_group(prefix_group_id=False)
136136
def standard_tasks_group():
137137
classic_bash = BashOperator(
138-
task_id="classic_bash", bash_command="echo Parameter is {{ params.mapping_count }}"
138+
task_id="classic_bash",
139+
bash_command='echo "Parameter is $MAPPING_COUNT"',
140+
env={"MAPPING_COUNT": "{{ params.mapping_count }}"},
141+
append_env=True,
139142
)
140143

141144
empty = EmptyOperator(task_id="not_executed")

‎providers/openlineage/tests/system/openlineage/example_openlineage_base_complex_dag.py‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,13 @@ def execute(self, context):
9797
task_2 = PythonOperator(
9898
task_id="task_2",
9999
python_callable=do_nothing,
100-
inlets=[Asset(uri="file://host1/dir2/file2.txt"), Asset(uri="unknown_scheme://host1/dir2/file3.txt")],
100+
inlets=[
101+
Asset(
102+
uri="file://host1/dir2/file2.txt",
103+
extra={"etag": '"0x8D0000000000001"', "owner": "O'Brien", "has_header": True, "note": None},
104+
),
105+
Asset(uri="unknown_scheme://host1/dir2/file3.txt"),
106+
],
101107
max_retry_delay=42,
102108
doc="text doc",
103109
doc_md="should be skipped",
@@ -107,7 +113,12 @@ def execute(self, context):
107113
)
108114
task_3 = EmptyOperator(
109115
task_id="task_3",
110-
outlets=[Asset(uri="file://host/dir/file.txt")],
116+
outlets=[
117+
Asset(
118+
uri="file://host/dir/file.txt",
119+
extra={"checks": [{"passed": True}], "ratio": 0.5},
120+
)
121+
],
111122
doc_md="MD doc",
112123
doc_json="should be skipped",
113124
doc_yaml="should be skipped",

0 commit comments

Comments
 (0)