apache-airflow-providers-google 10.12.0__tar.gz → 10.13.0__tar.gz

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-google
-Version: 10.12.0
+Version: 10.13.0
 Summary: Provider package apache-airflow-providers-google for Apache Airflow
 Keywords: airflow-provider,google,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -22,40 +22,40 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: PyOpenSSL
 Requires-Dist: apache-airflow-providers-common-sql>=1.7.2
-Requires-Dist: apache-airflow>=2.5.0
+Requires-Dist: apache-airflow>=2.6.0
 Requires-Dist: asgiref>=3.5.2
 Requires-Dist: gcloud-aio-auth>=4.0.0,<5.0.0
 Requires-Dist: gcloud-aio-bigquery>=6.1.2
-Requires-Dist: gcloud-aio-storage
-Requires-Dist: gcsfs>=2023.9.2
+Requires-Dist: gcloud-aio-storage>=9.0.0
+Requires-Dist: gcsfs>=2023.10.0
 Requires-Dist: google-ads>=22.1.0
 Requires-Dist: google-api-core>=2.11.0
 Requires-Dist: google-api-python-client>=1.6.0
 Requires-Dist: google-auth-httplib2>=0.0.1
 Requires-Dist: google-auth>=1.0.0
 Requires-Dist: google-cloud-aiplatform>=1.22.1
-Requires-Dist: google-cloud-automl>=2.11.0
+Requires-Dist: google-cloud-automl>=2.12.0
 Requires-Dist: google-cloud-batch>=0.13.0
-Requires-Dist: google-cloud-bigquery-datatransfer>=3.11.0
+Requires-Dist: google-cloud-bigquery-datatransfer>=3.13.0
 Requires-Dist: google-cloud-bigtable>=2.17.0
-Requires-Dist: google-cloud-build>=3.13.0
+Requires-Dist: google-cloud-build>=3.22.0
 Requires-Dist: google-cloud-compute>=1.10.0
 Requires-Dist: google-cloud-container>=2.17.4
 Requires-Dist: google-cloud-datacatalog>=3.11.1
-Requires-Dist: google-cloud-dataflow-client>=0.8.2
+Requires-Dist: google-cloud-dataflow-client>=0.8.6
 Requires-Dist: google-cloud-dataform>=0.5.0
-Requires-Dist: google-cloud-dataplex>=1.4.2
+Requires-Dist: google-cloud-dataplex>=1.10.0
 Requires-Dist: google-cloud-dataproc-metastore>=1.12.0
-Requires-Dist: google-cloud-dataproc>=5.5.0
+Requires-Dist: google-cloud-dataproc>=5.8.0
 Requires-Dist: google-cloud-dlp>=3.12.0
 Requires-Dist: google-cloud-kms>=2.15.0
 Requires-Dist: google-cloud-language>=2.9.0
 Requires-Dist: google-cloud-logging>=3.5.0
 Requires-Dist: google-cloud-memcache>=1.7.0
-Requires-Dist: google-cloud-monitoring>=2.14.1
-Requires-Dist: google-cloud-orchestration-airflow>=1.7.0
+Requires-Dist: google-cloud-monitoring>=2.18.0
+Requires-Dist: google-cloud-orchestration-airflow>=1.10.0
 Requires-Dist: google-cloud-os-login>=2.9.1
-Requires-Dist: google-cloud-pubsub>=2.15.0
+Requires-Dist: google-cloud-pubsub>=2.19.0
 Requires-Dist: google-cloud-redis>=2.12.0
 Requires-Dist: google-cloud-run>=0.9.0
 Requires-Dist: google-cloud-secret-manager>=2.16.0
@@ -98,8 +98,8 @@ Requires-Dist: apache-airflow-providers-sftp ; extra == "sftp"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Requires-Dist: apache-airflow-providers-trino ; extra == "trino"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-google/10.12.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/10.12.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-google/10.13.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/10.13.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -167,7 +167,7 @@ Provides-Extra: trino
 
 Package ``apache-airflow-providers-google``
 
-Release: ``10.12.0``
+Release: ``10.13.0``
 
 
 Google services including:
@@ -187,7 +187,7 @@ This is a provider package for ``google`` provider. All classes for this provide
 are in ``airflow.providers.google`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.12.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.13.0/>`_.
 
 Installation
 ------------
@@ -204,40 +204,40 @@ Requirements
 =======================================  ==================
 PIP package                              Version required
 =======================================  ==================
-``apache-airflow``                       ``>=2.5.0``
+``apache-airflow``                       ``>=2.6.0``
 ``apache-airflow-providers-common-sql``  ``>=1.7.2``
 ``asgiref``                              ``>=3.5.2``
 ``gcloud-aio-auth``                      ``>=4.0.0,<5.0.0``
 ``gcloud-aio-bigquery``                  ``>=6.1.2``
-``gcloud-aio-storage``
-``gcsfs``                                ``>=2023.9.2``
+``gcloud-aio-storage``                   ``>=9.0.0``
+``gcsfs``                                ``>=2023.10.0``
 ``google-ads``                           ``>=22.1.0``
 ``google-api-core``                      ``>=2.11.0``
 ``google-api-python-client``             ``>=1.6.0``
 ``google-auth``                          ``>=1.0.0``
 ``google-auth-httplib2``                 ``>=0.0.1``
 ``google-cloud-aiplatform``              ``>=1.22.1``
-``google-cloud-automl``                  ``>=2.11.0``
-``google-cloud-bigquery-datatransfer``   ``>=3.11.0``
+``google-cloud-automl``                  ``>=2.12.0``
+``google-cloud-bigquery-datatransfer``   ``>=3.13.0``
 ``google-cloud-bigtable``                ``>=2.17.0``
-``google-cloud-build``                   ``>=3.13.0``
+``google-cloud-build``                   ``>=3.22.0``
 ``google-cloud-compute``                 ``>=1.10.0``
 ``google-cloud-container``               ``>=2.17.4``
 ``google-cloud-datacatalog``             ``>=3.11.1``
-``google-cloud-dataflow-client``         ``>=0.8.2``
+``google-cloud-dataflow-client``         ``>=0.8.6``
 ``google-cloud-dataform``                ``>=0.5.0``
-``google-cloud-dataplex``                ``>=1.4.2``
-``google-cloud-dataproc``                ``>=5.5.0``
+``google-cloud-dataplex``                ``>=1.10.0``
+``google-cloud-dataproc``                ``>=5.8.0``
 ``google-cloud-dataproc-metastore``      ``>=1.12.0``
 ``google-cloud-dlp``                     ``>=3.12.0``
 ``google-cloud-kms``                     ``>=2.15.0``
 ``google-cloud-language``                ``>=2.9.0``
 ``google-cloud-logging``                 ``>=3.5.0``
 ``google-cloud-memcache``                ``>=1.7.0``
-``google-cloud-monitoring``              ``>=2.14.1``
-``google-cloud-orchestration-airflow``   ``>=1.7.0``
+``google-cloud-monitoring``              ``>=2.18.0``
+``google-cloud-orchestration-airflow``   ``>=1.10.0``
 ``google-cloud-os-login``                ``>=2.9.1``
-``google-cloud-pubsub``                  ``>=2.15.0``
+``google-cloud-pubsub``                  ``>=2.19.0``
 ``google-cloud-redis``                   ``>=2.12.0``
 ``google-cloud-secret-manager``          ``>=2.16.0``
 ``google-cloud-spanner``                 ``>=3.11.1``
@@ -300,4 +300,4 @@ Dependent package
 ========================================================================================================================  ====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.12.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.13.0/changelog.html>`_.

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-google``
 
-Release: ``10.12.0``
+Release: ``10.13.0``
 
 
 Google services including:
@@ -62,7 +62,7 @@ This is a provider package for ``google`` provider. All classes for this provide
 are in ``airflow.providers.google`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.12.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.13.0/>`_.
 
 Installation
 ------------
@@ -79,40 +79,40 @@ Requirements
 =======================================  ==================
 PIP package                              Version required
 =======================================  ==================
-``apache-airflow``                       ``>=2.5.0``
+``apache-airflow``                       ``>=2.6.0``
 ``apache-airflow-providers-common-sql``  ``>=1.7.2``
 ``asgiref``                              ``>=3.5.2``
 ``gcloud-aio-auth``                      ``>=4.0.0,<5.0.0``
 ``gcloud-aio-bigquery``                  ``>=6.1.2``
-``gcloud-aio-storage``
-``gcsfs``                                ``>=2023.9.2``
+``gcloud-aio-storage``                   ``>=9.0.0``
+``gcsfs``                                ``>=2023.10.0``
 ``google-ads``                           ``>=22.1.0``
 ``google-api-core``                      ``>=2.11.0``
 ``google-api-python-client``             ``>=1.6.0``
 ``google-auth``                          ``>=1.0.0``
 ``google-auth-httplib2``                 ``>=0.0.1``
 ``google-cloud-aiplatform``              ``>=1.22.1``
-``google-cloud-automl``                  ``>=2.11.0``
-``google-cloud-bigquery-datatransfer``   ``>=3.11.0``
+``google-cloud-automl``                  ``>=2.12.0``
+``google-cloud-bigquery-datatransfer``   ``>=3.13.0``
 ``google-cloud-bigtable``                ``>=2.17.0``
-``google-cloud-build``                   ``>=3.13.0``
+``google-cloud-build``                   ``>=3.22.0``
 ``google-cloud-compute``                 ``>=1.10.0``
 ``google-cloud-container``               ``>=2.17.4``
 ``google-cloud-datacatalog``             ``>=3.11.1``
-``google-cloud-dataflow-client``         ``>=0.8.2``
+``google-cloud-dataflow-client``         ``>=0.8.6``
 ``google-cloud-dataform``                ``>=0.5.0``
-``google-cloud-dataplex``                ``>=1.4.2``
-``google-cloud-dataproc``                ``>=5.5.0``
+``google-cloud-dataplex``                ``>=1.10.0``
+``google-cloud-dataproc``                ``>=5.8.0``
 ``google-cloud-dataproc-metastore``      ``>=1.12.0``
 ``google-cloud-dlp``                     ``>=3.12.0``
 ``google-cloud-kms``                     ``>=2.15.0``
 ``google-cloud-language``                ``>=2.9.0``
 ``google-cloud-logging``                 ``>=3.5.0``
 ``google-cloud-memcache``                ``>=1.7.0``
-``google-cloud-monitoring``              ``>=2.14.1``
-``google-cloud-orchestration-airflow``   ``>=1.7.0``
+``google-cloud-monitoring``              ``>=2.18.0``
+``google-cloud-orchestration-airflow``   ``>=1.10.0``
 ``google-cloud-os-login``                ``>=2.9.1``
-``google-cloud-pubsub``                  ``>=2.15.0``
+``google-cloud-pubsub``                  ``>=2.19.0``
 ``google-cloud-redis``                   ``>=2.12.0``
 ``google-cloud-secret-manager``          ``>=2.16.0``
 ``google-cloud-spanner``                 ``>=3.11.1``
@@ -175,4 +175,4 @@ Dependent package
 ========================================================================================================================  ====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.12.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.13.0/changelog.html>`_.

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/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+"
     )

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.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)

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/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]] = (),
     ):

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/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",

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/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:

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/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,

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/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

{apache_airflow_providers_google-10.12.0 → apache_airflow_providers_google-10.13.0}/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: