SqlOperator for database queries in Apache Airflow - Commands & Configuration

Introduction

Sometimes you need to run SQL commands automatically as part of your data workflows. SqlOperator in Airflow helps you run SQL queries on databases without writing extra code to connect and manage the database.

When you want to run a simple SQL query like inserting or updating data as part of a scheduled task.

When you need to automate database maintenance tasks such as cleaning old records regularly.

When you want to chain SQL queries with other tasks in your data pipeline.

When you want to run a query on a database and use the result in later steps of your workflow.

When you want to keep your SQL code organized inside Airflow DAGs for easy management.

Config File - example_sql_operator_dag.py

example_sql_operator_dag.py

from airflow import DAG
from airflow.providers.postgres.operators.postgres import PostgresOperator
from datetime import datetime

with DAG(
    dag_id='example_sql_operator_dag',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    catchup=False
) as dag:

    create_table = PostgresOperator(
        task_id='create_table',
        postgres_conn_id='my_postgres_conn',
        sql="""
        CREATE TABLE IF NOT EXISTS users (
            id SERIAL PRIMARY KEY,
            name VARCHAR(50),
            email VARCHAR(50)
        );
        """
    )

    insert_user = PostgresOperator(
        task_id='insert_user',
        postgres_conn_id='my_postgres_conn',
        sql="""
        INSERT INTO users (name, email) VALUES ('Alice', '[e-mail]');
        """
    )

    create_table >> insert_user

This Python file defines an Airflow DAG that runs two SQL commands on a Postgres database.

create_table: This task creates a table named users if it does not exist.

insert_user: This task inserts a new user record into the users table.

The postgres_conn_id refers to a connection configured in Airflow to connect to your Postgres database.

The tasks run in order: first create the table, then insert the user.

Commands

This command lists all the DAGs Airflow knows about, so you can check if your DAG is recognized.

Terminal

airflow dags list

Expected Output

example_sql_operator_dag

This command starts running the DAG named 'example_sql_operator_dag' immediately to execute the SQL tasks.

Terminal

airflow dags trigger example_sql_operator_dag

Expected Output

Created <DagRun example_sql_operator_dag @ 2024-06-01T12:00:00+00:00: manual__2024-06-01T12:00:00+00:00, externally triggered: True>

This command shows the tasks inside the DAG so you can verify the tasks are defined correctly.

Terminal

airflow tasks list example_sql_operator_dag

Expected Output

create_table

insert_user

This command runs the 'create_table' task for the DAG on the given date to test the SQL execution without scheduling.

Terminal

airflow tasks test example_sql_operator_dag create_table 2024-06-01

Expected Output

[2024-06-01 12:00:00,000] {taskinstance.py:xxxx} INFO - Executing task create_table

[2024-06-01 12:00:01,000] {taskinstance.py:xxxx} INFO - Task create_table succeeded

Key Concept

If you remember nothing else from this pattern, remember: SqlOperator lets you run SQL commands inside Airflow tasks easily by specifying the SQL and database connection.

Common Mistakes

Not setting up the database connection in Airflow before using postgres_conn_id.

Airflow cannot connect to the database and the task will fail with a connection error.

Create a connection in Airflow UI under Admin > Connections with the correct database credentials and use its connection ID in the operator.

Writing SQL commands with syntax errors inside the sql parameter.

The SQL query will fail and the task will not complete successfully.

Test your SQL queries separately in the database client before adding them to the operator.

Forgetting to set catchup=False in the DAG when testing to avoid running many past schedules.

Airflow will try to run the DAG for all past dates, causing unexpected load and confusion.

Set catchup=False in the DAG definition to run only the latest schedule.

Summary

Define SQL tasks in Airflow using SqlOperator by specifying the SQL query and database connection ID.

Use airflow CLI commands to list, trigger, and test your DAG and tasks.

Ensure your database connection is configured in Airflow and your SQL syntax is correct before running.

Practice

1. What is the primary purpose of the SqlOperator in Airflow?

easy

A. To run SQL commands as part of an Airflow task

B. To schedule Python scripts

C. To monitor server health

D. To send email notifications

Solution

1. Step 1: Understand the role of SqlOperator

   The SqlOperator is designed to execute SQL commands within Airflow tasks, automating database operations.

2. Step 2: Compare with other options

   Scheduling Python scripts, monitoring servers, and sending emails are handled by other operators or components, not SqlOperator.

3. Final Answer:

   To run SQL commands as part of an Airflow task -> Option A

4. Quick Check:

   SqlOperator runs SQL commands = D [OK]

Hint: SqlOperator always runs SQL queries inside tasks [OK]

Common Mistakes:

• Confusing SqlOperator with PythonOperator
• Thinking SqlOperator schedules scripts
• Assuming SqlOperator sends emails

2. Which of the following is the correct way to create a SqlOperator task in Airflow?

easy

A. SqlOperator(id='run_query', query='SELECT * FROM table', connection='my_db')

B. SqlOperator(name='run_query', sql_command='SELECT * FROM table', conn='my_db')

C. SqlOperator(task='run_query', command='SELECT * FROM table', connection_id='my_db')

D. SqlOperator(task_id='run_query', sql='SELECT * FROM table', conn_id='my_db')

Solution

1. Step 1: Check parameter names for SqlOperator

   The correct parameters are task_id, sql, and conn_id.

2. Step 2: Validate each option

   Only SqlOperator(task_id='run_query', sql='SELECT * FROM table', conn_id='my_db') uses the correct parameter names and syntax for SqlOperator.

3. Final Answer:

   SqlOperator(task_id='run_query', sql='SELECT * FROM table', conn_id='my_db') -> Option D

4. Quick Check:

   Correct parameters = B [OK]

Hint: Remember: task_id, sql, conn_id are required parameters [OK]

Common Mistakes:

• Using wrong parameter names like 'id' or 'query'
• Mixing connection parameter names
• Omitting task_id or sql

3. Given this Airflow task using SqlOperator:

task = SqlOperator(
    task_id='count_users',
    sql='SELECT COUNT(*) FROM users;',
    conn_id='my_postgres'
)

What will be the main effect when this task runs?

medium

A. It will create a new table named 'count_users'

B. It will return the count of users and store it in Airflow logs

C. It will delete all records from the users table

D. It will send an email with the user count

Solution

1. Step 1: Understand the SQL command

   The SQL command SELECT COUNT(*) FROM users; counts the number of rows in the users table.

2. Step 2: Understand SqlOperator behavior

   SqlOperator runs the query and logs the output; it does not create tables, delete data, or send emails by default.

3. Final Answer:

   It will return the count of users and store it in Airflow logs -> Option B

4. Quick Check:

   SELECT COUNT(*) logs result = A [OK]

Hint: SELECT queries log results; no table creation or deletion [OK]

Common Mistakes:

• Thinking SELECT creates or deletes tables
• Assuming SqlOperator sends emails automatically
• Confusing query effects with task side effects

4. You wrote this SqlOperator task:

task = SqlOperator(
    task_id='update_status',
    sql='UPDATE orders SET status = 'shipped' WHERE id = 10;',
    conn_id='my_db'
)

But Airflow throws a syntax error. What is the problem?

medium

A. conn_id should be connection_id

B. The task_id is missing

C. The SQL string uses single quotes inside single quotes without escaping

D. SqlOperator cannot run UPDATE statements

Solution

1. Step 1: Check the SQL string syntax

   The SQL string uses single quotes around the whole string and also inside for 'shipped', causing a conflict.

2. Step 2: Understand string escaping in Python

   Inner single quotes must be escaped or the string should use double quotes outside to avoid syntax errors.

3. Final Answer:

   The SQL string uses single quotes inside single quotes without escaping -> Option C

4. Quick Check:

   Escape inner quotes or use double quotes outside [OK]

Hint: Escape inner quotes or use double quotes outside SQL string [OK]

Common Mistakes:

• Ignoring string quoting rules in Python
• Changing parameter names incorrectly
• Thinking SqlOperator can't run UPDATE

5. You want to run multiple SQL commands in one SqlOperator task to reset a table:

DROP TABLE IF EXISTS temp_data;
CREATE TABLE temp_data (id INT, value TEXT);
INSERT INTO temp_data VALUES (1, 'A'), (2, 'B');

How should you write the sql parameter to run all commands correctly?

hard

A. Use a single string with commands separated by semicolons

B. Pass a list of strings, one per command

C. Use multiple SqlOperator tasks, one per command

D. Put all commands in a file and pass the filename to sql

Solution

1. Step 1: Understand SqlOperator's sql parameter

   SqlOperator accepts a single string with multiple SQL commands separated by semicolons.

2. Step 2: Compare other options

   Passing a list is invalid; multiple tasks or a file are alternatives but not the direct answer to the question.

3. Final Answer:

   Use a single string with commands separated by semicolons -> Option A

4. Quick Check:

   Multiple SQL commands in one string separated by semicolons [OK]

Hint: Separate SQL commands with semicolons in one string [OK]

Common Mistakes:

• Passing list instead of string to sql
• Trying to run multiple commands without semicolons
• Confusing file input with direct sql parameter