Airflow - 使用 COPY 導出 PostgreSQL 表

Question

在氣流中，我們有 Postgres 到 GCS 和 S3 運算符，這些都使用 SQL 查詢來獲取結果並將其導出到目標。

但是對於大表， COPY方法將是首選。 Airflow 中是否有這樣的操作員可以通過 COPY 導出數據並可以選擇提及 where 條件？

Answer 1

PostgresHook類有一個方法， bulk_dump ，它就是這樣做的，可用於將整個表導出到文件中。 如果您需要導出查詢而不是表，您將需要使用copy_expert方法。

您可以將它們與PythonOperator和一個小函數結合使用，以 (1) 從 Airflow 獲取PostgresHook和 (2) 運行下游任務的導出。

---

編輯：根據您的評論，聽起來您可能會從更明確的演示中受益。 這是一些（未經測試的）服務器代碼作為靈感：

import logging
from tempfile import NamedTemporaryFile

from airflow import models
from airflow.hooks.postgres_hook import PostgresHook
from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
from airflow.operators.python_operator import PythonOperator


# Change these to your identifiers, if needed.
GOOGLE_CONN_ID = "google_cloud_default"
POSTGRES_CONN_ID = "postgres_default"


def copy_to_gcs(copy_sql, file_name, bucket_name):
    gcs_hook = GoogleCloudStorageHook(GOOGLE_CONN_ID)
    pg_hook = PostgresHook.get_hook(POSTGRES_CONN_ID)

    with NamedTemporaryFile(suffix=".csv") as temp_file:
        temp_name = temp_file.name        

        logging.info("Exporting query to file '%s'", temp_name)
        pg_hook.copy_expert(copy_sql, filename=temp_name)

        logging.info("Uploading to %s/%s", bucket_name, file_name)
        gcs_hook.upload(bucket_name, file_name, temp_name)


with models.DAG(...) as dag:
    copy_to_gcs_task = python_operator.PythonOperator(
        task_id="copy_to_gcs",
        python_callable=copy_to_gcs,
        op_kwargs={
            "copy_sql": "SELECT * FROM the_source",
            "file_name": "name_for_gcs.csv",
            "bucket_name": "from_airflow"
        )