Postgres 使用 SQLExecuteQueryOperator 的操作指南¶
簡介¶
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 insert 陳述式
-- 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 屬性，以及如何透過在 hook_params 屬性中傳遞選項來控制工作階段參數的不同方式，將參數動態傳遞到 PostgresOperator 任務中。