From 919ab9921946e0332ead3ee77920f1140ed940f6 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 22 Oct 2025 16:20:09 +0800 Subject: [PATCH 1/7] Move Sentry integration to Task SDK The integration logic is moved from core into task sdk so it can be run at runtime. Things that used to rely on database access have been rewritten to use the execution API instead. Not yet applied to anything. Tests to come. --- airflow-core/docs/extra-packages-ref.rst | 16 +- airflow-core/pyproject.toml | 7 +- .../execution_api/datamodels/taskinstance.py | 7 + .../execution_api/routes/task_instances.py | 16 ++ airflow-core/src/airflow/sentry.py | 196 ------------------ .../unit/dag_processing/test_processor.py | 2 + .../tests/unit/jobs/test_triggerer_job.py | 2 + pyproject.toml | 11 +- .../ci/prek/update_airflow_pyproject_toml.py | 11 + task-sdk/pyproject.toml | 6 + task-sdk/src/airflow/sdk/api/client.py | 6 + .../airflow/sdk/api/datamodels/_generated.py | 8 + .../src/airflow/sdk/execution_time/comms.py | 24 +++ .../sdk/execution_time/sentry/__init__.py | 34 +++ .../sdk/execution_time/sentry/configured.py | 151 ++++++++++++++ .../airflow/sdk/execution_time/sentry/noop.py | 51 +++++ .../airflow/sdk/execution_time/supervisor.py | 5 + .../airflow/sdk/execution_time/task_runner.py | 10 + .../task_sdk/execution_time}/test_sentry.py | 92 ++++---- 19 files changed, 405 insertions(+), 250 deletions(-) delete mode 100644 airflow-core/src/airflow/sentry.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/sentry/__init__.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/sentry/configured.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/sentry/noop.py rename {airflow-core/tests/unit/core => task-sdk/tests/task_sdk/execution_time}/test_sentry.py (63%) diff --git a/airflow-core/docs/extra-packages-ref.rst b/airflow-core/docs/extra-packages-ref.rst index 0ba12979129f7..9f307fd95a7ce 100644 --- a/airflow-core/docs/extra-packages-ref.rst +++ b/airflow-core/docs/extra-packages-ref.rst @@ -426,10 +426,12 @@ Group extras The group extras are convenience extras. Such extra installs many optional dependencies together. It is not recommended to use it in production, but it is useful for CI, development and testing purposes. -+-----------+------------------------------------------+---------------------------------------------------+ -| extra | install command | enables | -+===========+==========================================+===================================================+ -| all | ``pip install apache-airflow[all]`` | All optional dependencies including all providers | -+-----------+------------------------------------------+---------------------------------------------------+ -| all-core | ``pip install apache-airflow[all-core]`` | All optional core dependencies | -+-----------+------------------------------------------+---------------------------------------------------+ ++--------------+----------------------------------------------+---------------------------------------------------+ +| extra | install command | enables | ++==============+==============================================+===================================================+ +| all | ``pip install apache-airflow[all]`` | All optional dependencies including all providers | ++--------------+----------------------------------------------+---------------------------------------------------+ +| all-core | ``pip install apache-airflow[all-core]`` | All optional core dependencies | ++--------------+----------------------------------------------+---------------------------------------------------+ +| all-task-sdk | ``pip install apache-airflow[all-task-sdk]`` | All optional task SDK dependencies | ++--------------+----------------------------------------------+---------------------------------------------------+ diff --git a/airflow-core/pyproject.toml b/airflow-core/pyproject.toml index b6e488e96519a..3b9fbe1763a4f 100644 --- a/airflow-core/pyproject.toml +++ b/airflow-core/pyproject.toml @@ -178,16 +178,11 @@ dependencies = [ "otel" = [ "opentelemetry-exporter-prometheus>=0.47b0", ] -"sentry" = [ - "blinker>=1.1", - # Apparently sentry needs flask to be installed to work properly - "sentry-sdk[flask]>=2.30.0", -] "statsd" = [ "statsd>=3.3.0", ] "all" = [ - "apache-airflow-core[graphviz,kerberos,otel,sentry,statsd]" + "apache-airflow-core[graphviz,kerberos,otel,statsd]" ] [project.scripts] diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index f2e7aa36b0d9b..4213af1b120d1 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -17,6 +17,7 @@ from __future__ import annotations import uuid +from collections.abc import Iterable from datetime import timedelta from enum import Enum from typing import Annotated, Any, Literal @@ -355,6 +356,12 @@ class TaskStatesResponse(BaseModel): task_states: dict[str, Any] +class TaskBreadcrumbsResponse(BaseModel): + """Response for task breadcrumbs.""" + + breadcrumbs: Iterable[dict[str, Any]] + + class InactiveAssetsResponse(BaseModel): """Response for inactive assets.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index dab8ac1842936..6963f7a8edf8a 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -43,6 +43,7 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( InactiveAssetsResponse, PrevSuccessfulDagRunResponse, + TaskBreadcrumbsResponse, TaskStatesResponse, TIDeferredStatePayload, TIEnterRunningPayload, @@ -917,6 +918,21 @@ def get_task_instance_states( return TaskStatesResponse(task_states=run_id_task_state_map) +@router.get("/breadcrumbs", status_code=status.HTTP_200_OK) +def get_task_instance_breadcrumbs(dag_id: str, run_id: str, session: SessionDep) -> TaskBreadcrumbsResponse: + result = session.execute( + select(TI.task_id, TI.map_index, TI.state, TI.operator, TI.duration).where( + TI.dag_id == dag_id, TI.run_id == run_id + ) + ).mappings() + + def _iter_breadcrumbs() -> Iterator[dict[str, Any]]: + for row in result: + yield {str(k): v for k, v in row.items()} + + return TaskBreadcrumbsResponse(breadcrumbs=_iter_breadcrumbs()) + + def _is_eligible_to_retry(state: str, try_number: int, max_tries: int) -> bool: """Is task instance is eligible for retry.""" if state == TaskInstanceState.RESTARTING: diff --git a/airflow-core/src/airflow/sentry.py b/airflow-core/src/airflow/sentry.py deleted file mode 100644 index 7a7b1df2e118e..0000000000000 --- a/airflow-core/src/airflow/sentry.py +++ /dev/null @@ -1,196 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""Sentry Integration.""" - -from __future__ import annotations - -import logging -from functools import wraps -from typing import TYPE_CHECKING - -from airflow.configuration import conf -from airflow.executors.executor_loader import ExecutorLoader -from airflow.utils.session import find_session_idx, provide_session -from airflow.utils.state import TaskInstanceState - -if TYPE_CHECKING: - from sqlalchemy.orm import Session - - from airflow.models.taskinstance import TaskInstance - -log = logging.getLogger(__name__) - - -class DummySentry: - """Blank class for Sentry.""" - - def add_tagging(self, task_instance): - """Blank function for tagging.""" - - def add_breadcrumbs(self, task_instance, session: Session | None = None): - """Blank function for breadcrumbs.""" - - def enrich_errors(self, run): - """Blank function for formatting a TaskInstance._run_raw_task.""" - return run - - def flush(self): - """Blank function for flushing errors.""" - - -Sentry: DummySentry = DummySentry() -if conf.getboolean("sentry", "sentry_on", fallback=False): - import sentry_sdk - from sentry_sdk.integrations.flask import FlaskIntegration - from sentry_sdk.integrations.logging import ignore_logger - - class ConfiguredSentry(DummySentry): - """Configure Sentry SDK.""" - - SCOPE_DAG_RUN_TAGS = frozenset(("data_interval_end", "data_interval_start", "logical_date")) - SCOPE_TASK_INSTANCE_TAGS = frozenset(("task_id", "dag_id", "try_number")) - SCOPE_CRUMBS = frozenset(("task_id", "state", "operator", "duration")) - - UNSUPPORTED_SENTRY_OPTIONS = frozenset( - ( - "integrations", - "in_app_include", - "in_app_exclude", - "ignore_errors", - "before_breadcrumb", - ) - ) - - def __init__(self): - """Initialize the Sentry SDK.""" - ignore_logger("airflow.task") - - sentry_flask = FlaskIntegration() - - # LoggingIntegration is set by default. - integrations = [sentry_flask] - - executor_class, _ = ExecutorLoader.import_default_executor_cls() - - if executor_class.supports_sentry: - from sentry_sdk.integrations.celery import CeleryIntegration - - sentry_celery = CeleryIntegration() - integrations.append(sentry_celery) - - dsn = None - sentry_config_opts = conf.getsection("sentry") or {} - if sentry_config_opts: - sentry_config_opts.pop("sentry_on") - old_way_dsn = sentry_config_opts.pop("sentry_dsn", None) - new_way_dsn = sentry_config_opts.pop("dsn", None) - # supported backward compatibility with old way dsn option - dsn = old_way_dsn or new_way_dsn - - unsupported_options = self.UNSUPPORTED_SENTRY_OPTIONS.intersection(sentry_config_opts.keys()) - if unsupported_options: - log.warning( - "There are unsupported options in [sentry] section: %s", - ", ".join(unsupported_options), - ) - - sentry_config_opts["before_send"] = conf.getimport("sentry", "before_send", fallback=None) - sentry_config_opts["transport"] = conf.getimport("sentry", "transport", fallback=None) - - if dsn: - sentry_sdk.init(dsn=dsn, integrations=integrations, **sentry_config_opts) - else: - # Setting up Sentry using environment variables. - log.debug("Defaulting to SENTRY_DSN in environment.") - sentry_sdk.init(integrations=integrations, **sentry_config_opts) - - def add_tagging(self, task_instance): - """Add tagging for a task_instance.""" - dag_run = task_instance.dag_run - task = task_instance.task - - with sentry_sdk.configure_scope() as scope: - for tag_name in self.SCOPE_TASK_INSTANCE_TAGS: - attribute = getattr(task_instance, tag_name) - scope.set_tag(tag_name, attribute) - for tag_name in self.SCOPE_DAG_RUN_TAGS: - attribute = getattr(dag_run, tag_name) - scope.set_tag(tag_name, attribute) - scope.set_tag("operator", task.__class__.__name__) - - @provide_session - def add_breadcrumbs( - self, - task_instance: TaskInstance, - session: Session | None = None, - ) -> None: - """Add breadcrumbs inside of a task_instance.""" - if session is None: - return - dr = task_instance.get_dagrun(session) - task_instances = dr.get_task_instances( - state={TaskInstanceState.SUCCESS, TaskInstanceState.FAILED}, - session=session, - ) - - for ti in task_instances: - data = {} - for crumb_tag in self.SCOPE_CRUMBS: - data[crumb_tag] = getattr(ti, crumb_tag) - - sentry_sdk.add_breadcrumb(category="completed_tasks", data=data, level="info") - - def enrich_errors(self, func): - """ - Decorate errors. - - Wrap TaskInstance._run_raw_task to support task specific tags and breadcrumbs. - """ - session_args_idx = find_session_idx(func) - - @wraps(func) - def wrapper(_self, *args, **kwargs): - # Wrapping the _run_raw_task function with push_scope to contain - # tags and breadcrumbs to a specific Task Instance - - try: - session = kwargs.get("session", args[session_args_idx]) - except IndexError: - session = None - - with sentry_sdk.push_scope(): - try: - # Is a LocalTaskJob get the task instance - if hasattr(_self, "task_instance"): - task_instance = _self.task_instance - else: - task_instance = _self - - self.add_tagging(task_instance) - self.add_breadcrumbs(task_instance, session=session) - return func(_self, *args, **kwargs) - except Exception as e: - sentry_sdk.capture_exception(e) - raise - - return wrapper - - def flush(self): - sentry_sdk.flush() - - Sentry = ConfiguredSentry() diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index ef1ee928bf706..35c6babe3a880 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -1822,6 +1822,7 @@ def get_type_names(union_type): "GetAssetEventByAssetAlias", "GetDagRunState", "GetDRCount", + "GetTaskBreadcrumbs", "GetTaskRescheduleStartDate", "GetTICount", "GetTaskStates", @@ -1848,6 +1849,7 @@ def get_type_names(union_type): "DRCount", "SentFDs", "StartupDetails", + "TaskBreadcrumbsResult", "TaskRescheduleStartDate", "TICount", "TaskStatesResult", diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 7ff1dac002bc1..ea1974637b90d 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -1199,6 +1199,7 @@ def get_type_names(union_type): "GetAssetEventByAssetAlias", "GetPrevSuccessfulDagRun", "GetPreviousDagRun", + "GetTaskBreadcrumbs", "GetTaskRescheduleStartDate", "GetXComCount", "GetXComSequenceItem", @@ -1221,6 +1222,7 @@ def get_type_names(union_type): "AssetEventsResult", "SentFDs", "StartupDetails", + "TaskBreadcrumbsResult", "TaskRescheduleStartDate", "InactiveAssetsResult", "CreateHITLDetailPayload", diff --git a/pyproject.toml b/pyproject.toml index 21277b1faee68..509897f4506f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,12 +92,15 @@ packages = [] "otel" = [ "apache-airflow-core[otel]" ] -"sentry" = [ - "apache-airflow-core[sentry]" -] "statsd" = [ "apache-airflow-core[statsd]" ] +"all-task-sdk" = [ + "apache-airflow-task-sdk[all]" +] +"sentry" = [ + "apache-airflow-task-sdk[sentry]" +] "airbyte" = [ "apache-airflow-providers-airbyte>=5.0.0" ] @@ -1284,7 +1287,7 @@ dev = [ "apache-airflow-task-sdk-tests", "apache-airflow-helm-tests", "apache-airflow-kubernetes-tests", - "apache-airflow-task-sdk", + "apache-airflow-task-sdk[all]", "apache-airflow-ctl", "apache-airflow-ctl-tests", "apache-airflow-shared-logging", diff --git a/scripts/ci/prek/update_airflow_pyproject_toml.py b/scripts/ci/prek/update_airflow_pyproject_toml.py index 1814fe3cc227c..962e9aa00d49c 100755 --- a/scripts/ci/prek/update_airflow_pyproject_toml.py +++ b/scripts/ci/prek/update_airflow_pyproject_toml.py @@ -45,6 +45,9 @@ AIRFLOW_CORE_ROOT_PATH = AIRFLOW_ROOT_PATH / "airflow-core" AIRFLOW_CORE_PYPROJECT_TOML_FILE = AIRFLOW_CORE_ROOT_PATH / "pyproject.toml" +AIRFLOW_TASK_SDK_ROOT_PATH = AIRFLOW_ROOT_PATH / "task-sdk" +AIRFLOW_TASK_SDK_PYPROJECT_TOML_FILE = AIRFLOW_TASK_SDK_ROOT_PATH / "pyproject.toml" + PROVIDERS_DIR = AIRFLOW_ROOT_PATH / "providers" START_OPTIONAL_DEPENDENCIES = ( @@ -178,6 +181,14 @@ def get_python_exclusion(provider_dependencies: dict[str, Any]) -> str: all_optional_dependencies.append('"all-core" = [\n "apache-airflow-core[all]"\n]\n') else: all_optional_dependencies.append(f'"{optional}" = [\n "apache-airflow-core[{optional}]"\n]\n') + optional_airflow_task_sdk_dependencies = get_optional_dependencies(AIRFLOW_TASK_SDK_PYPROJECT_TOML_FILE) + for optional in sorted(optional_airflow_task_sdk_dependencies): + if optional == "all": + all_optional_dependencies.append('"all-task-sdk" = [\n "apache-airflow-task-sdk[all]"\n]\n') + else: + all_optional_dependencies.append( + f'"{optional}" = [\n "apache-airflow-task-sdk[{optional}]"\n]\n' + ) all_providers = sorted(get_all_provider_ids()) all_provider_lines = [] for provider_id in all_providers: diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index a46b57cf69646..5601438b44490 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -75,6 +75,12 @@ dependencies = [ # End of shared logging dependencies ] +[project.optional-dependencies] +"sentry" = [ + "sentry-sdk>=2.30.0", +] +"all" = ["apache-airflow-task-sdk[sentry]"] + [project.urls] "Bug Tracker" = "https://github.com/apache/airflow/issues" Documentation = "https://airflow.apache.org/docs/" diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index c4cb0178c1a1f..9867c67535a7e 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -53,6 +53,7 @@ HITLUser, InactiveAssetsResponse, PrevSuccessfulDagRunResponse, + TaskBreadcrumbsResponse, TaskInstanceState, TaskStatesResponse, TerminalStateNonSuccess, @@ -348,6 +349,11 @@ def get_task_states( resp = self.client.get("task-instances/states", params=params) return TaskStatesResponse.model_validate_json(resp.read()) + def get_task_breakcrumbs(self, dag_id: str, run_id: str) -> TaskBreadcrumbsResponse: + params = {"dag_id": dag_id, "run_id": run_id} + resp = self.client.get("task-instances/breadcrumbs", params=params) + return TaskBreadcrumbsResponse.model_validate_json(resp.read()) + def validate_inlets_and_outlets(self, id: uuid.UUID) -> InactiveAssetsResponse: """Validate whether there're inactive assets in inlets and outlets of a given task instance.""" resp = self.client.get(f"task-instances/{id}/validate-inlets-and-outlets") diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index ea8fbf96de1c9..ca2dceec3b074 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -278,6 +278,14 @@ class TITargetStatePayload(BaseModel): state: IntermediateTIState +class TaskBreadcrumbsResponse(BaseModel): + """ + Response for task breadcrumbs. + """ + + breadcrumbs: Annotated[list[dict[str, Any]], Field(title="Breadcrumbs")] + + class TaskStatesResponse(BaseModel): """ Response for task states with run_id, task and state. diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 45986bd398fee..df24085ea66f4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -75,6 +75,7 @@ HITLDetailRequest, InactiveAssetsResponse, PrevSuccessfulDagRunResponse, + TaskBreadcrumbsResponse, TaskInstance, TaskInstanceState, TaskStatesResponse, @@ -550,6 +551,21 @@ def from_api_response(cls, task_states_response: TaskStatesResponse) -> TaskStat return cls(**task_states_response.model_dump(exclude_defaults=True), type="TaskStatesResult") +class TaskBreadcrumbsResult(TaskBreadcrumbsResponse): + type: Literal["TaskBreadcrumbsResult"] = "TaskBreadcrumbsResult" + + @classmethod + def from_api_response(cls, response: TaskBreadcrumbsResponse) -> TaskBreadcrumbsResult: + """ + Create result class from API Response. + + API Response is autogenerated from the API schema, so we need to convert + it to Result for communication between the Supervisor and the task + process since it needs a discriminator field. + """ + return cls(**response.model_dump(exclude_defaults=True), type="TaskBreadcrumbsResult") + + class DRCount(BaseModel): """Response containing count of Dag Runs matching certain filters.""" @@ -608,6 +624,7 @@ def from_api_response(cls, hitl_request: HITLDetailRequest) -> HITLDetailRequest | StartupDetails | TaskRescheduleStartDate | TICount + | TaskBreadcrumbsResult | TaskStatesResult | VariableResult | XComCountResponse @@ -895,6 +912,12 @@ class GetTaskStates(BaseModel): type: Literal["GetTaskStates"] = "GetTaskStates" +class GetTaskBreadcrumbs(BaseModel): + dag_id: str + run_id: str + type: Literal["GetTaskBreadcrumbs"] = "GetTaskBreadcrumbs" + + class GetDRCount(BaseModel): dag_id: str logical_dates: list[AwareDatetime] | None = None @@ -942,6 +965,7 @@ class MaskSecret(BaseModel): | GetPreviousDagRun | GetTaskRescheduleStartDate | GetTICount + | GetTaskBreadcrumbs | GetTaskStates | GetVariable | GetXCom diff --git a/task-sdk/src/airflow/sdk/execution_time/sentry/__init__.py b/task-sdk/src/airflow/sdk/execution_time/sentry/__init__.py new file mode 100644 index 0000000000000..547cb7076f811 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/sentry/__init__.py @@ -0,0 +1,34 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from airflow.configuration import conf + +__all__ = ["Sentry"] + +Sentry: NoopSentry + +if conf.getboolean("sentry", "sentry_on", fallback=False): + from airflow.sdk.execution_time.sentry.configured import ConfiguredSentry + + Sentry = ConfiguredSentry() +else: + from airflow.sdk.execution_time.sentry.noop import NoopSentry + + Sentry = NoopSentry() diff --git a/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py b/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py new file mode 100644 index 0000000000000..e9753ad9b706a --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py @@ -0,0 +1,151 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Configured Sentry integration. + +This module must only be imported conditionally since the Sentry SDK is NOT a +required dependency of the Airflow Task SDK. You shouldn't import this module +anyway, but use the parent ``airflow.sdk.execution_time.sentry`` path instead, +where things in this module are re-exported. +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING + +import sentry_sdk +import sentry_sdk.integrations.logging +import structlog + +from airflow.sdk.execution_time.sentry.noop import NoopSentry +from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + +if TYPE_CHECKING: + from structlog.typing import FilteringBoundLogger as Logger + + from airflow.sdk import Context + from airflow.sdk.execution_time.sentry.noop import Run, RunReturn + from airflow.sdk.types import DagRunProtocol, RuntimeTaskInstanceProtocol + +log = structlog.get_logger(logger_name=__name__) + + +class ConfiguredSentry(NoopSentry): + """Configure Sentry SDK.""" + + SCOPE_DAG_RUN_TAGS = frozenset(("data_interval_end", "data_interval_start", "logical_date")) + SCOPE_TASK_INSTANCE_TAGS = frozenset(("task_id", "dag_id", "try_number")) + + UNSUPPORTED_SENTRY_OPTIONS = frozenset( + ( + "integrations", + "in_app_include", + "in_app_exclude", + "ignore_errors", + "before_breadcrumb", + ) + ) + + def __init__(self): + """Initialize the Sentry SDK.""" + from airflow.configuration import conf + + sentry_sdk.integrations.logging.ignore_logger("airflow.task") + + # LoggingIntegration is set by default. + integrations = [] + + # TODO: How can we get executor info in the runner to support this? + # executor_class, _ = ExecutorLoader.import_default_executor_cls() + # if executor_class.supports_sentry: + # from sentry_sdk.integrations.celery import CeleryIntegration + + # sentry_celery = CeleryIntegration() + # integrations.append(sentry_celery) + + dsn = None + sentry_config_opts = conf.getsection("sentry") or {} + if sentry_config_opts: + sentry_config_opts.pop("sentry_on") + old_way_dsn = sentry_config_opts.pop("sentry_dsn", None) + new_way_dsn = sentry_config_opts.pop("dsn", None) + # supported backward compatibility with old way dsn option + dsn = old_way_dsn or new_way_dsn + + if unsupported_options := self.UNSUPPORTED_SENTRY_OPTIONS.intersection(sentry_config_opts): + log.warning( + "There are unsupported options in [sentry] section", + options=unsupported_options, + ) + + sentry_config_opts["before_send"] = conf.getimport("sentry", "before_send", fallback=None) + sentry_config_opts["transport"] = conf.getimport("sentry", "transport", fallback=None) + + if dsn: + sentry_sdk.init(dsn=dsn, integrations=integrations, **sentry_config_opts) + else: + # Setting up Sentry using environment variables. + log.debug("Defaulting to SENTRY_DSN in environment.") + sentry_sdk.init(integrations=integrations, **sentry_config_opts) + + def add_tagging(self, dag_run: DagRunProtocol, task_instance: RuntimeTaskInstanceProtocol) -> None: + """Add tagging for a task_instance.""" + task = task_instance.task + + with sentry_sdk.configure_scope() as scope: + for tag_name in self.SCOPE_TASK_INSTANCE_TAGS: + attribute = getattr(task_instance, tag_name) + scope.set_tag(tag_name, attribute) + for tag_name in self.SCOPE_DAG_RUN_TAGS: + attribute = getattr(dag_run, tag_name) + scope.set_tag(tag_name, attribute) + scope.set_tag("operator", task.__class__.__name__) + + def add_breadcrumbs(self, task_instance: RuntimeTaskInstanceProtocol) -> None: + """Add breadcrumbs inside of a task_instance.""" + breadcrumbs = RuntimeTaskInstance.get_task_breadcrumbs( + dag_id=task_instance.dag_id, + run_id=task_instance.run_id, + ) + for breadcrumb in breadcrumbs: + sentry_sdk.add_breadcrumb(category="completed_tasks", data=breadcrumb, level="info") + + def enrich_errors(self, run: Run) -> Run: + """ + Decorate errors. + + Wrap :func:`airflow.sdk.execution_time.task_runner.run` to support task + specific tags and breadcrumbs. + """ + + @functools.wraps(run) + def wrapped_run(ti: RuntimeTaskInstance, context: Context, log: Logger) -> RunReturn: + with sentry_sdk.push_scope(): + try: + self.add_tagging(context["dag_run"], ti) + self.add_breadcrumbs(ti) + return run(ti, context, log) + except Exception as e: + sentry_sdk.capture_exception(e) + raise + + return wrapped_run + + def flush(self): + sentry_sdk.flush() diff --git a/task-sdk/src/airflow/sdk/execution_time/sentry/noop.py b/task-sdk/src/airflow/sdk/execution_time/sentry/noop.py new file mode 100644 index 0000000000000..63281d32e6ed5 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/sentry/noop.py @@ -0,0 +1,51 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + from structlog.typing import FilteringBoundLogger as Logger + + from airflow.sdk import Context + from airflow.sdk.execution_time.comms import ToSupervisor + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + from airflow.sdk.types import DagRunProtocol, RuntimeTaskInstanceProtocol, TaskInstanceState + + RunReturn = tuple[TaskInstanceState, ToSupervisor | None, BaseException | None] + + class Run(Protocol): + def __call__(self, ti: RuntimeTaskInstance, context: Context, log: Logger) -> RunReturn: ... + + +class NoopSentry: + """Blank class for Sentry.""" + + def add_tagging(self, dag_run: DagRunProtocol, task_instance: RuntimeTaskInstanceProtocol) -> None: + """Blank function for tagging.""" + + def add_breadcrumbs(self, task_instance: RuntimeTaskInstanceProtocol) -> None: + """Blank function for breadcrumbs.""" + + def enrich_errors(self, run: Run) -> Run: + """Blank function for formatting a TaskInstance._run_raw_task.""" + return run + + def flush(self) -> None: + """Blank function for flushing errors.""" diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 48086a79f2f05..7211423807eb6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -86,6 +86,7 @@ GetDRCount, GetPreviousDagRun, GetPrevSuccessfulDagRun, + GetTaskBreadcrumbs, GetTaskRescheduleStartDate, GetTaskStates, GetTICount, @@ -109,6 +110,7 @@ SkipDownstreamTasks, StartupDetails, SucceedTask, + TaskBreadcrumbsResult, TaskState, TaskStatesResult, ToSupervisor, @@ -1358,6 +1360,9 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: resp = TaskStatesResult.from_api_response(task_states_map) else: resp = task_states_map + elif isinstance(msg, GetTaskBreadcrumbs): + api_resp = self.client.task_instances.get_task_breakcrumbs(dag_id=msg.dag_id, run_id=msg.run_id) + resp = TaskBreadcrumbsResult.from_api_response(api_resp) elif isinstance(msg, GetDRCount): resp = self.client.dag_runs.get_count( dag_id=msg.dag_id, diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index bd70135c95eb7..48307473838a2 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -70,6 +70,7 @@ GetDagRunState, GetDRCount, GetPreviousDagRun, + GetTaskBreadcrumbs, GetTaskRescheduleStartDate, GetTaskStates, GetTICount, @@ -84,6 +85,7 @@ SkipDownstreamTasks, StartupDetails, SucceedTask, + TaskBreadcrumbsResult, TaskRescheduleStartDate, TaskState, TaskStatesResult, @@ -524,6 +526,14 @@ def get_task_states( return response.task_states + @staticmethod + def get_task_breadcrumbs(dag_id: str, run_id: str) -> Iterable[dict[str, Any]]: + """Return task breadcrumbs for the given dag run.""" + response = SUPERVISOR_COMMS.send(GetTaskBreadcrumbs(dag_id=dag_id, run_id=run_id)) + if TYPE_CHECKING: + assert isinstance(response, TaskBreadcrumbsResult) + return response.breadcrumbs + @staticmethod def get_dr_count( dag_id: str, diff --git a/airflow-core/tests/unit/core/test_sentry.py b/task-sdk/tests/task_sdk/execution_time/test_sentry.py similarity index 63% rename from airflow-core/tests/unit/core/test_sentry.py rename to task-sdk/tests/task_sdk/execution_time/test_sentry.py index 723a9dc0ce510..d31dc92d73d59 100644 --- a/airflow-core/tests/unit/core/test_sentry.py +++ b/task-sdk/tests/task_sdk/execution_time/test_sentry.py @@ -23,11 +23,15 @@ import pytest import time_machine +import uuid6 from sentry_sdk import configure_scope from sentry_sdk.transport import Transport from airflow._shared.timezones import timezone from airflow.providers.standard.operators.python import PythonOperator +from airflow.sdk.api.datamodels._generated import DagRun, DagRunState, DagRunType +from airflow.sdk.execution_time.comms import GetTaskBreadcrumbs, TaskBreadcrumbsResult +from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance from airflow.utils.module_loading import import_string from airflow.utils.state import State @@ -38,6 +42,7 @@ DATA_INTERVAL = (LOGICAL_DATE, LOGICAL_DATE + SCHEDULE_INTERVAL) DAG_ID = "test_dag" TASK_ID = "test_task" +RUN_ID = "test_run" OPERATOR = "PythonOperator" TRY_NUMBER = 0 STATE = State.SUCCESS @@ -77,20 +82,36 @@ class CustomTransport(Transport): class TestSentryHook: @pytest.fixture - def task_instance(self, dag_maker): - # Mock the Dag - with dag_maker(DAG_ID, schedule=SCHEDULE_INTERVAL, serialized=True): - task = PythonOperator(task_id=TASK_ID, python_callable=int) + def dag_run(self): + return DagRun.model_construct( + dag_id=DAG_ID, + run_id=RUN_ID, + logical_date=LOGICAL_DATE, + data_interval_start=DATA_INTERVAL[0], + data_interval_end=DATA_INTERVAL[1], + run_after=max(DATA_INTERVAL), + start_date=max(DATA_INTERVAL), + run_type=DagRunType.MANUAL, + state=DagRunState.RUNNING, + consumed_asset_events=[], + ) - dr = dag_maker.create_dagrun(data_interval=DATA_INTERVAL, logical_date=LOGICAL_DATE) - ti = dr.task_instances[0] - ti.state = STATE - ti.task = task - dag_maker.session.commit() - - yield ti - - dag_maker.session.rollback() + @pytest.fixture + def task_instance(self, dag_run): + ti_date = timezone.utcnow() + return RuntimeTaskInstance.model_construct( + id=uuid6.uuid7(), + task_id=TASK_ID, + dag_id=dag_run.dag_id, + run_id=dag_run.run_id, + try_number=TRY_NUMBER, + dag_version_id=uuid6.uuid7(), + task=PythonOperator(task_id=TASK_ID, python_callable=bool), + bundle_instance=mock.Mock(), + start_date=ti_date, + end_date=ti_date, + state=STATE, + ) @pytest.fixture def sentry_sdk(self): @@ -103,10 +124,10 @@ def sentry(self): { ("sentry", "sentry_on"): "True", ("sentry", "default_integrations"): "False", - ("sentry", "before_send"): "unit.core.test_sentry.before_send", + ("sentry", "before_send"): "task_sdk.execution_time.test_sentry.before_send", }, ): - from airflow import sentry + from airflow.sdk.execution_time import sentry importlib.reload(sentry) yield sentry.Sentry @@ -119,10 +140,10 @@ def sentry_custom_transport(self): { ("sentry", "sentry_on"): "True", ("sentry", "default_integrations"): "False", - ("sentry", "transport"): "unit.core.test_sentry.CustomTransport", + ("sentry", "transport"): "task_sdk.execution_time.test_sentry.CustomTransport", }, ): - from airflow import sentry + from airflow.sdk.execution_time import sentry importlib.reload(sentry) yield sentry.Sentry @@ -135,41 +156,38 @@ def sentry_minimum(self): Minimum sentry config """ with conf_vars({("sentry", "sentry_on"): "True"}): - from airflow import sentry + from airflow.sdk.execution_time import sentry importlib.reload(sentry) yield sentry.Sentry importlib.reload(sentry) - @pytest.mark.db_test - def test_add_tagging(self, sentry, task_instance): + def test_add_tagging(self, sentry, dag_run, task_instance): """ Test adding tags. """ - sentry.add_tagging(task_instance=task_instance) + sentry.add_tagging(dag_run=dag_run, task_instance=task_instance) with configure_scope() as scope: - for key, value in scope._tags.items(): - assert value == TEST_SCOPE[key] + assert scope._tags == TEST_SCOPE - @pytest.mark.db_test @time_machine.travel(CRUMB_DATE) - def test_add_breadcrumbs(self, sentry, task_instance): + def test_add_breadcrumbs(self, mock_supervisor_comms, sentry, dag_run, task_instance): """ Test adding breadcrumbs. """ - sentry.add_tagging(task_instance=task_instance) - sentry.add_breadcrumbs(task_instance=task_instance) + mock_supervisor_comms.send.return_value = TaskBreadcrumbsResult.model_construct( + breadcrumbs=[TASK_DATA], + ) + sentry.add_breadcrumbs(task_instance=task_instance) with configure_scope() as scope: - test_crumb = scope._breadcrumbs.pop() - for item in CRUMB: - if item == "timestamp": - pass - elif item == "state": - assert str(CRUMB[item]) == str(test_crumb[item]) - else: - assert CRUMB[item] == test_crumb[item] + collected_crumb = scope._breadcrumbs.pop() + assert collected_crumb == CRUMB + + assert mock_supervisor_comms.send.mock_calls == [ + mock.call(GetTaskBreadcrumbs(dag_id=DAG_ID, run_id=RUN_ID)), + ] def test_before_send(self, sentry_sdk, sentry): """ @@ -177,7 +195,7 @@ def test_before_send(self, sentry_sdk, sentry): """ assert sentry called = sentry_sdk.call_args.kwargs["before_send"] - expected = import_string("unit.core.test_sentry.before_send") + expected = import_string("task_sdk.execution_time.test_sentry.before_send") assert called == expected def test_custom_transport(self, sentry_sdk, sentry_custom_transport): @@ -186,7 +204,7 @@ def test_custom_transport(self, sentry_sdk, sentry_custom_transport): """ assert sentry_custom_transport called = sentry_sdk.call_args.kwargs["transport"] - expected = import_string("unit.core.test_sentry.CustomTransport") + expected = import_string("task_sdk.execution_time.test_sentry.CustomTransport") assert called == expected def test_minimum_config(self, sentry_sdk, sentry_minimum): From f89fa74376e572411208ecc7461c5f6dd1227260 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 27 Oct 2025 17:16:17 +0800 Subject: [PATCH 2/7] Add test for server-side breadcrumb endpoint --- .../execution_api/routes/task_instances.py | 8 +-- .../versions/head/test_task_instances.py | 69 +++++++++++++++++++ 2 files changed, 73 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 6963f7a8edf8a..6dc4101d851ff 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -69,7 +69,7 @@ from airflow.sdk.definitions.asset import Asset, AssetUniqueKey from airflow.serialization.serialized_objects import SerializedDAG from airflow.task.trigger_rule import TriggerRule -from airflow.utils.state import DagRunState, TaskInstanceState +from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState if TYPE_CHECKING: from sqlalchemy.sql.dml import Update @@ -921,9 +921,9 @@ def get_task_instance_states( @router.get("/breadcrumbs", status_code=status.HTTP_200_OK) def get_task_instance_breadcrumbs(dag_id: str, run_id: str, session: SessionDep) -> TaskBreadcrumbsResponse: result = session.execute( - select(TI.task_id, TI.map_index, TI.state, TI.operator, TI.duration).where( - TI.dag_id == dag_id, TI.run_id == run_id - ) + select(TI.task_id, TI.map_index, TI.state, TI.operator, TI.duration) + .where(TI.dag_id == dag_id, TI.run_id == run_id, TI.state.in_(TerminalTIState)) + .order_by(TI.task_id, TI.map_index) ).mappings() def _iter_breadcrumbs() -> Iterator[dict[str, Any]]: diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 3efedecdbe90b..ee4dced435f9f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -2518,6 +2518,75 @@ def add_one(x): assert response.json() == {"task_states": {dr.run_id: expected}} +class TestGetTaskInstanceBreadcrumbs: + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + @pytest.fixture(autouse=True) + def dag_run(self, dag_maker, session): + with dag_maker(session=session): + for name in TaskInstanceState._member_names_: + EmptyOperator(task_id=name) + return dag_maker.create_dagrun(state="running") + + @pytest.fixture(autouse=True) + def task_instances(self, dag_run, session): + tis = {ti.task_id: ti for ti in dag_run.task_instances} + for name, value in TaskInstanceState._member_map_.items(): + tis[name].state = value + session.commit() + return tis + + def test_get_breadcrumbs(self, client, dag_run): + response = client.get( + "/execution/task-instances/breadcrumbs", + params={"dag_id": dag_run.dag_id, "run_id": dag_run.run_id}, + ) + assert response.status_code == 200 + assert response.json() == { # Should find tis with terminal states. + "breadcrumbs": [ + { + "duration": None, + "map_index": -1, + "operator": "EmptyOperator", + "state": "failed", + "task_id": "FAILED", + }, + { + "duration": None, + "map_index": -1, + "operator": "EmptyOperator", + "state": "removed", + "task_id": "REMOVED", + }, + { + "duration": None, + "map_index": -1, + "operator": "EmptyOperator", + "state": "skipped", + "task_id": "SKIPPED", + }, + { + "duration": None, + "map_index": -1, + "operator": "EmptyOperator", + "state": "success", + "task_id": "SUCCESS", + }, + { + "duration": None, + "map_index": -1, + "operator": "EmptyOperator", + "state": "upstream_failed", + "task_id": "UPSTREAM_FAILED", + }, + ] + } + + class TestInvactiveInletsAndOutlets: @pytest.mark.parametrize( "logical_date", From f3d98eb761f66574f4e34910be462fd6f634ceac Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 28 Oct 2025 11:15:42 +0800 Subject: [PATCH 3/7] Add supervisor test for breadcrumbs api --- .../execution_time/test_supervisor.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index cdb7d322c6f02..bef9c2c917ec5 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -86,6 +86,7 @@ GetHITLDetailResponse, GetPreviousDagRun, GetPrevSuccessfulDagRun, + GetTaskBreadcrumbs, GetTaskRescheduleStartDate, GetTaskStates, GetTICount, @@ -110,6 +111,7 @@ SetXCom, SkipDownstreamTasks, SucceedTask, + TaskBreadcrumbsResult, TaskRescheduleStartDate, TaskState, TaskStatesResult, @@ -2298,6 +2300,37 @@ class RequestTestCase: ), test_id="skip_downstream_tasks", ), + RequestTestCase( + message=GetTaskBreadcrumbs(dag_id="test_dag", run_id="test_run"), + client_mock=ClientMock( + method_path="task_instances.get_task_breakcrumbs", + kwargs={"dag_id": "test_dag", "run_id": "test_run"}, + response=TaskBreadcrumbsResult( + breadcrumbs=[ + { + "task_id": "test_task", + "map_index": 2, + "state": "success", + "operator": "PythonOperator", + "duration": 432.0, + }, + ], + ), + ), + expected_body={ + "breadcrumbs": [ + { + "task_id": "test_task", + "map_index": 2, + "state": "success", + "operator": "PythonOperator", + "duration": 432.0, + }, + ], + "type": "TaskBreadcrumbsResult", + }, + test_id="get_task_breadcrumbs", + ), ] From b069241dd0a03f68b36f3ad7a146cd5eb606fa1f Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 28 Oct 2025 12:06:30 +0800 Subject: [PATCH 4/7] Enable Sentry integration in task runner --- task-sdk/src/airflow/sdk/execution_time/task_runner.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 48307473838a2..e93fa3b68a8eb 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -107,6 +107,7 @@ get_previous_dagrun_success, set_current_context, ) +from airflow.sdk.execution_time.sentry import Sentry from airflow.sdk.execution_time.xcom import XCom from airflow.sdk.timezone import coerce_datetime from airflow.stats import Stats @@ -880,6 +881,7 @@ def _defer_task( return msg, state +@Sentry.enrich_errors def run( ti: RuntimeTaskInstance, context: Context, From d0877d1dc726250d400d3ff25aac572c4ac1d981 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 28 Oct 2025 13:40:23 +0800 Subject: [PATCH 5/7] Make Flask optional in tests This is needed to reflect core dep changes since we've removed Sentry (and its Flask integration) from core. --- .../tests/unit/plugins/test_plugin.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/airflow-core/tests/unit/plugins/test_plugin.py b/airflow-core/tests/unit/plugins/test_plugin.py index 62e642955ef6e..3d84c1a2e0930 100644 --- a/airflow-core/tests/unit/plugins/test_plugin.py +++ b/airflow-core/tests/unit/plugins/test_plugin.py @@ -18,7 +18,6 @@ from __future__ import annotations from fastapi import FastAPI -from flask import Blueprint try: # if flask_appbuilder is installed, we can use AppBuilderBaseView @@ -90,13 +89,20 @@ def plugin_macro(): # Creating a flask blueprint to integrate the templates and static folder -bp = Blueprint( - "test_plugin", - __name__, - template_folder="templates", # registers airflow/plugins/templates as a Jinja template folder - static_folder="static", - static_url_path="/static/test_plugin", -) +try: + from flask import Blueprint +except ImportError: + flask_blueprints = [] +else: + flask_blueprints = [ + Blueprint( + "test_plugin", + __name__, + template_folder="templates", # registers airflow/plugins/templates as a Jinja template folder + static_folder="static", + static_url_path="/static/test_plugin", + ) + ] app = FastAPI() @@ -149,7 +155,7 @@ def get_weight(self, ti): class AirflowTestPlugin(AirflowPlugin): name = "test_plugin" macros = [plugin_macro] - flask_blueprints = [bp] + flask_blueprints = flask_blueprints fastapi_apps = [app_with_metadata] fastapi_root_middlewares = [middleware_with_metadata] external_views = [external_view_with_metadata] From a5c986396adfc3afb9b2b731b390f91448dc9fb2 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Tue, 28 Oct 2025 18:30:49 +0800 Subject: [PATCH 6/7] Mock out all Sentry usages in tests Sentry tends to mess with logging (seems to be related to the debug helpers it runs AT IMPORT TIME), which in turn messes up out supervisor tests. The Sentry integration tests are now rewritten to not actually use the Sentry SDK at all to avoid this issue. The problem now is I have no real way to actually know if the integration actually works anymore. We'll need to figure that out in another way. --- .../sdk/execution_time/sentry/configured.py | 5 +- .../task_sdk/execution_time/test_sentry.py | 136 +++++++++++------- 2 files changed, 84 insertions(+), 57 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py b/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py index e9753ad9b706a..cfb2d916df80d 100644 --- a/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py +++ b/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py @@ -49,8 +49,8 @@ class ConfiguredSentry(NoopSentry): """Configure Sentry SDK.""" - SCOPE_DAG_RUN_TAGS = frozenset(("data_interval_end", "data_interval_start", "logical_date")) - SCOPE_TASK_INSTANCE_TAGS = frozenset(("task_id", "dag_id", "try_number")) + SCOPE_DAG_RUN_TAGS = ("data_interval_start", "data_interval_end", "logical_date") + SCOPE_TASK_INSTANCE_TAGS = ("task_id", "dag_id", "try_number") UNSUPPORTED_SENTRY_OPTIONS = frozenset( ( @@ -107,7 +107,6 @@ def __init__(self): def add_tagging(self, dag_run: DagRunProtocol, task_instance: RuntimeTaskInstanceProtocol) -> None: """Add tagging for a task_instance.""" task = task_instance.task - with sentry_sdk.configure_scope() as scope: for tag_name in self.SCOPE_TASK_INSTANCE_TAGS: attribute = getattr(task_instance, tag_name) diff --git a/task-sdk/tests/task_sdk/execution_time/test_sentry.py b/task-sdk/tests/task_sdk/execution_time/test_sentry.py index d31dc92d73d59..b68add79e294f 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_sentry.py +++ b/task-sdk/tests/task_sdk/execution_time/test_sentry.py @@ -19,13 +19,12 @@ import datetime import importlib +import sys +import types from unittest import mock import pytest -import time_machine import uuid6 -from sentry_sdk import configure_scope -from sentry_sdk.transport import Transport from airflow._shared.timezones import timezone from airflow.providers.standard.operators.python import PythonOperator @@ -46,15 +45,6 @@ OPERATOR = "PythonOperator" TRY_NUMBER = 0 STATE = State.SUCCESS -TEST_SCOPE = { - "dag_id": DAG_ID, - "task_id": TASK_ID, - "data_interval_start": DATA_INTERVAL[0], - "data_interval_end": DATA_INTERVAL[1], - "logical_date": LOGICAL_DATE, - "operator": OPERATOR, - "try_number": TRY_NUMBER, -} TASK_DATA = { "task_id": TASK_ID, "state": STATE, @@ -62,24 +52,21 @@ "duration": None, } -CRUMB_DATE = datetime.datetime(2019, 5, 15, tzinfo=datetime.timezone.utc) -CRUMB = { - "timestamp": CRUMB_DATE, - "type": "default", - "category": "completed_tasks", - "data": TASK_DATA, - "level": "info", -} - def before_send(_): pass -class CustomTransport(Transport): +class CustomTransport: pass +def is_configured(obj): + from airflow.sdk.execution_time.sentry.configured import ConfiguredSentry + + return isinstance(obj, ConfiguredSentry) + + class TestSentryHook: @pytest.fixture def dag_run(self): @@ -113,13 +100,33 @@ def task_instance(self, dag_run): state=STATE, ) - @pytest.fixture - def sentry_sdk(self): - with mock.patch("sentry_sdk.init") as sentry_sdk: - yield sentry_sdk + @pytest.fixture(scope="class", autouse=True) + def mock_sentry_sdk(self): + sentry_sdk_integrations_logging = types.ModuleType("sentry_sdk.integrations.logging") + sentry_sdk_integrations_logging.ignore_logger = mock.MagicMock() + + sentry_sdk = types.ModuleType("sentry_sdk") + sentry_sdk.init = mock.MagicMock() + sentry_sdk.integrations = mock.Mock(logging=sentry_sdk_integrations_logging) + sentry_sdk.configure_scope = mock.MagicMock() + sentry_sdk.add_breadcrumb = mock.MagicMock() + + sys.modules["sentry_sdk"] = sentry_sdk + sys.modules["sentry_sdk.integrations.logging"] = sentry_sdk_integrations_logging + yield sentry_sdk + del sys.modules["sentry_sdk"] + del sys.modules["sentry_sdk.integrations.logging"] + + @pytest.fixture(autouse=True) + def remove_mock_sentry_sdk(self, mock_sentry_sdk): + yield + mock_sentry_sdk.integrations.logging.ignore_logger.reset_mock() + mock_sentry_sdk.init.reset_mock() + mock_sentry_sdk.configure_scope.reset_mock() + mock_sentry_sdk.add_breadcrumb.reset_mock() @pytest.fixture - def sentry(self): + def sentry(self, mock_sentry_sdk): with conf_vars( { ("sentry", "sentry_on"): "True", @@ -135,7 +142,7 @@ def sentry(self): importlib.reload(sentry) @pytest.fixture - def sentry_custom_transport(self): + def sentry_custom_transport(self, mock_sentry_sdk): with conf_vars( { ("sentry", "sentry_on"): "True", @@ -151,7 +158,7 @@ def sentry_custom_transport(self): importlib.reload(sentry) @pytest.fixture - def sentry_minimum(self): + def sentry_minimum(self, mock_sentry_sdk): """ Minimum sentry config """ @@ -163,16 +170,37 @@ def sentry_minimum(self): importlib.reload(sentry) - def test_add_tagging(self, sentry, dag_run, task_instance): + def test_init(self, mock_sentry_sdk, sentry): + assert is_configured(sentry) + assert mock_sentry_sdk.integrations.logging.ignore_logger.mock_calls == [mock.call("airflow.task")] + assert mock_sentry_sdk.init.mock_calls == [ + mock.call( + integrations=[], + default_integrations=False, + before_send=import_string("task_sdk.execution_time.test_sentry.before_send"), + transport=None, + ), + ] + + def test_add_tagging(self, mock_sentry_sdk, sentry, dag_run, task_instance): """ Test adding tags. """ sentry.add_tagging(dag_run=dag_run, task_instance=task_instance) - with configure_scope() as scope: - assert scope._tags == TEST_SCOPE + assert mock_sentry_sdk.configure_scope.mock_calls == [ + mock.call.__call__(), + mock.call.__call__().__enter__(), + mock.call.__call__().__enter__().set_tag("task_id", TASK_ID), + mock.call.__call__().__enter__().set_tag("dag_id", DAG_ID), + mock.call.__call__().__enter__().set_tag("try_number", TRY_NUMBER), + mock.call.__call__().__enter__().set_tag("data_interval_start", DATA_INTERVAL[0]), + mock.call.__call__().__enter__().set_tag("data_interval_end", DATA_INTERVAL[1]), + mock.call.__call__().__enter__().set_tag("logical_date", LOGICAL_DATE), + mock.call.__call__().__enter__().set_tag("operator", OPERATOR), + mock.call.__call__().__exit__(None, None, None), + ] - @time_machine.travel(CRUMB_DATE) - def test_add_breadcrumbs(self, mock_supervisor_comms, sentry, dag_run, task_instance): + def test_add_breadcrumbs(self, mock_supervisor_comms, mock_sentry_sdk, sentry, task_instance): """ Test adding breadcrumbs. """ @@ -181,35 +209,35 @@ def test_add_breadcrumbs(self, mock_supervisor_comms, sentry, dag_run, task_inst ) sentry.add_breadcrumbs(task_instance=task_instance) - with configure_scope() as scope: - collected_crumb = scope._breadcrumbs.pop() - assert collected_crumb == CRUMB + assert mock_sentry_sdk.add_breadcrumb.mock_calls == [ + mock.call(category="completed_tasks", data=TASK_DATA, level="info"), + ] assert mock_supervisor_comms.send.mock_calls == [ mock.call(GetTaskBreadcrumbs(dag_id=DAG_ID, run_id=RUN_ID)), ] - def test_before_send(self, sentry_sdk, sentry): - """ - Test before send callable gets passed to the sentry SDK. - """ - assert sentry - called = sentry_sdk.call_args.kwargs["before_send"] - expected = import_string("task_sdk.execution_time.test_sentry.before_send") - assert called == expected - - def test_custom_transport(self, sentry_sdk, sentry_custom_transport): + def test_custom_transport(self, mock_sentry_sdk, sentry_custom_transport): """ Test transport gets passed to the sentry SDK """ - assert sentry_custom_transport - called = sentry_sdk.call_args.kwargs["transport"] - expected = import_string("task_sdk.execution_time.test_sentry.CustomTransport") - assert called == expected + assert is_configured(sentry_custom_transport) + assert mock_sentry_sdk.integrations.logging.ignore_logger.mock_calls == [mock.call("airflow.task")] + assert mock_sentry_sdk.init.mock_calls == [ + mock.call( + integrations=[], + default_integrations=False, + before_send=None, + transport=import_string("task_sdk.execution_time.test_sentry.CustomTransport"), + ), + ] - def test_minimum_config(self, sentry_sdk, sentry_minimum): + def test_minimum_config(self, mock_sentry_sdk, sentry_minimum): """ Test before_send doesn't raise an exception when not set """ - assert sentry_minimum - sentry_sdk.assert_called_once() + assert is_configured(sentry_minimum) + assert mock_sentry_sdk.integrations.logging.ignore_logger.mock_calls == [mock.call("airflow.task")] + assert mock_sentry_sdk.init.mock_calls == [ + mock.call(integrations=[], before_send=None, transport=None), + ] From 9282cf9ec0501cb694f9100d6974a00a38ff34d6 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 29 Oct 2025 12:08:18 +0800 Subject: [PATCH 7/7] Fix circular imports The Sentry integration need to not import other parts at top-level. --- task-sdk/src/airflow/sdk/execution_time/sentry/configured.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py b/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py index cfb2d916df80d..5826ae447c5bb 100644 --- a/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py +++ b/task-sdk/src/airflow/sdk/execution_time/sentry/configured.py @@ -34,13 +34,13 @@ import structlog from airflow.sdk.execution_time.sentry.noop import NoopSentry -from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance if TYPE_CHECKING: from structlog.typing import FilteringBoundLogger as Logger from airflow.sdk import Context from airflow.sdk.execution_time.sentry.noop import Run, RunReturn + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance from airflow.sdk.types import DagRunProtocol, RuntimeTaskInstanceProtocol log = structlog.get_logger(logger_name=__name__) @@ -118,6 +118,8 @@ def add_tagging(self, dag_run: DagRunProtocol, task_instance: RuntimeTaskInstanc def add_breadcrumbs(self, task_instance: RuntimeTaskInstanceProtocol) -> None: """Add breadcrumbs inside of a task_instance.""" + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + breadcrumbs = RuntimeTaskInstance.get_task_breadcrumbs( dag_id=task_instance.dag_id, run_id=task_instance.run_id,