Building a Running Pipeline

Lets look at another example: we need to get some data from a file which is hosted online and insert it into our local database. We also need to look at removing duplicate rows while inserting.

Initial setup

We need to have Docker installed as we will be using the Running Airflow in Docker procedure for this example. The steps below should be sufficient, but see the quick-start documentation for full instructions.

  1. # Download the docker-compose.yaml file
  2. curl -LfO 'https://airflow.apache.org/docs/apache-airflow/stable/docker-compose.yaml'
  4. # Make expected directories and set an expected environment variable
  5. mkdir -p ./dags ./logs ./plugins
  6. echo -e "AIRFLOW_UID=$(id -u)" > .env
  8. # Initialize the database
  9. docker-compose up airflow-init
  11. # Start up all services
  12. docker-compose up

After all services have started up, the web UI will be available at: http://localhost:8080. The default account has the username airflow and the password airflow.

We will also need to create a connection to the postgres db. To create one via the web UI, from the “Admin” menu, select “Connections”, then click the Plus sign to “Add a new record” to the list of connections.

Fill in the fields as shown below. Note the Connection Id value, which we’ll pass as a parameter for the postgres_conn_id kwarg.

  • Connection Id: tutorial_pg_conn

  • Connection Type: postgres

  • Host: postgres

  • Schema: airflow

  • Login: airflow

  • Password: airflow

  • Port: 5432

Test your connection and if the test is successful, save your connection.

Table Creation Tasks

We can use the PostgresOperator to define tasks that create tables in our postgres db.

We’ll create one table to facilitate data cleaning steps (employees_temp) and another table to store our cleaned data (employees).

  1. from airflow.providers.postgres.operators.postgres import PostgresOperator
  3. create_employees_table = PostgresOperator(
  4.     task_id="create_employees_table",
  5.     postgres_conn_id="tutorial_pg_conn",
  6.     sql="""
  7.         CREATE TABLE IF NOT EXISTS employees (
  8.             "Serial Number" NUMERIC PRIMARY KEY,
  9.             "Company Name" TEXT,
  10.             "Employee Markme" TEXT,
  11.             "Description" TEXT,
  12.             "Leave" INTEGER
  13.         );""",
  14. )
  16. create_employees_temp_table = PostgresOperator(
  17.     task_id="create_employees_temp_table",
  18.     postgres_conn_id="tutorial_pg_conn",
  19.     sql="""
  20.         DROP TABLE IF EXISTS employees_temp;
  21.         CREATE TABLE employees_temp (
  22.             "Serial Number" NUMERIC PRIMARY KEY,
  23.             "Company Name" TEXT,
  24.             "Employee Markme" TEXT,
  25.             "Description" TEXT,
  26.             "Leave" INTEGER
  27.         );""",
  28. )

Optional: Using SQL From Files

If you want to abstract these sql statements out of your DAG, you can move the statements sql files somewhere within the dags/ directory and pass the sql file_path (relative to dags/) to the sql kwarg. For employees for example, create a sql directory in dags/, put employees DDL in dags/sql/employees_schema.sql, and modify the PostgresOperator() to:

  1. create_employees_table = PostgresOperator(
  2.     task_id="create_employees_table",
  3.     postgres_conn_id="tutorial_pg_conn",
  4.     sql="sql/employees_schema.sql",
  5. )

and repeat for the employees_temp table.

Data Retrieval Task

Here we retrieve data, save it to a file on our Airflow instance, and load the data from that file into an intermediate table where we can execute data cleaning steps.

  1. import os
  2. import requests
  3. from airflow.decorators import task
  4. from airflow.providers.postgres.hooks.postgres import PostgresHook
  7. @task
  8. def get_data():
  9.     # NOTE: configure this as appropriate for your airflow environment
  10.     data_path = "/opt/airflow/dags/files/employees.csv"
  11.     os.makedirs(os.path.dirname(data_path), exist_ok=True)
  13.     url = "https://raw.githubusercontent.com/apache/airflow/main/docs/apache-airflow/tutorial/pipeline_example.csv"
  15.     response = requests.request("GET", url)
  17.     with open(data_path, "w") as file:
  18.         file.write(response.text)
  20.     postgres_hook = PostgresHook(postgres_conn_id="tutorial_pg_conn")
  21.     conn = postgres_hook.get_conn()
  22.     cur = conn.cursor()
  23.     with open(data_path, "r") as file:
  24.         cur.copy_expert(
  25.             "COPY employees_temp FROM STDIN WITH CSV HEADER DELIMITER AS ',' QUOTE '\"'",
  26.             file,
  27.         )
  28.     conn.commit()

Data Merge Task

Here we select completely unique records from the retrieved data, then we check to see if any employee Serial Numbers are already in the database (if they are, we update those records with the new data).

  1. from airflow.decorators import task
  2. from airflow.providers.postgres.hooks.postgres import PostgresHook
  5. @task
  6. def merge_data():
  7.     query = """
  8.         INSERT INTO employees
  9.         SELECT *
  10.         FROM (
  11.             SELECT DISTINCT *
  12.             FROM employees_temp
  13.         ) t
  14.         ON CONFLICT ("Serial Number") DO UPDATE
  15.         SET
  16.               "Employee Markme" = excluded."Employee Markme",
  17.               "Description" = excluded."Description",
  18.               "Leave" = excluded."Leave";
  19.     """
  20.     try:
  21.         postgres_hook = PostgresHook(postgres_conn_id="tutorial_pg_conn")
  22.         conn = postgres_hook.get_conn()
  23.         cur = conn.cursor()
  24.         cur.execute(query)
  25.         conn.commit()
  26.         return 0
  27.     except Exception as e:
  28.         return 1

Completing our DAG

We’ve developed our tasks, now we need to wrap them in a DAG, which enables us to define when and how tasks should run, and state any dependencies that tasks have on other tasks. The DAG below is configured to:

  • run every day at midnight starting on Jan 1, 2021,

  • only run once in the event that days are missed, and

  • timeout after 60 minutes

And from the last line in the definition of the process-employees DAG, we see:

  1. [create_employees_table, create_employees_temp_table] >> get_data() >> merge_data()

  • the merge_data() task depends on the get_data() task,

  • the get_data() depends on both the create_employees_table and create_employees_temp_table tasks, and

  • the create_employees_table and create_employees_temp_table tasks can run independently.

Putting all of the pieces together, we have our completed DAG.

  1. import datetime
  2. import pendulum
  3. import os
  5. import requests
  6. from airflow.decorators import dag, task
  7. from airflow.providers.postgres.hooks.postgres import PostgresHook
  8. from airflow.providers.postgres.operators.postgres import PostgresOperator
  11. @dag(
  12.     dag_id="process-employees",
  13.     schedule_interval="0 0 * * *",
  14.     start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
  15.     catchup=False,
  16.     dagrun_timeout=datetime.timedelta(minutes=60),
  17. )
  18. def ProcessEmployees():
  19.     create_employees_table = PostgresOperator(
  20.         task_id="create_employees_table",
  21.         postgres_conn_id="tutorial_pg_conn",
  22.         sql="""
  23.             CREATE TABLE IF NOT EXISTS employees (
  24.                 "Serial Number" NUMERIC PRIMARY KEY,
  25.                 "Company Name" TEXT,
  26.                 "Employee Markme" TEXT,
  27.                 "Description" TEXT,
  28.                 "Leave" INTEGER
  29.             );""",
  30.     )
  32.     create_employees_temp_table = PostgresOperator(
  33.         task_id="create_employees_temp_table",
  34.         postgres_conn_id="tutorial_pg_conn",
  35.         sql="""
  36.             DROP TABLE IF EXISTS employees_temp;
  37.             CREATE TABLE employees_temp (
  38.                 "Serial Number" NUMERIC PRIMARY KEY,
  39.                 "Company Name" TEXT,
  40.                 "Employee Markme" TEXT,
  41.                 "Description" TEXT,
  42.                 "Leave" INTEGER
  43.             );""",
  44.     )
  46.     @task
  47.     def get_data():
  48.         # NOTE: configure this as appropriate for your airflow environment
  49.         data_path = "/opt/airflow/dags/files/employees.csv"
  50.         os.makedirs(os.path.dirname(data_path), exist_ok=True)
  52.         url = "https://raw.githubusercontent.com/apache/airflow/main/docs/apache-airflow/tutorial/pipeline_example.csv"
  54.         response = requests.request("GET", url)
  56.         with open(data_path, "w") as file:
  57.             file.write(response.text)
  59.         postgres_hook = PostgresHook(postgres_conn_id="tutorial_pg_conn")
  60.         conn = postgres_hook.get_conn()
  61.         cur = conn.cursor()
  62.         with open(data_path, "r") as file:
  63.             cur.copy_expert(
  64.                 "COPY employees_temp FROM STDIN WITH CSV HEADER DELIMITER AS ',' QUOTE '\"'",
  65.                 file,
  66.             )
  67.         conn.commit()
  69.     @task
  70.     def merge_data():
  71.         query = """
  72.             INSERT INTO employees
  73.             SELECT *
  74.             FROM (
  75.                 SELECT DISTINCT *
  76.                 FROM employees_temp
  77.             ) t
  78.             ON CONFLICT ("Serial Number") DO UPDATE
  79.             SET "Serial Number" = excluded."Serial Number";
  80.         """
  81.         try:
  82.             postgres_hook = PostgresHook(postgres_conn_id="tutorial_pg_conn")
  83.             conn = postgres_hook.get_conn()
  84.             cur = conn.cursor()
  85.             cur.execute(query)
  86.             conn.commit()
  87.             return 0
  88.         except Exception as e:
  89.             return 1
  91.     [create_employees_table, create_employees_temp_table] >> get_data() >> merge_data()
  94. dag = ProcessEmployees()

Save this code to a python file in the /dags folder (e.g. dags/process-employees.py) and (after a brief delay), the process-employees DAG will be included in the list of available DAGs on the web UI.

../_images/tutorial-pipeline-1.png

You can trigger the process-employees DAG by unpausing it (via the slider on the left end) and running it (via the Run button under Actions).

../_images/tutorial-pipeline-2.png

In the process-employees DAG’s Grid view, we see all that all tasks ran successfully in all executed runs. Success!

What’s Next?

You now have a pipeline running inside Airflow using Docker Compose. Here are a few things you might want to do next:

See also