apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

airflow/providers/google/cloud/operators/vertex_ai/dataset.py

@@ -20,20 +20,23 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-from airflow.providers.google.cloud.hooks.vertex_ai.dataset import DatasetHook
-from airflow.providers.google.cloud.links.vertex_ai import VertexAIDatasetLink, VertexAIDatasetListLink
-from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1.types import Dataset, ExportDataConfig, ImportDataConfig
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.vertex_ai.dataset import DatasetHook
+from airflow.providers.google.cloud.links.vertex_ai import VertexAIDatasetLink, VertexAIDatasetListLink
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class CreateDatasetOperator(GoogleCloudBaseOperator):
     """
@@ -83,6 +86,13 @@ class CreateDatasetOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = DatasetHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -104,8 +114,8 @@ class CreateDatasetOperator(GoogleCloudBaseOperator):
         dataset_id = hook.extract_dataset_id(dataset)
         self.log.info("Dataset was created. Dataset id: %s", dataset_id)
 
-        self.xcom_push(context, key="dataset_id", value=dataset_id)
-        VertexAIDatasetLink.persist(context=context, task_instance=self, dataset_id=dataset_id)
+        context["ti"].xcom_push(key="dataset_id", value=dataset_id)
+        VertexAIDatasetLink.persist(context=context, dataset_id=dataset_id)
         return dataset
 
 
@@ -158,6 +168,13 @@ class GetDatasetOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = DatasetHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -175,7 +192,7 @@ class GetDatasetOperator(GoogleCloudBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            VertexAIDatasetLink.persist(context=context, task_instance=self, dataset_id=self.dataset_id)
+            VertexAIDatasetLink.persist(context=context, dataset_id=self.dataset_id)
             self.log.info("Dataset was gotten.")
             return Dataset.to_dict(dataset_obj)
         except NotFound:
@@ -319,7 +336,21 @@ class ExportDataOperator(GoogleCloudBaseOperator):
         self.log.info("Export was done successfully")
 
 
-class ImportDataOperator(GoogleCloudBaseOperator):
+class DatasetImportDataResultsCheckHelper:
+    """Helper utils to verify import dataset data results."""
+
+    @staticmethod
+    def _get_number_of_ds_items(dataset, total_key_name):
+        number_of_items = type(dataset).to_dict(dataset).get(total_key_name, 0)
+        return number_of_items
+
+    @staticmethod
+    def _raise_for_empty_import_result(dataset_id, initial_size, size_after_import):
+        if int(size_after_import) - int(initial_size) <= 0:
+            raise AirflowException(f"Empty results of data import for the dataset_id {dataset_id}.")
+
+
+class ImportDataOperator(GoogleCloudBaseOperator, DatasetImportDataResultsCheckHelper):
     """
     Imports data into a Dataset.
 
@@ -340,6 +371,7 @@ class ImportDataOperator(GoogleCloudBaseOperator):
         If set as a sequence, the identities from the list must grant
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
+    :param raise_for_empty_result: Raise an error if no additional data has been populated after the import.
     """
 
     template_fields = ("region", "dataset_id", "project_id", "impersonation_chain")
@@ -356,6 +388,7 @@ class ImportDataOperator(GoogleCloudBaseOperator):
         metadata: Sequence[tuple[str, str]] = (),
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
+        raise_for_empty_result: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -368,13 +401,24 @@ class ImportDataOperator(GoogleCloudBaseOperator):
         self.metadata = metadata
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
+        self.raise_for_empty_result = raise_for_empty_result
 
     def execute(self, context: Context):
         hook = DatasetHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-
+        initial_dataset_size = self._get_number_of_ds_items(
+            dataset=hook.get_dataset(
+                dataset=self.dataset_id,
+                project_id=self.project_id,
+                region=self.region,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            ),
+            total_key_name="data_item_count",
+        )
         self.log.info("Importing data: %s", self.dataset_id)
         operation = hook.import_data(
             project_id=self.project_id,
@@ -386,7 +430,21 @@ class ImportDataOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         hook.wait_for_operation(timeout=self.timeout, operation=operation)
+        result_dataset_size = self._get_number_of_ds_items(
+            dataset=hook.get_dataset(
+                dataset=self.dataset_id,
+                project_id=self.project_id,
+                region=self.region,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            ),
+            total_key_name="data_item_count",
+        )
+        if self.raise_for_empty_result:
+            self._raise_for_empty_import_result(self.dataset_id, initial_dataset_size, result_dataset_size)
         self.log.info("Import was done successfully")
+        return {"total_data_items_imported": int(result_dataset_size) - int(initial_dataset_size)}
 
 
 class ListDatasetsOperator(GoogleCloudBaseOperator):
@@ -449,6 +507,12 @@ class ListDatasetsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = DatasetHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -466,7 +530,7 @@ class ListDatasetsOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIDatasetListLink.persist(context=context, task_instance=self)
+        VertexAIDatasetListLink.persist(context=context)
         return [Dataset.to_dict(result) for result in results]
 
 

airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py

@@ -21,7 +21,11 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint, endpoint_service
 
 from airflow.providers.google.cloud.hooks.vertex_ai.endpoint_service import EndpointServiceHook
 from airflow.providers.google.cloud.links.vertex_ai import (
@@ -30,15 +34,13 @@ from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIModelLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint, endpoint_service
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class CreateEndpointOperator(GoogleCloudBaseOperator):
     """
@@ -91,6 +93,13 @@ class CreateEndpointOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = EndpointServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -113,8 +122,8 @@ class CreateEndpointOperator(GoogleCloudBaseOperator):
         endpoint_id = hook.extract_endpoint_id(endpoint)
         self.log.info("Endpoint was created. Endpoint ID: %s", endpoint_id)
 
-        self.xcom_push(context, key="endpoint_id", value=endpoint_id)
-        VertexAIEndpointLink.persist(context=context, task_instance=self, endpoint_id=endpoint_id)
+        context["ti"].xcom_push(key="endpoint_id", value=endpoint_id)
+        VertexAIEndpointLink.persist(context=context, endpoint_id=endpoint_id)
         return endpoint
 
 
@@ -253,6 +262,13 @@ class DeployModelOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = EndpointServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -276,8 +292,8 @@ class DeployModelOperator(GoogleCloudBaseOperator):
         deployed_model_id = hook.extract_deployed_model_id(deploy_model)
         self.log.info("Model was deployed. Deployed Model ID: %s", deployed_model_id)
 
-        self.xcom_push(context, key="deployed_model_id", value=deployed_model_id)
-        VertexAIModelLink.persist(context=context, task_instance=self, model_id=deployed_model_id)
+        context["ti"].xcom_push(key="deployed_model_id", value=deployed_model_id)
+        VertexAIModelLink.persist(context=context, model_id=deployed_model_id)
         return deploy_model
 
 
@@ -328,6 +344,13 @@ class GetEndpointOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = EndpointServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -344,7 +367,7 @@ class GetEndpointOperator(GoogleCloudBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            VertexAIEndpointLink.persist(context=context, task_instance=self, endpoint_id=self.endpoint_id)
+            VertexAIEndpointLink.persist(context=context, endpoint_id=self.endpoint_id)
             self.log.info("Endpoint was gotten.")
             return Endpoint.to_dict(endpoint_obj)
         except NotFound:
@@ -427,6 +450,12 @@ class ListEndpointsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = EndpointServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -444,7 +473,7 @@ class ListEndpointsOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIEndpointListLink.persist(context=context, task_instance=self)
+        VertexAIEndpointListLink.persist(context=context)
         return [Endpoint.to_dict(result) for result in results]
 
 
@@ -580,6 +609,13 @@ class UpdateEndpointOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = EndpointServiceHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -597,5 +633,5 @@ class UpdateEndpointOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         self.log.info("Endpoint was updated")
-        VertexAIEndpointLink.persist(context=context, task_instance=self, endpoint_id=self.endpoint_id)
+        VertexAIEndpointLink.persist(context=context, endpoint_id=self.endpoint_id)
         return Endpoint.to_dict(result)