Skip to content

Commit c499b6b

Browse files
authored
Small refactors in ClusterGenerator of dataproc (#30714)
1 parent 76ebc9b commit c499b6b

File tree

1 file changed

+15
-11
lines changed

1 file changed

+15
-11
lines changed

β€Žairflow/providers/google/cloud/operators/dataproc.py

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -236,46 +236,50 @@ def _get_init_action_timeout(self) -> dict:
236236
)
237237

238238
def _build_gce_cluster_config(self, cluster_data):
239+
# This variable is created since same string was being used multiple times
240+
config = "gce_cluster_config"
239241
if self.zone:
240242
zone_uri = f"https://www.googleapis.com/compute/v1/projects/{self.project_id}/zones/{self.zone}"
241-
cluster_data["gce_cluster_config"]["zone_uri"] = zone_uri
243+
cluster_data[config]["zone_uri"] = zone_uri
242244

243245
if self.metadata:
244-
cluster_data["gce_cluster_config"]["metadata"] = self.metadata
246+
cluster_data[config]["metadata"] = self.metadata
245247

246248
if self.network_uri:
247-
cluster_data["gce_cluster_config"]["network_uri"] = self.network_uri
249+
cluster_data[config]["network_uri"] = self.network_uri
248250

249251
if self.subnetwork_uri:
250-
cluster_data["gce_cluster_config"]["subnetwork_uri"] = self.subnetwork_uri
252+
cluster_data[config]["subnetwork_uri"] = self.subnetwork_uri
251253

252254
if self.internal_ip_only:
253255
if not self.subnetwork_uri:
254256
raise AirflowException("Set internal_ip_only to true only when you pass a subnetwork_uri.")
255-
cluster_data["gce_cluster_config"]["internal_ip_only"] = True
257+
cluster_data[config]["internal_ip_only"] = True
256258

257259
if self.tags:
258-
cluster_data["gce_cluster_config"]["tags"] = self.tags
260+
cluster_data[config]["tags"] = self.tags
259261

260262
if self.service_account:
261-
cluster_data["gce_cluster_config"]["service_account"] = self.service_account
263+
cluster_data[config]["service_account"] = self.service_account
262264

263265
if self.service_account_scopes:
264-
cluster_data["gce_cluster_config"]["service_account_scopes"] = self.service_account_scopes
266+
cluster_data[config]["service_account_scopes"] = self.service_account_scopes
265267

266268
return cluster_data
267269

268270
def _build_lifecycle_config(self, cluster_data):
271+
# This variable is created since same string was being used multiple times
272+
lifecycle_config = "lifecycle_config"
269273
if self.idle_delete_ttl:
270-
cluster_data["lifecycle_config"]["idle_delete_ttl"] = {"seconds": self.idle_delete_ttl}
274+
cluster_data[lifecycle_config]["idle_delete_ttl"] = {"seconds": self.idle_delete_ttl}
271275

272276
if self.auto_delete_time:
273277
utc_auto_delete_time = timezone.convert_to_utc(self.auto_delete_time)
274-
cluster_data["lifecycle_config"]["auto_delete_time"] = utc_auto_delete_time.strftime(
278+
cluster_data[lifecycle_config]["auto_delete_time"] = utc_auto_delete_time.strftime(
275279
"%Y-%m-%dT%H:%M:%S.%fZ"
276280
)
277281
elif self.auto_delete_ttl:
278-
cluster_data["lifecycle_config"]["auto_delete_ttl"] = {"seconds": int(self.auto_delete_ttl)}
282+
cluster_data[lifecycle_config]["auto_delete_ttl"] = {"seconds": int(self.auto_delete_ttl)}
279283

280284
return cluster_data
281285

0 commit comments

Comments
 (0)