feat: Support BigQueryToBigQueryOperator (#86)

* feat: Support BQ to BQ operator * feat: Added BigQueryToBigQueryOperator in pipeline.yaml sample
GoogleCloudPlatform · Jun 14, 2021 · fd26476 · fd26476
1 parent d832327
commit fd26476
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 1 deletion.
diff --git a/samples/pipeline.yaml b/samples/pipeline.yaml
@@ -278,6 +278,16 @@ dag:
           limit_memory: "250M"
           limit_cpu: "1"
 
+    - operator: "BigQueryToBigQueryOperator"
+      description: "Task to run a BQ to BQ operator"
+
+      args:
+        task_id: "sample_bq_to_bq_task"
+        source_project_dataset_tables: ["{{ var.json.DATASET_FOLDER_NAME.PIPELINE_NAME.source_project_dataset_table }}"]
+        destination_project_dataset_table: "{{ var.json.DATASET_FOLDER_NAME.PIPELINE_NAME.destination_project_dataset_table }}"
+        impersonation_chain: "{{ var.json.DATASET_FOLDER_NAME.service_account }}"
+        write_disposition: "WRITE_TRUNCATE"
+
   graph_paths:
     # This is where you specify the relationships (i.e. directed paths/edges)
     # among the tasks specified above. Use the bitshift operator to define the
@@ -286,5 +296,5 @@ dag:
     # For more info, see
     # https://airflow.apache.org/docs/apache-airflow/stable/tutorial.html#setting-up-dependencies
     - "sample_bash_task >> [sample_gcs_to_bq_task, sample_gcs_to_gcs_task]"
-    - "sample_gcs_to_bq_task >> sample_bq_sql_task"
+    - "sample_gcs_to_bq_task >> [sample_bq_sql_task, sample_bq_to_bq_task]"
     - "sample_bq_sql_task >> sample_gcs_delete_task"
diff --git a/scripts/dag_imports.json b/scripts/dag_imports.json
@@ -16,6 +16,10 @@
             "import": "from airflow.contrib.operators import bigquery_operator",
             "class": "bigquery_operator.BigQueryOperator"
         },
+        "BigQueryToBigQueryOperator": {
+            "import": "from airflow.contrib.operators import bigquery_to_bigquery",
+            "class": "bigquery_to_bigquery.BigQueryToBigQueryOperator"
+        },
         "KubernetesPodOperator": {
             "import": "from airflow.contrib.operators import kubernetes_pod_operator",
             "class": "kubernetes_pod_operator.KubernetesPodOperator"

diff --git a/scripts/generate_dag.py b/scripts/generate_dag.py
@@ -33,6 +33,7 @@
     "GoogleCloudStorageToGoogleCloudStorageOperator",
     "GoogleCloudStorageDeleteOperator",
     "BigQueryOperator",
+    "BigQueryToBigQueryOperator",
     "KubernetesPodOperator",
 }