Migration of System Tests: Cloud BigQuery Data Transfer (AIP-47) (#27312

)

docs/apache-airflow-providers-google/operators/cloud/bigquery_dts.rst

@@ -47,15 +47,15 @@ for example :class:`~airflow.providers.google.cloud.operators.bigquery_dts.BigQu
 scheduling option is present in passed configuration. If present then nothing is done, otherwise it's value is
 set to ``True``.
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_bigquery_dts.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/bigquery/example_bigquery_dts.py
     :language: python
     :start-after: [START howto_bigquery_dts_create_args]
     :end-before: [END howto_bigquery_dts_create_args]
 
 You can create the operator with or without project id. If project id is missing
 it will be retrieved from the Google Cloud connection used. Basic usage of the operator:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_bigquery_dts.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/bigquery/example_bigquery_dts.py
     :language: python
     :dedent: 4
     :start-after: [START howto_bigquery_create_data_transfer]
@@ -78,7 +78,7 @@ To delete DTS transfer configuration you can use
 
 Basic usage of the operator:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_bigquery_dts.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/bigquery/example_bigquery_dts.py
     :language: python
     :dedent: 4
     :start-after: [START howto_bigquery_delete_data_transfer]
@@ -99,7 +99,7 @@ Start manual transfer runs to be executed now with schedule_time equal to curren
 
 Basic usage of the operator:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_bigquery_dts.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/bigquery/example_bigquery_dts.py
     :language: python
     :dedent: 4
     :start-after: [START howto_bigquery_start_transfer]
@@ -112,7 +112,7 @@ parameters which allows you to dynamically determine values.
 To check if operation succeeded you can use
 :class:`~airflow.providers.google.cloud.sensors.bigquery_dts.BigQueryDataTransferServiceTransferRunSensor`.
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_bigquery_dts.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/bigquery/example_bigquery_dts.py
     :language: python
     :dedent: 4
     :start-after: [START howto_bigquery_dts_sensor]

airflow/providers/google/cloud/example_dags/example_bigquery_dts.py → tests/system/providers/google/cloud/bigquery/example_bigquery_dts.py

@@ -23,58 +23,97 @@
 import os
 import time
 from datetime import datetime
+from pathlib import Path
 from typing import cast
 
 from airflow import models
+from airflow.models.baseoperator import chain
 from airflow.models.xcom_arg import XComArg
+from airflow.providers.google.cloud.operators.bigquery import (
+    BigQueryCreateEmptyDatasetOperator,
+    BigQueryCreateEmptyTableOperator,
+    BigQueryDeleteDatasetOperator,
+)
 from airflow.providers.google.cloud.operators.bigquery_dts import (
     BigQueryCreateDataTransferOperator,
     BigQueryDataTransferServiceStartTransferRunsOperator,
     BigQueryDeleteDataTransferConfigOperator,
 )
+from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
 from airflow.providers.google.cloud.sensors.bigquery_dts import BigQueryDataTransferServiceTransferRunSensor
+from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator
+from airflow.utils.trigger_rule import TriggerRule
+
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT")
+
+DAG_ID = "example_gcp_bigquery_dts"
 
-GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
-BUCKET_URI = os.environ.get("GCP_DTS_BUCKET_URI", "gs://INVALID BUCKET NAME/bank-marketing.csv")
-GCP_DTS_BQ_DATASET = os.environ.get("GCP_DTS_BQ_DATASET", "test_dts")
-GCP_DTS_BQ_TABLE = os.environ.get("GCP_DTS_BQ_TABLE", "GCS_Test")
+BUCKET_NAME = f"bucket-{DAG_ID}-{ENV_ID}"
+
+FILE_NAME = "us-states.csv"
+CURRENT_FOLDER = Path(__file__).parent
+FILE_LOCAL_PATH = str(Path(CURRENT_FOLDER) / "resources" / FILE_NAME)
+BUCKET_URI = f"gs://{BUCKET_NAME}/{FILE_NAME}"
+
+DATASET_NAME = f"dataset_{DAG_ID}_{ENV_ID}"
+DTS_BQ_TABLE = "DTS_BQ_TABLE"
 
 # [START howto_bigquery_dts_create_args]
 
 # In the case of Airflow, the customer needs to create a transfer
 # config with the automatic scheduling disabled and then trigger
 # a transfer run using a specialized Airflow operator
-schedule_options = {"disable_auto_scheduling": True}
-
-PARAMS = {
-    "field_delimiter": ",",
-    "max_bad_records": "0",
-    "skip_leading_rows": "1",
-    "data_path_template": BUCKET_URI,
-    "destination_table_name_template": GCP_DTS_BQ_TABLE,
-    "file_format": "CSV",
-}
-
 TRANSFER_CONFIG = {
-    "destination_dataset_id": GCP_DTS_BQ_DATASET,
-    "display_name": "GCS Test Config",
+    "destination_dataset_id": DATASET_NAME,
+    "display_name": "test data transfer",
     "data_source_id": "google_cloud_storage",
-    "schedule_options": schedule_options,
-    "params": PARAMS,
+    "schedule_options": {"disable_auto_scheduling": True},
+    "params": {
+        "field_delimiter": ",",
+        "max_bad_records": "0",
+        "skip_leading_rows": "1",
+        "data_path_template": BUCKET_URI,
+        "destination_table_name_template": DTS_BQ_TABLE,
+        "file_format": "CSV",
+    },
 }
 
 # [END howto_bigquery_dts_create_args]
 
 with models.DAG(
-    "example_gcp_bigquery_dts",
+    DAG_ID,
+    schedule="@once",
     start_date=datetime(2021, 1, 1),
     catchup=False,
-    tags=["example"],
+    tags=["example", "bigquery"],
 ) as dag:
+
+    create_bucket = GCSCreateBucketOperator(
+        task_id="create_bucket", bucket_name=BUCKET_NAME, project_id=PROJECT_ID
+    )
+    upload_file = LocalFilesystemToGCSOperator(
+        task_id="upload_file",
+        src=FILE_LOCAL_PATH,
+        dst=FILE_NAME,
+        bucket=BUCKET_NAME,
+    )
+    create_dataset = BigQueryCreateEmptyDatasetOperator(task_id="create_dataset", dataset_id=DATASET_NAME)
+
+    create_table = BigQueryCreateEmptyTableOperator(
+        task_id="create_table",
+        dataset_id=DATASET_NAME,
+        table_id=DTS_BQ_TABLE,
+        schema_fields=[
+            {"name": "name", "type": "STRING", "mode": "REQUIRED"},
+            {"name": "post_abbr", "type": "STRING", "mode": "NULLABLE"},
+        ],
+    )
+
     # [START howto_bigquery_create_data_transfer]
     gcp_bigquery_create_transfer = BigQueryCreateDataTransferOperator(
         transfer_config=TRANSFER_CONFIG,
-        project_id=GCP_PROJECT_ID,
+        project_id=PROJECT_ID,
         task_id="gcp_bigquery_create_transfer",
     )
 
@@ -103,11 +142,49 @@
         transfer_config_id=transfer_config_id, task_id="gcp_bigquery_delete_transfer"
     )
     # [END howto_bigquery_delete_data_transfer]
+    gcp_bigquery_delete_transfer.trigger_rule = TriggerRule.ALL_DONE
+
+    delete_dataset = BigQueryDeleteDatasetOperator(
+        task_id="delete_dataset",
+        dataset_id=DATASET_NAME,
+        delete_contents=True,
+        trigger_rule=TriggerRule.ALL_DONE,
+    )
 
-    gcp_run_sensor >> gcp_bigquery_delete_transfer
+    delete_bucket = GCSDeleteBucketOperator(
+        task_id="delete_bucket", bucket_name=BUCKET_NAME, trigger_rule=TriggerRule.ALL_DONE
+    )
 
     # Task dependencies created via `XComArgs`:
     #   gcp_bigquery_create_transfer >> gcp_bigquery_start_transfer
     #   gcp_bigquery_create_transfer >> gcp_run_sensor
     #   gcp_bigquery_start_transfer >> gcp_run_sensor
     #   gcp_bigquery_create_transfer >> gcp_bigquery_delete_transfer
+
+    chain(
+        # TEST SETUP
+        create_bucket,
+        upload_file,
+        create_dataset,
+        create_table,
+        # TEST BODY
+        gcp_bigquery_create_transfer,
+        gcp_bigquery_start_transfer,
+        gcp_run_sensor,
+        gcp_bigquery_delete_transfer,
+        # TEST TEARDOWN
+        delete_dataset,
+        delete_bucket,
+    )
+
+    from tests.system.utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
+
+from tests.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)