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

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
@@ -29,6 +29,7 @@ from google.cloud.aiplatform import datasets
 from google.cloud.aiplatform.models import Model
 from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.auto_ml import AutoMLHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIModelLink,
@@ -36,11 +37,12 @@ from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAITrainingPipelinesLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.common.deprecated import deprecated
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
@@ -91,6 +93,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 +251,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 +343,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,15 +466,20 @@ 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
 
 
+@deprecated(
+    planned_removal_date="March 24, 2026",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create Auto ML Video Training job."""
 
@@ -531,12 +545,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 +654,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 +676,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

@@ -42,7 +42,7 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.aiplatform import BatchPredictionJob as BatchPredictionJobObject, Model, explain
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
@@ -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

@@ -51,8 +51,9 @@ if TYPE_CHECKING:
         CustomPythonPackageTrainingJob,
         CustomTrainingJob,
     )
+    from google.cloud.aiplatform_v1.types import PscInterfaceConfig
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
@@ -110,6 +111,7 @@ class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
         predefined_split_column_name: str | None = None,
         timestamp_split_column_name: str | None = None,
         tensorboard: str | None = None,
+        psc_interface_config: PscInterfaceConfig | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -166,21 +168,29 @@ class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
         self.predefined_split_column_name = predefined_split_column_name
         self.timestamp_split_column_name = timestamp_split_column_name
         self.tensorboard = tensorboard
+        self.psc_interface_config = psc_interface_config
         # END Run param
         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(
@@ -466,6 +476,8 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
             For more information on configuring your service account please visit:
             https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+    :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+        training.
     :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
     :param impersonation_chain: Optional service account to impersonate using short-term
         credentials, or chained list of accounts required to get the access_token
@@ -579,18 +591,19 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
             sync=True,
+            psc_interface_config=self.psc_interface_config,
         )
 
         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:
@@ -645,11 +658,12 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             predefined_split_column_name=self.predefined_split_column_name,
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
+            psc_interface_config=self.psc_interface_config,
         )
         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,
@@ -924,6 +938,8 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
             For more information on configuring your service account please visit:
             https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+    :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+        training.
     :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
     :param impersonation_chain: Optional service account to impersonate using short-term
         credentials, or chained list of accounts required to get the access_token
@@ -1036,18 +1052,19 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
             sync=True,
+            psc_interface_config=self.psc_interface_config,
         )
 
         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:
@@ -1103,11 +1120,12 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             predefined_split_column_name=self.predefined_split_column_name,
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
+            psc_interface_config=self.psc_interface_config,
         )
         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,
@@ -1382,6 +1400,8 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
             For more information on configuring your service account please visit:
             https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+    :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+        training.
     :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
     :param impersonation_chain: Optional service account to impersonate using short-term
         credentials, or chained list of accounts required to get the access_token
@@ -1499,18 +1519,19 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
             sync=True,
+            psc_interface_config=None,
         )
 
         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:
@@ -1566,11 +1587,12 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             predefined_split_column_name=self.predefined_split_column_name,
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
+            psc_interface_config=self.psc_interface_config,
         )
         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 +1770,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 +1792,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]

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

@@ -20,12 +20,13 @@
 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 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
@@ -34,7 +35,7 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CreateDatasetOperator(GoogleCloudBaseOperator):
@@ -85,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,
@@ -106,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
 
 
@@ -160,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,
@@ -177,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:
@@ -321,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.
 
@@ -342,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")
@@ -358,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)
@@ -370,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,
@@ -388,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):
@@ -451,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,
@@ -468,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]