Skip to content

Commit 66b3ca1

Browse files
authored
Migrate Google example bigquery_to_mssql to new design AIP-47 (#25174)
related: #22447, #22430
1 parent 911db9f commit 66b3ca1

File tree

1 file changed

+27
-6
lines changed

1 file changed

+27
-6
lines changed

‎airflow/providers/google/cloud/example_dags/example_bigquery_to_mssql.py renamed to ‎tests/system/providers/google/cloud/bigquery/example_bigquery_to_mssql.py

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,18 +30,21 @@
3030
)
3131
from airflow.providers.google.cloud.transfers.bigquery_to_mssql import BigQueryToMsSqlOperator
3232

33+
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
3334
PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
34-
DATASET_NAME = os.environ.get("GCP_BIGQUERY_DATASET_NAME", "test_dataset_transfer")
35+
DAG_ID = "example_bigquery_to_mssql"
36+
37+
DATASET_NAME = f"dataset_{DAG_ID}_{ENV_ID}"
3538
DATA_EXPORT_BUCKET_NAME = os.environ.get("GCP_BIGQUERY_EXPORT_BUCKET_NAME", "INVALID BUCKET NAME")
3639
TABLE = "table_42"
3740
destination_table = "mssql_table_test"
3841

3942
with models.DAG(
40-
"example_bigquery_to_mssql",
41-
schedule_interval=None, # Override to match your needs
43+
DAG_ID,
44+
schedule_interval="@once", # Override to match your needs
4245
start_date=datetime(2021, 1, 1),
4346
catchup=False,
44-
tags=["example"],
47+
tags=["example", "bigquery"],
4548
) as dag:
4649
bigquery_to_mssql = BigQueryToMsSqlOperator(
4750
task_id="bigquery_to_mssql",
@@ -61,10 +64,28 @@
6164
{"name": "salary", "type": "INTEGER", "mode": "NULLABLE"},
6265
],
6366
)
64-
create_dataset >> create_table >> bigquery_to_mssql
6567

6668
delete_dataset = BigQueryDeleteDatasetOperator(
6769
task_id="delete_dataset", dataset_id=DATASET_NAME, delete_contents=True
6870
)
6971

70-
bigquery_to_mssql >> delete_dataset
72+
(
73+
# TEST SETUP
74+
create_dataset
75+
>> create_table
76+
# TEST BODY
77+
>> bigquery_to_mssql
78+
# TEST TEARDOWN
79+
>> delete_dataset
80+
)
81+
82+
from tests.system.utils.watcher import watcher
83+
84+
# This test needs watcher in order to properly mark success/failure
85+
# when "tearDown" task with trigger rule is part of the DAG
86+
list(dag.tasks) >> watcher()
87+
88+
from tests.system.utils import get_test_run # noqa: E402
89+
90+
# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
91+
test_run = get_test_run(dag)

0 commit comments

Comments
 (0)