Trigger Airflow DAG with REST API using Snowflake

In an earlier article I wrote about How to Trigger Airflow DAG Using REST API. I have also extended that concept with the article where I highlighted how to Create An Event Driven Airflow Pipeline with SNS, SQS and Lambda.

Not a Medium member? Click here to read for free!

In this demo, I will create a connection between Snowflake and Airflow such that we can have a Snowflake object trigger an Airflow DAG.

Solution architecture

The steps are as follows:

1. Create a Lambda function
2. Create an API gateway in AWS
3. Create the required permissions in AWS
4. Create an API integration in Snowflake. Then update AWS role from (3) above with Snowflake details.
5. Create the Snowflake external function to call Lambda

Take note of the following:

• I am reusing Airflow REST API setup instructions from my article How to Trigger Airflow DAG Using REST
• I am reusing the DAG sns_receiver_dag from my article Create An Event Driven Airflow Pipeline with SNS, SQS and Lambda.
• I am also reusing NGROK setup steps from the event drive pipeline article above.
• I am also reusing the Lambda function code from the same article, although with a few tweaks to get the message from Snowflake.
• I will assume you have basic knowledge of AWS, Snowflake and Airflow so I will not dive too deep into certain setup steps.

Create Lambda Function

I’ll create a simple lambda function that will be used to call the airflow REST API.

Create Lambda function

Update Lambda code with the code below:

import json

import urllib3
from datetime import datetime

def lambda_handler(event, context):
    try:

        url = 'https://138f-2607-fea8-b5d-fe60-4c38-c877-8721-440e.ngrok-free.app/api/v1/dags/sns_receiver_dag/dagRuns'

        print('Retrieving message...')
        # grab message from Snowflake
        # snowflake payload looks like this: {'data': [[0, '{"config_id":99,"run_type":"regular"}']]}
        message = event['data'][0][1]

        body = {
            "dag_run_id":"lambda"+str(datetime.now()).replace(" ","T"),
            "conf":{"message":str(message)} # converting message to string so I don' t have to change the structure of existing DAG
        }


        # trigger DAG run
        print('Now calling airflow at ', url)
        http = urllib3.PoolManager()
        http.request('POST', url,
                headers={'Content-Type': 'application/json', 'Authorization':'Basic YWlyZmxvdzphaXJmbG93'},
                body=json.dumps(body))

        print('Completed')         

        return {
            'statusCode': 200,
            'data': [[0,{'description':'success'}]] 
            }
    except Exception as e:
        return {
            'statusCode': 500,
            'data': [[0,{'description':f'An error occurred: {e}'}]]
            } 

Create API Gateway in AWS

Next I create the API gateway in AWS. Snowflake will not call the external API directly but will go through a proxy.

Create REST API in API Gateway

Next I create a resource named callDag

Create resource

Within that resource, I create a method of type POST to call a Lambda function, and choose the Lambda function created earlier.

Create method of type POST

Next select Deploy API. I used a new stage I named dev and put a short description so I know why I made this resource.

Deploy API in AWS

Take note of the Invoke URL. This is what the Snowflake function will call.

Take note of Invoke URL

Once all this is done, you’ll notice in the Lambda console that a trigger has already been added:

API gateway as Lambda trigger

Add Permissions in AWS

Next I create a role to permit another AWS account gain access to resources in my AWS account. This is useful because the Snowflake account in AWS is an external entity to mine.

Create new AWS role

Configure role

Attach permissions for Lambda full access and API gateway invoke full access.

Create API Integration In Snowflake

Next I create an API integration object in Snowflake and configure it with the AWS role ARN and the gateway URL (including path to the resource configured).

CREATE OR REPLACE API INTEGRATION snf_to_airflow_int
  API_PROVIDER =  aws_api_gateway 
  API_AWS_ROLE_ARN = '<use role ARN>'
  API_ALLOWED_PREFIXES = ('<use gateway URL to resource>')
  ENABLED =  TRUE;

Take care to include the resource path in the invoke URL:

Full URL

Now run the SQL script describe integration snf_to_airflow_int;

describe integration result

Take the account ID within API_AWS_IAM_USER_ARN. Go to the AWS console and to the IAM role for the snowflake integration.

Go to Trust Relationships then Edit trust policy.

Replace AWS property of the principal to API_AWS_IAM_USER_ARN and replace the value of sts:ExternalId with API_AWS_EXTERNAL_ID.

Save and exit.

This change grants Snowflake the permission to use that role to call the API integration.

Create Snowflake External Function

In Snowflake, an external code allows you to call code that executes outside Snowflake.

CREATE OR REPLACE EXTERNAL FUNCTION call_airflow(message variant)
returns variant
api_integration=snf_to_airflow_int
as 'https://0ea87xxzs9.execute-api.us-east-1.amazonaws.com/dev/callDag';

Make the API call from Snowflake

I am reusing the DAG sns_receiver_dag. This DAG is able to accept a trigger parameter. I will use a simple DAG parameter message {“config_id”:99,”run_type”:”regular”}.

I call the DAG like this from Snowflake:

select call_airflow(parse_json('{"config_id":99,"run_type":"regular"}'));

In Cloudwatch, you can see some of the logs from the Lambda code:

Lambda logs in Cloudwatch

Ngrok showing that the Lambda requests are reaching my PC locally:

NGROK log

In airflow you can see that run_config contains the message sent as parameters for the DAG trigger:

Airflow result

Benefits:

• Integrating Snowflake with Airflow REST API is another way to create synergy between your tools of work
• It’s easy and straightforward to implement
• The Lambda function can be reused across many services, not just Snowflake. This is only one example

For another way to trigger Airflow DAGs with REST API and Lambda, read my article Create An Event Driven Airflow Pipeline with SNS, SQS and Lambda.

References

1. Creating an external function for AWS using the AWS Management Console | Snowflake Documentation
2. Remote service input and output data formats | Snowflake Documentation
3. External Functions Tutorial | Snowflake