|
28 | 28 | DataprocCreateClusterOperator,
|
29 | 29 | DataprocDeleteClusterOperator,
|
30 | 30 | DataprocSubmitJobOperator,
|
31 |
| - DataprocUpdateClusterOperator, |
32 | 31 | )
|
33 | 32 | from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
|
34 | 33 | from airflow.utils.trigger_rule import TriggerRule
|
|
53 | 52 | "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 1024},
|
54 | 53 | },
|
55 | 54 | "worker_config": {
|
56 |
| - "num_instances": 2, |
| 55 | + "num_instances": 3, |
57 | 56 | "machine_type_uri": "n1-standard-4",
|
58 | 57 | "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 1024},
|
59 | 58 | },
|
60 | 59 | }
|
61 | 60 |
|
62 |
| -# Update options |
63 |
| -CLUSTER_UPDATE = { |
64 |
| - "config": {"worker_config": {"num_instances": 3}, "secondary_worker_config": {"num_instances": 3}} |
65 |
| -} |
66 |
| -UPDATE_MASK = { |
67 |
| - "paths": ["config.worker_config.num_instances", "config.secondary_worker_config.num_instances"] |
68 |
| -} |
69 |
| - |
70 | 61 | TIMEOUT = {"seconds": 1 * 24 * 60 * 60}
|
71 | 62 |
|
72 | 63 | # Jobs definitions
|
|
87 | 78 | schedule="@once",
|
88 | 79 | start_date=datetime(2021, 1, 1),
|
89 | 80 | catchup=False,
|
90 |
| - tags=["example", "dataproc"], |
| 81 | + tags=["example", "dataproc", "hadoop"], |
91 | 82 | ) as dag:
|
92 | 83 | create_bucket = GCSCreateBucketOperator(
|
93 | 84 | task_id="create_bucket", bucket_name=BUCKET_NAME, project_id=PROJECT_ID
|
|
101 | 92 | cluster_name=CLUSTER_NAME,
|
102 | 93 | )
|
103 | 94 |
|
104 |
| - scale_cluster = DataprocUpdateClusterOperator( |
105 |
| - task_id="scale_cluster", |
106 |
| - cluster_name=CLUSTER_NAME, |
107 |
| - cluster=CLUSTER_UPDATE, |
108 |
| - update_mask=UPDATE_MASK, |
109 |
| - graceful_decommission_timeout=TIMEOUT, |
110 |
| - project_id=PROJECT_ID, |
111 |
| - region=REGION, |
112 |
| - ) |
113 |
| - |
114 | 95 | hadoop_task = DataprocSubmitJobOperator(
|
115 | 96 | task_id="hadoop_task", job=HADOOP_JOB, region=REGION, project_id=PROJECT_ID
|
116 | 97 | )
|
|
127 | 108 | task_id="delete_bucket", bucket_name=BUCKET_NAME, trigger_rule=TriggerRule.ALL_DONE
|
128 | 109 | )
|
129 | 110 |
|
130 |
| - create_bucket >> create_cluster >> scale_cluster >> hadoop_task >> delete_cluster >> delete_bucket |
| 111 | + ( |
| 112 | + # TEST SETUP |
| 113 | + [create_bucket, create_cluster] |
| 114 | + # TEST BODY |
| 115 | + >> hadoop_task |
| 116 | + # TEST TEARDOWN |
| 117 | + >> [delete_cluster, delete_bucket] |
| 118 | + ) |
131 | 119 |
|
132 | 120 | from tests.system.utils.watcher import watcher
|
133 | 121 |
|
|
0 commit comments