Introspecting Databases with Airflow and SQLAlchemy
By: L. D. Nicolas May
Tags: Sqlalchemy, Dag, Airflow, Introspect, Database

Airflow Setup

The following steps can be done via the Airflow CLI, but the web UI is easy and pretty.

Run Airflow

Run the Airflow scheduler and web server. How to run Airflow is covered in the Airflow getting started tutorial. Once Airflow’s started, point your browser to the web UI. If you installed Airflow via Docker or locally, the URL will likely be http://localhost:8080/ .

Create a Connection to your Postgres database

From the web UI’s navigation bar, click on Admin > Connections.

Create a Postgres database connection via Airflow web UI

Click on the blue + button to add a new connection.

Add a new connection via Airflow web UI

Enter your database host connection details and click the Save button. Note: The “schema” field is actually the database you’d like to connect to. I’m using the handy DVD Rental data available in my dvdrental database. Make note of your Conn Id, as we’ll need that in the next step.

Define your Postgres database connection via Airflow web UI

Create the DAG

In the dags folder you created during the Airflow installation, add a Python file called introspect_db.py . Copy-and-paste this DAG code into the file. If you have the Conn Id postgres_conn , leave the code as is. If not, change the pg_hook assignment line to reflect your Conn Id.

import logging

import re

from typing import List



from airflow import DAG

from airflow.operators.dummy import DummyOperator

from airflow.operators.python import PythonOperator

from airflow.providers.postgres.hooks.postgres import PostgresHook

from airflow.utils.dates import days_ago

from airflow.utils.task_group import TaskGroup



from sqlalchemy.inspection import inspect

from sqlalchemy.engine.reflection import Inspector

from sqlalchemy.engine import Engine





def log_info(info, *args, **kwargs):

logging.info(f"info={info}")

logging.info(f"args={args}")

logging.info(f"kwargs={kwargs}")





default_args = {

"owner": "airflow",

}



with DAG(

dag_id="introspect_db",

default_args=default_args,

schedule_interval="@once",

start_date=days_ago(2),

) as dag: pg_hook = PostgresHook(postgres_conn_id="postgres_conn")

pg_engine: Engine = pg_hook.get_sqlalchemy_engine()

pg_inspector: Inspector = inspect(pg_engine)



schema_names: List[str] = pg_inspector.get_schema_names()



schema_task_groups: List[TaskGroup] = [

TaskGroup(group_id=f"{s}") for s in schema_names

]



for schema_tg in schema_task_groups:

# Ex. of schema_tg.group_id: "public"

schema_name: str = schema_tg.group_id



table_names: List[str] = \

pg_inspector.get_table_names(schema=schema_name)



table_task_groups: List[TaskGroup] = [

TaskGroup(group_id=f"{t}", parent_group=schema_tg)

for t in table_names

]



for table_tg in table_task_groups:

# Ex. of table_tg.group_id: "public.actor"

table_name_p = re.compile(r"^\w+\.(\w+)$")

table_name: str = re.match(

table_name_p,

table_tg.group_id

).group(1)



start_table = PythonOperator(

task_id=f"start_{table_name}",

task_group=table_tg,

python_callable=log_info,

op_kwargs={

"info": [

f"table_tg.group_id={table_tg.group_id}",

f"table_name={table_name}",

]

}

) end_table = DummyOperator(

task_id=f"end_{table_name}",

task_group=table_tg

) start_table >> end_table start_schema = PythonOperator(

task_id=f"start_{schema_name}",

task_group=schema_tg,

python_callable=log_info,

op_kwargs={

"info": [

f"schema_tg.group_id={schema_tg.group_id}",

f"schema_name={schema_name}",

]

}

) end_schema = DummyOperator(

task_id=f"end_{schema_name}",

task_group=schema_tg

)



start_schema >> table_task_groups >> end_schema



start = DummyOperator(task_id="start")

end = DummyOperator(task_id="end")



start >> schema_task_groups >> end

There are a few snippets from the DAG code that are worth highlighting.

Snippet 1. The first snippet (below) illustrates how to get your database’s schema names. First, you use your Postgres connection (defined above) to get a Postgres hook. This hook is used to get the SQLAlchemy engine. The engine, in turn, is used to inspect the connected database for a list of all the schema names.

# Snippet 1 pg_hook = PostgresHook(postgres_conn_id="postgres_conn")

pg_engine: Engine = pg_hook.get_sqlalchemy_engine()

pg_inspector: Inspector = inspect(pg_engine) schema_names: List[str] = pg_inspector.get_schema_names()

Snippet 2. The next snippet illustrates how to create a list of TaskGroup objects for each of the introspected schemas. We iterate over each schema TaskGroup object, get each schema name as a string from the schema TaskGroup’s group ID, and then do some kind of work within the schema.

# Snippet 2 schema_task_groups: List[TaskGroup] = [

TaskGroup(group_id=f"{s}") for s in schema_names

] for schema_tg in schema_task_groups:

# Ex. of schema_tg.group_id: "public"

schema_name: str = schema_tg.group_id ...

Snippet 3. In this case, the “work” we do inside each schema is introspect all the schema’s tables. To do this, we take a similar approach as we did with the schemas. However, this time we need to do two things differently. First, for creating each table TaskGroup object, we need to pass the schema_tg TaskGroup object to the parent_group parameter. Second, we need to parse out the table name from table TaskGroup object’s group_id . The task group ID is of the format <schema>.<table> , so I’ve used the re regex module to parse out just the table name. For example, with a table task group ID of public.actor , the regex pattern ^\w+\.(\w+)$ will allow us to extract the actor portion of the task group ID.

# Snippet 3 table_names: List[str] = \

pg_inspector.get_table_names(schema=schema_name)



table_task_groups: List[TaskGroup] = [

TaskGroup(group_id=f"{t}", parent_group=schema_tg)

for t in table_names

]



for table_tg in table_task_groups:

# Ex. of table_tg.group_id: "public.actor"

table_name_p = re.compile(r"^\w+\.(\w+)$")

table_name: str = re.match(

table_name_p,

table_tg.group_id

).group(1) ...

Snippet 4. Finally, notice how the task dependencies are defined. The task pairs start and end , and start_schema and end_schema are single tasks. Between each task pair is a variable for a list of TaskGroup objects. It’s similar to defining task dependencies with a list like task_1 >> [task_2, task_3] >> task_4 because schema_task_groups and table_task_groups are lists of TaskGroup objects.

# Snippet 4 start_schema >> table_task_groups >> end_schema ... start >> schema_task_groups >> end

The code is ready. Now we can interact with the web UI.

Review the DAG’s Structure

Returning to the web UI at http://localhost:8080/ , click on the introspect_db DAG link, and then the Graph View button. You should see the schema task groups in the web UI.

Mine looks like this.

Schema Task Groups in Graph View

Click on one of the schema task groups to expand it. You should see the table task groups within the schema task group you selected.

Table Task Groups in Graph View

Finally, click on one of the table task groups to expand it, and you’ll reveal the tasks within the table group.