apache-airflow-providers-google 10.12.0rc1__py3-none-any.whl → 10.13.0rc2__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

@@ -55,7 +55,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.
 

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

@@ -821,6 +821,7 @@ class GCSHook(GoogleBaseHook):
                     delimiter=delimiter,
                     versions=versions,
                 )
+                list(blobs)
 
             if blobs.prefixes:
                 ids.extend(blobs.prefixes)
@@ -932,6 +933,7 @@ class GCSHook(GoogleBaseHook):
                     delimiter=delimiter,
                     versions=versions,
                 )
+                list(blobs)
 
             if blobs.prefixes:
                 ids.extend(blobs.prefixes)

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

@@ -112,7 +112,7 @@ class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
     :param metadata: Strings which should be sent along with the request as metadata.
     :param deferrable: Run operator in the deferrable mode
     :param pooling_period_seconds: Optional: Control the rate of the poll for the result of deferrable run.
-        By default the trigger will poll every 30 seconds.
+        By default, the trigger will poll every 30 seconds.
     """
 
     template_fields = (

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

@@ -264,7 +264,7 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
     :param deferrable: Run operator in the deferrable mode
     """
 
-    template_fields = ("project_id", "region", "gcp_conn_id", "impersonation_chain", "job_name")
+    template_fields = ("project_id", "region", "gcp_conn_id", "impersonation_chain", "job_name", "overrides")
 
     def __init__(
         self,
@@ -321,10 +321,10 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
     def execute_complete(self, context: Context, event: dict):
         status = event["status"]
 
-        if status == RunJobStatus.TIMEOUT:
+        if status == RunJobStatus.TIMEOUT.value:
             raise AirflowException("Operation timed out")
 
-        if status == RunJobStatus.FAIL:
+        if status == RunJobStatus.FAIL.value:
             error_code = event["operation_error_code"]
             error_message = event["operation_error_message"]
             raise AirflowException(