Setting Airflow context Environment variables for operators

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

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:
     """

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 <airflow_cluster> must be string, not <class 'list'>"
-
-        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,)

generated/provider_dependencies.json

@@ -179,6 +179,7 @@
     ],
     "cross-providers-deps": [
       "amazon",
+      "common.compat",
       "common.sql",
       "microsoft.mssql",
       "mysql",

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 <https://airflow.apache.org/docs/apache-airflow-providers-amazon>`_                    ``amazon``
+`apache-airflow-providers-common-compat <https://airflow.apache.org/docs/apache-airflow-providers-common-compat>`_      ``common.compat``
 `apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_            ``common.sql``
 `apache-airflow-providers-microsoft-mssql <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-mssql>`_  ``microsoft.mssql``
 `apache-airflow-providers-mysql <https://airflow.apache.org/docs/apache-airflow-providers-mysql>`_                      ``mysql``

providers/apache/hive/pyproject.toml

@@ -91,13 +91,17 @@ dependencies = [
 "vertica" = [
     "apache-airflow-providers-vertica"
 ]
+"common.compat" = [
+    "apache-airflow-providers-common-compat"
+]
 
 [dependency-groups]
 dev = [
     "apache-airflow",
     "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",

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": [],
     }

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"]
 

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)

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

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

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]

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)
 
 

providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_samba.py

@@ -24,15 +24,20 @@
 
 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,
     MockHiveServer2Hook,
     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):

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
 

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