def test_airflow_context_vars_patch(self):
patch_airflow_context_vars()
from airflow.utils.operator_helpers import context_to_airflow_vars
mock_task_instance = Mock(try_number=1)
context = dict(task_instance=mock_task_instance)
airflow_vars = context_to_airflow_vars(context)
assert airflow_vars["AIRFLOW_CTX_TRY_NUMBER"] == "1"
def dbnd_airflow_bootstrap():
"""
all relevant patches for airflow execution
"""
global _airflow_bootstrap_applied
if _airflow_bootstrap_applied:
return
_airflow_bootstrap_applied = True # prevent recursive call
from dbnd._core.configuration.dbnd_config import config as dbnd_config
if dbnd_config.getboolean("airflow", "enable_dbnd_context_vars"):
patch_airflow_context_vars()
if os.name == "nt" and dbnd_config.getboolean("airflow",
"enable_windows_support"):
from dbnd_airflow.airflow_override.dbnd_airflow_windows import (
patch_airflow_windows_support, )
patch_airflow_windows_support()
from airflow.plugins_manager import AirflowPlugin
from dbnd_airflow.tracking.airflow_patching import (
add_tracking_to_policy,
patch_airflow_context_vars,
)
add_tracking_to_policy()
patch_airflow_context_vars()
class DbndAutoTracking(AirflowPlugin):
name = "dbnd_airflow_auto_tracking"