使用 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 屬性中傳遞選項來控制會話引數。