Airflow - 如何從 BigQuery 表中獲取數據並將其用作列表？

Question

我正在嘗試獲取一列，然后使用值來創建文件名。

我嘗試了以下方法，它應該創建一個 csv，其名稱為指定列中的第一個值。 它說列表是空的，但當我嘗試使用它時

bq_data = []
get_data = BigQueryGetDataOperator(
    task_id='get_data_from_bq',
    dataset_id='SK22',
    table_id='current_times',
    max_results='100',
    selected_fields='current_timestamps',
)


def process_data_from_bq(**kwargs):
    ti = kwargs['ti']
    global bq_data
    bq_data = ti.xcom_pull(task_ids='get_data_from_bq')


process_data = PythonOperator(
        task_id='process_data_from_bq',
        python_callable=process_data_from_bq,
        provide_context=True)
run_export = BigQueryToCloudStorageOperator(
        task_id=f"save_data_on_storage{str(bq_data[0])}",
        source_project_dataset_table="a-data-set",
        destination_cloud_storage_uris=[f"gs://europe-west1-airflow-bucket/data/test{bq_data[0]}.csv"],
        export_format="CSV",
        field_delimiter=",",
        print_header=False,
        dag=dag,
    )

get_data >> process_data >> run_export

Answer 1

我認為不需要在BigQueryGetDataOperator和BigQueryToCloudStorageOperator之間使用PythonOperator ，您可以直接在BigQueryToCloudStorageOperator中使用xcom pull ：

get_data = BigQueryGetDataOperator(
    task_id='get_data_from_bq',
    dataset_id='SK22',
    table_id='current_times',
    max_results='100',
    selected_fields='current_timestamps',
)

run_export = BigQueryToCloudStorageOperator(
        task_id="save_data_on_storage",
        source_project_dataset_table="a-data-set",
        destination_cloud_storage_uris=[f"gs://europe-west1-airflow-bucket/data/test" + "{{ ti.xcom_pull(task_ids='get_data_from_bq')[0] }}" + ".csv"],
        export_format="CSV",
        field_delimiter=",",
        print_header=False,
        dag=dag,
    )

get_data >> run_export

destination_cloud_storage_uris是一個模板化參數，您可以在其中傳遞Jinja模板語法。

我沒有測試語法，但它應該可以工作。

我也不建議你使用像bq_data這樣的全局變量在operator之間傳遞數據，因為它不起作用，你需要找到一種方法直接在operator中使用xcom （ Jinja模板或訪問operator的當前Context ).

我還注意到您沒有使用最新的Airflow運算符：

• BigQueryToCloudStorageOperator -> BigQueryToGCSOperator

如果您想使用 BigQueryGetDataOperator 運算符提供的所有列表並從中計算目標 URI 列表，我建議您使用另一種解決方案：

from typing import List, Dict

from airflow.providers.google.cloud.transfers.bigquery_to_gcs import BigQueryToGCSOperator

class CustomBigQueryToGCSOperator(BigQueryToGCSOperator):

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

    def execute(self, context):
        task_instance = context['task_instance']
        data_from_bq: List[Dict] = task_instance.xcom_pull('get_data_from_bq')

        destination_cloud_storage_uris: List[str] = list(map(self.to_destination_cloud_storage_uris, data_from_bq))

        self.destination_cloud_storage_uris = destination_cloud_storage_uris

        super(CustomBigQueryToGCSOperator, self).execute(context)

    def to_destination_cloud_storage_uris(self, data_from_bq: Dict) -> str:
        return f"gs://europe-west1-airflow-bucket/data/test{data_from_bq['your_field']}.csv"

一些解釋：

• 我創建了一個擴展BigQueryToGCSOperator的自定義運算符
• 在execute方法中，我可以訪問操作員的當前上下文
• 從上下文中，我可以從BigQueryGetDataOperator提供的BQ中檢索列表。 我假設這是一個 Dict 列表，但你必須確認這一點
• 我從這個字典列表中計算出一個目標GCS URI 列表
• 我將計算出的目標GCS URI 分配給運算符中的相應字段

此解決方案的優點是，您可以更靈活地應用基於 xcom 值的邏輯。

缺點是它有點冗長。