Airflow db cleanup dag update (#8851)

* Fix and adjust to Airflow 2.3 db cleanup DAG

* Remove cleanup Task and TaskSet tables

* removing duplicated airflow_1 file + some minor fixes

* fix lint errors

* small fix

* small fix  in deletion

* separate dag logic between airflow versions

* move session closing to finally part

* lint fixes

* import order fix

Co-authored-by: Leah E. Cole <[email protected]>

Author: kubasieron

composer/airflow_1_samples/airflow_db_cleanup.py

@@ -53,11 +53,12 @@
 
 import airflow
 from airflow import settings
-from airflow.configuration import conf
 from airflow.jobs import BaseJob
 from airflow.models import DAG, DagModel, DagRun, Log, SlaMiss, \
     TaskInstance, Variable, XCom
 from airflow.operators.python_operator import PythonOperator
+from airflow.version import version as airflow_version
+
 import dateutil.parser
 from sqlalchemy import and_, func
 from sqlalchemy.exc import ProgrammingError
@@ -66,6 +67,7 @@
 try:
     # airflow.utils.timezone is available from v1.10 onwards
     from airflow.utils import timezone
+
     now = timezone.utcnow
 except ImportError:
     now = datetime.utcnow
@@ -79,9 +81,11 @@
 DAG_OWNER_NAME = "operations"
 # List of email address to send email alerts to if this job fails
 ALERT_EMAIL_ADDRESSES = []
+# Airflow version used by the environment in list form, value stored in
+# airflow_version is in format e.g "1.10.15+composer"
+AIRFLOW_VERSION = airflow_version[:-len("+composer")].split(".")
 # Length to retain the log files if not already provided in the conf. If this
 # is set to 30, the job will remove those files that arE 30 days old or older.
-
 DEFAULT_MAX_DB_ENTRY_AGE_IN_DAYS = int(
     Variable.get("airflow_db_cleanup__max_db_entry_age_in_days", 30))
 # Prints the database entries which will be getting deleted; set to False
@@ -139,6 +143,7 @@
 # Check for TaskReschedule model
 try:
     from airflow.models import TaskReschedule
+
     DATABASE_OBJECTS.append({
         "airflow_db_model": TaskReschedule,
         "age_check_column": TaskReschedule.execution_date,
@@ -153,6 +158,7 @@
 # Check for TaskFail model
 try:
     from airflow.models import TaskFail
+
     DATABASE_OBJECTS.append({
         "airflow_db_model": TaskFail,
         "age_check_column": TaskFail.execution_date,
@@ -164,23 +170,10 @@
 except Exception as e:
     logging.error(e)
 
-# Check for RenderedTaskInstanceFields model
-try:
-    from airflow.models import RenderedTaskInstanceFields
-    DATABASE_OBJECTS.append({
-        "airflow_db_model": RenderedTaskInstanceFields,
-        "age_check_column": RenderedTaskInstanceFields.execution_date,
-        "keep_last": False,
-        "keep_last_filters": None,
-        "keep_last_group_by": None
-    })
-
-except Exception as e:
-    logging.error(e)
-
 # Check for ImportError model
 try:
     from airflow.models import ImportError
+
     DATABASE_OBJECTS.append({
         "airflow_db_model": ImportError,
         "age_check_column": ImportError.timestamp,
@@ -193,34 +186,6 @@
 except Exception as e:
     logging.error(e)
 
-# Check for celery executor
-airflow_executor = str(conf.get("core", "executor"))
-logging.info("Airflow Executor: " + str(airflow_executor))
-if (airflow_executor == "CeleryExecutor"):
-    logging.info("Including Celery Modules")
-    try:
-        from celery.backends.database.models import Task, TaskSet
-        DATABASE_OBJECTS.extend(({
-            "airflow_db_model": Task,
-            "age_check_column": Task.date_done,
-            "keep_last": False,
-            "keep_last_filters": None,
-            "keep_last_group_by": None,
-            "do_not_delete_by_dag_id": True
-        }, {
-            "airflow_db_model": TaskSet,
-            "age_check_column": TaskSet.date_done,
-            "keep_last": False,
-            "keep_last_filters": None,
-            "keep_last_group_by": None,
-            "do_not_delete_by_dag_id": True
-        }))
-
-    except Exception as e:
-        logging.error(e)
-
-session = settings.Session()
-
 default_args = {
     "owner": DAG_OWNER_NAME,
     "depends_on_past": False,
@@ -252,7 +217,7 @@ def print_configuration_function(**context):
         max_db_entry_age_in_days = dag_run_conf.get(
             "maxDBEntryAgeInDays", None)
     logging.info("maxDBEntryAgeInDays from dag_run.conf: " + str(dag_run_conf))
-    if (max_db_entry_age_in_days is None or max_db_entry_age_in_days < 1):
+    if max_db_entry_age_in_days is None or max_db_entry_age_in_days < 1:
         logging.info(
             "maxDBEntryAgeInDays conf variable isn't included or Variable " +
             "value is less than 1. Using Default '" +
@@ -266,7 +231,6 @@ def print_configuration_function(**context):
     logging.info("max_db_entry_age_in_days: " + str(max_db_entry_age_in_days))
     logging.info("max_date:                 " + str(max_date))
     logging.info("enable_delete:            " + str(ENABLE_DELETE))
-    logging.info("session:                  " + str(session))
     logging.info("")
 
     logging.info("Setting max_execution_date to XCom for Downstream Processes")
@@ -280,7 +244,57 @@ def print_configuration_function(**context):
     dag=dag)
 
 
+def build_query(session, airflow_db_model, age_check_column, max_date,
+                keep_last, keep_last_filters=None, keep_last_group_by=None):
+    query = session.query(airflow_db_model).options(
+        load_only(age_check_column))
+
+    logging.info("INITIAL QUERY : " + str(query))
+
+    if not keep_last:
+        query = query.filter(age_check_column <= max_date, )
+    else:
+        subquery = session.query(func.max(DagRun.execution_date))
+        # workaround for MySQL "table specified twice" issue
+        # https://github.com/teamclairvoyant/airflow-maintenance-dags/issues/41
+        if keep_last_filters is not None:
+            for entry in keep_last_filters:
+                subquery = subquery.filter(entry)
+
+            logging.info("SUB QUERY [keep_last_filters]: " + str(subquery))
+
+        if keep_last_group_by is not None:
+            subquery = subquery.group_by(keep_last_group_by)
+            logging.info(
+                "SUB QUERY [keep_last_group_by]: " +
+                str(subquery))
+
+        subquery = subquery.from_self()
+
+        query = query.filter(
+            and_(age_check_column.notin_(subquery)),
+            and_(age_check_column <= max_date))
+
+    return query
+
+
+def print_query(query, airflow_db_model, age_check_column):
+    entries_to_delete = query.all()
+
+    logging.info("Query: " + str(query))
+    logging.info("Process will be Deleting the following " +
+                 str(airflow_db_model.__name__) + "(s):")
+    for entry in entries_to_delete:
+        date = str(entry.__dict__[str(age_check_column).split(".")[1]])
+        logging.info("\tEntry: " + str(entry) + ", Date: " + date)
+
+    logging.info("Process will be Deleting "
+                 + str(len(entries_to_delete)) + " "
+                 + str(airflow_db_model.__name__) + "(s)")
+
+
 def cleanup_function(**context):
+    session = settings.Session()
 
     logging.info("Retrieving max_execution_date from XCom")
     max_date = context["ti"].xcom_pull(
@@ -310,67 +324,34 @@ def cleanup_function(**context):
     logging.info("Running Cleanup Process...")
 
     try:
-        query = session.query(airflow_db_model).options(
-            load_only(age_check_column))
-
-        logging.info("INITIAL QUERY : " + str(query))
-
-        if keep_last:
-
-            subquery = session.query(func.max(DagRun.execution_date))
-            # workaround for MySQL "table specified twice" issue
-            # https://github.com/teamclairvoyant/airflow-maintenance-dags/issues/41
-            if keep_last_filters is not None:
-                for entry in keep_last_filters:
-                    subquery = subquery.filter(entry)
-
-                logging.info("SUB QUERY [keep_last_filters]: " + str(subquery))
-
-            if keep_last_group_by is not None:
-                subquery = subquery.group_by(keep_last_group_by)
-                logging.info(
-                    "SUB QUERY [keep_last_group_by]: " +
-                    str(subquery))
-
-            subquery = subquery.from_self()
-
-            query = query.filter(
-                and_(age_check_column.notin_(subquery)),
-                and_(age_check_column <= max_date))
-
-        else:
-            query = query.filter(age_check_column <= max_date,)
-
-        if PRINT_DELETES:
-            entries_to_delete = query.all()
-
-            logging.info("Query: " + str(query))
-            logging.info("Process will be Deleting the following " +
-                         str(airflow_db_model.__name__) + "(s):")
-            for entry in entries_to_delete:
-                date = str(entry.__dict__[str(age_check_column).split(".")[1]])
-                logging.info("\tEntry: " + str(entry) + ", Date: " + date)
-
-            logging.info("Process will be Deleting "
-                         + str(len(entries_to_delete)) + " "
-                         + str(airflow_db_model.__name__) + "(s)")
+        if context["params"].get("do_not_delete_by_dag_id"):
+            query = build_query(session, airflow_db_model, age_check_column,
+                                max_date, keep_last, keep_last_filters,
+                                keep_last_group_by)
+            if PRINT_DELETES:
+                print_query(query, airflow_db_model, age_check_column)
+            if ENABLE_DELETE:
+                logging.info("Performing Delete...")
+                query.delete(synchronize_session=False)
+            session.commit()
         else:
-            logging.warn(
-                "You've opted to skip printing the db entries to be deleted. "
-                "Set PRINT_DELETES to True to show entries!!!")
-
-        if ENABLE_DELETE:
-            logging.info("Performing Delete...")
-            if context["params"].get("do_not_delete_by_dag_id"):
-                query.filter(age_check_column <= max_date).delete(synchronize_session=False)
+            dags = session.query(airflow_db_model.dag_id).distinct()
+            session.commit()
+
+            list_dags = [str(list(dag)[0]) for dag in dags]
+            for dag in list_dags:
+                query = build_query(session, airflow_db_model, age_check_column,
+                                    max_date, keep_last, keep_last_filters,
+                                    keep_last_group_by)
+                query = query.filter(airflow_db_model.dag_id == dag)
+                if PRINT_DELETES:
+                    print_query(query, airflow_db_model, age_check_column)
+                if ENABLE_DELETE:
+                    logging.info("Performing Delete...")
+                    query.delete(synchronize_session=False)
                 session.commit()
-            else:
-                dags = session.query(airflow_db_model.dag_id).distinct()
-                list_dags = [str(list(dag)[0]) for dag in dags]
-                for dag in list_dags:
-                    query.filter(age_check_column <= max_date).filter(airflow_db_model.dag_id == dag).delete(synchronize_session=False)
-                    session.commit()
-        else:
+
+        if not ENABLE_DELETE:
             logging.warn("You've opted to skip deleting the db entries. "
                          "Set ENABLE_DELETE to True to delete entries!!!")
 
@@ -379,12 +360,13 @@ def cleanup_function(**context):
     except ProgrammingError as e:
         logging.error(e)
         logging.error(
-            str(airflow_db_model) + " is not present in the metadata."
-            "Skipping...")
+            str(airflow_db_model) + " is not present in the metadata. "
+                                    "Skipping...")
+    finally:
+        session.close()
 
 
 for db_object in DATABASE_OBJECTS:
-
     cleanup_op = PythonOperator(
         task_id="cleanup_" + str(db_object["airflow_db_model"].__name__),
         python_callable=cleanup_function,