apache-airflow-providers-google 10.17.0rc1__py3-none-any.whl → 10.18.0__py3-none-any.whl

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

@@ -59,6 +59,7 @@ from airflow.providers.google.cloud.links.data_loss_prevention import (
     CloudDLPPossibleInfoTypesListLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
@@ -108,7 +109,7 @@ class CloudDLPCancelDLPJobOperator(GoogleCloudBaseOperator):
         self,
         *,
         dlp_job_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -195,7 +196,7 @@ class CloudDLPCreateDeidentifyTemplateOperator(GoogleCloudBaseOperator):
         self,
         *,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         deidentify_template: dict | DeidentifyTemplate | None = None,
         template_id: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -304,7 +305,7 @@ class CloudDLPCreateDLPJobOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         inspect_job: dict | InspectJobConfig | None = None,
         risk_job: dict | RiskAnalysisJobConfig | None = None,
         job_id: str | None = None,
@@ -416,7 +417,7 @@ class CloudDLPCreateInspectTemplateOperator(GoogleCloudBaseOperator):
         self,
         *,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         inspect_template: InspectTemplate | None = None,
         template_id: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -522,7 +523,7 @@ class CloudDLPCreateJobTriggerOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         job_trigger: dict | JobTrigger | None = None,
         trigger_id: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -631,7 +632,7 @@ class CloudDLPCreateStoredInfoTypeOperator(GoogleCloudBaseOperator):
         self,
         *,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         config: StoredInfoTypeConfig | None = None,
         stored_info_type_id: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -752,7 +753,7 @@ class CloudDLPDeidentifyContentOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         deidentify_config: dict | DeidentifyConfig | None = None,
         inspect_config: dict | InspectConfig | None = None,
         item: dict | ContentItem | None = None,
@@ -842,7 +843,7 @@ class CloudDLPDeleteDeidentifyTemplateOperator(GoogleCloudBaseOperator):
         *,
         template_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -929,7 +930,7 @@ class CloudDLPDeleteDLPJobOperator(GoogleCloudBaseOperator):
         self,
         *,
         dlp_job_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1017,7 +1018,7 @@ class CloudDLPDeleteInspectTemplateOperator(GoogleCloudBaseOperator):
         *,
         template_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1103,7 +1104,7 @@ class CloudDLPDeleteJobTriggerOperator(GoogleCloudBaseOperator):
         self,
         *,
         job_trigger_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1191,7 +1192,7 @@ class CloudDLPDeleteStoredInfoTypeOperator(GoogleCloudBaseOperator):
         *,
         stored_info_type_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1281,7 +1282,7 @@ class CloudDLPGetDeidentifyTemplateOperator(GoogleCloudBaseOperator):
         *,
         template_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1364,7 +1365,7 @@ class CloudDLPGetDLPJobOperator(GoogleCloudBaseOperator):
         self,
         *,
         dlp_job_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1452,7 +1453,7 @@ class CloudDLPGetInspectTemplateOperator(GoogleCloudBaseOperator):
         *,
         template_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1538,7 +1539,7 @@ class CloudDLPGetDLPJobTriggerOperator(GoogleCloudBaseOperator):
         self,
         *,
         job_trigger_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1626,7 +1627,7 @@ class CloudDLPGetStoredInfoTypeOperator(GoogleCloudBaseOperator):
         *,
         stored_info_type_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1716,7 +1717,7 @@ class CloudDLPInspectContentOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         inspect_config: dict | InspectConfig | None = None,
         item: dict | ContentItem | None = None,
         inspect_template_name: str | None = None,
@@ -1802,7 +1803,7 @@ class CloudDLPListDeidentifyTemplatesOperator(GoogleCloudBaseOperator):
         self,
         *,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         page_size: int | None = None,
         order_by: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -1894,7 +1895,7 @@ class CloudDLPListDLPJobsOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         results_filter: str | None = None,
         page_size: int | None = None,
         job_type: str | None = None,
@@ -1986,7 +1987,7 @@ class CloudDLPListInfoTypesOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         language_code: str | None = None,
         results_filter: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -2077,7 +2078,7 @@ class CloudDLPListInspectTemplatesOperator(GoogleCloudBaseOperator):
         self,
         *,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         page_size: int | None = None,
         order_by: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -2168,7 +2169,7 @@ class CloudDLPListJobTriggersOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         page_size: int | None = None,
         order_by: str | None = None,
         results_filter: str | None = None,
@@ -2263,7 +2264,7 @@ class CloudDLPListStoredInfoTypesOperator(GoogleCloudBaseOperator):
         self,
         *,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         page_size: int | None = None,
         order_by: str | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -2359,7 +2360,7 @@ class CloudDLPRedactImageOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         inspect_config: dict | InspectConfig | None = None,
         image_redaction_configs: None | (list[dict] | list[RedactImageRequest.ImageRedactionConfig]) = None,
         include_findings: bool | None = None,
@@ -2453,7 +2454,7 @@ class CloudDLPReidentifyContentOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         reidentify_config: dict | DeidentifyConfig | None = None,
         inspect_config: dict | InspectConfig | None = None,
         item: dict | ContentItem | None = None,
@@ -2548,7 +2549,7 @@ class CloudDLPUpdateDeidentifyTemplateOperator(GoogleCloudBaseOperator):
         *,
         template_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         deidentify_template: dict | DeidentifyTemplate | None = None,
         update_mask: dict | FieldMask | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -2648,7 +2649,7 @@ class CloudDLPUpdateInspectTemplateOperator(GoogleCloudBaseOperator):
         *,
         template_id: str,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         inspect_template: dict | InspectTemplate | None = None,
         update_mask: dict | FieldMask | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -2744,7 +2745,7 @@ class CloudDLPUpdateJobTriggerOperator(GoogleCloudBaseOperator):
         self,
         *,
         job_trigger_id,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         job_trigger: dict | JobTrigger | None = None,
         update_mask: dict | FieldMask | None = None,
         retry: Retry | _MethodDefault = DEFAULT,
@@ -2843,7 +2844,7 @@ class CloudDLPUpdateStoredInfoTypeOperator(GoogleCloudBaseOperator):
         *,
         stored_info_type_id,
         organization_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         config: dict | StoredInfoTypeConfig | None = None,
         update_mask: dict | FieldMask | None = None,
         retry: Retry | _MethodDefault = DEFAULT,

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

@@ -35,6 +35,7 @@ from airflow.providers.google.cloud.utils.field_validator import (
     GcpBodyFieldValidator,
     GcpFieldValidationException,
 )
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.version import version
 
 if TYPE_CHECKING:
@@ -155,7 +156,7 @@ class CloudFunctionDeployFunctionOperator(GoogleCloudBaseOperator):
         *,
         location: str,
         body: dict,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
         zip_path: str | None = None,
@@ -363,7 +364,7 @@ class CloudFunctionDeleteFunctionOperator(GoogleCloudBaseOperator):
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
         impersonation_chain: str | Sequence[str] | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         **kwargs,
     ) -> None:
         self.name = name
@@ -447,7 +448,7 @@ class CloudFunctionInvokeFunctionOperator(GoogleCloudBaseOperator):
         function_id: str,
         input_data: dict,
         location: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
         impersonation_chain: str | Sequence[str] | None = None,

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

@@ -29,6 +29,8 @@ from typing import TYPE_CHECKING, Sequence
 
 import pendulum
 
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
@@ -119,7 +121,7 @@ class GCSCreateBucketOperator(GoogleCloudBaseOperator):
         resource: dict | None = None,
         storage_class: str = "MULTI_REGIONAL",
         location: str = "US",
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         labels: dict | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
@@ -297,7 +299,7 @@ class GCSDeleteObjectsOperator(GoogleCloudBaseOperator):
         *,
         bucket_name: str,
         objects: list[str] | None = None,
-        prefix: str | None = None,
+        prefix: str | list[str] | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -309,12 +311,14 @@ class GCSDeleteObjectsOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
 
         if objects is None and prefix is None:
-            err_message = "(Task {task_id}) Either object or prefix should be set. Both are None.".format(
+            err_message = "(Task {task_id}) Either objects or prefix should be set. Both are None.".format(
                 **kwargs
             )
             raise ValueError(err_message)
+        if objects is not None and prefix is not None:
+            err_message = "(Task {task_id}) Objects or prefix should be set. Both provided.".format(**kwargs)
+            raise ValueError(err_message)
 
-        self._objects: list[str] = []
         super().__init__(**kwargs)
 
     def execute(self, context: Context) -> None:
@@ -324,15 +328,14 @@ class GCSDeleteObjectsOperator(GoogleCloudBaseOperator):
         )
 
         if self.objects is not None:
-            self._objects = self.objects
+            objects = self.objects
         else:
-            self._objects = hook.list(bucket_name=self.bucket_name, prefix=self.prefix)
-        self.log.info("Deleting %s objects from %s", len(self._objects), self.bucket_name)
-        for object_name in self._objects:
+            objects = hook.list(bucket_name=self.bucket_name, prefix=self.prefix)
+        self.log.info("Deleting %s objects from %s", len(objects), self.bucket_name)
+        for object_name in objects:
             hook.delete(bucket_name=self.bucket_name, object_name=object_name)
 
-    def get_openlineage_facets_on_complete(self, task_instance):
-        """Implement on_complete as execute() resolves object names."""
+    def get_openlineage_facets_on_start(self):
         from openlineage.client.facet import (
             LifecycleStateChange,
             LifecycleStateChangeDatasetFacet,
@@ -342,8 +345,17 @@ class GCSDeleteObjectsOperator(GoogleCloudBaseOperator):
 
         from airflow.providers.openlineage.extractors import OperatorLineage
 
-        if not self._objects:
-            return OperatorLineage()
+        objects = []
+        if self.objects is not None:
+            objects = self.objects
+        elif self.prefix is not None:
+            prefixes = [self.prefix] if isinstance(self.prefix, str) else self.prefix
+            for pref in prefixes:
+                # Use parent if not a file (dot not in name) and not a dir (ends with slash)
+                if "." not in pref.split("/")[-1] and not pref.endswith("/"):
+                    pref = Path(pref).parent.as_posix()
+                pref = "/" if pref in (".", "", "/") else pref.rstrip("/")
+                objects.append(pref)
 
         bucket_url = f"gs://{self.bucket_name}"
         input_datasets = [
@@ -360,7 +372,7 @@ class GCSDeleteObjectsOperator(GoogleCloudBaseOperator):
                     )
                 },
             )
-            for object_name in self._objects
+            for object_name in objects
         ]
 
         return OperatorLineage(inputs=input_datasets)
@@ -774,8 +786,8 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
         self.upload_continue_on_fail = upload_continue_on_fail
         self.upload_num_attempts = upload_num_attempts
 
-        self._source_object_names: list[str] = []
-        self._destination_object_names: list[str] = []
+        self._source_prefix_interp: str | None = None
+        self._destination_prefix_interp: str | None = None
 
     def execute(self, context: Context) -> list[str]:
         # Define intervals and prefixes.
@@ -803,11 +815,11 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
         timespan_start = timespan_start.in_timezone(timezone.utc)
         timespan_end = timespan_end.in_timezone(timezone.utc)
 
-        source_prefix_interp = GCSTimeSpanFileTransformOperator.interpolate_prefix(
+        self._source_prefix_interp = GCSTimeSpanFileTransformOperator.interpolate_prefix(
             self.source_prefix,
             timespan_start,
         )
-        destination_prefix_interp = GCSTimeSpanFileTransformOperator.interpolate_prefix(
+        self._destination_prefix_interp = GCSTimeSpanFileTransformOperator.interpolate_prefix(
             self.destination_prefix,
             timespan_start,
         )
@@ -828,9 +840,9 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
         )
 
         # Fetch list of files.
-        self._source_object_names = source_hook.list_by_timespan(
+        blobs_to_transform = source_hook.list_by_timespan(
             bucket_name=self.source_bucket,
-            prefix=source_prefix_interp,
+            prefix=self._source_prefix_interp,
             timespan_start=timespan_start,
             timespan_end=timespan_end,
         )
@@ -840,7 +852,7 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
             temp_output_dir_path = Path(temp_output_dir)
 
             # TODO: download in parallel.
-            for blob_to_transform in self._source_object_names:
+            for blob_to_transform in blobs_to_transform:
                 destination_file = temp_input_dir_path / blob_to_transform
                 destination_file.parent.mkdir(parents=True, exist_ok=True)
                 try:
@@ -877,6 +889,8 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
 
             self.log.info("Transformation succeeded. Output temporarily located at %s", temp_output_dir_path)
 
+            files_uploaded = []
+
             # TODO: upload in parallel.
             for upload_file in temp_output_dir_path.glob("**/*"):
                 if upload_file.is_dir():
@@ -884,8 +898,8 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
 
                 upload_file_name = str(upload_file.relative_to(temp_output_dir_path))
 
-                if self.destination_prefix is not None:
-                    upload_file_name = f"{destination_prefix_interp}/{upload_file_name}"
+                if self._destination_prefix_interp is not None:
+                    upload_file_name = f"{self._destination_prefix_interp.rstrip('/')}/{upload_file_name}"
 
                 self.log.info("Uploading file %s to %s", upload_file, upload_file_name)
 
@@ -897,35 +911,46 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
                         chunk_size=self.chunk_size,
                         num_max_attempts=self.upload_num_attempts,
                     )
-                    self._destination_object_names.append(str(upload_file_name))
+                    files_uploaded.append(str(upload_file_name))
                 except GoogleCloudError:
                     if not self.upload_continue_on_fail:
                         raise
 
-            return self._destination_object_names
+            return files_uploaded
 
     def get_openlineage_facets_on_complete(self, task_instance):
-        """Implement on_complete as execute() resolves object names."""
+        """Implement on_complete as execute() resolves object prefixes."""
         from openlineage.client.run import Dataset
 
         from airflow.providers.openlineage.extractors import OperatorLineage
 
-        input_datasets = [
-            Dataset(
-                namespace=f"gs://{self.source_bucket}",
-                name=object_name,
-            )
-            for object_name in self._source_object_names
-        ]
-        output_datasets = [
-            Dataset(
-                namespace=f"gs://{self.destination_bucket}",
-                name=object_name,
-            )
-            for object_name in self._destination_object_names
-        ]
-
-        return OperatorLineage(inputs=input_datasets, outputs=output_datasets)
+        def _parse_prefix(pref):
+            # Use parent if not a file (dot not in name) and not a dir (ends with slash)
+            if "." not in pref.split("/")[-1] and not pref.endswith("/"):
+                pref = Path(pref).parent.as_posix()
+            return "/" if pref in (".", "/", "") else pref.rstrip("/")
+
+        input_prefix, output_prefix = "/", "/"
+        if self._source_prefix_interp is not None:
+            input_prefix = _parse_prefix(self._source_prefix_interp)
+
+        if self._destination_prefix_interp is not None:
+            output_prefix = _parse_prefix(self._destination_prefix_interp)
+
+        return OperatorLineage(
+            inputs=[
+                Dataset(
+                    namespace=f"gs://{self.source_bucket}",
+                    name=input_prefix,
+                )
+            ],
+            outputs=[
+                Dataset(
+                    namespace=f"gs://{self.destination_bucket}",
+                    name=output_prefix,
+                )
+            ],
+        )
 
 
 class GCSDeleteBucketOperator(GoogleCloudBaseOperator):