Introducing simplified interplay with the Airflow REST API in Amazon MWAA

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

At the moment, 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 techniques, and lets you work together along with your workflows programmatically. The Airflow REST API facilitates a variety of use instances, from centralizing and automating administrative duties to constructing event-driven, data-aware information pipelines.

On this put up, we focus on the enhancement and current a number of use instances that the enhancement unlocks in your Amazon MWAA setting.

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 comparable to invoking Directed Acyclic Graphs (DAGs), checking activity 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 in the present day, Amazon MWAA offered the muse 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 enhancing 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 performing 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 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 performing as an middleman, the improved API delivers Airflow REST API execution outcomes on to the consumer, lowering complexity and enhancing total usability.
• Automated administration – The simplified REST API entry permits automating numerous administrative and administration duties, comparable 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 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 permits the Amazon MWAA setting to handle the incoming workload and scale sources accordingly, enhancing the general reliability and effectivity of event-driven pipelines.

Within the following sections, we show methods to use the improved API in numerous use instances.

Tips on how to use the improved Amazon MWAA API

The next code snippet reveals the final request format for the improved REST API:

POST /restapi/Title HTTP/1.1
Content material-type: software/json

{
    Title: String,
    Technique: 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 Technique to make use of are the required parameters, whereas QueryParameters and Physique are non-compulsory and can be utilized as wanted within the API calls.

The next code snippet reveals the final 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 methods 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 consumer
mwaa_client = boto3.consumer("mwaa")

# Name the improved REST API utilizing boto3 consumer
# Utilizing QueryParameters, you'll be able to selectively specify the sphere to be up to date
# With out QueryParameters, all fields can be up to date
response = mwaa_client.invoke_rest_api(
    Title="",
    Technique="PATCH",
    Path=f"/variables/",
    Physique={
        "key": "",
        "worth": "",
        "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 consumer ought to have an AWS Identification and Entry Administration (IAM) principal of airflow:InvokeRestAPI hooked up to name the requisite setting. The permission might be scoped to particular Airflow roles (Admin, Op, Consumer, Viewer, or Public) to regulate entry ranges.

This easy but highly effective REST API helps numerous use instances in your Amazon MWAA environments. Let’s assessment some vital ones within the subsequent sections.

Automate administration and administration duties

Previous to this launch, to automate configurations and setup of sources comparable 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 may simplify this automation with the brand new enhanced REST API help.

For this instance, let’s assume you wish to automate sustaining your Amazon MWAA setting variables. It’s essential carry out API operations comparable to create, learn, replace, and delete on Airflow variables to realize this activity. The next is a straightforward 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 non-compulsory MWAA boto3 consumer
    def __init__(self, env_name, mwaa_client=None):
        self.env_name = env_name
        self.consumer = mwaa_client or boto3.consumer("mwaa")

    # Checklist all variables within the MWAA setting
    def record(self):
        response = self.consumer.invoke_rest_api(
            Title=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.consumer.invoke_rest_api(
            Title=self.env_name,
            Technique="GET",
            Path=f"/variables/{key}"
        )

        return response['RestApiResponse']

    # Create a brand new variable with key, worth, and non-compulsory description
    def create(self, key, worth, description=None):
        response = self.consumer.invoke_rest_api(
            Title=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.consumer.invoke_rest_api(
            Title=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.consumer.invoke_rest_api(
            Title=self.env_name,
            Technique="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="Take a look at worth",
        description="Take a look at description"
    )
    print(response)

    print("nListing all variables ...")
    variables = consumer.record()
    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="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 got 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 take a look at variable ...
{'description': 'Take a look at description', 'key': 'take a look at', 'worth': 'Take a look at worth'}

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

Getting the take a look at variable ...
{'key': 'take a look at', 'worth': 'Take a look at 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': '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', 'kind': '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 instances.

Construct event-driven information pipelines

The Airflow neighborhood has been actively innovating to reinforce the platform’s information consciousness, enabling you to construct extra dynamic and responsive workflows. Once we introduced help for model 2.9.2 in Amazon MWAA, we launched capabilities that enable pipelines to react to modifications in datasets, each inside Airflow environments and in exterior techniques. 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 you’ll want 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", kind="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 consumer that may invoke this DAG, supplying the bash command of your alternative 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 non-compulsory 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,
            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.consumer.invoke_rest_api(
            Title=self.env_name,
            Technique="POST",
            Path=f"/dags/{dag_id}/dagRuns",
            Physique=physique
        )
        return response['RestApiResponse']

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

The next snippet reveals a pattern run of the script:

$python mwaa_dag_run_client.py
Triggering a dag run ...
{'conf': {'command': "echo 'Hey 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 activity reveals 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 - 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 - Hey 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} ▶ Submit activity execution logs

You may additional increase this instance to extra helpful event-driven architectures. Let’s increase 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 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 numerous information in a brief timeframe.

One method 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 price limiting the API requests to the online server. You may 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 entails 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 setting 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 an inventory 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 = record(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"Working 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"Working 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 places of the newly added information within the S3 information lake. Upon receiving dataset occasions, the Amazon MWAA setting batches the dataset occasions based mostly on the load and schedules the DAG to deal with them appropriately. The modified structure to help the data-aware scheduling is introduced beneath.

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 non-compulsory 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,
            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

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
    outcome = consumer.create_event(
        dataset_uri="datalake",
        further={"uri": s3_file_uri}
    )

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

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 reveals the execution log of the extract activity from the pipeline. The log reveals 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 - Working information extraction for s3://example-bucket/path/to/file1.csv ...
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Working information extraction for s3://example-bucket/path/to/file2.csv ...
[2024-10-21, 16:47:15 UTC] {logging_mixin.py:190} INFO - Working information extraction for s3://example-bucket/path/to/file3.csv ...
[2024-10-21, 16:47:15 UTC] {python.py:240} INFO - Performed. Returned worth was: None
[2024-10-21, 16:47:16 UTC] {taskinstance.py:340} ▶ Submit activity execution logs

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

Concerns

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

• IAM permissions – Make sure 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, Consumer, Viewer, or Public).
• Error dealing with – Implement strong error dealing with mechanisms to deal with numerous 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 appropriate with any future modifications.
• Safety and compliance – Adhere to your group’s safety and compliance necessities when integrating exterior techniques 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 better, 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 instances, from centralizing and automating administrative duties, enhancing total usability, to constructing event-driven, data-aware information pipelines.

As you discover this new function, contemplate the assorted use instances and greatest practices outlined on this put up. Through the use of the brand new InvokeRestApi, you’ll be able to streamline your information administration processes, improve operational effectivity, and drive better worth out of your data-driven techniques.

In regards to the Authors

Chandan Rupakheti is a Senior Options Architect at AWS. His most important focus at AWS lies within the intersection of analytics, serverless, and AdTech providers. He’s a passionate technical chief, researcher, and mentor with a knack for constructing progressive options within the cloud. Outdoors of his skilled life, he loves spending time together with his household and mates, and listening to and taking part in music.

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