Repository: apache/airflow
File: providers/google/src/airflow/providers/google/cloud/operators/dataproc.py (Python)

Code before:
            if not self.subnetwork_uri and self.internal_ip_only:
                raise AirflowException("Set internal_ip_only to true only when you pass a subnetwork_uri.")
            cluster_data[config]["internal_ip_only"] = self.internal_ip_only

        if self.tags:
            cluster_data[config]["tags"] = self.tags

        if self.service_account:
            cluster_data[config]["service_account"] = self.service_account

        if self.service_account_scopes:
            cluster_data[config]["service_account_scopes"] = self.service_account_scopes

        return cluster_data

    def _build_lifecycle_config(self, cluster_data):
        # This variable is created since same string was being used multiple times
        lifecycle_config = "lifecycle_config"
        if self.idle_delete_ttl:
            cluster_data[lifecycle_config]["idle_delete_ttl"] = {"seconds": self.idle_delete_ttl}

        if self.auto_delete_time:
            try:
                from airflow.sdk import timezone
            except (ImportError, ModuleNotFoundError):
                from airflow.utils import timezone
            utc_auto_delete_time = timezone.convert_to_utc(self.auto_delete_time)
            cluster_data[lifecycle_config]["auto_delete_time"] = utc_auto_delete_time.strftime(
                "%Y-%m-%dT%H:%M:%S.%fZ"
            )
        elif self.auto_delete_ttl:
            cluster_data[lifecycle_config]["auto_delete_ttl"] = {"seconds": int(self.auto_delete_ttl)}

        return cluster_data

    def _build_driver_pool(self):
        driver_pool = {
            "node_group": {
                "roles": ["DRIVER"],
                "node_group_config": {"num_instances": self.driver_pool_size},
            },
        }
        if self.driver_pool_id:
            driver_pool["node_group_id"] = self.driver_pool_id
        return driver_pool

    def _build_cluster_data(self):
        if self.zone:
            master_type_uri = (
                f"projects/{self.project_id}/zones/{self.zone}/machineTypes/{self.master_machine_type}"
            )

Review comment (refactor):
Please use `from airflow.providers.common.compat.sdk import timezone` instead and move this to top-level.

Code after:
            if not self.subnetwork_uri and self.internal_ip_only:
                raise AirflowException("Set internal_ip_only to true only when you pass a subnetwork_uri.")
            cluster_data[config]["internal_ip_only"] = self.internal_ip_only

        if self.tags:
            cluster_data[config]["tags"] = self.tags

        if self.service_account:
            cluster_data[config]["service_account"] = self.service_account

        if self.service_account_scopes:
            cluster_data[config]["service_account_scopes"] = self.service_account_scopes

        return cluster_data

    def _build_lifecycle_config(self, cluster_data):
        # This variable is created since same string was being used multiple times
        lifecycle_config = "lifecycle_config"
        if self.idle_delete_ttl:
            cluster_data[lifecycle_config]["idle_delete_ttl"] = {"seconds": self.idle_delete_ttl}

        if self.auto_delete_time:
            utc_auto_delete_time = timezone.convert_to_utc(self.auto_delete_time)
            cluster_data[lifecycle_config]["auto_delete_time"] = utc_auto_delete_time.strftime(
                "%Y-%m-%dT%H:%M:%S.%fZ"
            )
        elif self.auto_delete_ttl:
            cluster_data[lifecycle_config]["auto_delete_ttl"] = {"seconds": int(self.auto_delete_ttl)}

        return cluster_data

    def _build_driver_pool(self):
        driver_pool = {
            "node_group": {
                "roles": ["DRIVER"],
                "node_group_config": {"num_instances": self.driver_pool_size},
            },
        }
        if self.driver_pool_id:
            driver_pool["node_group_id"] = self.driver_pool_id
        return driver_pool

    def _build_cluster_data(self):
        if self.zone:
            master_type_uri = (
                f"projects/{self.project_id}/zones/{self.zone}/machineTypes/{self.master_machine_type}"
            )
            worker_type_uri = (
                f"projects/{self.project_id}/zones/{self.zone}/machineTypes/{self.worker_machine_type}"
            )
        else: