apache-airflow-providers-google 14.1.0__py3-none-any.whl → 15.0.0rc1__py3-none-any.whl

airflow/providers/google/cloud/log/gcs_task_handler.py

@@ -25,6 +25,8 @@ from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+import attrs
+
 # not sure why but mypy complains on missing `storage` but it is clearly there and is importable
 from google.cloud import storage  # type: ignore[attr-defined]
 
@@ -42,6 +44,8 @@ from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
     from airflow.models.taskinstance import TaskInstance
+    from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
+    from airflow.utils.log.file_task_handler import LogMessages, LogSourceInfo
 
 _DEFAULT_SCOPESS = frozenset(
     [
@@ -52,6 +56,126 @@ _DEFAULT_SCOPESS = frozenset(
 logger = logging.getLogger(__name__)
 
 
+@attrs.define
+class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
+    remote_base: str
+    base_log_folder: Path = attrs.field(converter=Path)
+    delete_local_copy: bool
+
+    gcp_key_path: str | None
+    gcp_keyfile_dict: dict | None
+    scopes: Collection[str] | None
+    project_id: str
+
+    def upload(self, path: os.PathLike, ti: RuntimeTI):
+        """Upload the given log path to the remote storage."""
+        path = Path(path)
+        if path.is_absolute():
+            local_loc = path
+            remote_loc = os.path.join(self.remote_base, path.relative_to(self.base_log_folder))
+        else:
+            local_loc = self.base_log_folder.joinpath(path)
+            remote_loc = os.path.join(self.remote_base, path)
+
+        if local_loc.is_file():
+            # read log and remove old logs to get just the latest additions
+            log = local_loc.read_text()
+            has_uploaded = self.write(log, remote_loc)
+            if has_uploaded and self.delete_local_copy:
+                shutil.rmtree(os.path.dirname(local_loc))
+
+    @cached_property
+    def hook(self) -> GCSHook | None:
+        """Returns GCSHook if remote_log_conn_id configured."""
+        conn_id = conf.get("logging", "remote_log_conn_id", fallback=None)
+        if conn_id:
+            try:
+                return GCSHook(gcp_conn_id=conn_id)
+            except AirflowNotFoundException:
+                pass
+        return None
+
+    @cached_property
+    def client(self) -> storage.Client:
+        """Returns GCS Client."""
+        if self.hook:
+            credentials, project_id = self.hook.get_credentials_and_project_id()
+        else:
+            credentials, project_id = get_credentials_and_project_id(
+                key_path=self.gcp_key_path,
+                keyfile_dict=self.gcp_keyfile_dict,
+                scopes=self.scopes,
+                disable_logging=True,
+            )
+        return storage.Client(
+            credentials=credentials,
+            client_info=CLIENT_INFO,
+            project=self.project_id if self.project_id else project_id,
+        )
+
+    def write(self, log: str, remote_log_location: str) -> bool:
+        """
+        Write the log to the remote location and return `True`; fail silently and return `False` on error.
+
+        :param log: the log to write to the remote_log_location
+        :param remote_log_location: the log's location in remote storage
+        :return: whether the log is successfully written to remote location or not.
+        """
+        try:
+            blob = storage.Blob.from_string(remote_log_location, self.client)
+            old_log = blob.download_as_bytes().decode()
+            log = f"{old_log}\n{log}" if old_log else log
+        except Exception as e:
+            if not self.no_log_found(e):
+                self.log.warning("Error checking for previous log: %s", e)
+        try:
+            blob = storage.Blob.from_string(remote_log_location, self.client)
+            blob.upload_from_string(log, content_type="text/plain")
+        except Exception as e:
+            self.log.error("Could not write logs to %s: %s", remote_log_location, e)
+            return False
+        return True
+
+    @staticmethod
+    def no_log_found(exc):
+        """
+        Given exception, determine whether it is result of log not found.
+
+        :meta private:
+        """
+        return (exc.args and isinstance(exc.args[0], str) and "No such object" in exc.args[0]) or getattr(
+            exc, "resp", {}
+        ).get("status") == "404"
+
+    def read(self, relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages | None]:
+        messages = []
+        logs = []
+        remote_loc = os.path.join(self.remote_base, relative_path)
+        uris = []
+        bucket, prefix = _parse_gcs_url(remote_loc)
+        blobs = list(self.client.list_blobs(bucket_or_name=bucket, prefix=prefix))
+
+        if blobs:
+            uris = [f"gs://{bucket}/{b.name}" for b in blobs]
+            if AIRFLOW_V_3_0_PLUS:
+                messages = uris
+            else:
+                messages.extend(["Found remote logs:", *[f"  * {x}" for x in sorted(uris)]])
+        else:
+            return messages, None
+
+        try:
+            for key in sorted(uris):
+                blob = storage.Blob.from_string(key, self.client)
+                remote_log = blob.download_as_bytes().decode()
+                if remote_log:
+                    logs.append(remote_log)
+        except Exception as e:
+            if not AIRFLOW_V_3_0_PLUS:
+                messages.append(f"Unable to read remote log {e}")
+        return messages, logs
+
+
 class GCSTaskHandler(FileTaskHandler, LoggingMixin):
     """
     GCSTaskHandler is a python log handler that handles and reads task instance logs.
@@ -91,45 +215,19 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
     ):
         super().__init__(base_log_folder)
         self.handler: logging.FileHandler | None = None
-        self.remote_base = gcs_log_folder
         self.log_relative_path = ""
         self.closed = False
         self.upload_on_close = True
-        self.gcp_key_path = gcp_key_path
-        self.gcp_keyfile_dict = gcp_keyfile_dict
-        self.scopes = gcp_scopes
-        self.project_id = project_id
-        self.delete_local_copy = kwargs.get(
-            "delete_local_copy", conf.getboolean("logging", "delete_local_logs")
-        )
-
-    @cached_property
-    def hook(self) -> GCSHook | None:
-        """Returns GCSHook if remote_log_conn_id configured."""
-        conn_id = conf.get("logging", "remote_log_conn_id", fallback=None)
-        if conn_id:
-            try:
-                return GCSHook(gcp_conn_id=conn_id)
-            except AirflowNotFoundException:
-                pass
-        return None
-
-    @cached_property
-    def client(self) -> storage.Client:
-        """Returns GCS Client."""
-        if self.hook:
-            credentials, project_id = self.hook.get_credentials_and_project_id()
-        else:
-            credentials, project_id = get_credentials_and_project_id(
-                key_path=self.gcp_key_path,
-                keyfile_dict=self.gcp_keyfile_dict,
-                scopes=self.scopes,
-                disable_logging=True,
-            )
-        return storage.Client(
-            credentials=credentials,
-            client_info=CLIENT_INFO,
-            project=self.project_id if self.project_id else project_id,
+        self.io = GCSRemoteLogIO(
+            base_log_folder=base_log_folder,
+            remote_base=gcs_log_folder,
+            delete_local_copy=kwargs.get(
+                "delete_local_copy", conf.getboolean("logging", "delete_local_logs")
+            ),
+            gcp_key_path=gcp_key_path,
+            gcp_keyfile_dict=gcp_keyfile_dict,
+            scopes=gcp_scopes,
+            project_id=project_id,
         )
 
     def set_context(self, ti: TaskInstance, *, identifier: str | None = None) -> None:
@@ -140,6 +238,8 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         if TYPE_CHECKING:
             assert self.handler is not None
 
+        self.ti = ti
+
         full_path = self.handler.baseFilename
         self.log_relative_path = Path(full_path).relative_to(self.local_base).as_posix()
         is_trigger_log_context = getattr(ti, "is_trigger_log_context", False)
@@ -159,91 +259,23 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         if not self.upload_on_close:
             return
 
-        local_loc = os.path.join(self.local_base, self.log_relative_path)
-        remote_loc = os.path.join(self.remote_base, self.log_relative_path)
-        if os.path.exists(local_loc):
-            # read log and remove old logs to get just the latest additions
-            with open(local_loc) as logfile:
-                log = logfile.read()
-            gcs_write = self.gcs_write(log, remote_loc)
-            if gcs_write and self.delete_local_copy:
-                shutil.rmtree(os.path.dirname(local_loc))
+        if hasattr(self, "ti"):
+            self.io.upload(self.log_relative_path, self.ti)
 
         # Mark closed so we don't double write if close is called twice
         self.closed = True
 
-    def _add_message(self, msg):
-        filename, lineno, func, stackinfo = logger.findCaller()
-        record = logging.LogRecord("", logging.INFO, filename, lineno, msg + "\n", None, None, func=func)
-        return self.format(record)
+    def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[LogSourceInfo, LogMessages]:
+        # Explicitly getting log relative path is necessary as the given
+        # task instance might be different than task instance passed in
+        # in set_context method.
+        worker_log_rel_path = self._render_filename(ti, try_number)
 
-    def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[list[str], list[str]]:
-        # Explicitly getting log relative path is necessary because this method
-        # is called from webserver from TaskLogReader, where we don't call set_context
-        # and can read logs for different TIs in each request
-        messages = []
-        logs = []
-        worker_log_relative_path = self._render_filename(ti, try_number)
-        remote_loc = os.path.join(self.remote_base, worker_log_relative_path)
-        uris = []
-        bucket, prefix = _parse_gcs_url(remote_loc)
-        blobs = list(self.client.list_blobs(bucket_or_name=bucket, prefix=prefix))
+        messages, logs = self.io.read(worker_log_rel_path, ti)
 
-        if blobs:
-            uris = [f"gs://{bucket}/{b.name}" for b in blobs]
-            if AIRFLOW_V_3_0_PLUS:
-                messages = uris
-            else:
-                messages.extend(["Found remote logs:", *[f"  * {x}" for x in sorted(uris)]])
-        else:
+        if logs is None:
+            logs = []
             if not AIRFLOW_V_3_0_PLUS:
                 messages.append(f"No logs found in GCS; ti={ti}")
-        try:
-            for key in sorted(uris):
-                blob = storage.Blob.from_string(key, self.client)
-                remote_log = blob.download_as_bytes().decode()
-                if remote_log:
-                    logs.append(remote_log)
-        except Exception as e:
-            if not AIRFLOW_V_3_0_PLUS:
-                messages.append(f"Unable to read remote log {e}")
-        return messages, logs
-
-    def gcs_write(self, log, remote_log_location) -> bool:
-        """
-        Write the log to the remote location and return `True`; fail silently and return `False` on error.
-
-        :param log: the log to write to the remote_log_location
-        :param remote_log_location: the log's location in remote storage
-        :return: whether the log is successfully written to remote location or not.
-        """
-        try:
-            blob = storage.Blob.from_string(remote_log_location, self.client)
-            old_log = blob.download_as_bytes().decode()
-            log = f"{old_log}\n{log}" if old_log else log
-        except Exception as e:
-            if not self.no_log_found(e):
-                log += self._add_message(
-                    f"Error checking for previous log; if exists, may be overwritten: {e}"
-                )
-                self.log.warning("Error checking for previous log: %s", e)
-        try:
-            blob = storage.Blob.from_string(remote_log_location, self.client)
-            blob.upload_from_string(log, content_type="text/plain")
-        except Exception as e:
-            self.log.error("Could not write logs to %s: %s", remote_log_location, e)
-            return False
-        return True
 
-    @staticmethod
-    def no_log_found(exc):
-        """
-        Given exception, determine whether it is result of log not found.
-
-        :meta private:
-        """
-        if (exc.args and isinstance(exc.args[0], str) and "No such object" in exc.args[0]) or getattr(
-            exc, "resp", {}
-        ).get("status") == "404":
-            return True
-        return False
+        return messages, logs

airflow/providers/google/cloud/openlineage/facets.py

@@ -25,43 +25,43 @@ from airflow.providers.google import __version__ as provider_version
 
 if TYPE_CHECKING:
     from openlineage.client.generated.base import RunFacet
-else:
+
+try:
     try:
-        try:
-            from openlineage.client.generated.base import RunFacet
-        except ImportError:  # Old OpenLineage client is used
-            from openlineage.client.facet import BaseFacet as RunFacet
+        from openlineage.client.generated.base import RunFacet
+    except ImportError:  # Old OpenLineage client is used
+        from openlineage.client.facet import BaseFacet as RunFacet  # type: ignore[assignment]
 
-        @define
-        class BigQueryJobRunFacet(RunFacet):
-            """
-            Facet that represents relevant statistics of bigquery run.
+    @define
+    class BigQueryJobRunFacet(RunFacet):
+        """
+        Facet that represents relevant statistics of bigquery run.
 
-            :param cached: BigQuery caches query results. Rest of the statistics will not be provided for cached queries.
-            :param billedBytes: How many bytes BigQuery bills for.
-            :param properties: Full property tree of BigQUery run.
-            """
+        :param cached: BigQuery caches query results. Rest of the statistics will not be provided for cached queries.
+        :param billedBytes: How many bytes BigQuery bills for.
+        :param properties: Full property tree of BigQUery run.
+        """
 
-            cached: bool
-            billedBytes: int | None = field(default=None)
-            properties: str | None = field(default=None)
+        cached: bool
+        billedBytes: int | None = field(default=None)
+        properties: str | None = field(default=None)
 
-            @staticmethod
-            def _get_schema() -> str:
-                return (
-                    "https://raw.githubusercontent.com/apache/airflow/"
-                    f"providers-google/{provider_version}/airflow/providers/google/"
-                    "openlineage/BigQueryJobRunFacet.json"
-                )
-    except ImportError:  # OpenLineage is not available
+        @staticmethod
+        def _get_schema() -> str:
+            return (
+                "https://raw.githubusercontent.com/apache/airflow/"
+                f"providers-google/{provider_version}/airflow/providers/google/"
+                "openlineage/BigQueryJobRunFacet.json"
+            )
+except ImportError:  # OpenLineage is not available
 
-        def create_no_op(*_, **__) -> None:
-            """
-            Create a no-op placeholder.
+    def create_no_op(*_, **__) -> None:
+        """
+        Create a no-op placeholder.
 
-            This function creates and returns a None value, used as a placeholder when the OpenLineage client
-            library is available. It represents an action that has no effect.
-            """
-            return None
+        This function creates and returns a None value, used as a placeholder when the OpenLineage client
+        library is available. It represents an action that has no effect.
+        """
+        return None
 
-        BigQueryJobRunFacet = create_no_op
+    BigQueryJobRunFacet = create_no_op  # type: ignore[misc, assignment]

airflow/providers/google/cloud/openlineage/mixins.py

@@ -188,10 +188,10 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
         return list(final_outputs.values())
 
     def _get_input_dataset(self, table: dict) -> InputDataset:
-        return cast(InputDataset, self._get_dataset(table, "input"))
+        return cast("InputDataset", self._get_dataset(table, "input"))
 
     def _get_output_dataset(self, table: dict) -> OutputDataset:
-        return cast(OutputDataset, self._get_dataset(table, "output"))
+        return cast("OutputDataset", self._get_dataset(table, "output"))
 
     def _get_dataset(self, table: dict, dataset_type: str) -> Dataset:
         project = table.get("projectId")

airflow/providers/google/cloud/operators/automl.py

@@ -268,7 +268,7 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
     @cached_property
     def model(self) -> Model | None:
         if self.model_id:
-            hook = cast(CloudAutoMLHook, self.hook)
+            hook = cast("CloudAutoMLHook", self.hook)
             return hook.get_model(
                 model_id=self.model_id,
                 location=self.location,

airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py

@@ -280,7 +280,6 @@ class CloudDataTransferServiceCreateJobOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferJobLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 job_name=result[NAME],
             )
@@ -370,7 +369,6 @@ class CloudDataTransferServiceUpdateJobOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferJobLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 job_name=self.job_name,
             )
@@ -516,7 +514,6 @@ class CloudDataTransferServiceRunJobOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferJobLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 job_name=self.job_name,
             )

airflow/providers/google/cloud/operators/datafusion.py

@@ -26,7 +26,7 @@ from google.api_core.retry import exponential_sleep_generator
 from googleapiclient.errors import HttpError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.datafusion import SUCCESS_STATES, DataFusionHook, PipelineStates
 from airflow.providers.google.cloud.links.datafusion import (
     DataFusionInstanceLink,
@@ -37,33 +37,12 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 from airflow.providers.google.cloud.triggers.datafusion import DataFusionStartPipelineTrigger
 from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
 from airflow.providers.google.cloud.utils.helpers import resource_path_to_dict
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class DataFusionPipelineLinkHelper:
-    """
-    Helper class for Pipeline links.
-
-    .. warning::
-        This class is deprecated. Consider using ``resource_path_to_dict()`` instead.
-    """
-
-    @staticmethod
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="airflow.providers.google.cloud.utils.helpers.resource_path_to_dict",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_project_id(instance):
-        instance = instance["name"]
-        project_id = next(x for x in instance.split("/") if x.startswith("airflow"))
-        return project_id
-
-
 class CloudDataFusionRestartInstanceOperator(GoogleCloudBaseOperator):
     """
     Restart a single Data Fusion instance.

airflow/providers/google/cloud/operators/dataproc.py

@@ -44,7 +44,6 @@ from airflow.providers.google.cloud.hooks.dataproc import (
 )
 from airflow.providers.google.cloud.links.dataproc import (
     DATAPROC_BATCH_LINK,
-    DATAPROC_CLUSTER_LINK_DEPRECATED,
     DATAPROC_JOB_LINK_DEPRECATED,
     DataprocBatchesListLink,
     DataprocBatchLink,
@@ -63,7 +62,6 @@ from airflow.providers.google.cloud.triggers.dataproc import (
     DataprocSubmitTrigger,
 )
 from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.utils import timezone
 
@@ -919,145 +917,6 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
         return event["cluster"]
 
 
-# TODO: Remove one day
-@deprecated(
-    planned_removal_date="March 01, 2025",
-    use_instead="DataprocUpdateClusterOperator",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocScaleClusterOperator(GoogleCloudBaseOperator):
-    """
-    Scale, up or down, a cluster on Google Cloud Dataproc.
-
-    The operator will wait until the cluster is re-scaled.
-
-    Example usage:
-
-    .. code-block:: python
-
-        t1 = DataprocClusterScaleOperator(
-            task_id="dataproc_scale",
-            project_id="my-project",
-            cluster_name="cluster-1",
-            num_workers=10,
-            num_preemptible_workers=10,
-            graceful_decommission_timeout="1h",
-        )
-
-    .. seealso::
-        For more detail on about scaling clusters have a look at the reference:
-        https://cloud.google.com/dataproc/docs/concepts/configuring-clusters/scaling-clusters
-
-    :param cluster_name: The name of the cluster to scale. (templated)
-    :param project_id: The ID of the google cloud project in which
-        the cluster runs. (templated)
-    :param region: The region for the dataproc cluster. (templated)
-    :param num_workers: The new number of workers
-    :param num_preemptible_workers: The new number of preemptible workers
-    :param graceful_decommission_timeout: Timeout for graceful YARN decommissioning.
-        Maximum value is 1d
-    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
-    :param impersonation_chain: Optional service account to impersonate using short-term
-        credentials, or chained list of accounts required to get the access_token
-        of the last account in the list, which will be impersonated in the request.
-        If set as a string, the account must grant the originating account
-        the Service Account Token Creator IAM role.
-        If set as a sequence, the identities from the list must grant
-        Service Account Token Creator IAM role to the directly preceding identity, with first
-        account from the list granting this role to the originating account (templated).
-    """
-
-    template_fields: Sequence[str] = ("cluster_name", "project_id", "region", "impersonation_chain")
-
-    operator_extra_links = (DataprocLink(),)
-
-    def __init__(
-        self,
-        *,
-        cluster_name: str,
-        project_id: str = PROVIDE_PROJECT_ID,
-        region: str = "global",
-        num_workers: int = 2,
-        num_preemptible_workers: int = 0,
-        graceful_decommission_timeout: str | None = None,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        self.project_id = project_id
-        self.region = region
-        self.cluster_name = cluster_name
-        self.num_workers = num_workers
-        self.num_preemptible_workers = num_preemptible_workers
-        self.graceful_decommission_timeout = graceful_decommission_timeout
-        self.gcp_conn_id = gcp_conn_id
-        self.impersonation_chain = impersonation_chain
-
-    def _build_scale_cluster_data(self) -> dict:
-        scale_data = {
-            "config": {
-                "worker_config": {"num_instances": self.num_workers},
-                "secondary_worker_config": {"num_instances": self.num_preemptible_workers},
-            }
-        }
-        return scale_data
-
-    @property
-    def _graceful_decommission_timeout_object(self) -> dict[str, int] | None:
-        if not self.graceful_decommission_timeout:
-            return None
-
-        timeout = None
-        match = re.fullmatch(r"(\d+)([smdh])", self.graceful_decommission_timeout)
-        if match:
-            val = int(match.group(1))
-            unit = match.group(2)
-            if unit == "s":
-                timeout = val
-            elif unit == "m":
-                timeout = int(timedelta(minutes=val).total_seconds())
-            elif unit == "h":
-                timeout = int(timedelta(hours=val).total_seconds())
-            elif unit == "d":
-                timeout = int(timedelta(days=val).total_seconds())
-
-        if not timeout:
-            raise AirflowException(
-                "DataprocClusterScaleOperator "
-                " should be expressed in day, hours, minutes or seconds. "
-                " i.e. 1d, 4h, 10m, 30s"
-            )
-
-        return {"seconds": timeout}
-
-    def execute(self, context: Context) -> None:
-        """Scale, up or down, a cluster on Google Cloud Dataproc."""
-        self.log.info("Scaling cluster: %s", self.cluster_name)
-
-        scaling_cluster_data = self._build_scale_cluster_data()
-        update_mask = ["config.worker_config.num_instances", "config.secondary_worker_config.num_instances"]
-
-        hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
-        # Save data required to display extra link no matter what the cluster status will be
-        DataprocLink.persist(
-            context=context,
-            task_instance=self,
-            url=DATAPROC_CLUSTER_LINK_DEPRECATED,
-            resource=self.cluster_name,
-        )
-        operation = hook.update_cluster(
-            project_id=self.project_id,
-            region=self.region,
-            cluster_name=self.cluster_name,
-            cluster=scaling_cluster_data,
-            graceful_decommission_timeout=self._graceful_decommission_timeout_object,
-            update_mask={"paths": update_mask},
-        )
-        operation.result()
-        self.log.info("Cluster scaling finished")
-
-
 class DataprocDeleteClusterOperator(GoogleCloudBaseOperator):
     """
     Delete a cluster in a project.
@@ -1463,8 +1322,7 @@ class DataprocJobBaseOperator(GoogleCloudBaseOperator):
         """Initialize `self.job_template` with default values."""
         if self.project_id is None:
             raise AirflowException(
-                "project id should either be set via project_id "
-                "parameter or retrieved from the connection,"
+                "project id should either be set via project_id parameter or retrieved from the connection,"
             )
         job_template = DataProcJobBuilder(
             project_id=self.project_id,

airflow/providers/google/cloud/operators/dataproc_metastore.py

@@ -31,7 +31,6 @@ from google.cloud.metastore_v1.types import Backup, MetadataImport, Service
 from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator, BaseOperatorLink
 from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.links.storage import StorageLink