apache-airflow-providers-google 16.0.0a1__py3-none-any.whl → 16.1.0rc1__py3-none-any.whl

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

@@ -37,6 +37,7 @@ from airflow.providers.google.cloud.links.translate import (
     TranslationNativeDatasetLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.cloud.operators.vertex_ai.dataset import DatasetImportDataResultsCheckHelper
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
@@ -394,7 +395,6 @@ class TranslateTextBatchOperator(GoogleCloudBaseOperator):
         self.log.info("Translate text batch job started.")
         TranslateTextBatchLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
             output_config=self.output_config,
         )
@@ -480,15 +480,15 @@ class TranslateCreateDatasetOperator(GoogleCloudBaseOperator):
         result = hook.wait_for_operation_result(result_operation)
         result = type(result).to_dict(result)
         dataset_id = hook.extract_object_id(result)
-        self.xcom_push(context, key="dataset_id", value=dataset_id)
+        context["ti"].xcom_push(key="dataset_id", value=dataset_id)
         self.log.info("Dataset creation complete. The dataset_id: %s.", dataset_id)
 
         project_id = self.project_id or hook.project_id
         TranslationNativeDatasetLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dataset_id,
             project_id=project_id,
+            location=self.location,
         )
         return result
 
@@ -556,7 +556,6 @@ class TranslateDatasetsListOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationDatasetsListLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
         )
         self.log.info("Requesting datasets list")
@@ -577,7 +576,7 @@ class TranslateDatasetsListOperator(GoogleCloudBaseOperator):
         return result_ids
 
 
-class TranslateImportDataOperator(GoogleCloudBaseOperator):
+class TranslateImportDataOperator(GoogleCloudBaseOperator, DatasetImportDataResultsCheckHelper):
     """
     Import data to the translation dataset.
 
@@ -604,6 +603,7 @@ class TranslateImportDataOperator(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: Sequence[str] = (
@@ -629,6 +629,7 @@ class TranslateImportDataOperator(GoogleCloudBaseOperator):
         retry: Retry | _MethodDefault = DEFAULT,
         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)
@@ -641,9 +642,21 @@ class TranslateImportDataOperator(GoogleCloudBaseOperator):
         self.retry = retry
         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 = TranslateHook(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_id=self.dataset_id,
+                project_id=self.project_id,
+                location=self.location,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            ),
+            total_key_name="example_count",
+        )
         self.log.info("Importing data to dataset...")
         operation = hook.import_dataset_data(
             dataset_id=self.dataset_id,
@@ -657,12 +670,27 @@ class TranslateImportDataOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationNativeDatasetLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=self.dataset_id,
             project_id=project_id,
+            location=self.location,
         )
         hook.wait_for_operation_done(operation=operation, timeout=self.timeout)
+
+        result_dataset_size = self._get_number_of_ds_items(
+            dataset=hook.get_dataset(
+                dataset_id=self.dataset_id,
+                project_id=self.project_id,
+                location=self.location,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            ),
+            total_key_name="example_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("Importing data finished!")
+        return {"total_imported": int(result_dataset_size) - int(initial_dataset_size)}
 
 
 class TranslateDeleteDatasetOperator(GoogleCloudBaseOperator):
@@ -821,16 +849,16 @@ class TranslateCreateModelOperator(GoogleCloudBaseOperator):
         result = hook.wait_for_operation_result(operation=result_operation)
         result = type(result).to_dict(result)
         model_id = hook.extract_object_id(result)
-        self.xcom_push(context, key="model_id", value=model_id)
+        context["ti"].xcom_push(key="model_id", value=model_id)
         self.log.info("Model creation complete. The model_id: %s.", model_id)
 
         project_id = self.project_id or hook.project_id
         TranslationModelLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=self.dataset_id,
             model_id=model_id,
             project_id=project_id,
+            location=self.location,
         )
         return result
 
@@ -898,7 +926,6 @@ class TranslateModelsListOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationModelsListLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
         )
         self.log.info("Requesting models list")
@@ -1141,7 +1168,6 @@ class TranslateDocumentOperator(GoogleCloudBaseOperator):
         if self.document_output_config:
             TranslateResultByOutputConfigLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=self.project_id or hook.project_id,
                 output_config=self.document_output_config,
             )
@@ -1304,7 +1330,6 @@ class TranslateDocumentBatchOperator(GoogleCloudBaseOperator):
         self.log.info("Batch document translation job started.")
         TranslateResultByOutputConfigLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
             output_config=self.output_config,
         )
@@ -1411,7 +1436,7 @@ class TranslateCreateGlossaryOperator(GoogleCloudBaseOperator):
         result = type(result).to_dict(result)
 
         glossary_id = hook.extract_object_id(result)
-        self.xcom_push(context, key="glossary_id", value=glossary_id)
+        context["ti"].xcom_push(key="glossary_id", value=glossary_id)
         self.log.info("Glossary creation complete. The glossary_id: %s.", glossary_id)
         return result
 
@@ -1610,7 +1635,6 @@ class TranslateListGlossariesOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         TranslationGlossariesListLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
         )
         self.log.info("Requesting glossaries list")

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

@@ -173,7 +173,6 @@ class CloudTranslateSpeechOperator(GoogleCloudBaseOperator):
         if self.audio.uri:
             FileDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 # Slice from: "gs://{BUCKET_NAME}/{FILE_NAME}" to: "{BUCKET_NAME}/{FILE_NAME}"
                 uri=self.audio.uri[5:],
                 project_id=self.project_id or translate_hook.project_id,

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

@@ -21,7 +21,7 @@
 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
@@ -91,6 +91,13 @@ class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.hook: AutoMLHook | None = None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def on_kill(self) -> None:
         """Act as a callback called when the operator is killed; cancel any running job."""
         if self.hook:
@@ -242,12 +249,12 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
@@ -334,12 +341,12 @@ class CreateAutoMLImageTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
@@ -457,12 +464,12 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
@@ -531,12 +538,12 @@ class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
@@ -640,6 +647,12 @@ class ListAutoMLTrainingJobOperator(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 = AutoMLHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -656,5 +669,5 @@ class ListAutoMLTrainingJobOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAITrainingPipelinesLink.persist(context=context, task_instance=self)
+        VertexAITrainingPipelinesLink.persist(context=context)
         return [TrainingPipeline.to_dict(result) for result in results]

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

@@ -231,6 +231,13 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.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):
         self.log.info("Creating Batch prediction job")
         batch_prediction_job: BatchPredictionJobObject = self.hook.submit_batch_prediction_job(
@@ -262,9 +269,10 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
         batch_prediction_job_id = batch_prediction_job.name
         self.log.info("Batch prediction job was created. Job id: %s", batch_prediction_job_id)
 
-        self.xcom_push(context, key="batch_prediction_job_id", value=batch_prediction_job_id)
+        context["ti"].xcom_push(key="batch_prediction_job_id", value=batch_prediction_job_id)
         VertexAIBatchPredictionJobLink.persist(
-            context=context, task_instance=self, batch_prediction_job_id=batch_prediction_job_id
+            context=context,
+            batch_prediction_job_id=batch_prediction_job_id,
         )
 
         if self.deferrable:
@@ -295,13 +303,11 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
         job: dict[str, Any] = event["job"]
         self.log.info("Batch prediction job %s created and completed successfully.", job["name"])
         job_id = self.hook.extract_batch_prediction_job_id(job)
-        self.xcom_push(
-            context,
+        context["ti"].xcom_push(
             key="batch_prediction_job_id",
             value=job_id,
         )
-        self.xcom_push(
-            context,
+        context["ti"].xcom_push(
             key="training_conf",
             value={
                 "training_conf_id": job_id,
@@ -427,6 +433,13 @@ class GetBatchPredictionJobOperator(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 = BatchPredictionJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -445,7 +458,8 @@ class GetBatchPredictionJobOperator(GoogleCloudBaseOperator):
             )
             self.log.info("Batch prediction job was gotten.")
             VertexAIBatchPredictionJobLink.persist(
-                context=context, task_instance=self, batch_prediction_job_id=self.batch_prediction_job
+                context=context,
+                batch_prediction_job_id=self.batch_prediction_job,
             )
             return BatchPredictionJob.to_dict(result)
         except NotFound:
@@ -517,6 +531,12 @@ class ListBatchPredictionJobsOperator(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 = BatchPredictionJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -533,5 +553,5 @@ class ListBatchPredictionJobsOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIBatchPredictionJobListLink.persist(context=context, task_instance=self)
+        VertexAIBatchPredictionJobListLink.persist(context=context)
         return [BatchPredictionJob.to_dict(result) for result in results]

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

@@ -170,17 +170,24 @@ class CustomTrainingJobBaseOperator(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_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any] | None:
         if event["status"] == "error":
             raise AirflowException(event["message"])
         training_pipeline = event["job"]
         custom_job_id = self.hook.extract_custom_job_id_from_training_pipeline(training_pipeline)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
         try:
             model = training_pipeline["model_to_upload"]
             model_id = self.hook.extract_model_id(model)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
             return model
         except KeyError:
             self.log.warning(
@@ -584,13 +591,13 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
     def invoke_defer(self, context: Context) -> None:
@@ -648,8 +655,8 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
         )
         custom_container_training_job_obj.wait_for_resource_creation()
         training_pipeline_id: str = custom_container_training_job_obj.name
-        self.xcom_push(context, key="training_id", value=training_pipeline_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        context["ti"].xcom_push(key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_pipeline_id)
         self.defer(
             trigger=CustomContainerTrainingJobTrigger(
                 conn_id=self.gcp_conn_id,
@@ -1041,13 +1048,13 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
     def invoke_defer(self, context: Context) -> None:
@@ -1106,8 +1113,8 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
         )
         custom_python_training_job_obj.wait_for_resource_creation()
         training_pipeline_id: str = custom_python_training_job_obj.name
-        self.xcom_push(context, key="training_id", value=training_pipeline_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        context["ti"].xcom_push(key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_pipeline_id)
         self.defer(
             trigger=CustomPythonPackageTrainingJobTrigger(
                 conn_id=self.gcp_conn_id,
@@ -1504,13 +1511,13 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
     def invoke_defer(self, context: Context) -> None:
@@ -1569,8 +1576,8 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         )
         custom_training_job_obj.wait_for_resource_creation()
         training_pipeline_id: str = custom_training_job_obj.name
-        self.xcom_push(context, key="training_id", value=training_pipeline_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        context["ti"].xcom_push(key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_pipeline_id)
         self.defer(
             trigger=CustomTrainingJobTrigger(
                 conn_id=self.gcp_conn_id,
@@ -1748,6 +1755,12 @@ class ListCustomTrainingJobOperator(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 = CustomJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -1764,5 +1777,5 @@ class ListCustomTrainingJobOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAITrainingPipelinesLink.persist(context=context, task_instance=self)
+        VertexAITrainingPipelinesLink.persist(context=context)
         return [TrainingPipeline.to_dict(result) for result in results]