Introducing simplified interplay with the Airflow REST API in Amazon MWAA

Amazon Managed Workflows for Apache Airflow (Amazon MWAA) is a totally managed service that builds upon Apache Airflow, providing its advantages whereas eliminating the necessity so that you can arrange, function, and keep the underlying infrastructure, lowering operational overhead whereas growing safety and resilience.

At present, we’re excited to announce an enhancement to the Amazon MWAA integration with the Airflow REST API. This enchancment streamlines the power to entry and handle your Airflow environments and their integration with exterior programs, and lets you work together along with your workflows programmatically. The Airflow REST API facilitates a variety of use circumstances, from centralizing and automating administrative duties to constructing event-driven, data-aware information pipelines.

On this publish, we talk about the enhancement and current a number of use circumstances that the enhancement unlocks to your Amazon MWAA atmosphere.

Airflow REST API

The Airflow REST API is a programmatic interface that lets you work together with Airflow’s core functionalities. It’s a set of HTTP endpoints to carry out operations reminiscent of invoking Directed Acyclic Graphs (DAGs), checking activity statuses, retrieving metadata about workflows, managing connections and variables, and even initiating dataset-related occasions, with out immediately accessing the Airflow internet interface or command line instruments.

Earlier than immediately, Amazon MWAA offered the inspiration for interacting with the Airflow REST API. Although purposeful, the method of acquiring and managing entry tokens and session cookies added complexity to the workflow. Amazon MWAA now helps a simplified mechanism for interacting with the Airflow REST API utilizing AWS credentials, considerably lowering complexity and bettering total usability.

Enhancement overview

The brand new InvokeRestApi functionality lets you run Airflow REST API requests with a legitimate SigV4 signature utilizing your present AWS credentials. This function is now accessible to all Amazon MWAA environments (2.4.3+) in supported Amazon MWAA AWS Areas. By appearing as an middleman, this REST API processes requests on behalf of customers, requiring solely the atmosphere identify and API request payload as inputs.

Integrating with the Airflow REST API via the improved Amazon MWAA API supplies a number of key advantages:

• Simplified integration – The brand new InvokeRestApi functionality in Amazon MWAA removes the complexity of managing entry tokens and session cookies, making it simple to work together with the Airflow REST API.
• Improved usability – By appearing as an middleman, the improved API delivers Airflow REST API execution outcomes on to the shopper, lowering complexity and bettering total usability.
• Automated administration – The simplified REST API entry allows automating varied administrative and administration duties, reminiscent of managing Airflow variables, connections, slot swimming pools, and extra.
• Occasion-driven architectures – The improved API facilitates seamless integration with exterior occasions, enabling the triggering of Airflow DAGs based mostly on these occasions. This helps the rising emphasis on event-driven information pipelines.
• Information-aware scheduling – Utilizing the dataset-based scheduling function in Airflow, the improved API allows the Amazon MWAA atmosphere to handle the incoming workload and scale sources accordingly, bettering the general reliability and effectivity of event-driven pipelines.

Within the following sections, we reveal learn how to use the improved API in varied use circumstances.

Methods to use the improved Amazon MWAA API

The next code snippet exhibits the overall request format for the improved REST API:

POST /restapi/Identify HTTP/1.1
Content material-type: utility/json

{
    Identify: String,
    Technique: String,
    Path: String,
    QueryParameters: Json,
    Physique: Json
}

The Identify of the Amazon MWAA atmosphere, the Path of the Airflow REST API endpoint to be known as, and the HTTP Technique to make use of are the required parameters, whereas QueryParameters and Physique are non-obligatory and can be utilized as wanted within the API calls.

The next code snippet exhibits the overall response format:

{
    RestApiStatusCode: Quantity,
    RestApiResponse: Json
}

The RestApiStatusCode represents the HTTP standing code returned by the Airflow REST API name, and the RestApiResponse incorporates the response payload from the Airflow REST API.

The next pattern code snippet showcases learn how to replace the outline subject of an Airflow variable utilizing the improved integration. The decision makes use of the AWS Python SDK to invoke the Airflow REST API for the duty.

import boto3

# Create a boto3 shopper
mwaa_client = boto3.shopper("mwaa")

# Name the improved REST API utilizing boto3 shopper
# Utilizing QueryParameters, you'll be able to selectively specify the sector to be up to date
# With out QueryParameters, all fields shall be up to date
response = mwaa_client.invoke_rest_api(
    Identify="<your-environment-name>",
    Technique="PATCH",
    Path=f"/variables/<your-variable-key>",
    Physique={
        "key": "<key>",
        "worth": "<worth>",
        "description": "<description>"
    },
    QueryParameters={
        "update_mask": ["description"]
    }
)

# Entry the outputs of the REST name
status_code = response["RestApiStatusCode"]
outcome = response['RestApiResponse']

To make the invoke_rest_api SDK name, the calling shopper ought to have an AWS Id and Entry Administration (IAM) principal of airflow:InvokeRestAPI connected to name the requisite atmosphere. The permission might be scoped to particular Airflow roles (Admin, Op, Person, Viewer, or Public) to regulate entry ranges.

This easy but highly effective REST API helps varied use circumstances to your Amazon MWAA environments. Let’s overview some vital ones within the subsequent sections.

Automate administration and administration duties

Previous to this launch, to automate configurations and setup of sources reminiscent of variables, connections, slot swimming pools, and extra, you needed to develop a prolonged boilerplate code to make API requests to the Amazon MWAA internet servers. You needed to deal with the cookie and session administration within the course of. You’ll be able to simplify this automation with the brand new enhanced REST API assist.

For this instance, let’s assume you wish to automate sustaining your Amazon MWAA atmosphere variables. It is advisable to carry out API operations reminiscent of create, learn, replace, and delete on Airflow variables to realize this activity. The next is an easy Python shopper to take action (mwaa_variables_client.py):

import boto3

# Shopper for managing MWAA atmosphere variables
class MWAAVariablesClient:
    # Initialize the shopper with atmosphere identify and non-obligatory MWAA boto3 shopper
    def __init__(self, env_name, mwaa_client=None):
        self.env_name = env_name
        self.shopper = mwaa_client or boto3.shopper("mwaa")

    # Record all variables within the MWAA atmosphere
    def checklist(self):
        response = self.shopper.invoke_rest_api(
            Identify=self.env_name,
            Technique="GET",
            Path="/variables"
        )

        output = response['RestApiResponse']['variables']
        return output

    # Get a selected variable by key
    def get(self, key):
        response = self.shopper.invoke_rest_api(
            Identify=self.env_name,
            Technique="GET",
            Path=f"/variables/{key}"
        )

        return response['RestApiResponse']

    # Create a brand new variable with key, worth, and non-obligatory description
    def create(self, key, worth, description=None):
        response = self.shopper.invoke_rest_api(
            Identify=self.env_name,
            Technique="POST",
            Path="/variables",
            Physique={
                "key": key,
                "worth": worth,
                "description": description
            }
        )

        return response['RestApiResponse']
    
    # Replace an present variable's worth and outline
    def replace(self, key, worth, description, query_parameters=None):
        response = self.shopper.invoke_rest_api(
            Identify=self.env_name,
            Technique="PATCH",
            Path=f"/variables/{key}",
            Physique={
                "key": key,
                "worth": worth,
                "description": description
            },
            QueryParameters=query_parameters
        )

        return response['RestApiResponse']

    # Delete a variable by key
    def delete(self, key):
        response = self.shopper.invoke_rest_api(
            Identify=self.env_name,
            Technique="DELETE",
            Path=f"/variables/{key}"
        )
        return response['RestApiStatusCode']

if __name__ == "__main__":
    shopper = MWAAVariablesClient("<your-mwaa-environment-name>")

    print("nCreating a check variable ...")
    response = shopper.create(
        key="check",
        worth="Take a look at worth",
        description="Take a look at description"
    )
    print(response)

    print("nListing all variables ...")
    variables = shopper.checklist()
    print(variables)

    print("nGetting the check variable ...")
    response = shopper.get("check")
    print(response)

    print("nUpdating the worth and outline of check variable ...")
    response = shopper.replace(
        key="check",
        worth="Up to date Worth",
        description="Up to date description"
    )
    print(response)

    print("nUpdating solely description of check variable ...")
    response = shopper.replace(
        key="check", 
        worth="Up to date Worth", 
        description="Yet one more up to date description", 
        query_parameters={ "update_mask": ["description"] }
    )
    print(response)

    print("nDeleting the check variable ...")
    response_code = shopper.delete("check")
    print(f"Response code: {response_code}")

    print("nFinally, getting the deleted check variable ...")
    strive:
        response = shopper.get("check")
        print(response)
    besides Exception as e:
        print(e.response["RestApiResponse"])

Assuming that you’ve configured your terminal with acceptable AWS credentials, you’ll be able to run the previous Python script to realize the next outcomes:

$python mwaa_variables_client.py 

Making a check variable ...
{'description': 'Take a look at description', 'key': 'check', 'worth': 'Take a look at worth'}

Itemizing all variables ...
[{'key': 'test', 'value': 'Test value'}]

Getting the check variable ...
{'key': 'check', 'worth': 'Take a look at worth'}

Updating the worth and outline of check variable ...
{'description': 'Up to date description', 'key': 'check', 'worth': 'Up to date Worth'}

Updating solely description of check variable ...
{'description': 'Yet one more up to date description', 'key': 'check', 'worth': 'Up to date Worth'}

Deleting the check variable ...
Response code: 204

Lastly, getting the deleted check variable ...
{'element': 'Variable doesn't exist', 'standing': 404, 'title': 'Variable not discovered', 'sort': 'https://airflow.apache.org/docs/apache-airflow/2.8.1/stable-rest-api-ref.html#part/Errors/NotFound'}

Let’s additional discover different helpful use circumstances.

Construct event-driven information pipelines

The Airflow group has been actively innovating to reinforce the platform’s information consciousness, enabling you to construct extra dynamic and responsive workflows. After we introduced assist for model 2.9.2 in Amazon MWAA, we launched capabilities that permit pipelines to react to adjustments in datasets, each inside Airflow environments and in exterior programs. The brand new simplified integration with the Airflow REST API makes the implementation of data-driven pipelines extra simple.

Think about a use case the place it’s good to run a pipeline that makes use of enter from an exterior occasion. The next pattern DAG runs a bash command equipped as a parameter (any_bash_command.py):

"""
This DAG lets you execute a bash command equipped as a parameter to the DAG.
The command is handed as a parameter known as `command` within the DAG configuration.
"""

from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from airflow.fashions.param import Param

from datetime import datetime

with DAG(
    dag_id="any_bash_command", 
    schedule=None, 
    start_date=datetime(2022, 1, 1), 
    catchup=False,
    params={
        "command": Param("env", sort="string")
    },
) as dag:
    cli_command = BashOperator(
        task_id="triggered_bash_command",
        bash_command="{{ dag_run.conf['command'] }}"
    )

With the assistance of the improved REST API, you’ll be able to create a shopper that may invoke this DAG, supplying the bash command of your selection as follows (mwaa_dag_run_client.py):

import boto3

# Shopper for triggering DAG runs in Amazon MWAA
class MWAADagRunClient:
    # Initialize the shopper with MWAA atmosphere identify and non-obligatory MWAA boto3 shopper
    def __init__(self, env_name, mwaa_client=None):
        self.env_name = env_name
        self.shopper = mwaa_client or boto3.shopper("mwaa")

    # Set off a DAG run with specified parameters
    def trigger_run(self, 
            dag_id, 
            dag_run_id=None,
            logical_date=None,
            data_interval_start=None,
            data_interval_end=None,
            be aware=None,
            conf=None,
    ):
        physique = {}
        if dag_run_id:
            physique["dag_run_id"] = dag_run_id
        if logical_date:
            physique["logical_date"] = logical_date
        if data_interval_start:
            physique["data_interval_start"] = data_interval_start
        if data_interval_end:
            physique["data_interval_end"] = data_interval_end
        if be aware:
            physique["note"] = be aware
        physique["conf"] = conf or {}            

        response = self.shopper.invoke_rest_api(
            Identify=self.env_name,
            Technique="POST",
            Path=f"/dags/{dag_id}/dagRuns",
            Physique=physique
        )
        return response['RestApiResponse']

if __name__ == "__main__":
    shopper = MWAADagRunClient("<your-mwaa-environment-name>")
	
    print("nTriggering a dag run ...")
    outcome = shopper.trigger_run(
        dag_id="any_bash_command", 
        conf={
            "command": "echo 'Howdy from exterior set off!'"
        }
    )
    print(outcome)

The next snippet exhibits a pattern run of the script:

$python mwaa_dag_run_client.py
Triggering a dag run ...
{'conf': {'command': "echo 'Howdy from exterior set off!'"}, 'dag_id': 'any_bash_command', 'dag_run_id': 'manual__2024-10-21T16:56:09.852908+00:00', 'data_interval_end': '2024-10-21T16:56:09.852908+00:00', 'data_interval_start': '2024-10-21T16:56:09.852908+00:00', 'execution_date': '2024-10-21T16:56:09.852908+00:00', 'external_trigger': True, 'logical_date': '2024-10-21T16:56:09.852908+00:00', 'run_type': 'handbook', 'state': 'queued'}

On the Airflow UI, the trigger_bash_command activity exhibits the next execution log:

[2024-10-21, 16:56:12 UTC] {local_task_job_runner.py:123} ▶ Pre activity execution logs
[2024-10-21, 16:56:12 UTC] {subprocess.py:63} INFO - Tmp dir root location: /tmp
[2024-10-21, 16:56:12 UTC] {subprocess.py:75} INFO - Operating command: ['/usr/bin/bash', '-c', "echo 'Hello from external trigger!'"]
[2024-10-21, 16:56:12 UTC] {subprocess.py:86} INFO - Output:
[2024-10-21, 16:56:12 UTC] {subprocess.py:93} INFO - Howdy from exterior set off!
[2024-10-21, 16:56:12 UTC] {subprocess.py:97} INFO - Command exited with return code 0
[2024-10-21, 16:56:12 UTC] {taskinstance.py:340} ▶ Publish activity execution logs

You’ll be able to additional develop this instance to extra helpful event-driven architectures. Let’s develop the use case to run your information pipeline and carry out extract, remodel, and cargo (ETL) jobs when a brand new file lands in an Amazon Easy Storage Service (Amazon S3) bucket in your information lake. The next diagram illustrates one architectural method.

Within the context of invoking a DAG via an exterior enter, the AWS Lambda perform would haven’t any information of how busy the Amazon MWAA internet server is, probably resulting in the perform overwhelming the Amazon MWAA internet server by processing numerous recordsdata in a brief timeframe.

One solution to regulate the file processing throughput could be to introduce an Amazon Easy Queue Service (Amazon SQS) queue between the S3 bucket and the Lambda perform, which may also help with fee limiting the API requests to the net server. You’ll be able to obtain this by configuring most concurrency for Lambda for the SQS occasion supply. Nonetheless, the Lambda perform would nonetheless be unaware of the processing capability accessible within the Amazon MWAA atmosphere to run the DAGs.

Along with the SQS queue, to assist afford the Amazon MWAA atmosphere handle the load natively, you need to use the Airflow’s data-aware scheduling function utilizing datasets. This method includes utilizing the improved Amazon MWAA REST API to create dataset occasions, that are then utilized by the Airflow scheduler to schedule the DAG natively. This fashion, the Amazon MWAA atmosphere can successfully batch the dataset occasions and scale sources based mostly on the load. Let’s discover this method in additional element.

Configure data-aware scheduling

Think about the next DAG that showcases a framework for an ETL pipeline (data_aware_pipeline.py). It makes use of a dataset for scheduling.

"""
DAG to run the given ETL pipeline based mostly on the datalake dataset occasion.
"""
from airflow import DAG, Dataset
from airflow.decorators import activity

from datetime import datetime

# Create a dataset
datalake = Dataset("datalake")

# Return a listing of S3 file URIs from the equipped dataset occasions 
def get_resources(dataset_uri, triggering_dataset_events=None):
    occasions = triggering_dataset_events[dataset_uri] if triggering_dataset_events else []
    s3_uris = checklist(map(lambda e: e.further["uri"], occasions))
    return s3_uris

with DAG(
    dag_id="data_aware_pipeline",
    schedule=[datalake],
    start_date=datetime(2022, 1, 1), 
    catchup=False
):
    @activity
    def extract(triggering_dataset_events=None):
        sources = get_resources("datalake", triggering_dataset_events)
        for useful resource in sources:
            print(f"Operating information extraction for {useful resource} ...")

    @activity
    def remodel(triggering_dataset_events=None):
        sources = get_resources("datalake", triggering_dataset_events)
        for useful resource in sources:
            print(f"Operating information transformation for {useful resource} ...")

    @activity()
    def load(triggering_dataset_events=None):
        sources = get_resources("datalake", triggering_dataset_events)
        for useful resource in sources:
            print(f"Loading finalized information for {useful resource} ...")
    
    extract() >> remodel() >> load()

Within the previous code snippet, a Dataset object known as datalake is used to schedule the DAG. The get_resources perform extracts the further info that incorporates the areas of the newly added recordsdata within the S3 information lake. Upon receiving dataset occasions, the Amazon MWAA atmosphere batches the dataset occasions based mostly on the load and schedules the DAG to deal with them appropriately. The modified structure to assist the data-aware scheduling is offered under.

The next is a simplified shopper that may create a dataset occasion via the improved REST API (mwaa_dataset_client.py):

import boto3

# Shopper for interacting with MWAA datasets
class MWAADatasetClient:
    # Initialize the shopper with atmosphere identify and non-obligatory MWAA boto3 shopper
    def __init__(self, env_name, mwaa_client=None):
        self.env_name = env_name
        self.shopper = mwaa_client or boto3.shopper("mwaa")

    # Create a dataset occasion within the MWAA atmosphere
    def create_event(self, dataset_uri, further=None):
        physique = {
            "dataset_uri": dataset_uri
        }
        if further:
            physique["extra"] = further

        response = self.shopper.invoke_rest_api(
            Identify=self.env_name,
            Technique="POST",
            Path="/datasets/occasions",
            Physique=physique
        )
        return response['RestApiResponse']

The next is a code snippet for the Lambda perform within the previous structure to generate the dataset occasion, assuming the perform is configured to deal with one S3 PUT occasion at a time (dataset_event_lambda.py):

import os
import json

from mwaa_dataset_client import MWAADatasetClient

atmosphere = os.environ["MWAA_ENV_NAME"]
shopper = MWAADatasetClient(atmosphere)

def handler(occasion, context):
    # Extract S3 file URI from SQS document
    document = occasion["Records"][0]
    bucket = document["s3"]["bucket"]["name"]
    key = document["s3"]["object"]["key"]
    s3_file_uri = f"s3://{bucket}/{key}"

    # Create a dataset occasion
    outcome = shopper.create_event(
        dataset_uri="datalake",
        further={"uri": s3_file_uri}
    )

    return {
        "statusCode": 200,
        "physique": json.dumps(outcome)
    }

As new recordsdata get dropped into the S3 bucket, the Lambda perform will generate a dataset occasion per file, passing within the Amazon S3 location of the newly added recordsdata. The Amazon MWAA atmosphere will schedule the ETL pipeline upon receiving the dataset occasions. The next diagram illustrates a pattern run of the ETL pipeline on the Airflow UI.

The next snippet exhibits the execution log of the extract activity from the pipeline. The log exhibits how the Airflow scheduler batched three dataset occasions collectively to deal with the load.

[2024-10-21, 16:47:15 UTC] {local_task_job_runner.py:123} ▶ Pre activity execution logs
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Operating information extraction for s3://example-bucket/path/to/file1.csv ...
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Operating information extraction for s3://example-bucket/path/to/file2.csv ...
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Operating information extraction for s3://example-bucket/path/to/file3.csv ...
[2024-10-21, 16:47:15 UTC] {python.py:240} INFO - Executed. Returned worth was: None
[2024-10-21, 16:47:16 UTC] {taskinstance.py:340} ▶ Publish activity execution logs

On this approach, you need to use the improved REST API to create data-aware, event-driven pipelines.

Issues

When implementing options utilizing the improved Amazon MWAA REST API, it’s vital to contemplate the next:

• IAM permissions – Be certain that the IAM principal making the invoke_rest_api SDK name has the airflow:InvokeRestAPI permission on the Amazon MWAA useful resource. To regulate entry ranges, the permission might be scoped to particular Airflow roles (Admin, Op, Person, Viewer, or Public).
• Error dealing with – Implement strong error dealing with mechanisms to deal with varied HTTP standing codes and error responses from the Airflow REST API.
• Monitoring and logging – Arrange acceptable monitoring and logging to trace the efficiency and reliability of your API-based integrations and information pipelines.
• Versioning and compatibility – Monitor the versioning of the Airflow REST API and the Amazon MWAA service to ensure your integrations stay suitable with any future adjustments.
• Safety and compliance – Adhere to your group’s safety and compliance necessities when integrating exterior programs with Amazon MWAA and dealing with delicate information.

You can begin utilizing the simplified integration with the Airflow REST API in your Amazon MWAA environments with Airflow model 2.4.3 or higher, in all at present supported Areas.

Conclusion

The improved integration between Amazon MWAA and the Airflow REST API represents a big enchancment within the ease of interacting with Airflow’s core functionalities. This new functionality opens up a variety of use circumstances, from centralizing and automating administrative duties, bettering total usability, to constructing event-driven, data-aware information pipelines.

As you discover this new function, think about the assorted use circumstances and finest practices outlined on this publish. By utilizing the brand new InvokeRestApi, you’ll be able to streamline your information administration processes, improve operational effectivity, and drive higher worth out of your data-driven programs.

Concerning the Authors

Chandan Rupakheti is a Senior Options Architect at AWS. His primary focus at AWS lies within the intersection of analytics, serverless, and AdTech companies. He’s a passionate technical chief, researcher, and mentor with a knack for constructing modern options within the cloud. Outdoors of his skilled life, he loves spending time together with his household and buddies, and listening to and enjoying music.

Hernan Garcia is a Senior Options Architect at AWS based mostly out of Amsterdam. He has labored within the monetary companies business since 2018, specializing in utility modernization and supporting clients of their adoption of the cloud with a deal with serverless applied sciences.