Airflow DAG：如何使用 Python 运算符而不是 BigQuery 运算符将数据插入表中？

Question

我正在尝试使用简单的 Python 运算符而不是 BigQuery 运算符将一些数据插入表中，但我不确定如何实现它。我正在尝试以 Airflow DAG 的形式实现它。

我编写了一个简单的 DAG，并且设法使用以下方法将数据从 GCS Bucket 插入到 BigQuery，但我想使用 Python 运算符而不是 BigQuery 来执行此操作：

load_csv = gcs_to_bq.GoogleCloudStorageToBigQueryOperator(
    task_id='gcs_to_bq_example',
    bucket='cloud-samples-data',
    source_objects=['bigquery/us-states/us-states.csv'],
    destination_project_dataset_table='airflow_test.gcs_to_bq_table',
    schema_fields=[
        {'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'post_abbr', 'type': 'STRING', 'mode': 'NULLABLE'},
    ],
    write_disposition='WRITE_TRUNCATE',
    dag=dag)

我想使用简单的 Python 运算符而不是 BigQuery 来实现上述目标。

BQ 到 GCS： BigQuery 到 GCS：

# from google.cloud import bigquery
# client = bigquery.Client()
# bucket_name = 'my-bucket'
project = "bigquery-public-data"
dataset_id = "samples"
table_id = "shakespeare"

destination_uri = "gs://{}/{}".format(bucket_name, "shakespeare.csv")
dataset_ref = bigquery.DatasetReference(project, dataset_id)
table_ref = dataset_ref.table(table_id)

extract_job = client.extract_table(
    table_ref,
    destination_uri,
    # Location must match that of the source table.
    location="US",
)  # API request
extract_job.result()  # Waits for job to complete.

print(
    "Exported {}:{}.{} to {}".format(project, dataset_id, table_id, destination_uri)
)

Answer 1

您可以在PythonOperator中使用BigQueryPythonclient将GCS文件插入BigQuery，示例：

PythonOperator(
    task_id="gcs_to_bq",
    op_kwargs={
      'dataset': 'dataset',
      'table': 'table'
    },
    python_callable=load_gcs_files_to_bq
)

def load_gcs_files_to_bq(dataset, table):
   from google.cloud import bigquery

   # Construct a BigQuery client object.
   client = bigquery.Client()

   # TODO(developer): Set table_id to the ID of the table to create.
   table_id = f"your-project.{dataset}.{table}"

   job_config = bigquery.LoadJobConfig(
        schema=[
             bigquery.SchemaField("name", "STRING"),
             bigquery.SchemaField("post_abbr", "STRING"),
        ],
        skip_leading_rows=1,
        # The source format defaults to CSV, so the line below is optional.
        source_format=bigquery.SourceFormat.CSV,
    )
    
    uri = "gs://cloud-samples-data/bigquery/us-states/us-states.csv"

    load_job = client.load_table_from_uri(
        uri, table_id, job_config=job_config
    )  # Make an API request.

    load_job.result()  # Waits for the job to complete.

    destination_table = client.get_table(table_id)  # Make an API request.
    print("Loaded {} rows.".format(destination_table.num_rows))