apache-airflow-providers-google 12.0.0rc2__py3-none-any.whl → 14.0.0rc1__py3-none-any.whl

airflow/providers/google/cloud/transfers/postgres_to_gcs.py

@@ -24,6 +24,8 @@ import json
 import time
 import uuid
 from decimal import Decimal
+from functools import cached_property
+from typing import TYPE_CHECKING
 
 import pendulum
 from slugify import slugify
@@ -31,6 +33,9 @@ from slugify import slugify
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
 from airflow.providers.postgres.hooks.postgres import PostgresHook
 
+if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
+
 
 class _PostgresServerSideCursorDecorator:
     """
@@ -132,10 +137,13 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
             )
         return None
 
+    @cached_property
+    def db_hook(self) -> PostgresHook:
+        return PostgresHook(postgres_conn_id=self.postgres_conn_id)
+
     def query(self):
         """Query Postgres and returns a cursor to the results."""
-        hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
-        conn = hook.get_conn()
+        conn = self.db_hook.get_conn()
         cursor = conn.cursor(name=self._unique_name())
         cursor.execute(self.sql, self.parameters)
         if self.use_server_side_cursor:
@@ -180,3 +188,20 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
         if isinstance(value, Decimal):
             return float(value)
         return value
+
+    def get_openlineage_facets_on_start(self) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        sql_parsing_result = get_openlineage_facets_with_sql(
+            hook=self.db_hook,
+            sql=self.sql,
+            conn_id=self.postgres_conn_id,
+            database=self.db_hook.database,
+        )
+        gcs_output_datasets = self._get_openlineage_output_datasets()
+        if sql_parsing_result:
+            sql_parsing_result.outputs = gcs_output_datasets
+            return sql_parsing_result
+        return OperatorLineage(outputs=gcs_output_datasets, job_facets={"sql": SQLJobFacet(self.sql)})

airflow/providers/google/cloud/transfers/sftp_to_gcs.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 
 import os
 from collections.abc import Sequence
+from functools import cached_property
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
@@ -109,6 +110,10 @@ class SFTPToGCSOperator(BaseOperator):
         self.impersonation_chain = impersonation_chain
         self.sftp_prefetch = sftp_prefetch
 
+    @cached_property
+    def sftp_hook(self):
+        return SFTPHook(self.sftp_conn_id)
+
     def execute(self, context: Context):
         self.destination_path = self._set_destination_path(self.destination_path)
         self.destination_bucket = self._set_bucket_name(self.destination_bucket)
@@ -117,8 +122,6 @@ class SFTPToGCSOperator(BaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
 
-        sftp_hook = SFTPHook(self.sftp_conn_id)
-
         if WILDCARD in self.source_path:
             total_wildcards = self.source_path.count(WILDCARD)
             if total_wildcards > 1:
@@ -130,7 +133,7 @@ class SFTPToGCSOperator(BaseOperator):
             prefix, delimiter = self.source_path.split(WILDCARD, 1)
             base_path = os.path.dirname(prefix)
 
-            files, _, _ = sftp_hook.get_tree_map(base_path, prefix=prefix, delimiter=delimiter)
+            files, _, _ = self.sftp_hook.get_tree_map(base_path, prefix=prefix, delimiter=delimiter)
 
             for file in files:
                 destination_path = file.replace(base_path, self.destination_path, 1)
@@ -140,13 +143,13 @@ class SFTPToGCSOperator(BaseOperator):
                 # retain the "/" prefix, if it has.
                 if not self.destination_path:
                     destination_path = destination_path.lstrip("/")
-                self._copy_single_object(gcs_hook, sftp_hook, file, destination_path)
+                self._copy_single_object(gcs_hook, self.sftp_hook, file, destination_path)
 
         else:
             destination_object = (
                 self.destination_path if self.destination_path else self.source_path.rsplit("/", 1)[1]
             )
-            self._copy_single_object(gcs_hook, sftp_hook, self.source_path, destination_object)
+            self._copy_single_object(gcs_hook, self.sftp_hook, self.source_path, destination_object)
 
     def _copy_single_object(
         self,
@@ -188,3 +191,29 @@ class SFTPToGCSOperator(BaseOperator):
     def _set_bucket_name(name: str) -> str:
         bucket = name if not name.startswith("gs://") else name[5:]
         return bucket.strip("/")
+
+    def get_openlineage_facets_on_start(self):
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import extract_ds_name_from_gcs_path
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        source_name = extract_ds_name_from_gcs_path(self.source_path.split(WILDCARD, 1)[0])
+        if self.source_path.startswith("/") and source_name != "/":
+            source_name = "/" + source_name
+
+        if WILDCARD not in self.source_path and not self.destination_path:
+            dest_name = self.source_path.rsplit("/", 1)[1]
+        else:
+            dest_name = extract_ds_name_from_gcs_path(f"{self.destination_path}")
+
+        return OperatorLineage(
+            inputs=[
+                Dataset(
+                    namespace=f"file://{self.sftp_hook.remote_host}:{self.sftp_hook.port}",
+                    name=source_name,
+                )
+            ],
+            outputs=[
+                Dataset(namespace="gs://" + self._set_bucket_name(self.destination_bucket), name=dest_name)
+            ],
+        )

airflow/providers/google/cloud/transfers/sql_to_gcs.py

@@ -34,6 +34,7 @@ from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 
 if TYPE_CHECKING:
+    from airflow.providers.common.compat.openlineage.facet import OutputDataset
     from airflow.utils.context import Context
 
 
@@ -151,6 +152,7 @@ class BaseSQLToGCSOperator(BaseOperator):
         self.partition_columns = partition_columns
         self.write_on_empty = write_on_empty
         self.parquet_row_group_size = parquet_row_group_size
+        self._uploaded_file_names: list[str] = []
 
     def execute(self, context: Context):
         if self.partition_columns:
@@ -501,3 +503,16 @@ class BaseSQLToGCSOperator(BaseOperator):
             gzip=self.gzip if is_data_file else False,
             metadata=metadata,
         )
+        self._uploaded_file_names.append(object_name)
+
+    def _get_openlineage_output_datasets(self) -> list[OutputDataset]:
+        """Retrieve OpenLineage output datasets."""
+        from airflow.providers.common.compat.openlineage.facet import OutputDataset
+        from airflow.providers.google.cloud.openlineage.utils import extract_ds_name_from_gcs_path
+
+        return [
+            OutputDataset(
+                namespace=f"gs://{self.bucket}",
+                name=extract_ds_name_from_gcs_path(self.filename.split("{}", maxsplit=1)[0]),
+            )
+        ]

airflow/providers/google/cloud/transfers/trino_to_gcs.py

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
@@ -26,6 +27,8 @@ if TYPE_CHECKING:
     from trino.client import TrinoResult
     from trino.dbapi import Cursor as TrinoCursor
 
+    from airflow.providers.openlineage.extractors import OperatorLineage
+
 
 class _TrinoToGCSTrinoCursorAdapter:
     """
@@ -181,10 +184,13 @@ class TrinoToGCSOperator(BaseSQLToGCSOperator):
         super().__init__(**kwargs)
         self.trino_conn_id = trino_conn_id
 
+    @cached_property
+    def db_hook(self) -> TrinoHook:
+        return TrinoHook(trino_conn_id=self.trino_conn_id)
+
     def query(self):
         """Query trino and returns a cursor to the results."""
-        trino = TrinoHook(trino_conn_id=self.trino_conn_id)
-        conn = trino.get_conn()
+        conn = self.db_hook.get_conn()
         cursor = conn.cursor()
         self.log.info("Executing: %s", self.sql)
         cursor.execute(self.sql)
@@ -207,3 +213,20 @@ class TrinoToGCSOperator(BaseSQLToGCSOperator):
         :param schema_type: BigQuery data type
         """
         return value
+
+    def get_openlineage_facets_on_start(self) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        sql_parsing_result = get_openlineage_facets_with_sql(
+            hook=self.db_hook,
+            sql=self.sql,
+            conn_id=self.trino_conn_id,
+            database=None,
+        )
+        gcs_output_datasets = self._get_openlineage_output_datasets()
+        if sql_parsing_result:
+            sql_parsing_result.outputs = gcs_output_datasets
+            return sql_parsing_result
+        return OperatorLineage(outputs=gcs_output_datasets, job_facets={"sql": SQLJobFacet(self.sql)})

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

@@ -21,10 +21,9 @@ import asyncio
 from collections.abc import AsyncIterator, Sequence
 from typing import Any
 
-from google.cloud.bigquery_datatransfer_v1 import TransferRun, TransferState
-
 from airflow.providers.google.cloud.hooks.bigquery_dts import AsyncBiqQueryDataTransferServiceHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.bigquery_datatransfer_v1 import TransferRun, TransferState
 
 
 class BigQueryDataTransferRunTrigger(BaseTrigger):

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

@@ -20,10 +20,9 @@ import asyncio
 from collections.abc import AsyncIterator, Sequence
 from typing import Any
 
-from google.cloud.batch_v1 import Job, JobStatus
-
 from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.batch_v1 import Job, JobStatus
 
 DEFAULT_BATCH_LOCATION = "us-central1"
 

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

@@ -20,10 +20,9 @@ import asyncio
 from collections.abc import AsyncIterator, Sequence
 from typing import Any
 
-from google.cloud.devtools.cloudbuild_v1.types import Build
-
 from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.devtools.cloudbuild_v1.types import Build
 
 
 class CloudBuildCreateBuildTrigger(BaseTrigger):

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

@@ -25,11 +25,11 @@ from datetime import datetime
 from typing import Any
 
 from dateutil import parser
-from google.cloud.orchestration.airflow.service_v1.types import ExecuteAirflowCommandResponse
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.orchestration.airflow.service_v1.types import ExecuteAirflowCommandResponse
 
 
 class CloudComposerExecutionTrigger(BaseTrigger):
@@ -169,6 +169,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         poll_interval: int = 10,
+        composer_airflow_version: int = 2,
     ):
         super().__init__()
         self.project_id = project_id
@@ -181,6 +182,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
         self.poll_interval = poll_interval
+        self.composer_airflow_version = composer_airflow_version
 
         self.gcp_hook = CloudComposerAsyncHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -201,18 +203,24 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
                 "gcp_conn_id": self.gcp_conn_id,
                 "impersonation_chain": self.impersonation_chain,
                 "poll_interval": self.poll_interval,
+                "composer_airflow_version": self.composer_airflow_version,
             },
         )
 
     async def _pull_dag_runs(self) -> list[dict]:
         """Pull the list of dag runs."""
+        cmd_parameters = (
+            ["-d", self.composer_dag_id, "-o", "json"]
+            if self.composer_airflow_version < 3
+            else [self.composer_dag_id, "-o", "json"]
+        )
         dag_runs_cmd = await self.gcp_hook.execute_airflow_command(
             project_id=self.project_id,
             region=self.region,
             environment_id=self.environment_id,
             command="dags",
             subcommand="list-runs",
-            parameters=["-d", self.composer_dag_id, "-o", "json"],
+            parameters=cmd_parameters,
         )
         cmd_result = await self.gcp_hook.wait_command_execution_result(
             project_id=self.project_id,
@@ -232,7 +240,9 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         for dag_run in dag_runs:
             if (
                 start_date.timestamp()
-                < parser.parse(dag_run["logical_date"]).timestamp()
+                < parser.parse(
+                    dag_run["execution_date" if self.composer_airflow_version < 3 else "logical_date"]
+                ).timestamp()
                 < end_date.timestamp()
             ) and dag_run["state"] not in self.allowed_states:
                 return False

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

@@ -18,18 +18,18 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import AsyncIterator, Iterable, Sequence
 from typing import Any
 
-from google.api_core.exceptions import GoogleAPIError
-from google.cloud.storage_transfer_v1.types import TransferOperation
-
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     CloudDataTransferServiceAsyncHook,
+    GcpTransferOperationStatus,
 )
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.api_core.exceptions import GoogleAPIError
+from google.cloud.storage_transfer_v1.types import TransferOperation
 
 
 class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
@@ -132,3 +132,101 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
             project_id=self.project_id,
             gcp_conn_id=self.gcp_conn_id,
         )
+
+
+class CloudStorageTransferServiceCheckJobStatusTrigger(BaseTrigger):
+    """
+    CloudStorageTransferServiceCheckJobStatusTrigger run on the trigger worker to check Cloud Storage Transfer job.
+
+    :param job_name: The name of the transfer job
+    :param expected_statuses: The expected state of the operation.
+        See:
+        https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#Status
+    :param project_id: The ID of the project that owns the Transfer Job.
+    :param poke_interval: Polling period in seconds to check for the status
+    :param gcp_conn_id: The connection ID used to connect 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).
+    """
+
+    def __init__(
+        self,
+        job_name: str,
+        expected_statuses: set[str] | str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+        poke_interval: float = 10.0,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+    ):
+        super().__init__()
+        self.job_name = job_name
+        self.expected_statuses = expected_statuses
+        self.project_id = project_id
+        self.poke_interval = poke_interval
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize CloudStorageTransferServiceCheckJobStatusTrigger arguments and classpath."""
+        return (
+            f"{self.__class__.__module__ }.{self.__class__.__qualname__}",
+            {
+                "job_name": self.job_name,
+                "expected_statuses": self.expected_statuses,
+                "project_id": self.project_id,
+                "poke_interval": self.poke_interval,
+                "gcp_conn_id": self.gcp_conn_id,
+                "impersonation_chain": self.impersonation_chain,
+            },
+        )
+
+    def _get_async_hook(self) -> CloudDataTransferServiceAsyncHook:
+        return CloudDataTransferServiceAsyncHook(
+            project_id=self.project_id,
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        """Check the status of the transfer job and yield a TriggerEvent."""
+        hook = self._get_async_hook()
+        expected_statuses = (
+            {GcpTransferOperationStatus.SUCCESS} if not self.expected_statuses else self.expected_statuses
+        )
+
+        try:
+            while True:
+                operations = await hook.list_transfer_operations(
+                    request_filter={
+                        "project_id": self.project_id or hook.project_id,
+                        "job_names": [self.job_name],
+                    }
+                )
+                check = await CloudDataTransferServiceAsyncHook.operations_contain_expected_statuses(
+                    operations=operations,
+                    expected_statuses=expected_statuses,
+                )
+                if check:
+                    yield TriggerEvent(
+                        {
+                            "status": "success",
+                            "message": "Transfer operation completed successfully",
+                            "operations": operations,
+                        }
+                    )
+                    return
+
+                self.log.info(
+                    "Sleeping for %s seconds.",
+                    self.poke_interval,
+                )
+                await asyncio.sleep(self.poke_interval)
+        except Exception as e:
+            self.log.exception("Exception occurred while checking for query completion")
+            yield TriggerEvent({"status": "error", "message": str(e)})

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

@@ -22,6 +22,8 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
+from airflow.triggers.base import BaseTrigger, TriggerEvent
 from google.cloud.dataflow_v1beta3 import JobState
 from google.cloud.dataflow_v1beta3.types import (
     AutoscalingEvent,
@@ -32,9 +34,6 @@ from google.cloud.dataflow_v1beta3.types import (
     MetricUpdate,
 )
 
-from airflow.providers.google.cloud.hooks.dataflow import AsyncDataflowHook, DataflowJobStatus
-from airflow.triggers.base import BaseTrigger, TriggerEvent
-
 if TYPE_CHECKING:
     from google.cloud.dataflow_v1beta3.services.messages_v1_beta3.pagers import ListJobMessagesAsyncPager
 

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

@@ -22,10 +22,9 @@ from __future__ import annotations
 import asyncio
 from collections.abc import AsyncIterator, Sequence
 
-from google.cloud.dataplex_v1.types import DataScanJob
-
 from airflow.providers.google.cloud.hooks.dataplex import DataplexAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.dataplex_v1.types import DataScanJob
 
 
 class DataplexDataQualityJobTrigger(BaseTrigger):

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

@@ -25,9 +25,6 @@ import time
 from collections.abc import AsyncIterator, Sequence
 from typing import TYPE_CHECKING, Any
 
-from google.api_core.exceptions import NotFound
-from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
-
 from airflow.exceptions import AirflowException
 from airflow.models.taskinstance import TaskInstance
 from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook, DataprocHook
@@ -36,6 +33,8 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 from airflow.utils.session import provide_session
 from airflow.utils.state import TaskInstanceState
+from google.api_core.exceptions import NotFound
+from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
 
 if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session

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

@@ -23,7 +23,6 @@ from collections.abc import AsyncIterator, Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from google.cloud.container_v1.types import Operation
 from packaging.version import parse as parse_version
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
@@ -37,6 +36,7 @@ from airflow.providers.google.cloud.hooks.kubernetes_engine import (
 )
 from airflow.providers_manager import ProvidersManager
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.container_v1.types import Operation
 
 if TYPE_CHECKING:
     from datetime import datetime

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

@@ -23,10 +23,9 @@ from collections.abc import AsyncIterator, Sequence
 from functools import cached_property
 from typing import Any
 
-from google.cloud.pubsub_v1.types import ReceivedMessage
-
 from airflow.providers.google.cloud.hooks.pubsub import PubSubAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.pubsub_v1.types import ReceivedMessage
 
 
 class PubsubPullTrigger(BaseTrigger):

airflow/providers/google/cloud/triggers/vertex_ai.py

@@ -20,14 +20,6 @@ from collections.abc import AsyncIterator, Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from google.cloud.aiplatform_v1 import (
-    BatchPredictionJob,
-    HyperparameterTuningJob,
-    JobState,
-    PipelineState,
-    types,
-)
-
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobAsyncHook
 from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobAsyncHook
@@ -36,6 +28,13 @@ from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job im
 )
 from airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job import PipelineJobAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from google.cloud.aiplatform_v1 import (
+    BatchPredictionJob,
+    HyperparameterTuningJob,
+    JobState,
+    PipelineState,
+    types,
+)
 
 if TYPE_CHECKING:
     from proto import Message

airflow/providers/google/cloud/utils/credentials_provider.py

@@ -29,10 +29,6 @@ from urllib.parse import urlencode
 
 import google.auth
 import google.oauth2.service_account
-from google.auth import impersonated_credentials  # type: ignore[attr-defined]
-from google.auth.credentials import AnonymousCredentials, Credentials
-from google.auth.environment_vars import CREDENTIALS, LEGACY_PROJECT, PROJECT
-
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient
 from airflow.providers.google.cloud.utils.external_token_supplier import (
@@ -40,6 +36,9 @@ from airflow.providers.google.cloud.utils.external_token_supplier import (
 )
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.process_utils import patch_environ
+from google.auth import impersonated_credentials  # type: ignore[attr-defined]
+from google.auth.credentials import AnonymousCredentials, Credentials
+from google.auth.environment_vars import CREDENTIALS, LEGACY_PROJECT, PROJECT
 
 log = logging.getLogger(__name__)
 
@@ -219,11 +218,19 @@ class _CredentialProvider(LoggingMixin):
         idp_extra_params_dict: dict[str, str] | None = None,
     ) -> None:
         super().__init__()
-        key_options = [key_path, keyfile_dict, credential_config_file, key_secret_name, is_anonymous]
-        if len([x for x in key_options if x]) > 1:
+        key_options_map = {
+            "key_path": key_path,
+            "keyfile_dict": keyfile_dict,
+            "credential_config_file": credential_config_file,
+            "key_secret_name": key_secret_name,
+            "is_anonymous": is_anonymous,
+        }
+        key_options_label_provided = [label for label, credential in key_options_map.items() if credential]
+        if len(key_options_label_provided) > 1:
             raise AirflowException(
-                "The `keyfile_dict`, `key_path`, `credential_config_file`, `is_anonymous` and"
-                " `key_secret_name` fields are all mutually exclusive. Please provide only one value."
+                f"The `keyfile_dict`, `key_path`, `credential_config_file`, `is_anonymous` and"
+                f" `key_secret_name` fields are all mutually exclusive. "
+                f"Received options: {key_options_label_provided}. Please provide only one value."
             )
         self.key_path = key_path
         self.keyfile_dict = keyfile_dict

airflow/providers/google/cloud/utils/external_token_supplier.py

@@ -22,6 +22,7 @@ from functools import wraps
 from typing import TYPE_CHECKING, Any
 
 import requests
+
 from google.auth.exceptions import RefreshError
 from google.auth.identity_pool import SubjectTokenSupplier
 

airflow/providers/google/common/auth_backend/google_openid.py

@@ -23,17 +23,17 @@ import logging
 from functools import wraps
 from typing import Callable, TypeVar, cast
 
+from flask import Response, current_app, request as flask_request  # type: ignore
+
 import google
 import google.auth.transport.requests
 import google.oauth2.id_token
-from flask import Response, current_app, request as flask_request  # type: ignore
+from airflow.configuration import conf
+from airflow.providers.google.common.utils.id_token_credentials import get_default_id_token_credentials
 from google.auth import exceptions
 from google.auth.transport.requests import AuthorizedSession
 from google.oauth2 import service_account
 
-from airflow.configuration import conf
-from airflow.providers.google.common.utils.id_token_credentials import get_default_id_token_credentials
-
 log = logging.getLogger(__name__)
 
 _GOOGLE_ISSUERS = ("accounts.google.com", "https://accounts.google.com")

airflow/providers/google/common/consts.py

@@ -16,9 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from google.api_core.gapic_v1.client_info import ClientInfo
-
 from airflow import version
+from google.api_core.gapic_v1.client_info import ClientInfo
 
 GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME = "execute_complete"