Airflow

Run Airflow in Docker

Source: https://airflow.apache.org/docs/apache-airflow/stable/howto/docker-compose/index.html

Create directories

mkdir airflow-local
cd airflow-local
mkdir ./dags ./logs ./plugins

Fetch docker-compose.yaml

curl -LfO 'https://airflow.apache.org/docs/apache-airflow/2.5.0/docker-compose.yaml'

Before Airflow 2.2, the Docker Compose also had AIRFLOW_GID parameter, but it did not provide any additional functionality - only added confusion - so it has been removed.

Setting the Airflow user

echo -e "AIRFLOW_UID=$(id -u)" > .env
cat .env

Initialise the Airflow Database

docker-compose up airflow-init

Start Airflow services

docker-compose up

Enter the Airflow Worker container

docker exec -it <container-id> bash

Access to the CLI Application

docker-compose run airflow-worker airflow info

TLDR

mkdir airflow-local
cd airflow-local
mkdir ./dags ./logs ./plugins
curl -LfO 'https://airflow.apache.org/docs/apache-airflow/2.5.0/docker-compose.yaml'
echo -e "AIRFLOW_UID=$(id -u)" > .env
docker-compose up airflow-init

docker-compose down --volumes --remove-orphans
docker-compose up

Install additional packages

_PIP_ADDITIONAL_REQUIREMENTS

One option is to use the _PIP_ADDITIONAL_REQUIREMENTS,

Custome image

Another is to build a custom image.

  • Create a Dockerfile:
    FROM apache/airflow:2.5.0
    COPY requirements.txt .
    RUN pip install -r requirements.txt
    
  • Create a requirements.txt:
    airflow-code-editor
    black
    fs
    
  • Edit docker-compose.yaml, change ‘image: ${AIRFLOW_IMAGE_NAME:-apache/airflow:2.5.0}’ to ‘build: .’
  • Run ‘docker-compose build’ and then ‘docker-compose up’

Tips

DAG parameters:

max_active_runs - maximum number of active DAG runs

if you just want the DAGs to be able to execute two jobs in parallel (with no conditions between two distinct runs) then set max_active_runs=2

Operators parameters:

wait_for_downstream - when set to true, an instance of task X will wait for tasks immediately downstream of the previous instance of task X to finish successfully or be skipped before it runs

depends_on_past - when set to true, task instances will run sequentially and only if the previous instance has succeeded or has been skipped.

Example Dags

PythonVirtualenvOperator

from airflow.operators.python import PythonVirtualenvOperator
from airflow.utils.dates import days_ago
from airflow.utils.timezone import datetime

from airflow import DAG

# Define some default arguments for the DAG
default_args = {
    "owner": "airflow",
    "start_date": days_ago(1),
    # "start_date": "2022-12-19",
    "depends_on_past": False,
    "wait_for_downstream": False,
}

# Create a new DAG with the default arguments
dag = DAG(
    "m2k_tongji_hospital",
    default_args=default_args,
    # schedule_interval=None,
    schedule_interval="*/1 * * * *",
    max_active_runs=1,
    max_active_tasks=3,
)

# Define a Python function that transfer mqtt messages to kafka
def mqtt2kafka():

    import paho.mqtt.subscribe as subscribe
    from kafka import KafkaProducer
    import json
    import re
    import numpy as np
    import struct
    import pandas as pd
    import time
    import uuid

    # Subscribe to the MQTT server and the corresponding topic
    msg = subscribe.simple(
        "topic",
        hostname="127.0.0.1",
        qos=1,
        port=1883,
        msg_count=9,
        auth={"username": "admin", "password": "123456"},
    )

    # Define the producer
    producer = KafkaProducer(
        bootstrap_servers=[
            "127.0.0.3:9092",
            "127.0.0.4:9092",
            "127.0.0.2:9092",
        ],
    )

    for item in msg:
        # decode bytes to json
        encoded_message = item.payload
        decoded_message = json.loads(encoded_message.decode("utf-8"))

        # convert message value
        raw_msg_dict = decoded_message[0]
        value_char = raw_msg_dict["value"]
        start_reg = raw_msg_dict["start_reg"]
        end_reg = raw_msg_dict["end_reg"]
        reg_list = np.arange(start_reg, end_reg + 1)
        reg_list = [int(x) for x in reg_list]
        value_list = re.findall(r".{8}", value_char)

        value_list_decimal = []
        for value_hex_rev in value_list:
            value_hex = ""
            if len(value_hex_rev):
                byte_list = re.findall(r".{2}", value_hex_rev)
                byte_list.reverse()
                value_hex = "".join(byte_list)

            value_decimal = struct.unpack(">f", bytes.fromhex(value_hex))[0]
            value_decimal = round(value_decimal, 2)
            value_list_decimal.append(value_decimal)

        value_dict = dict(zip(reg_list, value_list_decimal))
        # norm_msg_dict = raw_msg_dict  # omitted as requested
        # norm_msg_dict["value"] = value_dict
        # message = json.dumps(norm_msg_dict, separators=(",", ":")).encode("utf-8")

        # Reformat message as requested
        new_struct = pd.DataFrame(
            list(value_dict.items()), columns=["tagCode", "value"]
        )
        new_struct["tagCode"] = new_struct["tagCode"].apply(lambda v: f"{v+1:04d}")
        if raw_msg_dict["deviceId"] == "xxx":
            new_struct["tagCode"] = "xxx" + new_struct["tagCode"].astype(str)
        else:
            new_struct["tagCode"] = "xxx" + new_struct["tagCode"].astype(str)


        # Send messages to Kafka topic
        message = new_struct.to_json(orient="records").encode("utf-8")
        key_uuid = str(uuid.uuid4()).encode("utf-8")

        # Topic for tests
        # producer.send('xxx', key=key_uuid, value=message)

        # Topic for production
        producer.send('xxx', key=key_uuid, value=message)

# Define the task
virtualenv_task = PythonVirtualenvOperator(
    task_id="m2k_tongji_hospital",
    python_callable=mqtt2kafka,
    requirements=["paho-mqtt", "kafka-python", "numpy", "pandas"],
    system_site_packages=False,
    dag=dag,
)

PythonVirtualenvOperator with XCom

from datetime import datetime, timedelta

import pendulum
from airflow.operators.python import PythonVirtualenvOperator

from airflow import DAG

local_tz = pendulum.timezone("Asia/Shanghai")

# Define some default arguments for the DAG
default_args = {
    "owner": "airflow",
    "depends_on_past": False,
    "wait_for_downstream": False,
    "start_date": datetime(2022, 12, 22, tzinfo=local_tz),
    "retries": 1,
    "retry_delay": timedelta(minutes=1),
}

# Create a new DAG with the default arguments
dag = DAG(
    "test_PythonVirtualenvOperator",
    default_args=default_args,
    schedule_interval=None,
    # schedule_interval="*/1 * * * *",
    max_active_runs=1,
    max_active_tasks=3,
)

# Define a Python function that subscribes mqtt messages
def subscribe_mqtt_messages():

    import json

    from paho.mqtt import subscribe

    # Subscribe to the MQTT server and the corresponding topic
    msg = subscribe.simple(
        "topic",
        hostname="127.0.0.1",
        qos=1,
        port=1883,
        msg_count=2,
        auth={"username": "admin", "password": "123456"},
    )

    payloads = []
    for item in msg:
        encoded_message = item.payload
        decoded_message = json.loads(encoded_message.decode("utf-8"))
        payloads.append(decoded_message)
    return payloads


# Define a function for converting MQTT messages to requested formats
def message_converter(msg):

    # import json
    import logging
    import re
    import struct
    import time

    import numpy as np
    import pandas as pd

    data_list = []
    # logging.info(print(msg))
    for item in eval(msg):

        # Convert message value
        # logging.info(print(item))
        raw_msg_dict = item[0]
        # logging.info(print(raw_msg_dict))
        value_char = raw_msg_dict["value"]
        start_reg = raw_msg_dict["start_reg"]
        end_reg = raw_msg_dict["end_reg"]
        reg_list = np.arange(start_reg, end_reg + 1)
        reg_list = [int(x) for x in reg_list]
        value_list = re.findall(r".{8}", value_char)

        value_list_decimal = []
        for value_hex_rev in value_list:
            value_hex = ""
            if len(value_hex_rev):
                byte_list = re.findall(r".{2}", value_hex_rev)
                byte_list.reverse()
                value_hex = "".join(byte_list)

            value_decimal = struct.unpack(">f", bytes.fromhex(value_hex))[0]
            value_decimal = round(value_decimal, 2)
            value_list_decimal.append(value_decimal)

        value_dict = dict(zip(reg_list, value_list_decimal))

        # Reformat message as requested
        new_struct = pd.DataFrame(
            list(value_dict.items()), columns=["tagCode", "value"]
        )
        message = new_struct.to_json(orient="records")
        data_list.append(message)

    return data_list


# Define a function to publish messages to Kafka topics
def publish_messages(messages):

    import logging
    import uuid

    from kafka import KafkaProducer

    # Define the producer
    producer = KafkaProducer(
        bootstrap_servers=[
            "172.1.1.3:9092",
            "172.1.1.4:9092",
            "172.1.1.2:9092",
        ],
    )

    # Send messages to Kafka topic

    key_uuid = str(uuid.uuid4()).encode("utf-8")
    messages = messages.encode("utf-8")
    logging.info(print(messages))
    producer.send("test_mbt_building_climate", key=key_uuid, value=messages)


# Define the task for fetching MQTT messages
subscribe_mqtt_messages_task = PythonVirtualenvOperator(
    task_id="subscribe_mqtt_messages",
    python_callable=subscribe_mqtt_messages,
    requirements=["paho-mqtt"],
    system_site_packages=False,
    dag=dag,
)

# Define the task for converting messages
message_converter_task = PythonVirtualenvOperator(
    task_id="message_converter",
    python_callable=message_converter,
    op_kwargs={"msg": '{{ti.xcom_pull(task_ids="subscribe_mqtt_messages")}}'},
    requirements=["numpy", "pandas"],
    system_site_packages=False,
    dag=dag,
)

# Define the task for publishing messages
publish_messages_task = PythonVirtualenvOperator(
    task_id="publish_messages",
    python_callable=publish_messages,
    op_kwargs={"messages": '{{ti.xcom_pull(task_ids="message_converter")}}'},
    requirements=["kafka-python"],
    system_site_packages=False,
    dag=dag,
)

subscribe_mqtt_messages_task >> message_converter_task >> publish_messages_task

DockerOperator

When using docker in airflow by docker-compose, add docker proxy in docker-compose.yaml file.

  docker-socket-proxy:
    image: tecnativa/docker-socket-proxy:latest
    environment:
      CONTAINERS: 1
      IMAGES: 1
      AUTH: 1
      POST: 1
    privileged: true
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock:ro
    restart: always

Then create task with DockerOperator.

from datetime import datetime, timedelta

import pendulum
import time
import logging

from airflow.operators.python_operator import BranchPythonOperator
from airflow.operators.python_operator import PythonOperator
from airflow.operators.docker_operator import DockerOperator
from airflow.providers.dingding.operators.dingding import DingdingOperator
from airflow.operators.dummy_operator import DummyOperator
from airflow.utils.dates import days_ago

from airflow import DAG

local_tz = pendulum.timezone("Asia/Shanghai")

# Define some default arguments for the DAG
default_args = {
    "owner": "airflow",
    "depends_on_past": False,
    "wait_for_downstream": False,
    "start_date": days_ago(1),
    "retries": 24,
    "retry_delay": timedelta(minutes=5),
}

# Create a new DAG with the default arguments
dag = DAG(
    "m2k_tongjihospital_v4_docker_schedule",
    default_args=default_args,
    # schedule_interval="*/1 * * * *",
    schedule_interval="0 */12 * * *",
    max_active_runs=1,
    max_active_tasks=30,
    catchup=False,
    dagrun_timeout=timedelta(hours=12),
)


def dag_start_time(**kwargs):
    time.sleep(30)
    start_time = time.time()
    task_instance = kwargs['ti']
    task_instance.xcom_push(key='start_time', value=start_time)


def dag_end_time(**kwargs):
    task_instance = kwargs['ti']
    starting_time = task_instance.xcom_pull(task_ids='start_dag', key="start_time")
    end_time = time.time()
    # logging.info(print(end_time, starting_time))
    duration = end_time - starting_time
    # logging.info(print(duration))
    if duration <= 14400:
        return 'dingtalk'
    else:
        return 'skip_task'


start_dag = PythonOperator(
    task_id='start_dag', python_callable=dag_start_time, provide_context=True, dag=dag
)

# Define the task for fetching MQTT messages
subscribe_mqtt_messages_task = DockerOperator(
    task_id="subscribe_mqtt_messages",
    api_version="auto",
    container_name="m2k_schedule",
    image="m2k:0.0.1",
    command="python ./m2k_TongjiHospital.py",
    # docker_url="unix://var/run/docker.sock",
    docker_url="TCP://docker-socket-proxy:2375",
    network_mode="bridge",
    xcom_all=True,
    # retrieve_output=True,
    # retrieve_output_path="/tmp/script.out",
    auto_remove=True,
    dag=dag,
)

end_dag = BranchPythonOperator(
    task_id='end_dag', python_callable=dag_end_time, provide_context=True, dag=dag
)

skip_task = DummyOperator(task_id='skip_task', trigger_rule='none_failed', dag=dag)

dingtalk = DingdingOperator(
    task_id='dingtalk',
    dingding_conn_id='dingding_default',
    message_type='text',
    message='DingTalk airflow m2k_tongjihospital ended within 4 hours.',
    at_mobiles=['17724559090'],
    # at_all=False,
    dag=dag,
)

start_dag >> subscribe_mqtt_messages_task >> end_dag >> [dingtalk, skip_task]

Auto-cleanup

"""
A maintenance workflow that you can deploy into Airflow to periodically clean
out the DagRun, TaskInstance, Log, XCom, Job DB and SlaMiss entries to avoid
having too much data in your Airflow MetaStore.

## Authors

The [DAG](https://cloud.google.com/composer/docs/composer-2/cleanup-airflow-database) is a fork of [teamclairvoyant repository.](https://github.com/teamclairvoyant/airflow-maintenance-dags/tree/master/db-cleanup)

## Usage

1. Update the global variables (SCHEDULE_INTERVAL, DAG_OWNER_NAME,
  ALERT_EMAIL_ADDRESSES and ENABLE_DELETE) in the DAG with the desired values

2. Modify the DATABASE_OBJECTS list to add/remove objects as needed. Each
   dictionary in the list features the following parameters:
    - airflow_db_model: Model imported from airflow.models corresponding to
      a table in the airflow metadata database
    - age_check_column: Column in the model/table to use for calculating max
      date of data deletion
    - keep_last: Boolean to specify whether to preserve last run instance
        - keep_last_filters: List of filters to preserve data from deleting
          during clean-up, such as DAG runs where the external trigger is set to 0.
        - keep_last_group_by: Option to specify column by which to group the
          database entries and perform aggregate functions.

3. Create and Set the following Variables in the Airflow Web Server
  (Admin -> Variables)
    - airflow_db_cleanup__max_db_entry_age_in_days - integer - 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.

4. Put the DAG in your gcs bucket.
"""
from datetime import datetime, timedelta
import logging
import os

import airflow
from airflow import settings
from airflow.configuration import conf
from airflow.jobs.base_job import BaseJob
from airflow.models import DAG, DagModel, DagRun, Log, SlaMiss, \
    TaskInstance, Variable, XCom
from airflow.operators.python import PythonOperator
import dateutil.parser
from sqlalchemy import and_, func
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.orm import load_only

try:
    # airflow.utils.timezone is available from v1.10 onwards
    from airflow.utils import timezone
    now = timezone.utcnow
except ImportError:
    now = datetime.utcnow

# airflow-db-cleanup
DAG_ID = os.path.basename(__file__).replace(".pyc", "").replace(".py", "")
START_DATE = airflow.utils.dates.days_ago(1)
# How often to Run. @daily - Once a day at Midnight (UTC)
SCHEDULE_INTERVAL = "@daily"
# Who is listed as the owner of this DAG in the Airflow Web Server
DAG_OWNER_NAME = "operations"
# List of email address to send email alerts to if this job fails
ALERT_EMAIL_ADDRESSES = []
# 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.
AIRFLOW_VERSION = os.getenv("MAJOR_VERSION", "2.5.0").split(".")
DEFAULT_MAX_DB_ENTRY_AGE_IN_DAYS = int(
    Variable.get("airflow_db_cleanup__max_db_entry_age_in_days", 1))
# Prints the database entries which will be getting deleted; set to False
# to avoid printing large lists and slowdown process
PRINT_DELETES = False
# Whether the job should delete the db entries or not. Included if you want to
# temporarily avoid deleting the db entries.
ENABLE_DELETE = True

# List of all the objects that will be deleted. Comment out the DB objects you
# want to skip.
DATABASE_OBJECTS = [{
    "airflow_db_model": BaseJob,
    "age_check_column": BaseJob.latest_heartbeat,
    "keep_last": False,
    "keep_last_filters": None,
    "keep_last_group_by": None
}, {
    "airflow_db_model": DagRun,
    "age_check_column": DagRun.execution_date,
    "keep_last": True,
    "keep_last_filters": [DagRun.external_trigger.is_(False)],
    "keep_last_group_by": DagRun.dag_id
}, {
    "airflow_db_model": TaskInstance,
    "age_check_column": TaskInstance.execution_date if AIRFLOW_VERSION < ['2', '2', '0'] else TaskInstance.start_date,
    "keep_last": False,
    "keep_last_filters": None,
    "keep_last_group_by": None
}, {
    "airflow_db_model": Log,
    "age_check_column": Log.dttm,
    "keep_last": False,
    "keep_last_filters": None,
    "keep_last_group_by": None
}, {
    "airflow_db_model": XCom,
    "age_check_column": XCom.execution_date if AIRFLOW_VERSION < ['2', '2', '0'] else XCom.timestamp,
    "keep_last": False,
    "keep_last_filters": None,
    "keep_last_group_by": None
}, {
    "airflow_db_model": SlaMiss,
    "age_check_column": SlaMiss.execution_date,
    "keep_last": False,
    "keep_last_filters": None,
    "keep_last_group_by": None
}, {
    "airflow_db_model": DagModel,
    "age_check_column": DagModel.last_scheduler_run if AIRFLOW_VERSION < ['2', '0', '2'] else DagModel.last_parsed_time,
    "keep_last": False,
    "keep_last_filters": None,
    "keep_last_group_by": None
}]

# Check for TaskReschedule model
try:
    from airflow.models import TaskReschedule
    DATABASE_OBJECTS.append({
        "airflow_db_model": TaskReschedule,
        "age_check_column": TaskReschedule.execution_date if AIRFLOW_VERSION < ['2', '2', '0'] else TaskReschedule.start_date,
        "keep_last": False,
        "keep_last_filters": None,
        "keep_last_group_by": None
    })

except Exception as e:
    logging.error(e)

# Check for TaskFail model
try:
    from airflow.models import TaskFail
    DATABASE_OBJECTS.append({
        "airflow_db_model": TaskFail,
        "age_check_column": TaskFail.execution_date,
        "keep_last": False,
        "keep_last_filters": None,
        "keep_last_group_by": None
    })

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 if AIRFLOW_VERSION < ['2', '2', '0'] else RenderedTaskInstanceFields.run_id,
        "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,
        "keep_last": False,
        "keep_last_filters": None,
        "keep_last_group_by": None
    })

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
        }, {
            "airflow_db_model": TaskSet,
            "age_check_column": TaskSet.date_done,
            "keep_last": False,
            "keep_last_filters": None,
            "keep_last_group_by": None
        }))

    except Exception as e:
        logging.error(e)

session = settings.Session()

default_args = {
    "owner": DAG_OWNER_NAME,
    "depends_on_past": False,
    "email": ALERT_EMAIL_ADDRESSES,
    "email_on_failure": True,
    "email_on_retry": False,
    "start_date": START_DATE,
    "retries": 1,
    "retry_delay": timedelta(minutes=1)
}

dag = DAG(
    DAG_ID,
    default_args=default_args,
    schedule_interval=SCHEDULE_INTERVAL,
    start_date=START_DATE)
if hasattr(dag, "doc_md"):
    dag.doc_md = __doc__
if hasattr(dag, "catchup"):
    dag.catchup = False


def print_configuration_function(**context):
    logging.info("Loading Configurations...")
    dag_run_conf = context.get("dag_run").conf
    logging.info("dag_run.conf: " + str(dag_run_conf))
    max_db_entry_age_in_days = None
    if dag_run_conf:
        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):
        logging.info(
            "maxDBEntryAgeInDays conf variable isn't included or Variable " +
            "value is less than 1. Using Default '" +
            str(DEFAULT_MAX_DB_ENTRY_AGE_IN_DAYS) + "'")
        max_db_entry_age_in_days = DEFAULT_MAX_DB_ENTRY_AGE_IN_DAYS
    max_date = now() + timedelta(-max_db_entry_age_in_days)
    logging.info("Finished Loading Configurations")
    logging.info("")

    logging.info("Configurations:")
    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")
    context["ti"].xcom_push(key="max_date", value=max_date.isoformat())


print_configuration = PythonOperator(
    task_id="print_configuration",
    python_callable=print_configuration_function,
    provide_context=True,
    dag=dag)


def cleanup_function(**context):

    if AIRFLOW_VERSION > ['3', '2', '5']:
        logging.error("The script is not supported for this Airflow Version. "
                      "Skipped deleting the data.")
        return
    logging.info("Retrieving max_execution_date from XCom")
    max_date = context["ti"].xcom_pull(
        task_ids=print_configuration.task_id, key="max_date")
    max_date = dateutil.parser.parse(max_date)  # stored as iso8601 str in xcom

    airflow_db_model = context["params"].get("airflow_db_model")
    state = context["params"].get("state")
    age_check_column = context["params"].get("age_check_column")
    keep_last = context["params"].get("keep_last")
    keep_last_filters = context["params"].get("keep_last_filters")
    keep_last_group_by = context["params"].get("keep_last_group_by")

    logging.info("Configurations:")
    logging.info("max_date:                 " + str(max_date))
    logging.info("enable_delete:            " + str(ENABLE_DELETE))
    logging.info("session:                  " + str(session))
    logging.info("airflow_db_model:         " + str(airflow_db_model))
    logging.info("state:                    " + str(state))
    logging.info("age_check_column:         " + str(age_check_column))
    logging.info("keep_last:                " + str(keep_last))
    logging.info("keep_last_filters:        " + str(keep_last_filters))
    logging.info("keep_last_group_by:       " + str(keep_last_group_by))

    logging.info("")

    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)")
        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...")
            # using bulk delete
            query.delete(synchronize_session=False)
            session.commit()
            logging.info("Finished Performing Delete")
        else:
            logging.warn("You've opted to skip deleting the db entries. "
                         "Set ENABLE_DELETE to True to delete entries!!!")

        logging.info("Finished Running Cleanup Process")

    except ProgrammingError as e:
        logging.error(e)
        logging.error(
            str(airflow_db_model) + " is not present in the metadata."
            "Skipping...")


for db_object in DATABASE_OBJECTS:

    cleanup_op = PythonOperator(
        task_id="cleanup_" + str(db_object["airflow_db_model"].__name__),
        python_callable=cleanup_function,
        params=db_object,
        provide_context=True,
        dag=dag)

    print_configuration.set_downstream(cleanup_op)
Previous
Next