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, decreasing operational overhead whereas growing safety and resilience.

Right now, we’re excited to announce an enhancement to the Amazon MWAA integration with the Airflow REST API. This enchancment streamlines the flexibility to entry and handle your Airflow environments and their integration with exterior methods, and means that you can work together together 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 knowledge pipelines.

On this publish, we focus on the enhancement and current a number of use circumstances that the enhancement unlocks to your Amazon MWAA setting.

Airflow REST API

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

Earlier than at this time, Amazon MWAA offered the inspiration for interacting with the Airflow REST API. Although useful, 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 decreasing complexity and bettering total usability.

Enhancement overview

The brand new InvokeRestApi functionality means that you can 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 setting title and API request payload as inputs.

Integrating with the Airflow REST API by way of the improved Amazon MWAA API gives 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 easy 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 consumer, decreasing complexity and bettering total usability.
• Automated administration – The simplified REST API entry permits automating varied administrative and administration duties, similar to 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 primarily based on these occasions. This helps the rising emphasis on event-driven knowledge pipelines.
• Information-aware scheduling – Utilizing the dataset-based scheduling function in Airflow, the improved API permits the Amazon MWAA setting to handle the incoming workload and scale sources accordingly, bettering the general reliability and effectivity of event-driven pipelines.

Within the following sections, we display easy methods 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/Title HTTP/1.1
Content material-type: utility/json

{
    Title: String,
    Methodology: String,
    Path: String,
    QueryParameters: Json,
    Physique: Json
}

The Title of the Amazon MWAA setting, the Path of the Airflow REST API endpoint to be known as, and the HTTP Methodology to make use of are the required parameters, whereas QueryParameters and Physique are elective 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 accommodates the response payload from the Airflow REST API.

The next pattern code snippet showcases easy methods to replace the outline discipline 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 consumer
mwaa_client = boto3.consumer("mwaa")

# Name the improved REST API utilizing boto3 consumer
# Utilizing QueryParameters, you may selectively specify the sphere to be up to date
# With out QueryParameters, all fields shall be up to date
response = mwaa_client.invoke_rest_api(
    Title="",
    Methodology="PATCH",
    Path=f"/variables/",
    Physique={
        "key": "",
        "worth": "",
        "description": ""
    },
    QueryParameters={
        "update_mask": ["description"]
    }
)

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

To make the invoke_rest_api SDK name, the calling consumer ought to have an AWS Identification and Entry Administration (IAM) principal of airflow:InvokeRestAPI connected to name the requisite setting. The permission will be scoped to particular Airflow roles (Admin, Op, Consumer, Viewer, or Public) to manage entry ranges.

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

Automate administration and administration duties

Previous to this launch, to automate configurations and setup of sources similar to variables, connections, slot swimming pools, and extra, you needed to develop a prolonged boilerplate code to make API requests to the Amazon MWAA net servers. You needed to deal with the cookie and session administration within the course of. You possibly can 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 setting variables. You might want to carry out API operations similar to create, learn, replace, and delete on Airflow variables to attain this process. The next is an easy Python consumer to take action (mwaa_variables_client.py):

import boto3

# Consumer for managing MWAA setting variables
class MWAAVariablesClient:
    # Initialize the consumer with setting title and elective MWAA boto3 consumer
    def __init__(self, env_name, mwaa_client=None):
        self.env_name = env_name
        self.consumer = mwaa_client or boto3.consumer("mwaa")

    # Record all variables within the MWAA setting
    def listing(self):
        response = self.consumer.invoke_rest_api(
            Title=self.env_name,
            Methodology="GET",
            Path="/variables"
        )

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

    # Get a particular variable by key
    def get(self, key):
        response = self.consumer.invoke_rest_api(
            Title=self.env_name,
            Methodology="GET",
            Path=f"/variables/{key}"
        )

        return response['RestApiResponse']

    # Create a brand new variable with key, worth, and elective description
    def create(self, key, worth, description=None):
        response = self.consumer.invoke_rest_api(
            Title=self.env_name,
            Methodology="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.consumer.invoke_rest_api(
            Title=self.env_name,
            Methodology="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.consumer.invoke_rest_api(
            Title=self.env_name,
            Methodology="DELETE",
            Path=f"/variables/{key}"
        )
        return response['RestApiStatusCode']

if __name__ == "__main__":
    consumer = MWAAVariablesClient("")

    print("nCreating a take a look at variable ...")
    response = consumer.create(
        key="take a look at",
        worth="Check worth",
        description="Check description"
    )
    print(response)

    print("nListing all variables ...")
    variables = consumer.listing()
    print(variables)

    print("nGetting the take a look at variable ...")
    response = consumer.get("take a look at")
    print(response)

    print("nUpdating the worth and outline of take a look at variable ...")
    response = consumer.replace(
        key="take a look at",
        worth="Up to date Worth",
        description="Up to date description"
    )
    print(response)

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

    print("nDeleting the take a look at variable ...")
    response_code = consumer.delete("take a look at")
    print(f"Response code: {response_code}")

    print("nFinally, getting the deleted take a look at variable ...")
    strive:
        response = consumer.get("take a look at")
        print(response)
    besides Exception as e:
        print(e.response["RestApiResponse"])

Assuming that you’ve configured your terminal with acceptable AWS credentials, you may run the previous Python script to attain the next outcomes:

$python mwaa_variables_client.py 

Making a take a look at variable ...
{'description': 'Check description', 'key': 'take a look at', 'worth': 'Check worth'}

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

Getting the take a look at variable ...
{'key': 'take a look at', 'worth': 'Check worth'}

Updating the worth and outline of take a look at variable ...
{'description': 'Up to date description', 'key': 'take a look at', 'worth': 'Up to date Worth'}

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

Deleting the take a look at variable ...
Response code: 204

Lastly, getting the deleted take a look at 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 knowledge pipelines

The Airflow neighborhood has been actively innovating to reinforce the platform’s knowledge consciousness, enabling you to construct extra dynamic and responsive workflows. Once 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 methods. The brand new simplified integration with the Airflow REST API makes the implementation of data-driven pipelines extra easy.

Take into account a use case the place it’s essential 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 means that you can 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 may create a consumer that may invoke this DAG, supplying the bash command of your selection as follows (mwaa_dag_run_client.py):

import boto3

# Consumer for triggering DAG runs in Amazon MWAA
class MWAADagRunClient:
    # Initialize the consumer with MWAA setting title and elective MWAA boto3 consumer
    def __init__(self, env_name, mwaa_client=None):
        self.env_name = env_name
        self.consumer = mwaa_client or boto3.consumer("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,
            notice=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 notice:
            physique["note"] = notice
        physique["conf"] = conf or {}            

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

if __name__ == "__main__":
    consumer = MWAADagRunClient("")
	
    print("nTriggering a dag run ...")
    consequence = consumer.trigger_run(
        dag_id="any_bash_command", 
        conf={
            "command": "echo 'Good day from exterior set off!'"
        }
    )
    print(consequence)

The next snippet exhibits a pattern run of the script:

$python mwaa_dag_run_client.py
Triggering a dag run ...
{'conf': {'command': "echo 'Good day 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': 'guide', 'state': 'queued'}

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

[2024-10-21, 16:56:12 UTC] {local_task_job_runner.py:123} ▶ Pre process 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 - Working 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 - Good day 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} ▶ Put up process execution logs

You possibly can additional increase this instance to extra helpful event-driven architectures. Let’s increase the use case to run your knowledge 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 knowledge lake. The next diagram illustrates one architectural method.

Within the context of invoking a DAG by way of an exterior enter, the AWS Lambda perform would don’t have any data of how busy the Amazon MWAA net server is, doubtlessly resulting in the perform overwhelming the Amazon MWAA net server by processing a lot of information in a brief timeframe.

One technique to regulate the file processing throughput can be to introduce an Amazon Easy Queue Service (Amazon SQS) queue between the S3 bucket and the Lambda perform, which might help with price limiting the API requests to the net server. You possibly can 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 setting to run the DAGs.

Along with the SQS queue, to assist afford the Amazon MWAA setting handle the load natively, you should 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 manner, the Amazon MWAA setting can successfully batch the dataset occasions and scale sources primarily based on the load. Let’s discover this method in additional element.

Configure data-aware scheduling

Take into account 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 primarily based on the datalake dataset occasion.
"""
from airflow import DAG, Dataset
from airflow.decorators import process

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 = listing(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
):
    @process
    def extract(triggering_dataset_events=None):
        sources = get_resources("datalake", triggering_dataset_events)
        for useful resource in sources:
            print(f"Working knowledge extraction for {useful resource} ...")

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

    @process()
    def load(triggering_dataset_events=None):
        sources = get_resources("datalake", triggering_dataset_events)
        for useful resource in sources:
            print(f"Loading finalized knowledge 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 accommodates the areas of the newly added information within the S3 knowledge lake. Upon receiving dataset occasions, the Amazon MWAA setting batches the dataset occasions primarily based on the load and schedules the DAG to deal with them appropriately. The modified structure to assist the data-aware scheduling is introduced under.

The next is a simplified consumer that may create a dataset occasion by way of the improved REST API (mwaa_dataset_client.py):

import boto3

# Consumer for interacting with MWAA datasets
class MWAADatasetClient:
    # Initialize the consumer with setting title and elective MWAA boto3 consumer
    def __init__(self, env_name, mwaa_client=None):
        self.env_name = env_name
        self.consumer = mwaa_client or boto3.consumer("mwaa")

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

        response = self.consumer.invoke_rest_api(
            Title=self.env_name,
            Methodology="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

setting = os.environ["MWAA_ENV_NAME"]
consumer = MWAADatasetClient(setting)

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

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

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

As new information 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 information. The Amazon MWAA setting 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 process 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 process execution logs
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Working knowledge extraction for s3://example-bucket/path/to/file1.csv ...
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Working knowledge extraction for s3://example-bucket/path/to/file2.csv ...
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Working knowledge extraction for s3://example-bucket/path/to/file3.csv ...
[2024-10-21, 16:47:15 UTC] {python.py:240} INFO - Accomplished. Returned worth was: None
[2024-10-21, 16:47:16 UTC] {taskinstance.py:340} ▶ Put up process execution logs

On this manner, you should 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 essential to think about the next:

• IAM permissions – Be sure that the IAM principal making the invoke_rest_api SDK name has the airflow:InvokeRestAPI permission on the Amazon MWAA useful resource. To manage entry ranges, the permission will be scoped to particular Airflow roles (Admin, Op, Consumer, Viewer, or Public).
• Error dealing with – Implement sturdy 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 knowledge pipelines.
• Versioning and compatibility – Monitor the versioning of the Airflow REST API and the Amazon MWAA service to ensure your integrations stay appropriate with any future adjustments.
• Safety and compliance – Adhere to your group’s safety and compliance necessities when integrating exterior methods with Amazon MWAA and dealing with delicate knowledge.

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

Conclusion

The improved integration between Amazon MWAA and the Airflow REST API represents a major 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 knowledge pipelines.

As you discover this new function, take into account the assorted use circumstances and finest practices outlined on this publish. Through the use of the brand new InvokeRestApi, you may streamline your knowledge administration processes, improve operational effectivity, and drive larger worth out of your data-driven methods.

In regards to the Authors

Chandan Rupakheti is a Senior Options Architect at AWS. His foremost 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 progressive options within the cloud. Exterior of his skilled life, he loves spending time along with his household and buddies, and listening to and enjoying music.

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