diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 259184a2febd8..73d3ab9a57e8d 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -104,6 +104,7 @@ from airflow.models.taskreschedule import TaskReschedule from airflow.models.xcom import LazyXComSelectSequence, XComModel from airflow.plugins_manager import integrate_macros_plugins +from airflow.sdk.execution_time.context import context_to_airflow_vars from airflow.sentry import Sentry from airflow.settings import task_instance_mutation_hook from airflow.stats import Stats @@ -115,7 +116,7 @@ from airflow.utils.helpers import prune_dict, render_template_to_string from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname -from airflow.utils.operator_helpers import ExecutionCallableRunner, context_to_airflow_vars +from airflow.utils.operator_helpers import ExecutionCallableRunner from airflow.utils.platform import getuser from airflow.utils.retries import run_with_db_retries from airflow.utils.session import NEW_SESSION, create_session, provide_session diff --git a/airflow-core/src/airflow/utils/operator_helpers.py b/airflow-core/src/airflow/utils/operator_helpers.py index d2239b21d673f..ef7cd35203902 100644 --- a/airflow-core/src/airflow/utils/operator_helpers.py +++ b/airflow-core/src/airflow/utils/operator_helpers.py @@ -20,10 +20,8 @@ import inspect import logging from collections.abc import Collection, Mapping -from datetime import datetime from typing import TYPE_CHECKING, Any, Callable, Protocol, TypeVar -from airflow import settings from airflow.typing_compat import ParamSpec from airflow.utils.types import NOTSET @@ -33,105 +31,6 @@ P = ParamSpec("P") R = TypeVar("R") -DEFAULT_FORMAT_PREFIX = "airflow.ctx." -ENV_VAR_FORMAT_PREFIX = "AIRFLOW_CTX_" - -AIRFLOW_VAR_NAME_FORMAT_MAPPING = { - "AIRFLOW_CONTEXT_DAG_ID": { - "default": f"{DEFAULT_FORMAT_PREFIX}dag_id", - "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_ID", - }, - "AIRFLOW_CONTEXT_TASK_ID": { - "default": f"{DEFAULT_FORMAT_PREFIX}task_id", - "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}TASK_ID", - }, - "AIRFLOW_CONTEXT_LOGICAL_DATE": { - "default": f"{DEFAULT_FORMAT_PREFIX}logical_date", - "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}LOGICAL_DATE", - }, - "AIRFLOW_CONTEXT_TRY_NUMBER": { - "default": f"{DEFAULT_FORMAT_PREFIX}try_number", - "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}TRY_NUMBER", - }, - "AIRFLOW_CONTEXT_DAG_RUN_ID": { - "default": f"{DEFAULT_FORMAT_PREFIX}dag_run_id", - "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_RUN_ID", - }, - "AIRFLOW_CONTEXT_DAG_OWNER": { - "default": f"{DEFAULT_FORMAT_PREFIX}dag_owner", - "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_OWNER", - }, - "AIRFLOW_CONTEXT_DAG_EMAIL": { - "default": f"{DEFAULT_FORMAT_PREFIX}dag_email", - "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_EMAIL", - }, -} - - -def context_to_airflow_vars(context: Mapping[str, Any], in_env_var_format: bool = False) -> dict[str, str]: - """ - Return values used to externally reconstruct relations between dags, dag_runs, tasks and task_instances. - - Given a context, this function provides a dictionary of values that can be used to - externally reconstruct relations between dags, dag_runs, tasks and task_instances. - Default to abc.def.ghi format and can be made to ABC_DEF_GHI format if - in_env_var_format is set to True. - - :param context: The context for the task_instance of interest. - :param in_env_var_format: If returned vars should be in ABC_DEF_GHI format. - :return: task_instance context as dict. - """ - params = {} - if in_env_var_format: - name_format = "env_var_format" - else: - name_format = "default" - - task = context.get("task") - task_instance = context.get("task_instance") - dag_run = context.get("dag_run") - - ops = [ - (task, "email", "AIRFLOW_CONTEXT_DAG_EMAIL"), - (task, "owner", "AIRFLOW_CONTEXT_DAG_OWNER"), - (task_instance, "dag_id", "AIRFLOW_CONTEXT_DAG_ID"), - (task_instance, "task_id", "AIRFLOW_CONTEXT_TASK_ID"), - (task_instance, "logical_date", "AIRFLOW_CONTEXT_LOGICAL_DATE"), - (task_instance, "try_number", "AIRFLOW_CONTEXT_TRY_NUMBER"), - (dag_run, "run_id", "AIRFLOW_CONTEXT_DAG_RUN_ID"), - ] - - context_params = settings.get_airflow_context_vars(context) - for key, value in context_params.items(): - if not isinstance(key, str): - raise TypeError(f"key <{key}> must be string") - if not isinstance(value, str): - raise TypeError(f"value of key <{key}> must be string, not {type(value)}") - - if in_env_var_format: - if not key.startswith(ENV_VAR_FORMAT_PREFIX): - key = ENV_VAR_FORMAT_PREFIX + key.upper() - else: - if not key.startswith(DEFAULT_FORMAT_PREFIX): - key = DEFAULT_FORMAT_PREFIX + key - params[key] = value - - for subject, attr, mapping_key in ops: - _attr = getattr(subject, attr, None) - if subject and _attr: - mapping_value = AIRFLOW_VAR_NAME_FORMAT_MAPPING[mapping_key][name_format] - if isinstance(_attr, str): - params[mapping_value] = _attr - elif isinstance(_attr, datetime): - params[mapping_value] = _attr.isoformat() - elif isinstance(_attr, list): - # os env variable value needs to be string - params[mapping_value] = ",".join(_attr) - else: - params[mapping_value] = str(_attr) - - return params - class KeywordParameters: """ diff --git a/airflow-core/tests/unit/utils/test_operator_helpers.py b/airflow-core/tests/unit/utils/test_operator_helpers.py index 1f6b6df4c6198..9d4db63a3aaec 100644 --- a/airflow-core/tests/unit/utils/test_operator_helpers.py +++ b/airflow-core/tests/unit/utils/test_operator_helpers.py @@ -50,53 +50,6 @@ def setup_method(self): "task": mock.MagicMock(name="task", owner=self.owner, email=self.email), } - def test_context_to_airflow_vars_empty_context(self): - assert operator_helpers.context_to_airflow_vars({}) == {} - - def test_context_to_airflow_vars_all_context(self): - assert operator_helpers.context_to_airflow_vars(self.context) == { - "airflow.ctx.dag_id": self.dag_id, - "airflow.ctx.logical_date": self.logical_date, - "airflow.ctx.task_id": self.task_id, - "airflow.ctx.dag_run_id": self.dag_run_id, - "airflow.ctx.try_number": str(self.try_number), - "airflow.ctx.dag_owner": "owner1,owner2", - "airflow.ctx.dag_email": "email1@test.com", - } - - assert operator_helpers.context_to_airflow_vars(self.context, in_env_var_format=True) == { - "AIRFLOW_CTX_DAG_ID": self.dag_id, - "AIRFLOW_CTX_LOGICAL_DATE": self.logical_date, - "AIRFLOW_CTX_TASK_ID": self.task_id, - "AIRFLOW_CTX_TRY_NUMBER": str(self.try_number), - "AIRFLOW_CTX_DAG_RUN_ID": self.dag_run_id, - "AIRFLOW_CTX_DAG_OWNER": "owner1,owner2", - "AIRFLOW_CTX_DAG_EMAIL": "email1@test.com", - } - - def test_context_to_airflow_vars_with_default_context_vars(self): - with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: - airflow_cluster = "cluster-a" - mock_method.return_value = {"airflow_cluster": airflow_cluster} - - context_vars = operator_helpers.context_to_airflow_vars(self.context) - assert context_vars["airflow.ctx.airflow_cluster"] == airflow_cluster - - context_vars = operator_helpers.context_to_airflow_vars(self.context, in_env_var_format=True) - assert context_vars["AIRFLOW_CTX_AIRFLOW_CLUSTER"] == airflow_cluster - - with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: - mock_method.return_value = {"airflow_cluster": [1, 2]} - with pytest.raises(TypeError) as error: - operator_helpers.context_to_airflow_vars(self.context) - assert str(error.value) == "value of key must be string, not " - - with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: - mock_method.return_value = {1: "value"} - with pytest.raises(TypeError) as error: - operator_helpers.context_to_airflow_vars(self.context) - assert str(error.value) == "key <1> must be string" - def callable1(ds_nodash): return (ds_nodash,) diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 893a70209ca03..33beef3a1c2fa 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -179,6 +179,7 @@ ], "cross-providers-deps": [ "amazon", + "common.compat", "common.sql", "microsoft.mssql", "mysql", diff --git a/providers/apache/hive/README.rst b/providers/apache/hive/README.rst index b887e980d6ea1..d46a9b901d2ee 100644 --- a/providers/apache/hive/README.rst +++ b/providers/apache/hive/README.rst @@ -79,6 +79,7 @@ You can install such cross-provider dependencies when installing from PyPI. For Dependent package Extra ====================================================================================================================== =================== `apache-airflow-providers-amazon `_ ``amazon`` +`apache-airflow-providers-common-compat `_ ``common.compat`` `apache-airflow-providers-common-sql `_ ``common.sql`` `apache-airflow-providers-microsoft-mssql `_ ``microsoft.mssql`` `apache-airflow-providers-mysql `_ ``mysql`` diff --git a/providers/apache/hive/pyproject.toml b/providers/apache/hive/pyproject.toml index 0dacdb1d073b4..c35281060034a 100644 --- a/providers/apache/hive/pyproject.toml +++ b/providers/apache/hive/pyproject.toml @@ -91,6 +91,9 @@ dependencies = [ "vertica" = [ "apache-airflow-providers-vertica" ] +"common.compat" = [ + "apache-airflow-providers-common-compat" +] [dependency-groups] dev = [ @@ -98,6 +101,7 @@ dev = [ "apache-airflow-task-sdk", "apache-airflow-devel-common", "apache-airflow-providers-amazon", + "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", "apache-airflow-providers-microsoft-mssql", "apache-airflow-providers-mysql", diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/get_provider_info.py b/providers/apache/hive/src/airflow/providers/apache/hive/get_provider_info.py index 85738ae924885..574d7a433ba6d 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/get_provider_info.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/get_provider_info.py @@ -199,6 +199,7 @@ def get_provider_info(): "presto": ["apache-airflow-providers-presto"], "samba": ["apache-airflow-providers-samba"], "vertica": ["apache-airflow-providers-vertica"], + "common.compat": ["apache-airflow-providers-common-compat"], }, "devel-dependencies": [], } diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py index dde421e01e69b..a944b9841a231 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py @@ -35,10 +35,17 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.security import utils from airflow.utils.helpers import as_flattened_list -from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import AIRFLOW_VAR_NAME_FORMAT_MAPPING +else: + from airflow.utils.operator_helpers import ( # type: ignore[no-redef, attr-defined] + AIRFLOW_VAR_NAME_FORMAT_MAPPING, + ) HIVE_QUEUE_PRIORITIES = ["VERY_HIGH", "HIGH", "NORMAL", "LOW", "VERY_LOW"] diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py b/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py index c131f89ee7c60..f702a17e781b6 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py @@ -26,8 +26,15 @@ from airflow.configuration import conf from airflow.models import BaseOperator from airflow.providers.apache.hive.hooks.hive import HiveCliHook -from airflow.utils import operator_helpers -from airflow.utils.operator_helpers import context_to_airflow_vars +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import AIRFLOW_VAR_NAME_FORMAT_MAPPING, context_to_airflow_vars +else: + from airflow.utils.operator_helpers import ( # type: ignore[no-redef, attr-defined] + AIRFLOW_VAR_NAME_FORMAT_MAPPING, + context_to_airflow_vars, + ) if TYPE_CHECKING: from airflow.utils.context import Context @@ -173,7 +180,5 @@ def on_kill(self) -> None: def clear_airflow_vars(self) -> None: """Reset airflow environment variables to prevent existing ones from impacting behavior.""" - blank_env_vars = { - value["env_var_format"]: "" for value in operator_helpers.AIRFLOW_VAR_NAME_FORMAT_MAPPING.values() - } + blank_env_vars = {value["env_var_format"]: "" for value in AIRFLOW_VAR_NAME_FORMAT_MAPPING.values()} os.environ.update(blank_env_vars) diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py index e2d80dc0701d2..3ee6b03e2e39e 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py @@ -25,8 +25,14 @@ from airflow.models import BaseOperator from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS from airflow.providers.mysql.hooks.mysql import MySqlHook -from airflow.utils.operator_helpers import context_to_airflow_vars + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import context_to_airflow_vars +else: + from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] + if TYPE_CHECKING: from airflow.utils.context import Context diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py index 50b1f669bc8dc..91e17e16c0e01 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py @@ -25,8 +25,13 @@ from airflow.models import BaseOperator from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS from airflow.providers.samba.hooks.samba import SambaHook -from airflow.utils.operator_helpers import context_to_airflow_vars + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import context_to_airflow_vars +else: + from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] if TYPE_CHECKING: from airflow.utils.context import Context diff --git a/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py b/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py index bbfee1d1d0c61..0f239bbe9892a 100644 --- a/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py +++ b/providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py @@ -32,7 +32,6 @@ from airflow.providers.apache.hive.hooks.hive import HiveCliHook, HiveMetastoreHook, HiveServer2Hook from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.utils import timezone -from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING from tests_common.test_utils.asserts import assert_equal_ignore_multiple_spaces from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS @@ -44,6 +43,13 @@ MockSubProcess, ) +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import AIRFLOW_VAR_NAME_FORMAT_MAPPING +else: + from airflow.utils.operator_helpers import ( # type: ignore[no-redef, attr-defined] + AIRFLOW_VAR_NAME_FORMAT_MAPPING, + ) + DEFAULT_DATE = timezone.datetime(2015, 1, 1) DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat() DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10] diff --git a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py index a3fb324c32a94..3c11dd6682f32 100644 --- a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py +++ b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py @@ -25,10 +25,15 @@ from airflow.providers.apache.hive.transfers.hive_to_mysql import HiveToMySqlOperator from airflow.utils import timezone -from airflow.utils.operator_helpers import context_to_airflow_vars +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS from unit.apache.hive import MockHiveServer2Hook, MockMySqlHook, TestHiveEnvironment +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import context_to_airflow_vars +else: + from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] + DEFAULT_DATE = timezone.datetime(2015, 1, 1) diff --git a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py index abf94365cb802..ec621bb59f7fb 100644 --- a/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py +++ b/providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py @@ -24,8 +24,8 @@ from airflow.providers.apache.hive.transfers.hive_to_samba import HiveToSambaOperator from airflow.providers.samba.hooks.samba import SambaHook -from airflow.utils.operator_helpers import context_to_airflow_vars +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS from unit.apache.hive import ( DEFAULT_DATE, MockConnectionCursor, @@ -33,6 +33,11 @@ TestHiveEnvironment, ) +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import context_to_airflow_vars +else: + from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] + class MockSambaHook(SambaHook): def __init__(self, *args, **kwargs): diff --git a/providers/edge/src/airflow/providers/edge/example_dags/win_test.py b/providers/edge/src/airflow/providers/edge/example_dags/win_test.py index 5bda7739a7b6f..78f40095e9b5e 100644 --- a/providers/edge/src/airflow/providers/edge/example_dags/win_test.py +++ b/providers/edge/src/airflow/providers/edge/example_dags/win_test.py @@ -40,7 +40,7 @@ from airflow.models.variable import Variable from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import Param -from airflow.utils.operator_helpers import context_to_airflow_vars +from airflow.sdk.execution_time.context import context_to_airflow_vars from airflow.utils.trigger_rule import TriggerRule from airflow.utils.types import ArgNotSet diff --git a/providers/presto/src/airflow/providers/presto/hooks/presto.py b/providers/presto/src/airflow/providers/presto/hooks/presto.py index 7d98a3ad872fe..d012510c96d1d 100644 --- a/providers/presto/src/airflow/providers/presto/hooks/presto.py +++ b/providers/presto/src/airflow/providers/presto/hooks/presto.py @@ -30,7 +30,14 @@ from airflow.exceptions import AirflowException from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.presto.version_compat import AIRFLOW_V_3_0_PLUS -from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING, DEFAULT_FORMAT_PREFIX + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import AIRFLOW_VAR_NAME_FORMAT_MAPPING, DEFAULT_FORMAT_PREFIX +else: + from airflow.utils.operator_helpers import ( # type: ignore[no-redef, attr-defined] + AIRFLOW_VAR_NAME_FORMAT_MAPPING, + DEFAULT_FORMAT_PREFIX, + ) if TYPE_CHECKING: from airflow.models import Connection diff --git a/providers/standard/src/airflow/providers/standard/operators/bash.py b/providers/standard/src/airflow/providers/standard/operators/bash.py index 9b4852a151b42..02d5373758890 100644 --- a/providers/standard/src/airflow/providers/standard/operators/bash.py +++ b/providers/standard/src/airflow/providers/standard/operators/bash.py @@ -27,10 +27,15 @@ from airflow.exceptions import AirflowException, AirflowSkipException from airflow.models.baseoperator import BaseOperator from airflow.providers.standard.hooks.subprocess import SubprocessHook, SubprocessResult, working_directory -from airflow.utils.operator_helpers import context_to_airflow_vars +from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.types import ArgNotSet +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import context_to_airflow_vars +else: + from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] + if TYPE_CHECKING: from sqlalchemy.orm import Session as SASession diff --git a/providers/trino/src/airflow/providers/trino/hooks/trino.py b/providers/trino/src/airflow/providers/trino/hooks/trino.py index af1e1b4264b20..728d993472fb1 100644 --- a/providers/trino/src/airflow/providers/trino/hooks/trino.py +++ b/providers/trino/src/airflow/providers/trino/hooks/trino.py @@ -32,7 +32,14 @@ from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.trino.version_compat import AIRFLOW_V_3_0_PLUS from airflow.utils.helpers import exactly_one -from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING, DEFAULT_FORMAT_PREFIX + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.execution_time.context import AIRFLOW_VAR_NAME_FORMAT_MAPPING, DEFAULT_FORMAT_PREFIX +else: + from airflow.utils.operator_helpers import ( # type: ignore[no-redef, attr-defined] + AIRFLOW_VAR_NAME_FORMAT_MAPPING, + DEFAULT_FORMAT_PREFIX, + ) if TYPE_CHECKING: from airflow.models import Connection diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 398a03e731764..c4d6d9c439ce1 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -54,6 +54,42 @@ VariableResult, ) + +DEFAULT_FORMAT_PREFIX = "airflow.ctx." +ENV_VAR_FORMAT_PREFIX = "AIRFLOW_CTX_" + +AIRFLOW_VAR_NAME_FORMAT_MAPPING = { + "AIRFLOW_CONTEXT_DAG_ID": { + "default": f"{DEFAULT_FORMAT_PREFIX}dag_id", + "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_ID", + }, + "AIRFLOW_CONTEXT_TASK_ID": { + "default": f"{DEFAULT_FORMAT_PREFIX}task_id", + "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}TASK_ID", + }, + "AIRFLOW_CONTEXT_LOGICAL_DATE": { + "default": f"{DEFAULT_FORMAT_PREFIX}logical_date", + "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}LOGICAL_DATE", + }, + "AIRFLOW_CONTEXT_TRY_NUMBER": { + "default": f"{DEFAULT_FORMAT_PREFIX}try_number", + "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}TRY_NUMBER", + }, + "AIRFLOW_CONTEXT_DAG_RUN_ID": { + "default": f"{DEFAULT_FORMAT_PREFIX}dag_run_id", + "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_RUN_ID", + }, + "AIRFLOW_CONTEXT_DAG_OWNER": { + "default": f"{DEFAULT_FORMAT_PREFIX}dag_owner", + "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_OWNER", + }, + "AIRFLOW_CONTEXT_DAG_EMAIL": { + "default": f"{DEFAULT_FORMAT_PREFIX}dag_email", + "env_var_format": f"{ENV_VAR_FORMAT_PREFIX}DAG_EMAIL", + }, +} + + log = structlog.get_logger(logger_name="task") @@ -456,3 +492,72 @@ def context_update_for_unmapped(context: Context, task: BaseOperator) -> None: context["params"] = process_params( context["dag"], task, context["dag_run"].conf, suppress_exception=False ) + + +def context_to_airflow_vars(context: Mapping[str, Any], in_env_var_format: bool = False) -> dict[str, str]: + """ + Return values used to externally reconstruct relations between dags, dag_runs, tasks and task_instances. + + Given a context, this function provides a dictionary of values that can be used to + externally reconstruct relations between dags, dag_runs, tasks and task_instances. + Default to abc.def.ghi format and can be made to ABC_DEF_GHI format if + in_env_var_format is set to True. + + :param context: The context for the task_instance of interest. + :param in_env_var_format: If returned vars should be in ABC_DEF_GHI format. + :return: task_instance context as dict. + """ + from datetime import datetime + + from airflow import settings + + params = {} + if in_env_var_format: + name_format = "env_var_format" + else: + name_format = "default" + + task = context.get("task") + task_instance = context.get("task_instance") + dag_run = context.get("dag_run") + + ops = [ + (task, "email", "AIRFLOW_CONTEXT_DAG_EMAIL"), + (task, "owner", "AIRFLOW_CONTEXT_DAG_OWNER"), + (task_instance, "dag_id", "AIRFLOW_CONTEXT_DAG_ID"), + (task_instance, "task_id", "AIRFLOW_CONTEXT_TASK_ID"), + (task_instance, "logical_date", "AIRFLOW_CONTEXT_LOGICAL_DATE"), + (task_instance, "try_number", "AIRFLOW_CONTEXT_TRY_NUMBER"), + (dag_run, "run_id", "AIRFLOW_CONTEXT_DAG_RUN_ID"), + ] + + context_params = settings.get_airflow_context_vars(context) + for key, value in context_params.items(): + if not isinstance(key, str): + raise TypeError(f"key <{key}> must be string") + if not isinstance(value, str): + raise TypeError(f"value of key <{key}> must be string, not {type(value)}") + + if in_env_var_format: + if not key.startswith(ENV_VAR_FORMAT_PREFIX): + key = ENV_VAR_FORMAT_PREFIX + key.upper() + else: + if not key.startswith(DEFAULT_FORMAT_PREFIX): + key = DEFAULT_FORMAT_PREFIX + key + params[key] = value + + for subject, attr, mapping_key in ops: + _attr = getattr(subject, attr, None) + if subject and _attr: + mapping_value = AIRFLOW_VAR_NAME_FORMAT_MAPPING[mapping_key][name_format] + if isinstance(_attr, str): + params[mapping_value] = _attr + elif isinstance(_attr, datetime): + params[mapping_value] = _attr.isoformat() + elif isinstance(_attr, list): + # os env variable value needs to be string + params[mapping_value] = ",".join(_attr) + else: + params[mapping_value] = str(_attr) + + return params 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 fcd401ab5db20..57b0107b9b2bc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -76,6 +76,7 @@ MacrosAccessor, OutletEventAccessors, VariableAccessor, + context_to_airflow_vars, get_previous_dagrun_success, set_current_context, ) @@ -807,6 +808,10 @@ def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger): # Populate the context var so ExecutorSafeguard doesn't complain ctx.run(ExecutorSafeguard.tracker.set, task) + # Export context in os.environ to make it available for operators to use. + airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True) + os.environ.update(airflow_context_vars) + for i, callback in enumerate(task.on_execute_callback): try: callback(context) diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py index 947d4a318d761..c81b991788a51 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -18,6 +18,7 @@ from __future__ import annotations from datetime import datetime +from unittest import mock from unittest.mock import MagicMock, patch import pytest @@ -49,6 +50,7 @@ VariableAccessor, _convert_connection_result_conn, _convert_variable_result_to_variable, + context_to_airflow_vars, set_current_context, ) @@ -103,6 +105,79 @@ def test_convert_variable_result_to_variable_with_deserialize_json(): ) +class TestAirflowContextHelpers: + def setup_method(self): + self.dag_id = "dag_id" + self.task_id = "task_id" + self.try_number = 1 + self.logical_date = "2017-05-21T00:00:00" + self.dag_run_id = "dag_run_id" + self.owner = ["owner1", "owner2"] + self.email = ["email1@test.com"] + self.context = { + "dag_run": mock.MagicMock( + name="dag_run", + run_id=self.dag_run_id, + logical_date=datetime.strptime(self.logical_date, "%Y-%m-%dT%H:%M:%S"), + ), + "task_instance": mock.MagicMock( + name="task_instance", + task_id=self.task_id, + dag_id=self.dag_id, + try_number=self.try_number, + logical_date=datetime.strptime(self.logical_date, "%Y-%m-%dT%H:%M:%S"), + ), + "task": mock.MagicMock(name="task", owner=self.owner, email=self.email), + } + + def test_context_to_airflow_vars_empty_context(self): + assert context_to_airflow_vars({}) == {} + + def test_context_to_airflow_vars_all_context(self): + assert context_to_airflow_vars(self.context) == { + "airflow.ctx.dag_id": self.dag_id, + "airflow.ctx.logical_date": self.logical_date, + "airflow.ctx.task_id": self.task_id, + "airflow.ctx.dag_run_id": self.dag_run_id, + "airflow.ctx.try_number": str(self.try_number), + "airflow.ctx.dag_owner": "owner1,owner2", + "airflow.ctx.dag_email": "email1@test.com", + } + + assert context_to_airflow_vars(self.context, in_env_var_format=True) == { + "AIRFLOW_CTX_DAG_ID": self.dag_id, + "AIRFLOW_CTX_LOGICAL_DATE": self.logical_date, + "AIRFLOW_CTX_TASK_ID": self.task_id, + "AIRFLOW_CTX_TRY_NUMBER": str(self.try_number), + "AIRFLOW_CTX_DAG_RUN_ID": self.dag_run_id, + "AIRFLOW_CTX_DAG_OWNER": "owner1,owner2", + "AIRFLOW_CTX_DAG_EMAIL": "email1@test.com", + } + + def test_context_to_airflow_vars_with_default_context_vars(self): + with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: + airflow_cluster = "cluster-a" + mock_method.return_value = {"airflow_cluster": airflow_cluster} + + context_vars = context_to_airflow_vars(self.context) + assert context_vars["airflow.ctx.airflow_cluster"] == airflow_cluster + + context_vars = context_to_airflow_vars(self.context, in_env_var_format=True) + assert context_vars["AIRFLOW_CTX_AIRFLOW_CLUSTER"] == airflow_cluster + + with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: + mock_method.return_value = {"airflow_cluster": [1, 2]} + with pytest.raises(TypeError) as error: + context_to_airflow_vars(self.context) + assert str(error.value) == "value of key must be string, not " + + with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: + mock_method.return_value = {1: "value"} + with pytest.raises(TypeError) as error: + context_to_airflow_vars(self.context) + assert str(error.value) == "key <1> must be string" + + class TestConnectionAccessor: def test_getattr_connection(self, mock_supervisor_comms): """ diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 10eb2bdf719bf..ed638fc3885f6 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -923,6 +923,31 @@ def test_run_with_inlets_and_outlets( mock_supervisor_comms.send_request.assert_any_call(msg=last_expected_msg, log=mock.ANY) +@mock.patch("airflow.sdk.execution_time.task_runner.context_to_airflow_vars") +@mock.patch.dict(os.environ, {}, clear=True) +def test_execute_task_exports_env_vars( + mock_context_to_airflow_vars, create_runtime_ti, mock_supervisor_comms +): + """Test that _execute_task exports airflow context to environment variables.""" + + def test_function(): + return "test function" + + task = PythonOperator( + task_id="test_task", + python_callable=test_function, + ) + + ti = create_runtime_ti(task=task, dag_id="dag_with_env_vars") + + mock_env_vars = {"AIRFLOW_CTX_DAG_ID": "test_dag_env_vars", "AIRFLOW_CTX_TASK_ID": "test_env_task"} + mock_context_to_airflow_vars.return_value = mock_env_vars + run(ti, log=mock.MagicMock()) + + assert os.environ["AIRFLOW_CTX_DAG_ID"] == "test_dag_env_vars" + assert os.environ["AIRFLOW_CTX_TASK_ID"] == "test_env_task" + + class TestRuntimeTaskInstance: def test_get_context_without_ti_context_from_server(self, mocked_parse, make_ti_context): """Test get_template_context without ti_context_from_server."""