使用 SQLExecuteQueryOperator 操作 Postgres 的操作指南¶
简介¶
Apache Airflow 拥有大量的运算符，可用于实现构成工作流的各种任务。Airflow 本质上是一个由任务（节点）和依赖项（边）组成的图（有向无环图）。
由运算符定义或实现的任务是数据管道中的一个工作单元。
本指南的目的是定义使用 SQLExecuteQueryOperator 与 PostgreSQL 数据库交互的任务。
注意
以前，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');
然后，我们可以创建一个 SQLExecuteQueryOperator 任务来填充 pet 表。
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 都使得可以以许多有趣的方式动态传入参数，但它们的用法略有不同，如下例所示。
要查找两个日期之间所有宠物的出生日期，当我们在代码中直接使用 SQL 语句时，我们将使用 parameters 属性
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 文件来整理一下。这次我们将使用从父 BaseOperator 类中免费获得的 params 属性。
-- dags/sql/birth_date.sql
SELECT * FROM pet WHERE birth_date BETWEEN SYMMETRIC {{ params.begin_date }} AND {{ params.end_date }};
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"},
)
启用记录发送到客户端的数据库消息¶
SQLExecuteQueryOperator 提供了 hook_params 属性，该属性允许您将参数添加到 DbApiHook。您可以使用 enable_log_db_messages 来记录 RAISE 语句发出的数据库消息或错误。
call_proc = SQLExecuteQueryOperator(
task_id="call_proc",
conn_id="postgres_default",
sql="call proc();",
hook_params={"enable_log_db_messages": True},
)
将服务器配置参数传递到 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 属性中传递选项来控制会话参数。