DRY — “Don’t Repeat Yourself” — is a fundamental coding principle that aims to reduce repetitive patterns and duplication, emphasising the development of modular and reusable code. When building your DAGs, you often encounter such situations, where certain tasks need to be executed sequentially and repeatedly. These repetitive tasks are ideal candidates for pattern implementation. In Airflow, we can effectively implement these patterns using TaskGroups.
✏️ A quick note: if you’re unfamiliar with TaskGroups, our friends over at Astronomer have written an excellent post about Airflow TaskGroups, what they are and how they work.
This post explores how we can implement such a pattern by the means of a simplified real world example.
Our use-case
Every day, we need to snapshot and ingest data objects from Salesforce to our datalake on S3. We’re fully on the AWS Stack, and are using AWS AppFlow to manage this integration for us.
With AppFlow, we’re running into a problem, though. AppFlow gives the resulting snapshots a random name. This is not ideal for downstream processes that need to further process the snapshot. As a solution, every time we snapshot using AppFlow, we need to rename the resulting file on our datalake to something constant. We have two sequential and repetitive tasks, time for a pattern!
The DAG
For the example DAG with an implemented pattern, see the code below. Take your time to read and understand the code. The next section will explain how it works.
It’s good to know that this DAG contains the custom operators AppFlow & S3Rename that trigger an AppFlow and rename a file, respectively. If you want an article on how these custom operators work, let me know!
"""DAG that ingests Salesforce data"""
import os
from datetime import date, timedelta
import pendulum
from airflow import DAG
from airflow.models import Variable
from airflow.sensors.external_task import ExternalTaskMarker
from airflow.utils.task_group import TaskGroup
from operators.aws.appflow import AppFlow
from operators.aws.s3 import S3Rename
default_args = {
'owner': 'data',
'depends_on_past': False,
'email': ['john@example.com'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 2,
'retry_delay': timedelta(minutes=5),
}
# Get path from execution date
PATH = '{{ execution_date.strftime("%Y/%m/%d") }}'
bucket = Variable.get('datalake-production')
def ingest_data(dag_obj: DAG, tables: list) -> TaskGroup:
"""
This pattern consists of 2 tasks:
1. Trigger a Salesforce AppFlow
2. Rename the file to a specified filename
"""
with TaskGroup(group_id='ingestion', dag=dag_obj) as paths:
previous = None
for table in tables:
with TaskGroup(group_id=table) as path:
appflow_name = f"Salesforce{table}"
trigger_appflow = AppFlow(
task_id=appflow_name,
flow_name=appflow_name,
)
rename_file = S3Rename(
task_id=f"rename_{appflow_name}_file",
bucket=bucket,
old_prefix=f"snapshots/salesforce/{appflow_name}/{PATH}",
new_prefix=f"snapshots/salesforce/{appflow_name}/{PATH}/{table.lower()}.csv"
)
trigger_appflow >> rename_file
if previous:
previous >> path
previous = path
return paths
with DAG(
os.path.basename(__file__).replace(".py", ""),
default_args=default_args,
description='DAG that ingests all data from Salesforce',
schedule_interval="30 6 * * *",
# Make a timezone aware DAG
start_date=pendulum.datetime(2024, 10, 1, tz="Europe/Amsterdam"),
catchup=False,
tags=['ingest'],
) as dag:
ingest_data_group = ingest_data(dag, [
"Accounts",
"Contacts",
"Opportunities",
"Products",
"Assets",
])
# Bidirectional coupling of the ingest & load
marker = ExternalTaskMarker(
task_id='load_salesforce_trigger',
external_dag_id='load_salesforce',
external_task_id='wait_for_salesforce_ingest'
)
ingest_data_group >> marker
This creates the following DAG:
Implementation of an Airflow pattern through TaskGroups
So, how does it work?
The magic of patterns resides in the ingest_date() function definition. Notice how this function contains nested TaskGroups. This is done so that we can repeat the pattern as much as we’d like, whilst making sure we always return a single TaskGroup. This makes it easy to set up the dependencies.
Let’s zoom in.
with TaskGroup(group_id='ingestion', dag=dag_obj) as paths:
previous = None
for table in tables:
with TaskGroup(group_id=table) as path:
appflow_name = f"Salesforce{table}"
trigger_appflow = AppFlow(
task_id=appflow_name,
flow_name=appflow_name,
)
rename_file = S3Rename(
task_id=f"rename_{appflow_name}_file",
bucket=bucket,
old_prefix=f"snapshots/salesforce/{appflow_name}/{PATH}",
new_prefix=f"snapshots/salesforce/{appflow_name}/{PATH}/{table.lower()}.csv"
)
trigger_appflow >> rename_file
if previous:
previous >> path
previous = path
The key to this function is the way it adds dependencies of the nested TaskGroups within the main ingestion group.
As we can see in the snippet above, it first sets the variable previous to None which prevents entering the if previous conditional. The function then loops through the provided tables. At the end of the iteration it sets previous to path, with path containing the two tasks and their dependencies. In the next iteration, previous is set, meaning it appends the new task group to the previous task group in the if previous conditional, and so on until the loop is completed.
This logic allows you to dynamically define your dependencies, and work with patterns within Airflow!
Next step
Defining your pattern within your DAG is fine as an example, but not very scalable. For further abstraction, you could opt to store all your patterns in a patterns.py utils file in your project’s root folder. This allows you to create cross-DAG patterns.
Hi, I’m Bastiaan 👋🏼 Data Lead at a scale-up. I write about the Modern Data Workflow, where I explore tools & processes to supercharge your data capabilities. Follow me for more!