def test_clean_airflow_import_if_environment_is_set(monkeypatch):
temp_home = tempfile.mkdtemp()
temp_airflow_home = path.join(temp_home, "airflow")
temp_airflow_cfg = path.join(temp_airflow_home, "airflow.cfg")
monkeypatch.setenv("AIRFLOW_HOME", temp_airflow_home)
monkeypatch.setenv("AIRFLOW_CONFIG", temp_airflow_cfg)
with CleanAirflowImport():
import airflow
temp_home_content = listdir(temp_home)
rmtree(temp_home)
assert len(temp_home_content) == 0, \
"Writing into the wrong Airflow home directory"
assert environ.get("AIRFLOW_HOME") == temp_airflow_home, \
"AIRFLOW_HOME is not restored after exit from CleanAirflowImport"
assert environ.get("AIRFLOW_CONFIG") == temp_airflow_cfg, \
"AIRFLOW_CONFIG is not restored after exit from CleanAirflowImport"
def test_clean_airflow_import_if_environment_is_not_set(monkeypatch):
temp_home = tempfile.mkdtemp()
monkeypatch.delenv("AIRFLOW_HOME", raising=False)
monkeypatch.delenv("AIRFLOW_CONFIG", raising=False)
monkeypatch.setattr(path, "expanduser",
lambda x: x.replace("~", temp_home))
with CleanAirflowImport():
import airflow
temp_home_content = listdir(temp_home)
rmtree(temp_home)
assert len(temp_home_content) == 0, \
"Writing into the wrong Airflow home directory"
assert environ.get("AIRFLOW_HOME") is None, \
"AIRFLOW_HOME is not cleaned after exit from CleanAirflowImport"
assert environ.get("AIRFLOW_CONFIG") is None, \
"AIRFLOW_CONFIG is not cleaned after exit from CleanAirflowImport"
import logging
import psutil
import shutil
from cwl_airflow.utilities.report import dag_on_success, dag_on_failure
from cwl_airflow.utilities.helpers import CleanAirflowImport
with CleanAirflowImport():
from airflow import configuration
from airflow.models import DAG, DagRun, TaskInstance
from airflow.operators.python_operator import PythonOperator
from airflow.utils.dates import days_ago
from airflow.utils.db import provide_session
from airflow.utils.state import State
logger = logging.getLogger(__name__)
TIMEOUT = configuration.conf.getint('core', 'KILLED_TASK_CLEANUP_TIME')
@provide_session
def clean_db(dr, session=None):
logger.debug(f"""Clean DB for {dr.dag_id} - {dr.run_id}""")
for ti in dr.get_task_instances():
logger.debug(
f"""process {ti.dag_id} - {ti.task_id} - {ti.execution_date}""")
ti.clear_xcom_data()
logger.debug(" - clean Xcom table")
session.query(TaskInstance).filter(
TaskInstance.task_id == ti.task_id,
TaskInstance.dag_id == ti.dag_id,