Airflow: Connections, Hooks, and Providers

What is a Connection?

A connection is a saved set of credentials stored inside of Airflow with a name attached to it. Instead of putting the host and password inside the code, you create a connection and store everything there. This is more secure because the password or any other credentials do not appear anywhere.

What is a Hook?

A connection by itself is just raw data, so you still need something to actually use it. This is called a hook, which is generally a Python class. You pass the connection name to it, and it returns a working client for that specific external system. Each database or system has a separate hook because Postgres can only talk to Postgres, HTTP can only interact with APIs, and so on.

What are Providers?

The Postgres Hook you want to use is not bundled with core Airflow. You have to download it separately. Providers are Python packages that ship hooks and operators for specific external systems. There are providers for almost all database systems and tools, such as Snowflake, MongoDB, Slack, Google Cloud, MySQL, and so on. All you need to do is install the packages.

Setting Up the Environment

Run the environment. Once it is up, go to the Airflow UI to set up the connection with the Postgres database.

On the left, click the admin section and go to connections. Click the blue button on the right. Fill in these fields:

• Connection ID: The name you will reference from the DAG. Call it postgres.
• Connection Type: Find Postgres from the dropdown.
• Description: Postgres database
• Host: postgres (the name of the service inside your Docker Compose setup)
• Login: airflow
• Password: airflow
• Port: 5432 (standard for Postgres)
• Database: airflow

Click save.

Building the DAG

Go to VS Code, navigate to the Airflow directory and DAG folder, and create a new file called postgres_dag.py.

Import the necessary components:

from datetime import datetime
from airflow.decorators import dag, task
from airflow.providers.postgres.hooks.postgres import PostgresHook

• DAG and TaskDecorators from TaskFlow
• PostgresHook – the class to talk to the database
• datetime to set the starting date in the DAG block

Creating the Table

Within the PostgresHook, indicate the name of your connection. Use hook.run to execute a SQL statement:

@task
def create_table():
    hook = PostgresHook(postgres_conn_id="postgres")
    hook.run("""
        CREATE TABLE IF NOT EXISTS users (
            id SERIAL PRIMARY KEY,
            name VARCHAR(100) NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    print("Table ready")

Inserting Data

@task
def insert_users():
    hook = PostgresHook(postgres_conn_id="postgres")
    hook.run("INSERT INTO users (name) VALUES ('Alice'), ('Bob'), ('Charlie')")
    print("Users inserted")

Reading Data

@task
def read_users():
    hook = PostgresHook(postgres_conn_id="postgres")
    records = hook.get_records("SELECT id, name FROM users")
    for row in records:
        print(f"ID: {row[0]}, Name: {row[1]}")

Defining the DAG

Define the DAG using the @dag decorator:

@dag(
    dag_id="postgres_dag",
    start_date=datetime(2023, 1, 1),
    schedule=None,
    catchup=False
)
def postgres_dag():
    create_table_task = create_table()
    insert_users_task = insert_users()
    read_users_task = read_users()

    create_table_task >> insert_users_task >> read_users_task

postgres_dag()

Running the DAG

Save the file and go back to the Airflow UI. Go to the DAG section, search for postgres_dag, and trigger it.

Check the logs for create_table - you should see that the connection is retrieved and "Table ready" is printed. For insert_users, you should see that all three names were added. For read_users, you should see users with ID 1, ID 2, and ID 3.

Verifying in the Database

To verify everything, go back to VS Code's terminal and run this command:

docker exec -it airflow-postgres1 psql -U airflow -d airflow -c "SELECT * FROM users"