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

airflow/providers/google/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "19.1.0"
+__version__ = "19.3.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.11.0"

airflow/providers/google/ads/hooks/ads.py

@@ -27,8 +27,7 @@ from google.ads.googleads.client import GoogleAdsClient
 from google.ads.googleads.errors import GoogleAdsException
 from google.auth.exceptions import GoogleAuthError
 
-from airflow.exceptions import AirflowException
-from airflow.providers.common.compat.sdk import BaseHook
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
 from airflow.providers.google.common.hooks.base_google import get_field
 
 if TYPE_CHECKING:

airflow/providers/google/cloud/bundles/gcs.py

@@ -23,7 +23,7 @@ import structlog
 from google.api_core.exceptions import NotFound
 
 from airflow.dag_processing.bundles.base import BaseDagBundle
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 

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

@@ -27,7 +27,7 @@ import tenacity
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud import alloydb_v1
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 

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

@@ -59,12 +59,9 @@ from pandas_gbq import read_gbq
 from pandas_gbq.gbq import GbqConnector  # noqa: F401 used in ``airflow.contrib.hooks.bigquery``
 from sqlalchemy import create_engine
 
-from airflow.exceptions import (
-    AirflowException,
-    AirflowOptionalProviderFeatureException,
-    AirflowProviderDeprecationWarning,
-)
+from airflow.exceptions import AirflowOptionalProviderFeatureException, AirflowProviderDeprecationWarning
 from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.cloud.utils.bigquery import bq_cast
 from airflow.providers.google.cloud.utils.credentials_provider import _get_scopes
@@ -1993,18 +1990,19 @@ def _format_schema_for_description(schema: dict) -> list:
     internal_size, precision, scale, null_ok.
     """
     description = []
-    for field in schema["fields"]:
-        mode = field.get("mode", "NULLABLE")
-        field_description = (
-            field["name"],
-            field["type"],
-            None,
-            None,
-            None,
-            None,
-            mode == "NULLABLE",
-        )
-        description.append(field_description)
+    if "fields" in schema:
+        for field in schema["fields"]:
+            mode = field.get("mode", "NULLABLE")
+            field_description = (
+                field["name"],
+                field["type"],
+                None,
+                None,
+                None,
+                None,
+                mode == "NULLABLE",
+            )
+            description.append(field_description)
     return description
 
 

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

@@ -33,7 +33,7 @@ from google.cloud.batch_v1 import (
     Task,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 

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

@@ -27,7 +27,7 @@ from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.devtools.cloudbuild_v1 import CloudBuildAsyncClient, CloudBuildClient, GetBuildRequest
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk 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 airflow.providers.google.common.hooks.operation_helpers import OperationHelper

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

@@ -35,7 +35,7 @@ from google.cloud.orchestration.airflow.service_v1 import (
     PollAirflowCommandResponse,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 

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

@@ -45,7 +45,7 @@ from google.cloud.redis_v1 import (
 )
 
 from airflow import version
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 if TYPE_CHECKING:

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

@@ -40,7 +40,7 @@ from google.cloud.run_v2 import (
 )
 from google.longrunning import operations_pb2
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,

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

@@ -49,7 +49,7 @@ from googleapiclient.errors import HttpError
 
 # Number of retries - used by googleapiclient method calls to perform retries
 # For requests that are "retriable"
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.version_compat import AIRFLOW_V_3_1_PLUS
 
 if AIRFLOW_V_3_1_PLUS:

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

@@ -47,7 +47,8 @@ from google.protobuf.json_format import MessageToDict
 from googleapiclient.discovery import Resource, build
 from googleapiclient.errors import HttpError
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,

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

@@ -28,7 +28,7 @@ from google.cloud.compute_v1.services.instance_templates import InstanceTemplate
 from google.cloud.compute_v1.services.instances import InstancesClient
 from googleapiclient.discovery import build
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 

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

@@ -26,7 +26,7 @@ from typing import Any
 from googleapiclient.errors import HttpError
 from paramiko.ssh_exception import SSHException
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 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

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

@@ -32,7 +32,8 @@ from google.cloud.datacatalog import (
     TagTemplateField,
 )
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException
 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

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

@@ -49,9 +49,9 @@ from google.cloud.dataflow_v1beta3.types import (
 from google.cloud.dataflow_v1beta3.types.jobs import ListJobsRequest
 from googleapiclient.discovery import Resource, build
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import 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.common.compat.sdk import AirflowException, timeout
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,
     GoogleBaseAsyncHook,

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

@@ -31,7 +31,7 @@ from google.cloud.dataform_v1beta1.types import (
     WriteFileResponse,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 if TYPE_CHECKING:

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

@@ -32,7 +32,7 @@ from gcloud.aio.auth import AioSession, Token
 from google.api_core.retry import exponential_sleep_generator
 from googleapiclient.discovery import Resource, build
 
-from airflow.exceptions import AirflowException, AirflowNotFoundException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,
@@ -469,31 +469,33 @@ class DataFusionHook(GoogleBaseHook):
             is always default. If your pipeline belongs to an Enterprise edition instance, you
             can create a namespace.
         """
-        # TODO: This API endpoint starts multiple pipelines. There will eventually be a fix
-        #  return the run Id as part of the API request to run a single pipeline.
-        #  https://github.com/apache/airflow/pull/8954#discussion_r438223116
+        # Use the single-program start endpoint for better error handling
+        # https://cdap.atlassian.net/wiki/spaces/DOCS/pages/477560983/Lifecycle+Microservices#Start-a-Program
+        program_type = self.cdap_program_type(pipeline_type=pipeline_type)
+        program_id = self.cdap_program_id(pipeline_type=pipeline_type)
         url = os.path.join(
-            instance_url,
-            "v3",
-            "namespaces",
-            quote(namespace),
+            self._base_url(instance_url, namespace),
+            quote(pipeline_name),
+            f"{program_type}s",
+            program_id,
             "start",
         )
         runtime_args = runtime_args or {}
-        body = [
-            {
-                "appId": pipeline_name,
-                "runtimeargs": runtime_args,
-                "programType": self.cdap_program_type(pipeline_type=pipeline_type),
-                "programId": self.cdap_program_id(pipeline_type=pipeline_type),
-            }
-        ]
-        response = self._cdap_request(url=url, method="POST", body=body)
+        response = self._cdap_request(url=url, method="POST", body=runtime_args)
         self._check_response_status_and_data(
             response, f"Starting a pipeline failed with code {response.status}"
         )
         response_json = json.loads(response.data)
-        return response_json[0]["runId"]
+
+        # Extract and validate runId from response
+        if "runId" not in response_json:
+            error_message = response_json.get("error", "Unknown error")
+            raise AirflowException(
+                f"Failed to start pipeline '{pipeline_name}'. "
+                f"The response does not contain a runId. Error: {error_message}"
+            )
+
+        return str(response_json["runId"])
 
     def stop_pipeline(self, pipeline_name: str, instance_url: str, namespace: str = "default") -> None:
         """

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

@@ -46,7 +46,7 @@ from google.cloud.dataplex_v1.types import (
 )
 from google.protobuf.field_mask_pb2 import FieldMask
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,

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

@@ -45,7 +45,7 @@ from google.cloud.dataproc_v1 import (
     WorkflowTemplateServiceClient,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.version import version as airflow_version

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

@@ -26,7 +26,7 @@ from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.metastore_v1 import DataprocMetastoreClient
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 

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

@@ -55,7 +55,7 @@ from google.cloud.dlp_v2.types import (
 )
 from google.protobuf.field_mask_pb2 import FieldMask
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 

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

@@ -25,7 +25,7 @@ from collections.abc import Sequence
 import requests
 from googleapiclient.discovery import build
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 # Time to sleep between active checks of the operation results

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

@@ -43,9 +43,9 @@ from google.api_core.exceptions import GoogleAPICallError, NotFound
 from google.cloud.exceptions import GoogleCloudError
 from google.cloud.storage.retry import DEFAULT_RETRY
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
-from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.common.compat.sdk import AirflowException, timezone
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (

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

@@ -21,7 +21,7 @@ from typing import Any
 
 from googleapiclient.discovery import Resource, build
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 

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

@@ -38,9 +38,9 @@ from kubernetes_asyncio import client as async_client
 from kubernetes_asyncio.config.kube_config import FileOrData
 
 from airflow import version
-from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook, KubernetesHook
 from airflow.providers.cncf.kubernetes.kube_client import _enable_tcp_keepalive
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,

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

@@ -28,8 +28,7 @@ from looker_sdk.rtl import api_settings, auth_session, requests_transport, seria
 from looker_sdk.sdk.api40 import methods as methods40
 from packaging.version import parse as parse_version
 
-from airflow.exceptions import AirflowException
-from airflow.providers.common.compat.sdk import BaseHook
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
 from airflow.version import version
 
 if TYPE_CHECKING:

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

@@ -31,7 +31,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.auth.transport import requests as google_requests
 from google.cloud.managedkafka_v1 import Cluster, ConsumerGroup, ManagedKafkaClient, Topic, types
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 

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

@@ -31,7 +31,7 @@ from gcloud.aio.auth import AioSession, Token
 from googleapiclient.discovery import Resource, build
 from googleapiclient.errors import HttpError
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import (
     PROVIDE_PROJECT_ID,
     GoogleBaseAsyncHook,

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

@@ -82,6 +82,7 @@ class PubSubHook(GoogleBaseHook):
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         enable_message_ordering: bool = False,
+        enable_open_telemetry_tracing: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -90,6 +91,7 @@ class PubSubHook(GoogleBaseHook):
             **kwargs,
         )
         self.enable_message_ordering = enable_message_ordering
+        self.enable_open_telemetry_tracing = enable_open_telemetry_tracing
         self._client = None
 
     def get_conn(self) -> PublisherClient:
@@ -104,6 +106,7 @@ class PubSubHook(GoogleBaseHook):
                 client_info=CLIENT_INFO,
                 publisher_options=PublisherOptions(
                     enable_message_ordering=self.enable_message_ordering,
+                    enable_open_telemetry_tracing=self.enable_open_telemetry_tracing,
                 ),
             )
         return self._client

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

@@ -27,7 +27,7 @@ from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
 from google.cloud.spanner_v1.client import Client
 from sqlalchemy import create_engine
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field

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

@@ -31,7 +31,7 @@ from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
 from google.protobuf.field_mask_pb2 import FieldMask
 from googleapiclient.errors import HttpError
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 if TYPE_CHECKING:

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

@@ -26,7 +26,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.tasks_v2 import CloudTasksClient
 from google.cloud.tasks_v2.types import Queue, Task
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 

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

@@ -32,7 +32,7 @@ from google.cloud.translate_v2 import Client
 from google.cloud.translate_v3 import TranslationServiceClient
 from google.cloud.translate_v3.types.translation_service import GlossaryInputConfig
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk 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 airflow.providers.google.common.hooks.operation_helpers import OperationHelper

airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py

@@ -36,7 +36,8 @@ from google.cloud.aiplatform import (
 )
 from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py

@@ -28,7 +28,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import BatchPredictionJob, Model, explain
 from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py

@@ -42,7 +42,7 @@ from google.cloud.aiplatform_v1 import (
     types,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py

@@ -30,7 +30,7 @@ from google.cloud.aiplatform_v1beta1 import (
     FeatureOnlineStoreServiceClient,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 

airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py

@@ -34,7 +34,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import CustomJob, HyperparameterTuningJob, gapic, hyperparameter_tuning
 from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

airflow/providers/google/cloud/hooks/vertex_ai/model_service.py

@@ -27,7 +27,7 @@ from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import ModelServiceClient
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 

airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py

@@ -39,7 +39,7 @@ from google.cloud.aiplatform_v1 import (
     types,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

airflow/providers/google/cloud/hooks/vertex_ai/ray.py

@@ -52,7 +52,7 @@ class RayHook(GoogleBaseHook):
                 return [__encode_value(nested_value) for nested_value in value]
             if not isinstance(value, dict) and isinstance(value, MutableMapping):
                 return {key: __encode_value(nested_value) for key, nested_value in dict(value).items()}
-            if dataclasses.is_dataclass(value):
+            if dataclasses.is_dataclass(value) and not isinstance(value, type):
                 return dataclasses.asdict(value)
             return value
 

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

@@ -36,7 +36,7 @@ from google.cloud.vision_v1 import (
 )
 from google.protobuf.json_format import MessageToDict
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook