apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

airflow/providers/google/cloud/hooks/compute.py

@@ -23,13 +23,14 @@ import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from google.cloud.compute_v1.services.instance_group_managers import InstanceGroupManagersClient
+from google.cloud.compute_v1.services.instance_templates import InstanceTemplatesClient
+from google.cloud.compute_v1.services.instances import InstancesClient
 from googleapiclient.discovery import build
 
 from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
-from google.cloud.compute_v1.services.instance_group_managers import InstanceGroupManagersClient
-from google.cloud.compute_v1.services.instance_templates import InstanceTemplatesClient
-from google.cloud.compute_v1.services.instances import InstancesClient
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
@@ -85,15 +86,15 @@ class ComputeEngineHook(GoogleBaseHook):
 
     def get_compute_instance_template_client(self):
         """Return Compute Engine Instance Template Client."""
-        return InstanceTemplatesClient(credentials=self.get_credentials(), client_info=self.client_info)
+        return InstanceTemplatesClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
 
     def get_compute_instance_client(self):
         """Return Compute Engine Instance Client."""
-        return InstancesClient(credentials=self.get_credentials(), client_info=self.client_info)
+        return InstancesClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
 
     def get_compute_instance_group_managers_client(self):
         """Return Compute Engine Instance Group Managers Client."""
-        return InstanceGroupManagersClient(credentials=self.get_credentials(), client_info=self.client_info)
+        return InstanceGroupManagersClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
 
     @GoogleBaseHook.fallback_to_default_project_id
     def insert_instance_template(

airflow/providers/google/cloud/hooks/compute_ssh.py

@@ -31,7 +31,11 @@ from airflow.providers.google.cloud.hooks.compute import ComputeEngineHook
 from airflow.providers.google.cloud.hooks.os_login import OSLoginHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.providers.ssh.hooks.ssh import SSHHook
-from airflow.utils.types import NOTSET, ArgNotSet
+
+try:
+    from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
+except ImportError:
+    from airflow.utils.types import NOTSET, ArgNotSet  # type: ignore[attr-defined,no-redef]
 
 # Paramiko should be imported after airflow.providers.ssh. Then the import will fail with
 # cannot import "airflow.providers.ssh" and will be correctly discovered as optional feature
@@ -148,8 +152,7 @@ class ComputeEngineSSHHook(SSHHook):
             return ComputeEngineHook(
                 gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
             )
-        else:
-            return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id)
+        return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id)
 
     def _load_connection_config(self):
         def _boolify(value):
@@ -158,7 +161,7 @@ class ComputeEngineSSHHook(SSHHook):
             if isinstance(value, str):
                 if value.lower() == "false":
                     return False
-                elif value.lower() == "true":
+                if value.lower() == "true":
                     return True
             return False
 

airflow/providers/google/cloud/hooks/datacatalog.py

@@ -19,9 +19,6 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud import datacatalog
 from google.cloud.datacatalog import (
@@ -35,11 +32,23 @@ from google.cloud.datacatalog import (
     TagTemplateField,
 )
 
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.deprecated import deprecated
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
 
+@deprecated(
+    planned_removal_date="January 30, 2026",
+    use_instead="airflow.providers.google.cloud.hooks.dataplex.DataplexHook",
+    reason="The Data Catalog will be discontinued on January 30, 2026 "
+    "in favor of Dataplex Universal Catalog.",
+    category=AirflowProviderDeprecationWarning,
+)
 class CloudDataCatalogHook(GoogleBaseHook):
     """
     Hook for Google Cloud Data Catalog Service.

airflow/providers/google/cloud/hooks/dataflow.py

@@ -27,22 +27,10 @@ import subprocess
 import time
 import uuid
 import warnings
-from collections.abc import Generator, Sequence
+from collections.abc import Callable, Generator, Sequence
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
-from googleapiclient.discovery import Resource, build
-
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType, beam_options_to_args
-from airflow.providers.google.common.deprecated import deprecated
-from airflow.providers.google.common.hooks.base_google import (
-    PROVIDE_PROJECT_ID,
-    GoogleBaseAsyncHook,
-    GoogleBaseHook,
-)
-from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.timeout import timeout
 from google.cloud.dataflow_v1beta3 import (
     GetJobRequest,
     Job,
@@ -59,6 +47,17 @@ from google.cloud.dataflow_v1beta3.types import (
     JobMetrics,
 )
 from google.cloud.dataflow_v1beta3.types.jobs import ListJobsRequest
+from googleapiclient.discovery import Resource, build
+
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType, beam_options_to_args
+from airflow.providers.common.compat.sdk import timeout
+from airflow.providers.google.common.hooks.base_google import (
+    PROVIDE_PROJECT_ID,
+    GoogleBaseAsyncHook,
+    GoogleBaseHook,
+)
+from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
     from google.cloud.dataflow_v1beta3.services.jobs_v1_beta3.pagers import ListJobsAsyncPager
@@ -136,7 +135,7 @@ def _fallback_variable_parameter(parameter_name: str, variable_key_name: str) ->
 
             return func(self, *args, **kwargs)
 
-        return cast(T, inner_wrapper)
+        return cast("T", inner_wrapper)
 
     return _wrapper
 
@@ -185,7 +184,67 @@ class DataflowJobType:
     JOB_TYPE_STREAMING = "JOB_TYPE_STREAMING"
 
 
-class _DataflowJobsController(LoggingMixin):
+class DataflowJobTerminalStateHelper(LoggingMixin):
+    """Helper to define and validate the dataflow job terminal state."""
+
+    @staticmethod
+    def expected_terminal_state_is_allowed(expected_terminal_state):
+        job_allowed_terminal_states = DataflowJobStatus.TERMINAL_STATES | {
+            DataflowJobStatus.JOB_STATE_RUNNING
+        }
+        if expected_terminal_state not in job_allowed_terminal_states:
+            raise AirflowException(
+                f"Google Cloud Dataflow job's expected terminal state "
+                f"'{expected_terminal_state}' is invalid."
+                f" The value should be any of the following: {job_allowed_terminal_states}"
+            )
+        return True
+
+    @staticmethod
+    def expected_terminal_state_is_valid_for_job_type(expected_terminal_state, is_streaming: bool):
+        if is_streaming:
+            invalid_terminal_state = DataflowJobStatus.JOB_STATE_DONE
+            job_type = "streaming"
+        else:
+            invalid_terminal_state = DataflowJobStatus.JOB_STATE_DRAINED
+            job_type = "batch"
+
+        if expected_terminal_state == invalid_terminal_state:
+            raise AirflowException(
+                f"Google Cloud Dataflow job's expected terminal state cannot be {invalid_terminal_state} while it is a {job_type} job"
+            )
+        return True
+
+    def job_reached_terminal_state(self, job, wait_until_finished=None, custom_terminal_state=None) -> bool:
+        """
+        Check the job reached terminal state, if job failed raise exception.
+
+        :return: True if job is done.
+        :raise: Exception
+        """
+        current_state = job["currentState"]
+        is_streaming = job.get("type") == DataflowJobType.JOB_TYPE_STREAMING
+        expected_terminal_state = (
+            DataflowJobStatus.JOB_STATE_RUNNING if is_streaming else DataflowJobStatus.JOB_STATE_DONE
+        )
+        if custom_terminal_state is not None:
+            expected_terminal_state = custom_terminal_state
+        self.expected_terminal_state_is_allowed(expected_terminal_state)
+        self.expected_terminal_state_is_valid_for_job_type(expected_terminal_state, is_streaming=is_streaming)
+        if current_state == expected_terminal_state:
+            if expected_terminal_state == DataflowJobStatus.JOB_STATE_RUNNING and wait_until_finished:
+                return False
+            return True
+        if current_state in DataflowJobStatus.AWAITING_STATES:
+            return wait_until_finished is False
+        self.log.debug("Current job: %s", job)
+        raise AirflowException(
+            f"Google Cloud Dataflow job {job['name']} is in an unexpected terminal state: {current_state}, "
+            f"expected terminal state: {expected_terminal_state}"
+        )
+
+
+class _DataflowJobsController(DataflowJobTerminalStateHelper):
     """
     Interface for communication with Google Cloud Dataflow API.
 
@@ -262,15 +321,14 @@ class _DataflowJobsController(LoggingMixin):
         """
         if not self._multiple_jobs and self._job_id:
             return [self.fetch_job_by_id(self._job_id)]
-        elif self._jobs:
+        if self._jobs:
             return [self.fetch_job_by_id(job["id"]) for job in self._jobs]
-        elif self._job_name:
+        if self._job_name:
             jobs = self._fetch_jobs_by_prefix_name(self._job_name.lower())
             if len(jobs) == 1:
                 self._job_id = jobs[0]["id"]
             return jobs
-        else:
-            raise ValueError("Missing both dataflow job ID and name.")
+        raise ValueError("Missing both dataflow job ID and name.")
 
     def fetch_job_by_id(self, job_id: str) -> dict[str, str]:
         """
@@ -435,12 +493,12 @@ class _DataflowJobsController(LoggingMixin):
                 f"'{current_expected_state}' is invalid."
                 f" The value should be any of the following: {terminal_states}"
             )
-        elif is_streaming and current_expected_state == DataflowJobStatus.JOB_STATE_DONE:
+        if is_streaming and current_expected_state == DataflowJobStatus.JOB_STATE_DONE:
             raise AirflowException(
                 "Google Cloud Dataflow job's expected terminal state cannot be "
                 "JOB_STATE_DONE while it is a streaming job"
             )
-        elif not is_streaming and current_expected_state == DataflowJobStatus.JOB_STATE_DRAINED:
+        if not is_streaming and current_expected_state == DataflowJobStatus.JOB_STATE_DRAINED:
             raise AirflowException(
                 "Google Cloud Dataflow job's expected terminal state cannot be "
                 "JOB_STATE_DRAINED while it is a batch job"
@@ -463,7 +521,10 @@ class _DataflowJobsController(LoggingMixin):
         """Wait for result of submitted job."""
         self.log.info("Start waiting for done.")
         self._refresh_jobs()
-        while self._jobs and not all(self._check_dataflow_job_state(job) for job in self._jobs):
+        while self._jobs and not all(
+            self.job_reached_terminal_state(job, self._wait_until_finished, self._expected_terminal_state)
+            for job in self._jobs
+        ):
             self.log.info("Waiting for done. Sleep %s s", self._poll_sleep)
             time.sleep(self._poll_sleep)
             self._refresh_jobs()
@@ -586,66 +647,6 @@ class DataflowHook(GoogleBaseHook):
         http_authorized = self._authorize()
         return build("datapipelines", "v1", http=http_authorized, cache_discovery=False)
 
-    @_fallback_to_location_from_variables
-    @_fallback_to_project_id_from_variables
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="airflow.providers.apache.beam.hooks.beam.start.start_java_pipeline, "
-        "providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done",
-        instructions="Please use airflow.providers.apache.beam.hooks.beam.start.start_java_pipeline "
-        "to start pipeline and providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done method "
-        "to wait for the required pipeline state instead.",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def start_java_dataflow(
-        self,
-        job_name: str,
-        variables: dict,
-        jar: str,
-        project_id: str,
-        job_class: str | None = None,
-        append_job_name: bool = True,
-        multiple_jobs: bool = False,
-        on_new_job_id_callback: Callable[[str], None] | None = None,
-        location: str = DEFAULT_DATAFLOW_LOCATION,
-    ) -> None:
-        """
-        Start Dataflow java job.
-
-        :param job_name: The name of the job.
-        :param variables: Variables passed to the job.
-        :param project_id: Optional, the Google Cloud project ID in which to start a job.
-            If set to None or missing, the default project_id from the Google Cloud connection is used.
-        :param jar: Name of the jar for the job
-        :param job_class: Name of the java class for the job.
-        :param append_job_name: True if unique suffix has to be appended to job name.
-        :param multiple_jobs: True if to check for multiple job in dataflow
-        :param on_new_job_id_callback: Callback called when the job ID is known.
-        :param location: Job location.
-        """
-        name = self.build_dataflow_job_name(job_name, append_job_name)
-
-        variables["jobName"] = name
-        variables["region"] = location
-        variables["project"] = project_id
-
-        if "labels" in variables:
-            variables["labels"] = json.dumps(variables["labels"], separators=(",", ":"))
-
-        self.beam_hook.start_java_pipeline(
-            variables=variables,
-            jar=jar,
-            job_class=job_class,
-            process_line_callback=process_line_and_extract_dataflow_job_id_callback(on_new_job_id_callback),
-        )
-        self.wait_for_done(
-            job_name=name,
-            location=location,
-            job_id=self.job_id,
-            multiple_jobs=multiple_jobs,
-        )
-
     @_fallback_to_location_from_variables
     @_fallback_to_project_id_from_variables
     @GoogleBaseHook.fallback_to_default_project_id
@@ -1004,7 +1005,7 @@ class DataflowHook(GoogleBaseHook):
         success_code = 0
 
         with self.provide_authorized_gcloud():
-            proc = subprocess.run(cmd, capture_output=True)
+            proc = subprocess.run(cmd, check=False, capture_output=True)
 
         if proc.returncode != success_code:
             stderr_last_20_lines = "\n".join(proc.stderr.decode().strip().splitlines()[-20:])
@@ -1027,82 +1028,6 @@ class DataflowHook(GoogleBaseHook):
                 "While reading job object after template execution error occurred. Job object has no id."
             )
 
-    @_fallback_to_location_from_variables
-    @_fallback_to_project_id_from_variables
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="airflow.providers.apache.beam.hooks.beam.start.start_python_pipeline method, "
-        "providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done",
-        instructions="Please use airflow.providers.apache.beam.hooks.beam.start.start_python_pipeline method "
-        "to start pipeline and providers.google.cloud.hooks.dataflow.DataflowHook.wait_for_done method "
-        "to wait for the required pipeline state instead.",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def start_python_dataflow(
-        self,
-        job_name: str,
-        variables: dict,
-        dataflow: str,
-        py_options: list[str],
-        project_id: str,
-        py_interpreter: str = "python3",
-        py_requirements: list[str] | None = None,
-        py_system_site_packages: bool = False,
-        append_job_name: bool = True,
-        on_new_job_id_callback: Callable[[str], None] | None = None,
-        location: str = DEFAULT_DATAFLOW_LOCATION,
-    ):
-        """
-        Start Dataflow job.
-
-        :param job_name: The name of the job.
-        :param variables: Variables passed to the job.
-        :param dataflow: Name of the Dataflow process.
-        :param py_options: Additional options.
-        :param project_id: The ID of the GCP project that owns the job.
-            If set to ``None`` or missing, the default project_id from the GCP connection is used.
-        :param py_interpreter: Python version of the beam pipeline.
-            If None, this defaults to the python3.
-            To track python versions supported by beam and related
-            issues check: https://issues.apache.org/jira/browse/BEAM-1251
-        :param py_requirements: Additional python package(s) to install.
-            If a value is passed to this parameter, a new virtual environment has been created with
-            additional packages installed.
-
-            You could also install the apache-beam package if it is not installed on your system or you want
-            to use a different version.
-        :param py_system_site_packages: Whether to include system_site_packages in your virtualenv.
-            See virtualenv documentation for more information.
-
-            This option is only relevant if the ``py_requirements`` parameter is not None.
-        :param append_job_name: True if unique suffix has to be appended to job name.
-        :param project_id: Optional, the Google Cloud project ID in which to start a job.
-            If set to None or missing, the default project_id from the Google Cloud connection is used.
-        :param on_new_job_id_callback: Callback called when the job ID is known.
-        :param location: Job location.
-        """
-        name = self.build_dataflow_job_name(job_name, append_job_name)
-        variables["job_name"] = name
-        variables["region"] = location
-        variables["project"] = project_id
-
-        self.beam_hook.start_python_pipeline(
-            variables=variables,
-            py_file=dataflow,
-            py_options=py_options,
-            py_interpreter=py_interpreter,
-            py_requirements=py_requirements,
-            py_system_site_packages=py_system_site_packages,
-            process_line_callback=process_line_and_extract_dataflow_job_id_callback(on_new_job_id_callback),
-        )
-
-        self.wait_for_done(
-            job_name=name,
-            location=location,
-            job_id=self.job_id,
-        )
-
     @staticmethod
     def build_dataflow_job_name(job_name: str, append_job_name: bool = True) -> str:
         """Build Dataflow job name."""
@@ -1200,85 +1125,6 @@ class DataflowHook(GoogleBaseHook):
         )
         jobs_controller.cancel()
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    def start_sql_job(
-        self,
-        job_name: str,
-        query: str,
-        options: dict[str, Any],
-        project_id: str,
-        location: str = DEFAULT_DATAFLOW_LOCATION,
-        on_new_job_id_callback: Callable[[str], None] | None = None,
-        on_new_job_callback: Callable[[dict], None] | None = None,
-    ):
-        """
-        Start Dataflow SQL query.
-
-        :param job_name: The unique name to assign to the Cloud Dataflow job.
-        :param query: The SQL query to execute.
-        :param options: Job parameters to be executed.
-            For more information, look at:
-            `https://cloud.google.com/sdk/gcloud/reference/beta/dataflow/sql/query
-            <gcloud beta dataflow sql query>`__
-            command reference
-        :param location: The location of the Dataflow job (for example europe-west1)
-        :param project_id: The ID of the GCP project that owns the job.
-            If set to ``None`` or missing, the default project_id from the GCP connection is used.
-        :param on_new_job_id_callback: (Deprecated) Callback called when the job ID is known.
-        :param on_new_job_callback: Callback called when the job is known.
-        :return: the new job object
-        """
-        gcp_options = {
-            "project": project_id,
-            "format": "value(job.id)",
-            "job-name": job_name,
-            "region": location,
-        }
-        cmd = self._build_gcloud_command(
-            command=["gcloud", "dataflow", "sql", "query", query], parameters={**gcp_options, **options}
-        )
-        self.log.info("Executing command: %s", " ".join(shlex.quote(c) for c in cmd))
-        with self.provide_authorized_gcloud():
-            proc = subprocess.run(cmd, capture_output=True)
-        self.log.info("Output: %s", proc.stdout.decode())
-        self.log.warning("Stderr: %s", proc.stderr.decode())
-        self.log.info("Exit code %d", proc.returncode)
-        stderr_last_20_lines = "\n".join(proc.stderr.decode().strip().splitlines()[-20:])
-        if proc.returncode != 0:
-            raise AirflowException(
-                f"Process exit with non-zero exit code. Exit code: {proc.returncode} Error Details : "
-                f"{stderr_last_20_lines}"
-            )
-        job_id = proc.stdout.decode().strip()
-
-        self.log.info("Created job ID: %s", job_id)
-
-        jobs_controller = _DataflowJobsController(
-            dataflow=self.get_conn(),
-            project_number=project_id,
-            job_id=job_id,
-            location=location,
-            poll_sleep=self.poll_sleep,
-            num_retries=self.num_retries,
-            drain_pipeline=self.drain_pipeline,
-            wait_until_finished=self.wait_until_finished,
-        )
-        job = jobs_controller.get_jobs(refresh=True)[0]
-
-        if on_new_job_id_callback:
-            warnings.warn(
-                "on_new_job_id_callback is Deprecated. Please start using on_new_job_callback",
-                AirflowProviderDeprecationWarning,
-                stacklevel=3,
-            )
-            on_new_job_id_callback(cast(str, job.get("id")))
-
-        if on_new_job_callback:
-            on_new_job_callback(job)
-
-        jobs_controller.wait_for_done()
-        return jobs_controller.get_jobs(refresh=True)[0]
-
     @GoogleBaseHook.fallback_to_default_project_id
     def get_job(
         self,
@@ -1427,8 +1273,7 @@ class DataflowHook(GoogleBaseHook):
             location=location,
         )
         job = job_controller.fetch_job_by_id(job_id)
-
-        return job_controller._check_dataflow_job_state(job)
+        return job_controller.job_reached_terminal_state(job)
 
     @GoogleBaseHook.fallback_to_default_project_id
     def create_data_pipeline(
@@ -1557,7 +1402,7 @@ class DataflowHook(GoogleBaseHook):
         return f"projects/{project_id}/locations/{location}"
 
 
-class AsyncDataflowHook(GoogleBaseAsyncHook):
+class AsyncDataflowHook(GoogleBaseAsyncHook, DataflowJobTerminalStateHelper):
     """Async hook class for dataflow service."""
 
     sync_hook_class = DataflowHook

airflow/providers/google/cloud/hooks/dataform.py

@@ -20,8 +20,6 @@ import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.dataform_v1beta1 import DataformClient
 from google.cloud.dataform_v1beta1.types import (
@@ -33,6 +31,9 @@ from google.cloud.dataform_v1beta1.types import (
     WriteFileResponse,
 )
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.dataform_v1beta1.services.dataform.pagers import QueryWorkflowInvocationActionsPager
@@ -452,8 +453,7 @@ class DataformHook(GoogleBaseHook):
         """
         client = self.get_dataform_client()
         workspace_path = (
-            f"projects/{project_id}/locations/{region}/"
-            f"repositories/{repository_id}/workspaces/{workspace_id}"
+            f"projects/{project_id}/locations/{region}/repositories/{repository_id}/workspaces/{workspace_id}"
         )
         request = {
             "name": workspace_path,
@@ -495,8 +495,7 @@ class DataformHook(GoogleBaseHook):
         """
         client = self.get_dataform_client()
         workspace_path = (
-            f"projects/{project_id}/locations/{region}/"
-            f"repositories/{repository_id}/workspaces/{workspace_id}"
+            f"projects/{project_id}/locations/{region}/repositories/{repository_id}/workspaces/{workspace_id}"
         )
         request = {
             "workspace": workspace_path,
@@ -541,8 +540,7 @@ class DataformHook(GoogleBaseHook):
         """
         client = self.get_dataform_client()
         workspace_path = (
-            f"projects/{project_id}/locations/{region}/"
-            f"repositories/{repository_id}/workspaces/{workspace_id}"
+            f"projects/{project_id}/locations/{region}/repositories/{repository_id}/workspaces/{workspace_id}"
         )
         request = {
             "workspace": workspace_path,
@@ -586,8 +584,7 @@ class DataformHook(GoogleBaseHook):
         """
         client = self.get_dataform_client()
         workspace_path = (
-            f"projects/{project_id}/locations/{region}/"
-            f"repositories/{repository_id}/workspaces/{workspace_id}"
+            f"projects/{project_id}/locations/{region}/repositories/{repository_id}/workspaces/{workspace_id}"
         )
         request = {
             "workspace": workspace_path,
@@ -628,8 +625,7 @@ class DataformHook(GoogleBaseHook):
         """
         client = self.get_dataform_client()
         workspace_path = (
-            f"projects/{project_id}/locations/{region}/"
-            f"repositories/{repository_id}/workspaces/{workspace_id}"
+            f"projects/{project_id}/locations/{region}/repositories/{repository_id}/workspaces/{workspace_id}"
         )
         request = {
             "workspace": workspace_path,
@@ -670,8 +666,7 @@ class DataformHook(GoogleBaseHook):
         """
         client = self.get_dataform_client()
         workspace_path = (
-            f"projects/{project_id}/locations/{region}/"
-            f"repositories/{repository_id}/workspaces/{workspace_id}"
+            f"projects/{project_id}/locations/{region}/repositories/{repository_id}/workspaces/{workspace_id}"
         )
         request = {
             "workspace": workspace_path,

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

@@ -26,11 +26,12 @@ from collections.abc import Sequence
 from typing import Any
 from urllib.parse import quote, urlencode, urljoin
 
+import google.auth
 from aiohttp import ClientSession
 from gcloud.aio.auth import AioSession, Token
+from google.api_core.retry import exponential_sleep_generator
 from googleapiclient.discovery import Resource, build
 
-import google.auth
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
 from airflow.providers.google.common.hooks.base_google import (
@@ -38,7 +39,6 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseAsyncHook,
     GoogleBaseHook,
 )
-from google.api_core.retry import exponential_sleep_generator
 
 Operation = dict[str, Any]
 
@@ -137,8 +137,7 @@ class DataFusionHook(GoogleBaseHook):
 
         # Time is up!
         raise AirflowException(
-            f"Pipeline {pipeline_name} state {current_state} is not "
-            f"one of {success_states} after {timeout}s"
+            f"Pipeline {pipeline_name} state {current_state} is not one of {success_states} after {timeout}s"
         )
 
     @staticmethod
@@ -171,9 +170,9 @@ class DataFusionHook(GoogleBaseHook):
     def _check_response_status_and_data(response, message: str) -> None:
         if response.status == 404:
             raise AirflowNotFoundException(message)
-        elif response.status == 409:
+        if response.status == 409:
             raise ConflictException("Conflict: Resource is still in use.")
-        elif response.status != 200:
+        if response.status != 200:
             raise AirflowException(message)
         if response.data is None:
             raise AirflowException(
@@ -435,7 +434,7 @@ class DataFusionHook(GoogleBaseHook):
         pipeline_id: str,
         pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH,
         namespace: str = "default",
-    ) -> Any:
+    ) -> dict:
         url = os.path.join(
             self._base_url(instance_url, namespace),
             quote(pipeline_name),
@@ -573,8 +572,7 @@ class DataFusionAsyncHook(GoogleBaseAsyncHook):
                         raise
         if pipeline:
             return pipeline
-        else:
-            raise AirflowException("Could not retrieve pipeline. Aborting.")
+        raise AirflowException("Could not retrieve pipeline. Aborting.")
 
     async def get_pipeline(
         self,