apache-airflow-providers-google 10.12.0__py3-none-any.whl → 10.13.0__py3-none-any.whl

airflow/providers/google/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "10.12.0"
+__version__ = "10.13.0"
 
 try:
     from airflow import __version__ as airflow_version
@@ -35,8 +35,8 @@ except ImportError:
     from airflow.version import version as airflow_version
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.5.0"
+    "2.6.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 2.5.0+"
+        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 2.6.0+"
     )

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

@@ -39,7 +39,7 @@ GCS_VERSION_AWARE = "gcs.version-aware"
 schemes = ["gs", "gcs"]
 
 
-def get_fs(conn_id: str | None) -> AbstractFileSystem:
+def get_fs(conn_id: str | None, storage_options: dict[str, str] | None = None) -> AbstractFileSystem:
     # https://gcsfs.readthedocs.io/en/latest/api.html#gcsfs.core.GCSFileSystem
     from gcsfs import GCSFileSystem
 
@@ -49,15 +49,18 @@ def get_fs(conn_id: str | None) -> AbstractFileSystem:
     g = GoogleBaseHook(gcp_conn_id=conn_id)
     creds = g.get_credentials()
 
-    return GCSFileSystem(
-        project=g.project_id,
-        access=g.extras.get(GCS_ACCESS, "full_control"),
-        token=creds.token,
-        consistency=g.extras.get(GCS_CONSISTENCY, "none"),
-        cache_timeout=g.extras.get(GCS_CACHE_TIMEOUT),
-        requester_pays=g.extras.get(GCS_REQUESTER_PAYS, False),
-        session_kwargs=g.extras.get(GCS_SESSION_KWARGS, {}),
-        endpoint_url=g.extras.get(GCS_ENDPOINT),
-        default_location=g.extras.get(GCS_DEFAULT_LOCATION),
-        version_aware=g.extras.get(GCS_VERSION_AWARE, "false").lower() == "true",
-    )
+    options = {
+        "project": g.project_id,
+        "access": g.extras.get(GCS_ACCESS, "full_control"),
+        "token": creds.token,
+        "consistency": g.extras.get(GCS_CONSISTENCY, "none"),
+        "cache_timeout": g.extras.get(GCS_CACHE_TIMEOUT),
+        "requester_pays": g.extras.get(GCS_REQUESTER_PAYS, False),
+        "session_kwargs": g.extras.get(GCS_SESSION_KWARGS, {}),
+        "endpoint_url": g.extras.get(GCS_ENDPOINT),
+        "default_location": g.extras.get(GCS_DEFAULT_LOCATION),
+        "version_aware": g.extras.get(GCS_VERSION_AWARE, "false").lower() == "true",
+    }
+    options.update(storage_options or {})
+
+    return GCSFileSystem(**options)

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

@@ -38,6 +38,7 @@ from airflow.providers.google.common.hooks.base_google import (
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
+    from google.api_core.retry_async import AsyncRetry
     from googleapiclient.discovery import Resource
 
 
@@ -321,7 +322,7 @@ class AsyncBiqQueryDataTransferServiceHook(GoogleBaseAsyncHook):
         run_id: str,
         project_id: str | None,
         location: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ):

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

@@ -33,6 +33,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
+    from google.api_core.retry_async import AsyncRetry
     from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
 
 # Time to sleep between active checks of the operation results
@@ -645,7 +646,7 @@ class CloudBuildAsyncHook(GoogleBaseHook):
         self,
         id_: str,
         project_id: str = PROVIDE_PROJECT_ID,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
         location: str = "global",

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

@@ -35,6 +35,7 @@ if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.operation_async import AsyncOperation
     from google.api_core.retry import Retry
+    from google.api_core.retry_async import AsyncRetry
     from google.cloud.orchestration.airflow.service_v1.services.environments.pagers import (
         ListEnvironmentsPager,
     )
@@ -332,7 +333,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         project_id: str,
         region: str,
         environment: Environment | dict,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -361,7 +362,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         project_id: str,
         region: str,
         environment_id: str,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -389,7 +390,7 @@ class CloudComposerAsyncHook(GoogleBaseHook):
         environment_id: str,
         environment: Environment | dict,
         update_mask: dict | FieldMask,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:

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

@@ -86,6 +86,9 @@ class ComputeEngineSSHHook(SSHHook):
     :param gcp_conn_id: The connection id to use when fetching connection information
     :param max_retries: Maximum number of retries the process will try to establish connection to instance.
         Could be decreased/increased by user based on the amount of parallel SSH connections to the instance.
+    :param impersonation_chain: Optional. The service account email to impersonate using short-term
+        credentials. The provided service account must grant the originating account
+        the Service Account Token Creator IAM role and have the sufficient rights to perform the request
     """
 
     conn_name_attr = "gcp_conn_id"
@@ -93,8 +96,8 @@ class ComputeEngineSSHHook(SSHHook):
     conn_type = "gcpssh"
     hook_name = "Google Cloud SSH"
 
-    @staticmethod
-    def get_ui_field_behaviour() -> dict[str, Any]:
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
         return {
             "hidden_fields": ["host", "schema", "login", "password", "port", "extra"],
             "relabeling": {},
@@ -114,15 +117,17 @@ class ComputeEngineSSHHook(SSHHook):
         expire_time: int = 300,
         cmd_timeout: int | ArgNotSet = NOTSET,
         max_retries: int = 10,
+        impersonation_chain: str | None = None,
         **kwargs,
     ) -> None:
         if kwargs.get("delegate_to") is not None:
             raise RuntimeError(
                 "The `delegate_to` parameter has been deprecated before and finally removed in this version"
-                " of Google Provider. You MUST convert it to `impersonate_chain`"
+                " of Google Provider. You MUST convert it to `impersonation_chain`"
             )
         # Ignore original constructor
         # super().__init__()
+        self.gcp_conn_id = gcp_conn_id
         self.instance_name = instance_name
         self.zone = zone
         self.user = user
@@ -132,9 +137,9 @@ class ComputeEngineSSHHook(SSHHook):
         self.use_iap_tunnel = use_iap_tunnel
         self.use_oslogin = use_oslogin
         self.expire_time = expire_time
-        self.gcp_conn_id = gcp_conn_id
         self.cmd_timeout = cmd_timeout
         self.max_retries = max_retries
+        self.impersonation_chain = impersonation_chain
         self._conn: Any | None = None
 
     @cached_property
@@ -143,7 +148,12 @@ class ComputeEngineSSHHook(SSHHook):
 
     @cached_property
     def _compute_hook(self) -> ComputeEngineHook:
-        return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id)
+        if self.impersonation_chain:
+            return ComputeEngineHook(
+                gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
+            )
+        else:
+            return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id)
 
     def _load_connection_config(self):
         def _boolify(value):
@@ -254,6 +264,8 @@ class ComputeEngineSSHHook(SSHHook):
                         f"--zone={self.zone}",
                         "--verbosity=warning",
                     ]
+                    if self.impersonation_chain:
+                        proxy_command_args.append(f"--impersonate-service-account={self.impersonation_chain}")
                     proxy_command = " ".join(shlex.quote(arg) for arg in proxy_command_args)
                 sshclient = self._connect_to_instance(user, hostname, privkey, proxy_command)
                 break
@@ -283,7 +295,7 @@ class ComputeEngineSSHHook(SSHHook):
                 client = _GCloudAuthorizedSSHClient(self._compute_hook)
                 # Default is RejectPolicy
                 # No known host checking since we are not storing privatekey
-                client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+                client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # nosec B507
                 client.connect(
                     hostname=hostname,
                     username=user,

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

@@ -27,9 +27,10 @@ import time
 import uuid
 import warnings
 from copy import deepcopy
-from typing import Any, Callable, Generator, Sequence, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Callable, Generator, Sequence, TypeVar, cast
 
 from google.cloud.dataflow_v1beta3 import GetJobRequest, Job, JobState, JobsV1Beta3AsyncClient, JobView
+from google.cloud.dataflow_v1beta3.types.jobs import ListJobsRequest
 from googleapiclient.discovery import build
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
@@ -42,6 +43,10 @@ from airflow.providers.google.common.hooks.base_google import (
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.timeout import timeout
 
+if TYPE_CHECKING:
+    from google.cloud.dataflow_v1beta3.services.jobs_v1_beta3.pagers import ListJobsAsyncPager
+
+
 # This is the default location
 # https://cloud.google.com/dataflow/pipelines/specifying-exec-params
 DEFAULT_DATAFLOW_LOCATION = "us-central1"
@@ -55,7 +60,7 @@ T = TypeVar("T", bound=Callable)
 
 
 def process_line_and_extract_dataflow_job_id_callback(
-    on_new_job_id_callback: Callable[[str], None] | None
+    on_new_job_id_callback: Callable[[str], None] | None,
 ) -> Callable[[str], None]:
     """Build callback that triggers the specified function.
 
@@ -219,7 +224,7 @@ class _DataflowJobsController(LoggingMixin):
 
     def is_job_running(self) -> bool:
         """
-        Helper method to check if jos is still running in dataflow.
+        Helper method to check if job is still running in dataflow.
 
         :return: True if job is running.
         """
@@ -1203,6 +1208,24 @@ class DataflowHook(GoogleBaseHook):
         )
         job_controller.wait_for_done()
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def is_job_done(self, location: str, project_id: str, job_id: str) -> bool:
+        """
+        Check that Dataflow job is started(for streaming job) or finished(for batch job).
+
+        :param location: location the job is running
+        :param project_id: Google Cloud project ID in which to start a job
+        :param job_id: Dataflow job ID
+        """
+        job_controller = _DataflowJobsController(
+            dataflow=self.get_conn(),
+            project_number=project_id,
+            location=location,
+        )
+        job = job_controller.fetch_job_by_id(job_id)
+
+        return job_controller._check_dataflow_job_state(job)
+
 
 class AsyncDataflowHook(GoogleBaseAsyncHook):
     """Async hook class for dataflow service."""
@@ -1295,3 +1318,38 @@ class AsyncDataflowHook(GoogleBaseAsyncHook):
         )
         state = job.current_state
         return state
+
+    async def list_jobs(
+        self,
+        jobs_filter: int | None = None,
+        project_id: str | None = PROVIDE_PROJECT_ID,
+        location: str | None = DEFAULT_DATAFLOW_LOCATION,
+        page_size: int | None = None,
+        page_token: str | None = None,
+    ) -> ListJobsAsyncPager:
+        """List jobs.
+
+        For detail see:
+        https://cloud.google.com/python/docs/reference/dataflow/latest/google.cloud.dataflow_v1beta3.types.ListJobsRequest
+
+        :param jobs_filter: Optional. This field filters out and returns jobs in the specified job state.
+        :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 location: Optional. The location of the Dataflow job (for example europe-west1).
+        :param page_size: Optional. If there are many jobs, limit response to at most this many.
+        :param page_token: Optional. Set this to the 'next_page_token' field of a previous response to request
+            additional results in a long list.
+        """
+        project_id = project_id or (await self.get_project_id())
+        client = await self.initialize_client(JobsV1Beta3AsyncClient)
+        request: ListJobsRequest = ListJobsRequest(
+            {
+                "project_id": project_id,
+                "location": location,
+                "filter": jobs_filter,
+                "page_size": page_size,
+                "page_token": page_token,
+            }
+        )
+        page_result: ListJobsAsyncPager = await client.list_jobs(request=request)
+        return page_result

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

@@ -40,6 +40,7 @@ from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHoo
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
+    from google.api_core.retry_async import AsyncRetry
     from googleapiclient.discovery import Resource
 
 PATH_DATA_SCAN = "projects/{project_id}/locations/{region}/dataScans/{data_scan_id}"
@@ -896,7 +897,7 @@ class DataplexAsyncHook(GoogleBaseAsyncHook):
         region: str,
         data_scan_id: str | None = None,
         job_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Any:

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

@@ -51,6 +51,7 @@ if TYPE_CHECKING:
     from google.api_core.operation_async import AsyncOperation
     from google.api_core.operations_v1.operations_client import OperationsClient
     from google.api_core.retry import Retry
+    from google.api_core.retry_async import AsyncRetry
     from google.protobuf.duration_pb2 import Duration
     from google.protobuf.field_mask_pb2 import FieldMask
 
@@ -256,7 +257,7 @@ class DataprocHook(GoogleBaseHook):
         self,
         operation: Operation,
         timeout: float | None = None,
-        result_retry: Retry | _MethodDefault = DEFAULT,
+        result_retry: AsyncRetry | _MethodDefault = DEFAULT,
     ) -> Any:
         """Wait for a long-lasting operation to complete."""
         try:
@@ -997,7 +998,7 @@ class DataprocHook(GoogleBaseHook):
         region: str,
         project_id: str,
         wait_check_interval: int = 10,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Batch:
@@ -1132,7 +1133,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         virtual_cluster_config: dict | None = None,
         labels: dict[str, str] | None = None,
         request_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -1199,7 +1200,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         project_id: str,
         cluster_uuid: str | None = None,
         request_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -1242,7 +1243,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         region: str,
         cluster_name: str,
         project_id: str,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> str:
@@ -1277,7 +1278,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         region: str,
         cluster_name: str,
         project_id: str,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Cluster:
@@ -1309,7 +1310,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         filter_: str,
         project_id: str,
         page_size: int | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ):
@@ -1349,7 +1350,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         region: str,
         graceful_decommission_timeout: dict | Duration | None = None,
         request_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -1429,7 +1430,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         template: dict | WorkflowTemplate,
         project_id: str,
         region: str,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> WorkflowTemplate:
@@ -1465,7 +1466,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         version: int | None = None,
         request_id: str | None = None,
         parameters: dict[str, str] | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -1511,7 +1512,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         project_id: str,
         region: str,
         request_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -1554,7 +1555,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         job_id: str,
         project_id: str,
         region: str,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Job:
@@ -1588,7 +1589,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         project_id: str,
         region: str,
         request_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Job:
@@ -1624,7 +1625,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         job_id: str,
         project_id: str,
         region: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Job:
@@ -1658,7 +1659,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         batch: dict | Batch,
         batch_id: str | None = None,
         request_id: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AsyncOperation:
@@ -1703,7 +1704,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         batch_id: str,
         region: str,
         project_id: str,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> None:
@@ -1737,7 +1738,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         batch_id: str,
         region: str,
         project_id: str,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Batch:
@@ -1773,7 +1774,7 @@ class DataprocAsyncHook(GoogleBaseHook):
         project_id: str,
         page_size: int | None = None,
         page_token: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
         filter: str | None = None,

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

@@ -822,10 +822,12 @@ class GCSHook(GoogleBaseHook):
                     versions=versions,
                 )
 
+            blob_names = [blob.name for blob in blobs]
+
             if blobs.prefixes:
                 ids.extend(blobs.prefixes)
             else:
-                ids.extend(blob.name for blob in blobs)
+                ids.extend(blob_names)
 
             page_token = blobs.next_page_token
             if page_token is None:
@@ -933,14 +935,16 @@ class GCSHook(GoogleBaseHook):
                     versions=versions,
                 )
 
+            blob_names = [
+                blob.name
+                for blob in blobs
+                if timespan_start <= blob.updated.replace(tzinfo=timezone.utc) < timespan_end
+            ]
+
             if blobs.prefixes:
                 ids.extend(blobs.prefixes)
             else:
-                ids.extend(
-                    blob.name
-                    for blob in blobs
-                    if timespan_start <= blob.updated.replace(tzinfo=timezone.utc) < timespan_end
-                )
+                ids.extend(blob_names)
 
             page_token = blobs.next_page_token
             if page_token is None:

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

@@ -49,6 +49,7 @@ from airflow.version import version
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
+    from google.api_core.retry_async import AsyncRetry
     from google.cloud.pubsub_v1.types import (
         DeadLetterPolicy,
         Duration,
@@ -611,7 +612,7 @@ class PubSubAsyncHook(GoogleBaseAsyncHook):
         project_id: str,
         ack_ids: list[str] | None = None,
         messages: list[ReceivedMessage] | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> None:
@@ -665,7 +666,7 @@ class PubSubAsyncHook(GoogleBaseAsyncHook):
         max_messages: int,
         project_id: str = PROVIDE_PROJECT_ID,
         return_immediately: bool = False,
-        retry: Retry | _MethodDefault = DEFAULT,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> list[ReceivedMessage]:

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

@@ -26,7 +26,6 @@ from typing import TYPE_CHECKING, Collection
 
 # not sure why but mypy complains on missing `storage` but it is clearly there and is importable
 from google.cloud import storage  # type: ignore[attr-defined]
-from packaging.version import Version
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowNotFoundException
@@ -48,18 +47,6 @@ _DEFAULT_SCOPESS = frozenset(
 logger = logging.getLogger(__name__)
 
 
-def get_default_delete_local_copy():
-    """Load delete_local_logs conf if Airflow version > 2.6 and return False if not.
-
-    TODO: delete this function when min airflow version >= 2.6.
-    """
-    from airflow.version import version
-
-    if Version(version) < Version("2.6"):
-        return False
-    return conf.getboolean("logging", "delete_local_logs")
-
-
 class GCSTaskHandler(FileTaskHandler, LoggingMixin):
     """
     GCSTaskHandler is a python log handler that handles and reads task instance logs.
@@ -108,8 +95,8 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         self.gcp_keyfile_dict = gcp_keyfile_dict
         self.scopes = gcp_scopes
         self.project_id = project_id
-        self.delete_local_copy = (
-            kwargs["delete_local_copy"] if "delete_local_copy" in kwargs else get_default_delete_local_copy()
+        self.delete_local_copy = kwargs.get(
+            "delete_local_copy", conf.getboolean("logging", "delete_local_logs")
         )
 
     @cached_property
@@ -218,30 +205,6 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
             messages.append(f"Unable to read remote log {e}")
         return messages, logs
 
-    def _read(self, ti, try_number, metadata=None):
-        """
-        Read logs of given task instance and try_number from GCS.
-
-        If failed, read the log from task instance host machine.
-
-        todo: when min airflow version >= 2.6, remove this method
-
-        :param ti: task instance object
-        :param try_number: task instance try_number to read logs from
-        :param metadata: log metadata,
-                         can be used for steaming log reading and auto-tailing.
-        """
-        if hasattr(super(), "_read_remote_logs"):
-            # from Airflow 2.6, we don't implement the `_read` method.
-            # if parent has _read_remote_logs, we're >= 2.6
-            return super()._read(ti, try_number, metadata)
-
-        messages, logs = self._read_remote_logs(ti, try_number, metadata)
-        if not logs:
-            return super()._read(ti, try_number, metadata)
-
-        return "".join([f"*** {x}\n" for x in messages]) + "\n".join(logs), {"end_of_log": True}
-
     def gcs_write(self, log, remote_log_location) -> bool:
         """
         Write the log to the remote location and return `True`; fail silently and return `False` on error.

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

@@ -30,21 +30,13 @@ from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesR
 
 from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id
 from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.utils.log.trigger_handler import ctx_indiv_trigger
 
 if TYPE_CHECKING:
-    from contextvars import ContextVar
-
     from google.auth.credentials import Credentials
 
     from airflow.models import TaskInstance
 
-try:
-    # todo: remove this conditional import when min airflow version >= 2.6
-    ctx_indiv_trigger: ContextVar | None
-    from airflow.utils.log.trigger_handler import ctx_indiv_trigger
-except ImportError:
-    ctx_indiv_trigger = None
-
 DEFAULT_LOGGER_NAME = "airflow"
 _GLOBAL_RESOURCE = Resource(type="global", labels={})
 
@@ -174,8 +166,7 @@ class StackdriverTaskHandler(logging.Handler):
         """
         message = self.format(record)
         ti = None
-        # todo: remove ctx_indiv_trigger is not None check when min airflow version >= 2.6
-        if ctx_indiv_trigger is not None and getattr(record, ctx_indiv_trigger.name, None):
+        if getattr(record, ctx_indiv_trigger.name, None):
             ti = getattr(record, "task_instance", None)  # trigger context
         labels = self._get_labels(ti)
         self._transport.send(record, message, resource=self.resource, labels=labels)