使用 SQLExecuteQueryOperator 操作 Postgres 的指南¶
引言¶
Apache Airflow 拥有丰富的 operator 宝库，可用于实现构成您工作流的各种任务。Airflow 本质上是一个图 (有向无环图)，由任务 (节点) 和依赖关系 (边) 组成。
由 operator 定义或实现的任务是数据管道中的一个工作单元。
本指南的目的是定义使用 SQLExecuteQueryOperator 与 PostgreSQL 数据库交互的任务。
注意
之前，使用 PostgresOperator 来执行此类操作。弃用后，该 operator 已被移除。请改用 SQLExecuteQueryOperator。
使用 SQLExecuteQueryOperator 的常见数据库操作¶
要使用 SQLExecuteQueryOperator 执行 PostgreSQL 请求，需要两个参数：sql 和 conn_id。这两个参数最终会被传递给直接与 Postgres 数据库交互的 DbApiHook 对象。
创建 Postgres 数据库表¶
下面的代码片段基于 Airflow-2.0
tests/system/postgres/example_postgres.py
# 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 语句直接写在 operator 中不太优雅，并且将来会带来维护上的困难。为防止这种情况，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 都使得以许多有趣的方式动态传递参数成为可能，但它们的用法略有不同，如下面的示例所示。
要查找两个日期之间所有宠物的出生日期，当我们在代码中直接使用 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 参数，以便指定 命令行选项，这些选项在连接开始时发送到服务器。
tests/system/postgres/example_postgres.py
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 应如下所示：
tests/system/postgres/example_postgres.py
# 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 属性中传递选项来控制会话参数。