apache-airflow-providers-google 10.7.0__py3-none-any.whl → 10.8.0__py3-none-any.whl

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

@@ -27,10 +27,9 @@ from __future__ import annotations
 
 import re
 import time
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.dlp import DlpServiceClient
 from google.cloud.dlp_v2.types import (
     ByteContentItem,
@@ -58,6 +57,9 @@ from airflow.exceptions import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
+if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
 DLP_JOB_PATH_PATTERN = "^projects/[^/]+/dlpJobs/(?P<job>.*?)$"
 
 
@@ -148,9 +150,9 @@ class CloudDLPHook(GoogleBaseHook):
 
         name = DlpServiceClient.dlp_job_path(project_id, dlp_job_id)
         client.cancel_dlp_job(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -195,11 +197,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.create_deidentify_template(
-            request=dict(
-                parent=parent,
-                deidentify_template=deidentify_template,
-                template_id=template_id,
-            ),
+            request={
+                "parent": parent,
+                "deidentify_template": deidentify_template,
+                "template_id": template_id,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -242,12 +244,12 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         job = client.create_dlp_job(
-            request=dict(
-                parent=parent,
-                inspect_job=inspect_job,
-                risk_job=risk_job,
-                job_id=job_id,
-            ),
+            request={
+                "parent": parent,
+                "inspect_job": inspect_job,
+                "risk_job": risk_job,
+                "job_id": job_id,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -321,11 +323,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.create_inspect_template(
-            request=dict(
-                parent=parent,
-                inspect_template=inspect_template,
-                template_id=template_id,
-            ),
+            request={
+                "parent": parent,
+                "inspect_template": inspect_template,
+                "template_id": template_id,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -360,11 +362,11 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         return client.create_job_trigger(
-            request=dict(
-                parent=parent,
-                job_trigger=job_trigger,
-                trigger_id=trigger_id,
-            ),
+            request={
+                "parent": parent,
+                "job_trigger": job_trigger,
+                "trigger_id": trigger_id,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -410,11 +412,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.create_stored_info_type(
-            request=dict(
-                parent=parent,
-                config=config,
-                stored_info_type_id=stored_info_type_id,
-            ),
+            request={
+                "parent": parent,
+                "config": config,
+                "stored_info_type_id": stored_info_type_id,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -461,14 +463,14 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         return client.deidentify_content(
-            request=dict(
-                parent=parent,
-                deidentify_config=deidentify_config,
-                inspect_config=inspect_config,
-                item=item,
-                inspect_template_name=inspect_template_name,
-                deidentify_template_name=deidentify_template_name,
-            ),
+            request={
+                "parent": parent,
+                "deidentify_config": deidentify_config,
+                "inspect_config": inspect_config,
+                "item": item,
+                "inspect_template_name": inspect_template_name,
+                "deidentify_template_name": deidentify_template_name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -509,9 +511,9 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         client.delete_deidentify_template(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -550,9 +552,9 @@ class CloudDLPHook(GoogleBaseHook):
 
         name = DlpServiceClient.dlp_job_path(project_id, dlp_job_id)
         client.delete_dlp_job(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -599,9 +601,9 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         client.delete_inspect_template(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -637,9 +639,9 @@ class CloudDLPHook(GoogleBaseHook):
 
         name = DlpServiceClient.job_trigger_path(project_id, job_trigger_id)
         client.delete_job_trigger(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -686,9 +688,9 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         client.delete_stored_info_type(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -735,9 +737,9 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.get_deidentify_template(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -773,9 +775,9 @@ class CloudDLPHook(GoogleBaseHook):
 
         name = DlpServiceClient.dlp_job_path(project_id, dlp_job_id)
         return client.get_dlp_job(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -822,9 +824,9 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.get_inspect_template(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -860,9 +862,9 @@ class CloudDLPHook(GoogleBaseHook):
 
         name = DlpServiceClient.job_trigger_path(project_id, job_trigger_id)
         return client.get_job_trigger(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -909,9 +911,9 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.get_stored_info_type(
-            request=dict(
-                name=name,
-            ),
+            request={
+                "name": name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -950,12 +952,12 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         return client.inspect_content(
-            request=dict(
-                parent=parent,
-                inspect_config=inspect_config,
-                item=item,
-                inspect_template_name=inspect_template_name,
-            ),
+            request={
+                "parent": parent,
+                "inspect_config": inspect_config,
+                "item": item,
+                "inspect_template_name": inspect_template_name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1003,11 +1005,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         results = client.list_deidentify_templates(
-            request=dict(
-                parent=parent,
-                page_size=page_size,
-                order_by=order_by,
-            ),
+            request={
+                "parent": parent,
+                "page_size": page_size,
+                "order_by": order_by,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1050,13 +1052,13 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         results = client.list_dlp_jobs(
-            request=dict(
-                parent=parent,
-                filter=results_filter,
-                page_size=page_size,
-                type_=job_type,
-                order_by=order_by,
-            ),
+            request={
+                "parent": parent,
+                "filter": results_filter,
+                "page_size": page_size,
+                "type_": job_type,
+                "order_by": order_by,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1088,10 +1090,10 @@ class CloudDLPHook(GoogleBaseHook):
         client = self.get_conn()
 
         return client.list_info_types(
-            request=dict(
-                language_code=language_code,
-                filter=results_filter,
-            ),
+            request={
+                "language_code": language_code,
+                "filter": results_filter,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1139,11 +1141,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         results = client.list_inspect_templates(
-            request=dict(
-                parent=parent,
-                page_size=page_size,
-                order_by=order_by,
-            ),
+            request={
+                "parent": parent,
+                "page_size": page_size,
+                "order_by": order_by,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1183,12 +1185,12 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         results = client.list_job_triggers(
-            request=dict(
-                parent=parent,
-                page_size=page_size,
-                order_by=order_by,
-                filter=results_filter,
-            ),
+            request={
+                "parent": parent,
+                "page_size": page_size,
+                "order_by": order_by,
+                "filter": results_filter,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1237,11 +1239,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         results = client.list_stored_info_types(
-            request=dict(
-                parent=parent,
-                page_size=page_size,
-                order_by=order_by,
-            ),
+            request={
+                "parent": parent,
+                "page_size": page_size,
+                "order_by": order_by,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1285,13 +1287,13 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         return client.redact_image(
-            request=dict(
-                parent=parent,
-                inspect_config=inspect_config,
-                image_redaction_configs=image_redaction_configs,
-                include_findings=include_findings,
-                byte_item=byte_item,
-            ),
+            request={
+                "parent": parent,
+                "inspect_config": inspect_config,
+                "image_redaction_configs": image_redaction_configs,
+                "include_findings": include_findings,
+                "byte_item": byte_item,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1336,14 +1338,14 @@ class CloudDLPHook(GoogleBaseHook):
 
         parent = DlpServiceClient.common_project_path(project_id)
         return client.reidentify_content(
-            request=dict(
-                parent=parent,
-                reidentify_config=reidentify_config,
-                inspect_config=inspect_config,
-                item=item,
-                inspect_template_name=inspect_template_name,
-                reidentify_template_name=reidentify_template_name,
-            ),
+            request={
+                "parent": parent,
+                "reidentify_config": reidentify_config,
+                "inspect_config": inspect_config,
+                "item": item,
+                "inspect_template_name": inspect_template_name,
+                "reidentify_template_name": reidentify_template_name,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1394,11 +1396,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.update_deidentify_template(
-            request=dict(
-                name=name,
-                deidentify_template=deidentify_template,
-                update_mask=update_mask,
-            ),
+            request={
+                "name": name,
+                "deidentify_template": deidentify_template,
+                "update_mask": update_mask,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1448,11 +1450,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.update_inspect_template(
-            request=dict(
-                name=name,
-                inspect_template=inspect_template,
-                update_mask=update_mask,
-            ),
+            request={
+                "name": name,
+                "inspect_template": inspect_template,
+                "update_mask": update_mask,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,
@@ -1552,11 +1554,11 @@ class CloudDLPHook(GoogleBaseHook):
             raise AirflowException("Please provide either organization_id or project_id.")
 
         return client.update_stored_info_type(
-            request=dict(
-                name=name,
-                config=config,
-                update_mask=update_mask,
-            ),
+            request={
+                "name": name,
+                "config": config,
+                "update_mask": update_mask,
+            },
             retry=retry,
             timeout=timeout,
             metadata=metadata,

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

@@ -26,18 +26,14 @@ import shutil
 import time
 import warnings
 from contextlib import contextmanager
-from datetime import datetime
 from functools import partial
 from io import BytesIO
-from os import path
 from tempfile import NamedTemporaryFile
-from typing import IO, Any, Callable, Generator, Sequence, TypeVar, cast, overload
+from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Sequence, TypeVar, cast, overload
 from urllib.parse import urlsplit
 
-from aiohttp import ClientSession
 from gcloud.aio.storage import Storage
 from google.api_core.exceptions import GoogleAPICallError, NotFound
-from google.api_core.retry import Retry
 
 # 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]
@@ -52,6 +48,12 @@ from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHoo
 from airflow.utils import timezone
 from airflow.version import version
 
+if TYPE_CHECKING:
+    from datetime import datetime
+
+    from aiohttp import ClientSession
+    from google.api_core.retry import Retry
+
 try:
     # Airflow 2.3 doesn't have this yet
     from airflow.typing_compat import ParamSpec
@@ -358,7 +360,6 @@ class GCSHook(GoogleBaseHook):
                 # Wait with exponential backoff scheme before retrying.
                 timeout_seconds = 2 ** (num_file_attempts - 1)
                 time.sleep(timeout_seconds)
-                continue
 
     def download_as_byte_array(
         self,
@@ -505,28 +506,23 @@ class GCSHook(GoogleBaseHook):
 
             :param f: Callable that should be retried.
             """
-            num_file_attempts = 0
-
-            while num_file_attempts < num_max_attempts:
+            for attempt in range(1, 1 + num_max_attempts):
                 try:
-                    num_file_attempts += 1
                     f()
-
                 except GoogleCloudError as e:
-                    if num_file_attempts == num_max_attempts:
+                    if attempt == num_max_attempts:
                         self.log.error(
                             "Upload attempt of object: %s from %s has failed. Attempt: %s, max %s.",
                             object_name,
                             object_name,
-                            num_file_attempts,
+                            attempt,
                             num_max_attempts,
                         )
                         raise e
 
                     # Wait with exponential backoff scheme before retrying.
-                    timeout_seconds = 2 ** (num_file_attempts - 1)
+                    timeout_seconds = 2 ** (attempt - 1)
                     time.sleep(timeout_seconds)
-                    continue
 
         client = self.get_conn()
         bucket = client.bucket(bucket_name, user_project=user_project)
@@ -550,10 +546,9 @@ class GCSHook(GoogleBaseHook):
             if gzip:
                 filename_gz = filename + ".gz"
 
-                with open(filename, "rb") as f_in:
-                    with gz.open(filename_gz, "wb") as f_out:
-                        shutil.copyfileobj(f_in, f_out)
-                        filename = filename_gz
+                with open(filename, "rb") as f_in, gz.open(filename_gz, "wb") as f_out:
+                    shutil.copyfileobj(f_in, f_out)
+                    filename = filename_gz
 
             _call_with_retry(
                 partial(blob.upload_from_filename, filename=filename, content_type=mime_type, timeout=timeout)
@@ -1294,7 +1289,7 @@ class GCSHook(GoogleBaseHook):
         self, blob: storage.Blob, destination_object: str | None, source_object_prefix_len: int
     ) -> str:
         return (
-            path.join(destination_object, blob.name[source_object_prefix_len:])
+            os.path.join(destination_object, blob.name[source_object_prefix_len:])
             if destination_object
             else blob.name[source_object_prefix_len:]
         )

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

@@ -19,15 +19,17 @@
 from __future__ import annotations
 
 import base64
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.kms_v1 import KeyManagementServiceClient
 
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
+if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
 
 def _b64encode(s: bytes) -> str:
     """Base 64 encodes a bytes object to a string."""