使用 SQLExecuteQueryOperator 操作 Postgres 的方法指南¶
简介¶
Apache Airflow 拥有丰富的操作符库，可用于实现构成工作流的各种任务。Airflow 本质上是一个由任务（节点）和依赖关系（边）组成的图（有向无环图）。
由操作符定义或实现的任务是数据管道中的一个工作单元。
本指南的目的是使用 SQLExecuteQueryOperator 定义涉及与 PostgreSQL 数据库交互的任务。
警告
以前，PostgresOperator 用于执行此类操作。但目前 PostgresOperator 已弃用，并将在该提供程序的未来版本中删除。请考虑尽快切换到 SQLExecuteQueryOperator。
使用 SQLExecuteQueryOperator 进行常见的数据库操作¶
要使用 SQLExecuteQueryOperator 执行 PostgreSQL 请求，需要两个参数：sql 和 conn_id。这两个参数最终会被提供给直接与 Postgres 数据库交互的 DbApiHook 对象。
创建 Postgres 数据库表¶
以下代码片段基于 Airflow-2.0
# create_pet_table, populate_pet_table, get_all_pets, and get_birth_date are examples of tasks created by
# instantiating the Postgres Operator
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
DAG_ID = "postgres_operator_dag"
with DAG(
dag_id=DAG_ID,
start_date=datetime.datetime(2020, 2, 2),
schedule="@once",
catchup=False,
) as dag:
create_pet_table = SQLExecuteQueryOperator(
task_id="create_pet_table",
sql="""
CREATE TABLE IF NOT EXISTS pet (
pet_id SERIAL PRIMARY KEY,
name VARCHAR NOT NULL,
pet_type VARCHAR NOT NULL,
birth_date DATE NOT NULL,
OWNER VARCHAR NOT NULL);
""",
)
将 SQL 语句直接写入操作符中并不十分吸引人，并且会在以后的某个时候造成可维护性问题。为了防止这种情况，Airflow 提供了一个优雅的解决方案。其工作原理如下：您只需在 DAG 文件夹内创建一个名为 sql 的目录，然后将包含 SQL 查询的所有 SQL 文件放入其中。
您的 dags/sql/pet_schema.sql 应该如下所示
-- create pet table
CREATE TABLE IF NOT EXISTS pet (
pet_id SERIAL PRIMARY KEY,
name VARCHAR NOT NULL,
pet_type VARCHAR NOT NULL,
birth_date DATE NOT NULL,
OWNER VARCHAR NOT NULL);
现在让我们重构 DAG 中的 create_pet_table
create_pet_table = SQLExecuteQueryOperator(
task_id="create_pet_table",
conn_id="postgres_default",
sql="sql/pet_schema.sql",
)
将数据插入 Postgres 数据库表¶
假设我们在 dags/sql/pet_schema.sql 文件中已经有了以下 SQL 插入语句
-- populate pet table
INSERT INTO pet VALUES ( 'Max', 'Dog', '2018-07-05', 'Jane');
INSERT INTO pet VALUES ( 'Susie', 'Cat', '2019-05-01', 'Phil');
INSERT INTO pet VALUES ( 'Lester', 'Hamster', '2020-06-23', 'Lily');
INSERT INTO pet VALUES ( 'Quincy', 'Parrot', '2013-08-11', 'Anne');
然后，我们可以创建一个填充 pet 表的 SQLExecuteQueryOperator 任务。
populate_pet_table = SQLExecuteQueryOperator(
task_id="populate_pet_table",
conn_id="postgres_default",
sql="sql/pet_schema.sql",
)
从 Postgres 数据库表中获取记录¶
从 Postgres 数据库表中获取记录可以像下面这样简单
get_all_pets = SQLExecuteQueryOperator(
task_id="get_all_pets",
conn_id="postgres_default",
sql="SELECT * FROM pet;",
)
将参数传递给 Postgres 的 SQLExecuteQueryOperator¶
SQLExecuteQueryOperator 提供了 parameters 属性，可以在运行时将值动态注入到 SQL 请求中。BaseOperator 类具有 params 属性，SQLExecuteQueryOperator 可以通过继承获得该属性。parameters 和 params 都可以通过多种有趣的方式动态传递参数。
要查找名为“Lester”的宠物的主人
get_birth_date = SQLExecuteQueryOperator(
task_id="get_birth_date",
conn_id="postgres_default",
sql="SELECT * FROM pet WHERE birth_date BETWEEN SYMMETRIC %(begin_date)s AND %(end_date)s",
parameters={"begin_date": "2020-01-01", "end_date": "2020-12-31"},
)
现在让我们重构 get_birth_date 任务。我们不直接将 SQL 语句写入代码中，而是通过创建一个 sql 文件来整理一下。
-- dags/sql/birth_date.sql
SELECT * FROM pet WHERE birth_date BETWEEN SYMMETRIC {{ params.begin_date }} AND {{ params.end_date }};
这一次，我们将使用从父类 BaseOperator 中免费获得的 params 属性。
get_birth_date = SQLExecuteQueryOperator(
task_id="get_birth_date",
conn_id="postgres_default",
sql="sql/birth_date.sql",
params={"begin_date": "2020-01-01", "end_date": "2020-12-31"},
)
将服务器配置参数传递给 PostgresOperator¶
SQLExecuteQueryOperator 提供了 hook_params 属性，允许您将参数传递给 DbApiHook。您可以通过这种方式传递 options 参数，以便指定在连接启动时发送到服务器的命令行选项。
get_birth_date = SQLExecuteQueryOperator(
task_id="get_birth_date",
sql="SELECT * FROM pet WHERE birth_date BETWEEN SYMMETRIC %(begin_date)s AND %(end_date)s",
parameters={"begin_date": "2020-01-01", "end_date": "2020-12-31"},
hook_params={"options": "-c statement_timeout=3000ms"},
)
完整的 Postgres Operator DAG¶
当我们把所有东西放在一起时，我们的 DAG 应该是这样的
# create_pet_table, populate_pet_table, get_all_pets, and get_birth_date are examples of tasks created by
# instantiating the Postgres Operator
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
DAG_ID = "postgres_operator_dag"
with DAG(
dag_id=DAG_ID,
start_date=datetime.datetime(2020, 2, 2),
schedule="@once",
catchup=False,
) as dag:
create_pet_table = SQLExecuteQueryOperator(
task_id="create_pet_table",
sql="""
CREATE TABLE IF NOT EXISTS pet (
pet_id SERIAL PRIMARY KEY,
name VARCHAR NOT NULL,
pet_type VARCHAR NOT NULL,
birth_date DATE NOT NULL,
OWNER VARCHAR NOT NULL);
""",
)
populate_pet_table = SQLExecuteQueryOperator(
task_id="populate_pet_table",
sql="""
INSERT INTO pet (name, pet_type, birth_date, OWNER)
VALUES ( 'Max', 'Dog', '2018-07-05', 'Jane');
INSERT INTO pet (name, pet_type, birth_date, OWNER)
VALUES ( 'Susie', 'Cat', '2019-05-01', 'Phil');
INSERT INTO pet (name, pet_type, birth_date, OWNER)
VALUES ( 'Lester', 'Hamster', '2020-06-23', 'Lily');
INSERT INTO pet (name, pet_type, birth_date, OWNER)
VALUES ( 'Quincy', 'Parrot', '2013-08-11', 'Anne');
""",
)
get_all_pets = SQLExecuteQueryOperator(task_id="get_all_pets", sql="SELECT * FROM pet;")
get_birth_date = SQLExecuteQueryOperator(
task_id="get_birth_date",
sql="SELECT * FROM pet WHERE birth_date BETWEEN SYMMETRIC %(begin_date)s AND %(end_date)s",
parameters={"begin_date": "2020-01-01", "end_date": "2020-12-31"},
hook_params={"options": "-c statement_timeout=3000ms"},
)
create_pet_table >> populate_pet_table >> get_all_pets >> get_birth_date
结论¶
在本操作指南中，我们探讨了如何使用 Apache Airflow SQLExecuteQueryOperator 连接到 PostgreSQL 数据库。让我们快速回顾一下要点。最佳实践是在 dags 目录中创建一个名为 sql 的子目录，用于存储 sql 文件。这将使您的代码更优雅、更易于维护。最后，我们还了解了如何使用 parameters 或 params 属性将参数动态传递给 PostgresOperator 任务，以及如何通过在 hook_params 属性中传递选项来控制会话参数。