Goal

Programmatically pause and unpause DAGs.

Steps

Prerequisites

See the KB article on using the Airflow API to generate an API key for the respective Astronomer product.

Python + Airflow REST API

Astro

"""
Programmatically pause and unpause DAGs
"""
import json
from typing import List

import requests

class PauseUnpauseDags:
    def __init__(
    self,
    url: str,
    token: str,
    ) -> None:
        self.url = f"{url}/api/v1/dags"
        self.token = token

    def _get_all_dag_ids(self) -> List[str]:
        dags = requests.get(
            f'{self.url}?limit=999999',
            headers={'Authorization': f"Bearer {self.token}"}
        )
        dags = dags.json()['dags']
        return [dag['dag_id'] for dag in dags]

    def pause_dags(self, dag_ids: List[str] = []) -> None:
        if not dag_ids:
            dag_ids = self._get_all_dag_ids()
        for dag_id in dag_ids:
            response = requests.patch(
                url=f"{self.url}/{dag_id}",
                headers={'Content-Type': 'application/json', 'Authorization': f"Bearer {self.token}"},
                data=json.dumps({"is_paused": True})
            )

    def unpause_dags(self, dag_ids: List[str] = []) -> None:
        if not dag_ids:
            dag_ids = self._get_all_dag_ids()
        for dag_id in dag_ids:
            response = requests.patch(
                url=f"{self.url}/{dag_id}",
                headers={'Content-Type': 'application/json', 'Authorization': f"Bearer {self.token}"},
                data=json.dumps({"is_paused": False})
            )

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--action",
        choices=["pause", "unpause"],
        help="Choose to pause or unpause a list of DAGs. Choices: ['pause', 'unpause']",
        required=True
    )
    parser.add_argument(
        "--airflow_url",
        help="The URL of your airflow instance, e.g. `airflow.astronomer.run/abc123`",
        required=True
    )
    parser.add_argument(
        "--dag_ids",
        nargs="+",
        help="Optional: Space separated list of DAG IDs that you'd like to pause or unpause, e.g. dag1 dag2 dag3"
    )
    parser.add_argument(
        "--token",
        help="Token generated by Astro API key/secret. To generate a token, see https://docs.astronomer.io/astro/airflow-api"
    )

    args = parser.parse_args()
    if not args.dag_ids:
        args.dag_ids = []

    pause_unpause_dags = PauseUnpauseDags(
        token=args.token,
        url=args.airflow_url
    )
    if args.action == "pause":
        print("Pausing DAG(s)...")
        pause_unpause_dags.pause_dags(args.dag_ids)
    else:
        print("Unpausing DAG(s)...")
        pause_unpause_dags.unpause_dags(args.dag_ids)

This script can be used as follows:

# pause all dags
# to unpause, just change the action arg to "unpause"
python3 pause_unpause_dags_cloud.py \
  --action pause \
  --airflow_url <Airflow URL> \
  --token <token generated from api key/secret>

# pause list of dags
# to unpause, just change the action arg to "unpause"
python3 pause_unpause_dags_astronomer_software.py \
  --action pause \
  --airflow_url <Airflow URL> \
  --token <token generated from api key/secret> \
  --dag_ids dag_id_1 dag_id_2 dag_id_3

Astronomer Software

The below script is also attached as `pause_unpause_dags_astronomer_software.py`

"""
Programmatically pause and unpause DAGs
"""
import json
from typing import List

import requests

class PauseUnpauseDags:
    def __init__(
        self,
        api_key: str,
        release_name: str,
        domain: str,
    ) -> None:
        self.api_key = api_key
        self.url = f"https://deployments.{domain}/{release_name}/airflow/api/v1/dags"

    def _get_all_dag_ids(self) -> List[str]:
        dags = requests.get(
            f'{self.url}?limit=999999',
            headers={'Authorization': self.api_key}
        )
        dags = dags.json()['dags']
        return [dag['dag_id'] for dag in dags]

    def pause_dags(self, dag_ids: List[str] = []) -> None:
        if not dag_ids:
            dag_ids = self._get_all_dag_ids()
        for dag_id in dag_ids:
            response = requests.patch(
                url=f"{self.url}/{dag_id}",
                headers={'Content-Type': 'application/json', 'Authorization': self.api_key},
                data=json.dumps({"is_paused": True})
            )

    def unpause_dags(self, dag_ids: List[str] = []) -> None:
        if not dag_ids:
            dag_ids = self._get_all_dag_ids()
        for dag_id in dag_ids:
            response = requests.patch(
                url=f"{self.url}/{dag_id}",
                headers={'Content-Type': 'application/json', 'Authorization': self.api_key},
                data=json.dumps({"is_paused": False})
            )

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--action",
        choices=["pause", "unpause"],
        help="Choose to pause or unpause a list of DAGs. Choices: ['pause', 'unpause']",
        required=True
    )
    parser.add_argument(
        "--api_key",
        help="Astronomer service account API key",
        required=True
    )
    parser.add_argument(
        "--release_name",
        help="The release name of the Astronomer deployment whose DAGs you wish to pause or unpause",
        required=True
    )
    parser.add_argument(
        "--domain",
        help="The base domain for your Astronomer installation, e.g. `astronomer.mycompany.com`",
        required=True
    )
    parser.add_argument(
        "--dag_ids",
        nargs="+",
        help="Optional: Space separated list of DAG IDs that you'd like to pause or unpause, e.g. dag1 dag2 dag3"
    )
    args = parser.parse_args()
    if not args.dag_ids:
        args.dag_ids = []

    pause_unpause_dags = PauseUnpauseDags(
        api_key=args.api_key,
        release_name=args.release_name,
        domain=args.domain
    )
    if args.action == "pause":
        print("Pausing DAG(s)...")
        pause_unpause_dags.pause_dags(args.dag_ids)
    else:
        print("Unpausing DAG(s)...")
        pause_unpause_dags.unpause_dags(args.dag_ids)

This script can be used as follows:

# pause all dags
# to unpause, just change the action arg to "unpause"
python3 pause_unpause_dags_astronomer_software.py \
  --action pause \
  --api_key <api key> \
  --release_name <release name> \
  --domain <astronomer software base domain>

# pause list of dags
# to unpause, just change the action arg to "unpause"
python3 pause_unpause_dags_astronomer_software.py \
  --action pause \
  --api_key <api key> \
  --release_name <release name> \
  --domain <astronomer software base domain> \
  --dag_ids dag_id_1 dag_id_2 dag_id_3

Custom Airflow Operator

from airflow.models.baseoperator import BaseOperator
from airflow.settings import Session
from airflow.models import DagModel
from enum import Enum
from typing import Union, List
from airflow.exceptions import AirflowBadRequest


class Operation(Enum):
    GET = 0
    PAUSE = 1


class PauseDagsOperator(BaseOperator):
    # required to be set in execute
    context = None

    def __init__(self, operation: Operation, pause: Union[bool, None] = None, dag_list: List[DagModel] = [], *args, **kwargs):
        super().__init__(*args, **kwargs)
        # do some checks to make sure that we get the right parameters
        if operation == Operation.GET and (pause is not None or len(dag_list) != 0):
            raise AirflowBadRequest(
                "If GET operation is desire, pause and dag_list parameters should not be set"
            )
        elif operation == Operation.PAUSE and (pause is None or len(dag_list) == 0):
            raise AirflowBadRequest(
                "If PAUSE or UNPAUSE operation is desired, pause and dag_list parameters should be set"
            )
        self.template_fields = ("dag_list",)
    
        self.operation = operation
        self.pause = pause
        self.dag_list = dag_list

    def _get_paused_dags_list(self):
        self.context["ti"].xcom_push(key="paused_dags", value=Session().query(DagModel).filter(
            DagModel.is_active,
            ~DagModel.is_paused,
            DagModel.dag_id != self.context["dag"].dag_id
        ).all())
    
    def _pause_dags_in_list(self):
        for dag in self.dag_list:
            dag.set_is_paused(self.pause)

    def execute(self, context):
        self.context = context
        operations_map = {
            Operation.GET: self._get_paused_dags_list,
            Operation.PAUSE: self._pause_dags_in_list,
        }
        operations_map[self.operation]()

The above operator can be used in a DAG as follows:

from datetime import datetime
from time import sleep
from airflow.operators.python import PythonOperator
from airflow.models import DAG
from plugins.pause_controller import Operation, PauseDagsOperator


default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
}


with DAG(
    dag_id="test_pausing_dags",
    default_args=default_args,
    catchup=False,
    schedule_interval=None,
    start_date=datetime(1970, 1, 1),
    render_template_as_native_obj=True,
) as dag:
    get_unpaused_dags = PauseDagsOperator(
        task_id="get_unpaused_dags",
        operation=Operation.GET
    )

    pause_dags = PauseDagsOperator(
        task_id="pause_running_dags",
        operation=Operation.PAUSE,
        pause=True,
        dag_list="{{ ti.xcom_pull(task_ids='get_unpaused_dags', key='paused_dags') }}"
    )

    just_sleep = PythonOperator(
        task_id="just_sleep_and_wait",
        python_callable=lambda: sleep(15)
    )

    unpause_the_dags_that_were_paused = PauseDagsOperator(
        task_id="unpause_the_dags_that_were_paused",
        operation=Operation.PAUSE,
        pause=False,
        dag_list="{{ ti.xcom_pull(task_ids='get_unpaused_dags', key='paused_dags') }}"
    )

     get_unpaused_dags >> pause_dags >> just_sleep >> unpause_the_dags_that_were_paused