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

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

@@ -31,7 +31,7 @@ import attrs
 import google.cloud.storage as storage
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowNotFoundException
+from airflow.providers.common.compat.sdk import AirflowNotFoundException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
 from airflow.providers.google.cloud.utils.credentials_provider import (
     get_credentials_and_project_id,
@@ -43,9 +43,11 @@ from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
+    from io import TextIOWrapper
+
     from airflow.models.taskinstance import TaskInstance
     from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
-    from airflow.utils.log.file_task_handler import LogMessages, LogSourceInfo
+    from airflow.utils.log.file_task_handler import LogResponse, RawLogStream, StreamingLogResponse
 
 _DEFAULT_SCOPESS = frozenset(
     [
@@ -149,11 +151,26 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
             exc, "resp", {}
         ).get("status") == "404"
 
-    def read(self, relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages | None]:
-        messages = []
-        logs = []
+    def read(self, relative_path: str, ti: RuntimeTI) -> LogResponse:
+        messages, log_streams = self.stream(relative_path, ti)
+        if not log_streams:
+            return messages, None
+
+        logs: list[str] = []
+        try:
+            # for each log_stream, exhaust the generator into a string
+            logs = ["".join(line for line in log_stream) for log_stream in log_streams]
+        except Exception as e:
+            if not AIRFLOW_V_3_0_PLUS:
+                messages.append(f"Unable to read remote log {e}")
+
+        return messages, logs
+
+    def stream(self, relative_path: str, ti: RuntimeTI) -> StreamingLogResponse:
+        messages: list[str] = []
+        log_streams: list[RawLogStream] = []
         remote_loc = os.path.join(self.remote_base, relative_path)
-        uris = []
+        uris: list[str] = []
         bucket, prefix = _parse_gcs_url(remote_loc)
         blobs = list(self.client.list_blobs(bucket_or_name=bucket, prefix=prefix))
 
@@ -164,18 +181,29 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
             else:
                 messages.extend(["Found remote logs:", *[f"  * {x}" for x in sorted(uris)]])
         else:
-            return messages, None
+            return messages, []
 
         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)
+                stream = blob.open("r")
+                log_streams.append(self._get_log_stream(stream))
         except Exception as e:
             if not AIRFLOW_V_3_0_PLUS:
                 messages.append(f"Unable to read remote log {e}")
-        return messages, logs
+        return messages, log_streams
+
+    def _get_log_stream(self, stream: TextIOWrapper) -> RawLogStream:
+        """
+        Yield lines from the given stream.
+
+        :param stream: The opened stream to read from.
+        :yield: Lines of the log file.
+        """
+        try:
+            yield from stream
+        finally:
+            stream.close()
 
 
 class GCSTaskHandler(FileTaskHandler, LoggingMixin):
@@ -273,7 +301,7 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         # Mark closed so we don't double write if close is called twice
         self.closed = True
 
-    def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[LogSourceInfo, LogMessages]:
+    def _read_remote_logs(self, ti, try_number, metadata=None) -> LogResponse:
         # Explicitly getting log relative path is necessary as the given
         # task instance might be different than task instance passed in
         # in set_context method.
@@ -283,7 +311,7 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
 
         if logs is None:
             logs = []
-            if not AIRFLOW_V_3_0_PLUS:
+            if not AIRFLOW_V_3_0_PLUS and not messages:
                 messages.append(f"No logs found in GCS; ti={ti}")
 
         return messages, logs

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

@@ -27,7 +27,7 @@ from google.api_core.exceptions import NotFound
 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.cloud.hooks.alloy_db import AlloyDbHook
 from airflow.providers.google.cloud.links.alloy_db import (
     AlloyDBBackupsLink,

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

@@ -33,8 +33,8 @@ from google.cloud.bigquery import DEFAULT_RETRY, CopyJob, ExtractJob, LoadJob, Q
 from google.cloud.bigquery.table import RowIterator, Table, TableListItem, TableReference
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import AirflowSkipException
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException, AirflowSkipException
 from airflow.providers.common.sql.operators.sql import (  # for _parse_boolean
     SQLCheckOperator,
     SQLColumnCheckOperator,

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

@@ -33,7 +33,7 @@ from google.cloud.bigquery_datatransfer_v1 import (
 )
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook, get_object_id
 from airflow.providers.google.cloud.links.bigquery_dts import BigQueryDataTransferConfigLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any
 
 import google.api_core.exceptions
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
 from airflow.providers.google.cloud.links.bigtable import (
     BigtableClusterLink,

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

@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING
 from google.cloud.batch_v1 import Job, Task
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.cloud_batch import CloudBatchJobFinishedTrigger

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

@@ -30,7 +30,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildHook
 from airflow.providers.google.cloud.links.cloud_build import (
     CloudBuildLink,

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

@@ -27,7 +27,7 @@ from google.cloud.orchestration.airflow.service_v1 import ImageVersion
 from google.cloud.orchestration.airflow.service_v1.types import Environment, ExecuteAirflowCommandResponse
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -25,7 +25,7 @@ import google.cloud.exceptions
 from google.api_core.exceptions import AlreadyExists
 from google.cloud.logging_v2.types import LogSink
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_logging import CloudLoggingHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 

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

@@ -25,7 +25,7 @@ from google.api_core.exceptions import AlreadyExists
 from google.cloud.run_v2 import Job, Service
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_run import CloudRunHook, CloudRunServiceHook
 from airflow.providers.google.cloud.links.cloud_run import CloudRunJobLoggingLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -27,8 +27,7 @@ from typing import TYPE_CHECKING, Any
 from googleapiclient.errors import HttpError
 
 from airflow.configuration import conf
-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.cloud.hooks.cloud_sql import CloudSQLDatabaseHook, CloudSQLHook
 from airflow.providers.google.cloud.links.cloud_sql import CloudSQLInstanceDatabaseLink, CloudSQLInstanceLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -25,8 +25,8 @@ from datetime import date, time, timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     ACCESS_KEY_ID,
     AWS_ACCESS_KEY,

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

@@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Any
 from google.api_core import exceptions
 from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google._vendor.json_merge_patch import merge
 from airflow.providers.google.cloud.hooks.compute import ComputeEngineHook
 from airflow.providers.google.cloud.links.compute import (

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

@@ -28,7 +28,7 @@ from typing import TYPE_CHECKING, Any
 from googleapiclient.errors import HttpError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.dataflow import (
     DEFAULT_DATAFLOW_LOCATION,
     DataflowHook,
@@ -583,6 +583,7 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
     def hook(self) -> DataflowHook:
         hook = DataflowHook(
             gcp_conn_id=self.gcp_conn_id,
+            poll_sleep=self.poll_sleep,
             drain_pipeline=self.drain_pipeline,
             cancel_timeout=self.cancel_timeout,
             wait_until_finished=self.wait_until_finished,

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
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.datafusion import SUCCESS_STATES, DataFusionHook, PipelineStates
 from airflow.providers.google.cloud.links.datafusion import (
     DataFusionInstanceLink,

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

@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, Any
 
 from google.protobuf.json_format import MessageToDict
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.triggers.dataplex import (
     DataplexDataProfileJobTrigger,
     DataplexDataQualityJobTrigger,

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

@@ -36,7 +36,8 @@ from google.api_core.retry import Retry, exponential_sleep_generator
 from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.dataproc import (
     DataprocHook,
     DataProcJobBuilder,

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

@@ -30,7 +30,7 @@ from google.cloud.metastore_v1 import MetadataExport, MetadataManagementActivity
 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.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -22,7 +22,7 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.datastore import DatastoreHook
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.links.datastore import (

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

@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, Any
 
 from googleapiclient.errors import HttpError
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.functions import CloudFunctionsHook
 from airflow.providers.google.cloud.links.cloud_functions import (
     CloudFunctionsDetailsLink,

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

@@ -38,7 +38,8 @@ if TYPE_CHECKING:
 from google.api_core.exceptions import Conflict
 from google.cloud.exceptions import GoogleCloudError
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.links.storage import FileDetailsLink, StorageLink

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

@@ -29,7 +29,7 @@ from kubernetes.client import V1JobList, models as k8s
 from packaging.version import parse as parse_version
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.cncf.kubernetes.operators.job import KubernetesJobOperator
 from airflow.providers.cncf.kubernetes.operators.kueue import (
     KubernetesInstallKueueOperator,
@@ -41,6 +41,7 @@ from airflow.providers.cncf.kubernetes.operators.resource import (
     KubernetesDeleteResourceOperator,
 )
 from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.kubernetes_engine import (
     GKEHook,
     GKEKubernetesHook,

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

@@ -21,7 +21,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.looker import LookerHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 

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

@@ -27,7 +27,7 @@ from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.managedkafka_v1 import types
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.managed_kafka import ManagedKafkaHook
 from airflow.providers.google.cloud.links.managed_kafka import (
     ApacheKafkaClusterLink,

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

@@ -42,7 +42,7 @@ from google.cloud.pubsub_v1.types import (
 )
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
 from airflow.providers.google.cloud.links.pubsub import PubSubSubscriptionLink, PubSubTopicLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
@@ -680,6 +680,9 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         ordering_key in PubsubMessage will be delivered to the subscribers in the order
         in which they are received by the Pub/Sub system. Otherwise, they may be
         delivered in any order. Default is False.
+    :param enable_open_telemetry_tracing: If true, enables OpenTelemetry tracing for
+        published messages. This allows distributed tracing of messages as they flow
+        through Pub/Sub topics. Default is False.
     :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.
@@ -695,6 +698,7 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         "topic",
         "messages",
         "enable_message_ordering",
+        "enable_open_telemetry_tracing",
         "impersonation_chain",
     )
     ui_color = "#0273d4"
@@ -707,6 +711,7 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         enable_message_ordering: bool = False,
+        enable_open_telemetry_tracing: bool = False,
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
@@ -716,6 +721,7 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         self.messages = messages
         self.gcp_conn_id = gcp_conn_id
         self.enable_message_ordering = enable_message_ordering
+        self.enable_open_telemetry_tracing = enable_open_telemetry_tracing
         self.impersonation_chain = impersonation_chain
 
     @cached_property
@@ -724,6 +730,7 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
             enable_message_ordering=self.enable_message_ordering,
+            enable_open_telemetry_tracing=self.enable_open_telemetry_tracing,
         )
 
     def execute(self, context: Context) -> None:

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

@@ -23,7 +23,7 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.spanner import SpannerHook
 from airflow.providers.google.cloud.links.spanner import SpannerDatabaseLink, SpannerInstanceLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.protobuf.json_format import MessageToDict
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook, RecognitionAudio
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID

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

@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 
-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.cloud.hooks.text_to_speech import CloudTextToSpeechHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, cast
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook, TranslateHook
 from airflow.providers.google.cloud.links.translate import (
     TranslateResultByOutputConfigLink,

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

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING
 
 from google.protobuf.json_format import MessageToDict
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

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

@@ -29,7 +29,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1.types import BatchPredictionJob
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIBatchPredictionJobLink,

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

@@ -30,7 +30,7 @@ from google.cloud.aiplatform_v1.types.dataset import Dataset
 from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIModelLink,

airflow/providers/google/cloud/operators/vertex_ai/dataset.py

@@ -26,7 +26,7 @@ from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1.types import Dataset, ExportDataConfig, ImportDataConfig
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.dataset import DatasetHook
 from airflow.providers.google.cloud.links.vertex_ai import VertexAIDatasetLink, VertexAIDatasetListLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator

airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py

@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING
 
 from google.api_core import exceptions
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.experiment_service import (
     ExperimentHook,
     ExperimentRunHook,

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

@@ -26,7 +26,7 @@ from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1beta1.types import FeatureViewDataFormat
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

airflow/providers/google/cloud/operators/vertex_ai/generative_model.py

@@ -24,7 +24,8 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from google.api_core import exceptions
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import (
     ExperimentRunHook,
     GenerativeModelHook,

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

@@ -28,7 +28,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import types
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
     HyperparameterTuningJobHook,
 )

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

@@ -28,7 +28,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import types
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job import PipelineJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIPipelineJobLink,

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

@@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Literal
 from google.api_core.exceptions import NotFound
 from google.cloud.aiplatform.vertex_ray.util import resources
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.ray import RayHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIRayClusterLink,

airflow/providers/google/cloud/secrets/secret_manager.py

@@ -23,7 +23,7 @@ from collections.abc import Sequence
 
 from google.auth.exceptions import DefaultCredentialsError
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient
 from airflow.providers.google.cloud.utils.credentials_provider import (
     _get_target_principal_and_delegates,
@@ -160,11 +160,12 @@ class CloudSecretManagerBackend(BaseSecretsBackend, LoggingMixin):
 
         return self._get_secret(self.connections_prefix, conn_id)
 
-    def get_variable(self, key: str) -> str | None:
+    def get_variable(self, key: str, team_name: str | None = None) -> str | None:
         """
         Get Airflow Variable from Environment Variable.
 
         :param key: Variable Key
+        :param team_name: Team name associated to the task trying to access the variable (if any)
         :return: Variable Value
         """
         if self.variables_prefix is None:

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

@@ -25,8 +25,8 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import BaseSensorOperator
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.triggers.bigquery import (
     BigQueryTableExistenceTrigger,

airflow/providers/google/cloud/sensors/bigquery_dts.py

@@ -25,8 +25,7 @@ from typing import TYPE_CHECKING
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.bigquery_datatransfer_v1 import TransferState
 
-from airflow.exceptions import AirflowException
-from airflow.providers.common.compat.sdk import BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 

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

@@ -30,8 +30,7 @@ from google.api_core.exceptions import NotFound
 from google.cloud.orchestration.airflow.service_v1.types import Environment, ExecuteAirflowCommandResponse
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
-from airflow.providers.common.compat.sdk import AirflowSkipException, BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowException, AirflowSkipException, BaseSensorOperator
 from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
 from airflow.providers.google.cloud.triggers.cloud_composer import (
     CloudComposerDAGRunTrigger,

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

@@ -23,8 +23,7 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
-from airflow.providers.common.compat.sdk import BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     COUNTERS,
     METADATA,