Apache Airflow - A Python Hands-On Guide
Apache Airflow - A Python Hands-On Guide
Available Provider
Available Provider
1. Writing a DAG in Python
A Directed Acyclic Graph (DAG) is the core abstraction in Airflow. It defines the
workflow and task dependencies.
Basic DAG Structure
from datetime import datetime
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
# Define default arguments
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
}
# Initialize DAG
with DAG(
dag_id='example_dag',
default_args=default_args,
description='A simple example DAG',
schedule_interval='@daily',
start_date=datetime(2024, 1, 1),
catchup=False,
) as dag:
def print_hello():
print("Hello, Airflow!")
task = PythonOperator(
1/9
Apache Airflow - A Python Hands-On Guide
task_id='print_hello',
python_callable=print_hello,
)
Key Notes:
DAG : The container for your workflow.
PythonOperator : Executes Python functions.
schedule_interval : Defines the schedule (e.g., @daily , @hourly ).
catchup : Prevents past executions from running.
2. Options for Python Operators
Airflow offers various operators that use Python extensively:
PythonOperator
Executes Python callables.
PythonOperator(
task_id='process_data',
python_callable=process_data_function,
op_kwargs={'param': 'value'}, # Pass arguments to the callable
)
BranchPythonOperator
Allows branching based on a condition.
from airflow.operators.python_operator import BranchPythonOperator
def choose_branch(**kwargs):
return 'branch_1' if kwargs['some_condition'] else 'branch_2'
branch_task = BranchPythonOperator(
task_id='branching',
python_callable=choose_branch,
provide_context=True,
)
2/9
Apache Airflow - A Python Hands-On Guide
PythonVirtualenvOperator
Executes Python code within a virtual environment.
from airflow.operators.python_operator import
PythonVirtualenvOperator
virtualenv_task = PythonVirtualenvOperator(
task_id='venv_task',
python_callable=lambda: print("Running in a virtualenv!"),
requirements=["numpy", "pandas"],
system_site_packages=False,
)
3. Working with Providers
Providers are integrations for various platforms. Here's how to use five popular ones
with Python:
1. AWS Provider
Install: pip install apache-airflow-providers-amazon
Example: S3 File Upload
from airflow.providers.amazon.aws.operators.s3 import
S3CreateObjectOperator
upload_task = S3CreateObjectOperator(
task_id='upload_to_s3',
aws_conn_id='my_aws_conn',
s3_bucket='my_bucket',
s3_key='path/to/file.txt',
data="Sample Data",
)
2. Google Cloud Provider
Install: pip install apache-airflow-providers-google
3/9
Apache Airflow - A Python Hands-On Guide
Example: BigQuery Query Execution
from airflow.providers.google.cloud.operators.bigquery import
BigQueryExecuteQueryOperator
bq_task = BigQueryExecuteQueryOperator(
task_id='bq_query',
sql='SELECT * FROM my_dataset.my_table',
gcp_conn_id='my_gcp_conn',
use_legacy_sql=False,
)
3. PostgreSQL Provider
Install: pip install apache-airflow-providers-postgres
Example: Run SQL on PostgreSQL
from airflow.providers.postgres.operators.postgres import
PostgresOperator
sql_task = PostgresOperator(
task_id='run_postgres_query',
postgres_conn_id='my_postgres_conn',
sql='SELECT * FROM my_table;',
)
4. Slack Provider
Install: pip install apache-airflow-providers-slack
Example: Send Slack Notification
from airflow.providers.slack.operators.slack_webhook import
SlackWebhookOperator
slack_task = SlackWebhookOperator(
task_id='send_slack_message',
4/9
Apache Airflow - A Python Hands-On Guide
http_conn_id='slack_conn',
message="Workflow completed successfully!",
channel="#alerts",
)
5. MySQL Provider
Install: pip install apache-airflow-providers-mysql
Example: Execute SQL in MySQL
from airflow.providers.mysql.operators.mysql import MySqlOperator
mysql_task = MySqlOperator(
task_id='mysql_query',
mysql_conn_id='my_mysql_conn',
sql='INSERT INTO my_table (id, value) VALUES (1, "test");',
)
5. Python Cheatsheet for Apache Airflow
Component Python Example
DAG DAG(dag_id='my_dag', schedule_interval='@daily', ...)
Initialization
PythonOperator PythonOperator(task_id='task',
python_callable=my_func)
Branching BranchPythonOperator(task_id='branch',
python_callable=my_func)
S3 Upload S3CreateObjectOperator(...,
s3_key='path/to/file.txt')
SQL Execution PostgresOperator(sql='SELECT * FROM table;')
BigQuery BigQueryExecuteQueryOperator(sql='SELECT * FROM
table')
Slack SlackWebhookOperator(message="Job done!")
Notification
5/9
Apache Airflow - A Python Hands-On Guide
Component Python Example
Virtualenv PythonVirtualenvOperator(python_callable=my_func,
...)
Working with Apache Spark in Airflow
Apache Spark is a distributed data processing framework widely used for big data
tasks. In Airflow, we can manage and orchestrate Spark jobs using operators such as:
SparkSubmitOperator : Submits a Spark job directly to a cluster.
EmrAddStepsOperator : Submits a Spark job to an Amazon EMR cluster.
DataprocSubmitJobOperator : Submits a Spark job to Google Dataproc.
These operators allow us to control Spark jobs programmatically within Airflow
workflows.
Complex Workflow: Conditional Spark Job Execution
Workflow Logic:
1. Execute spark_job_1 .
2. Execute spark_job_2 .
3. If spark_job_2 fails, run spark_job_3 .
4. If spark_job_2 succeeds, run spark_job_4 .
DAG Implementation
Step 1: Import Required Modules
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.providers.apache.spark.operators.spark_submit import
SparkSubmitOperator
6/9
Apache Airflow - A Python Hands-On Guide
from airflow.operators.python_operator import BranchPythonOperator
from datetime import datetime
Step 2: Define the DAG and Default Arguments
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'email_on_failure': False,
'retries': 1,
}
dag = DAG(
dag_id='spark_conditional_jobs',
default_args=default_args,
description='A DAG with conditional Spark job execution',
schedule_interval=None,
start_date=datetime(2023, 12, 1),
catchup=False,
)
Step 3: Define the SparkSubmitOperator Jobs
# Spark Job 1
spark_job_1 = SparkSubmitOperator(
task_id='spark_job_1',
application='/path/to/spark_job_1.py',
conn_id='spark_default', # Connection to your Spark cluster
application_args=['arg1', 'arg2'],
dag=dag,
)
# Spark Job 2
spark_job_2 = SparkSubmitOperator(
task_id='spark_job_2',
application='/path/to/spark_job_2.py',
conn_id='spark_default',
application_args=['arg1', 'arg2'],
dag=dag,
)
# Spark Job 3
spark_job_3 = SparkSubmitOperator(
task_id='spark_job_3',
7/9
Apache Airflow - A Python Hands-On Guide
application='/path/to/spark_job_3.py',
conn_id='spark_default',
application_args=['arg1', 'arg2'],
dag=dag,
)
# Spark Job 4
spark_job_4 = SparkSubmitOperator(
task_id='spark_job_4',
application='/path/to/spark_job_4.py',
conn_id='spark_default',
application_args=['arg1', 'arg2'],
dag=dag,
)
Step 4: Branch Logic Using BranchPythonOperator
def choose_next_task(**kwargs):
# Check the state of spark_job_2
task_instance = kwargs['ti']
spark_job_2_state =
task_instance.xcom_pull(task_ids='spark_job_2', key='return_value')
# Return the task to execute next
if spark_job_2_state == 'failed':
return 'spark_job_3'
return 'spark_job_4'
branch_task = BranchPythonOperator(
task_id='branch_task',
python_callable=choose_next_task,
provide_context=True,
dag=dag,
)
Step 5: Define the Task Dependencies
start = DummyOperator(task_id='start', dag=dag)
end = DummyOperator(task_id='end', dag=dag)
# Define dependencies
start >> spark_job_1
spark_job_1 >> spark_job_2
spark_job_2 >> branch_task
8/9
Apache Airflow - A Python Hands-On Guide
branch_task >> spark_job_3 >> end
branch_task >> spark_job_4 >> end
Explanation of Components
1. SparkSubmitOperator :
Used to submit Spark jobs to a cluster.
Specify the application path, connection ID, and arguments for the Spark job.
2. BranchPythonOperator :
Dynamically determines the next task based on the state of a previous task.
In this case, it checks if spark_job_2 succeeded or failed.
3. Dependencies:
The workflow ensures sequential execution from spark_job_1 to
spark_job_2 and conditional branching to spark_job_3 or spark_job_4 .
Python Cheatsheet for Spark in Airflow
Component Python Example
SparkSubmitOperator SparkSubmitOperator(application='/path/app.py',
...)
BranchPythonOperator BranchPythonOperator(python_callable=my_func,
...)
Conditional Execution branch_task >> task_1 >> end or branch_task >>
task_2
Task Dependency task_1 >> task_2 >> task_3
9/9
