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

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

@@ -21,8 +21,10 @@ from __future__ import annotations
 
 import ast
 import warnings
+from functools import cached_property
 from typing import TYPE_CHECKING, Sequence, Tuple
 
+from deprecated import deprecated
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.automl_v1beta1 import (
     BatchPredictResult,
@@ -33,8 +35,9 @@ from google.cloud.automl_v1beta1 import (
     TableSpec,
 )
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.automl import CloudAutoMLHook
+from airflow.providers.google.cloud.hooks.vertex_ai.prediction_service import PredictionServiceHook
 from airflow.providers.google.cloud.links.automl import (
     AutoMLDatasetLink,
     AutoMLDatasetListLink,
@@ -53,12 +56,36 @@ if TYPE_CHECKING:
 MetaData = Sequence[Tuple[str, str]]
 
 
+def _raise_exception_for_deprecated_operator(
+    deprecated_class_name: str, alternative_class_names: str | list[str]
+):
+    if isinstance(alternative_class_names, str):
+        alternative_class_name_str = alternative_class_names
+    elif len(alternative_class_names) == 1:
+        alternative_class_name_str = alternative_class_names[0]
+    else:
+        alternative_class_name_str = ", ".join(f"`{cls_name}`" for cls_name in alternative_class_names[:-1])
+        alternative_class_name_str += f" or `{alternative_class_names[-1]}`"
+
+    raise AirflowException(
+        f"{deprecated_class_name} for text, image, and video prediction has been "
+        f"deprecated and no longer available. All the functionality of "
+        f"legacy AutoML Natural Language, Vision, Video Intelligence and Tables "
+        f"and new features are available on the Vertex AI platform. "
+        f"Please use {alternative_class_name_str} from Vertex AI."
+    )
+
+
 class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
     """
     Creates Google Cloud AutoML model.
 
-    AutoMLTrainModelOperator for text prediction is deprecated. Please use
-    :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTextTrainingJobOperator`
+    AutoMLTrainModelOperator for tables, video intelligence, vision and natural language has been deprecated
+    and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTabularTrainingJobOperator`,
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLVideoTrainingJobOperator`,
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLImageTrainingJobOperator`,
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTextTrainingJobOperator`,
     instead.
 
     .. seealso::
@@ -120,17 +147,16 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
 
     def execute(self, context: Context):
-        # Output warning if running not AutoML Translation prediction job
+        # Raise exception if running not AutoML Translation prediction job
         if "translation_model_metadata" not in self.model:
-            warnings.warn(
-                "AutoMLTrainModelOperator for text, image and video prediction is deprecated. "
-                "All the functionality of legacy "
-                "AutoML Natural Language, Vision and Video Intelligence and new features are available "
-                "on the Vertex AI platform. "
-                "Please use `CreateAutoMLTextTrainingJobOperator`, `CreateAutoMLImageTrainingJobOperator` or"
-                " `CreateAutoMLVideoTrainingJobOperator` from VertexAI.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=3,
+            _raise_exception_for_deprecated_operator(
+                self.__class__.__name__,
+                [
+                    "CreateAutoMLTabularTrainingJobOperator",
+                    "CreateAutoMLVideoTrainingJobOperator",
+                    "CreateAutoMLImageTrainingJobOperator",
+                    "CreateAutoMLTextTrainingJobOperator",
+                ],
             )
         hook = CloudAutoMLHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -174,7 +200,8 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
         :ref:`howto/operator:AutoMLPredictOperator`
 
     :param model_id: Name of the model requested to serve the batch prediction.
-    :param payload: Name od the model used for the prediction.
+    :param endpoint_id: Name of the endpoint used for the prediction.
+    :param payload: Name of the model used for the prediction.
     :param project_id: ID of the Google Cloud project where model is located if None then
         default project_id is used.
     :param location: The location of the project.
@@ -206,10 +233,12 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        model_id: str,
+        model_id: str | None = None,
+        endpoint_id: str | None = None,
         location: str,
         payload: dict,
         operation_params: dict[str, str] | None = None,
+        instances: list[str] | None = None,
         project_id: str = PROVIDE_PROJECT_ID,
         metadata: MetaData = (),
         timeout: float | None = None,
@@ -221,7 +250,9 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
         super().__init__(**kwargs)
 
         self.model_id = model_id
+        self.endpoint_id = endpoint_id
         self.operation_params = operation_params  # type: ignore
+        self.instances = instances
         self.location = location
         self.project_id = project_id
         self.metadata = metadata
@@ -231,23 +262,69 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
-    def execute(self, context: Context):
-        hook = CloudAutoMLHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-        result = hook.predict(
+    @cached_property
+    def hook(self) -> CloudAutoMLHook | PredictionServiceHook:
+        if self.model_id:
+            return CloudAutoMLHook(
+                gcp_conn_id=self.gcp_conn_id,
+                impersonation_chain=self.impersonation_chain,
+            )
+        else:  # endpoint_id defined
+            return PredictionServiceHook(
+                gcp_conn_id=self.gcp_conn_id,
+                impersonation_chain=self.impersonation_chain,
+            )
+
+    def _check_model_type(self):
+        hook = self.hook
+        model = hook.get_model(
             model_id=self.model_id,
-            payload=self.payload,
             location=self.location,
             project_id=self.project_id,
-            params=self.operation_params,
             retry=self.retry,
             timeout=self.timeout,
             metadata=self.metadata,
         )
+        if not hasattr(model, "translation_model_metadata"):
+            raise AirflowException(
+                "AutoMLPredictOperator for text, image, and video prediction has been deprecated. "
+                "Please use endpoint_id param instead of model_id param."
+            )
+
+    def execute(self, context: Context):
+        if self.model_id is None and self.endpoint_id is None:
+            raise AirflowException("You must specify model_id or endpoint_id!")
+
+        if self.model_id:
+            self._check_model_type()
+
+        hook = self.hook
+        if self.model_id:
+            result = hook.predict(
+                model_id=self.model_id,
+                payload=self.payload,
+                location=self.location,
+                project_id=self.project_id,
+                params=self.operation_params,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+        else:  # self.endpoint_id is defined
+            result = hook.predict(
+                endpoint_id=self.endpoint_id,
+                instances=self.instances,
+                payload=self.payload,
+                location=self.location,
+                project_id=self.project_id,
+                parameters=self.operation_params,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+
         project_id = self.project_id or hook.project_id
-        if project_id:
+        if project_id and self.model_id:
             AutoMLModelPredictLink.persist(
                 context=context,
                 task_instance=self,
@@ -261,6 +338,14 @@ class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
     """
     Perform a batch prediction on Google Cloud AutoML.
 
+    AutoMLBatchPredictOperator for tables, video intelligence, vision and natural language has been deprecated
+    and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job.CreateBatchPredictionJobOperator`,
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job.GetBatchPredictionJobOperator`,
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job.ListBatchPredictionJobsOperator`,
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job.DeleteBatchPredictionJobOperator`,
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLBatchPredictOperator`
@@ -341,6 +426,25 @@ class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        model: Model = hook.get_model(
+            model_id=self.model_id,
+            location=self.location,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+
+        if not hasattr(model, "translation_model_metadata"):
+            _raise_exception_for_deprecated_operator(
+                self.__class__.__name__,
+                [
+                    "CreateBatchPredictionJobOperator",
+                    "GetBatchPredictionJobOperator",
+                    "ListBatchPredictionJobsOperator",
+                    "DeleteBatchPredictionJobOperator",
+                ],
+            )
         self.log.info("Fetch batch prediction.")
         operation = hook.batch_predict(
             model_id=self.model_id,
@@ -371,6 +475,10 @@ class AutoMLCreateDatasetOperator(GoogleCloudBaseOperator):
     """
     Creates a Google Cloud AutoML dataset.
 
+    AutoMLCreateDatasetOperator for tables, video intelligence, vision and natural language has been
+    deprecated and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.CreateDatasetOperator` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLCreateDatasetOperator`
@@ -430,6 +538,8 @@ class AutoMLCreateDatasetOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
 
     def execute(self, context: Context):
+        if "translation_dataset_metadata" not in self.dataset:
+            _raise_exception_for_deprecated_operator(self.__class__.__name__, "CreateDatasetOperator")
         hook = CloudAutoMLHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
@@ -463,6 +573,10 @@ class AutoMLImportDataOperator(GoogleCloudBaseOperator):
     """
     Imports data to a Google Cloud AutoML dataset.
 
+    AutoMLImportDataOperator for tables, video intelligence, vision and natural language has been deprecated
+    and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.ImportDataOperator` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLImportDataOperator`
@@ -530,6 +644,16 @@ class AutoMLImportDataOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        dataset: Dataset = hook.get_dataset(
+            dataset_id=self.dataset_id,
+            location=self.location,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        if not hasattr(dataset, "translation_dataset_metadata"):
+            _raise_exception_for_deprecated_operator(self.__class__.__name__, "ImportDataOperator")
         self.log.info("Importing data to dataset...")
         operation = hook.import_data(
             dataset_id=self.dataset_id,
@@ -662,10 +786,22 @@ class AutoMLTablesListColumnSpecsOperator(GoogleCloudBaseOperator):
         return result
 
 
+@deprecated(
+    reason=(
+        "Class `AutoMLTablesUpdateDatasetOperator` has been deprecated and no longer available. "
+        "Please use `UpdateDatasetOperator` instead"
+    ),
+    category=AirflowProviderDeprecationWarning,
+    action="error",
+)
 class AutoMLTablesUpdateDatasetOperator(GoogleCloudBaseOperator):
     """
     Updates a dataset.
 
+    AutoMLTablesUpdateDatasetOperator has been deprecated and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.UpdateDatasetOperator`
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLTablesUpdateDatasetOperator`
@@ -753,6 +889,10 @@ class AutoMLGetModelOperator(GoogleCloudBaseOperator):
     """
     Get Google Cloud AutoML model.
 
+    AutoMLGetModelOperator for tables, video intelligence, vision and natural language has been deprecated
+    and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLGetModelOperator`
@@ -823,6 +963,8 @@ class AutoMLGetModelOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
+        if not hasattr(result, "translation_model_metadata"):
+            _raise_exception_for_deprecated_operator(self.__class__.__name__, "GetModelOperator")
         model = Model.to_dict(result)
         project_id = self.project_id or hook.project_id
         if project_id:
@@ -840,6 +982,10 @@ class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
     """
     Delete Google Cloud AutoML model.
 
+    AutoMLDeleteModelOperator for tables, video intelligence, vision and natural language has been deprecated
+    and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLDeleteModelOperator`
@@ -901,6 +1047,16 @@ class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        model: Model = hook.get_model(
+            model_id=self.model_id,
+            location=self.location,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        if not hasattr(model, "translation_model_metadata"):
+            _raise_exception_for_deprecated_operator(self.__class__.__name__, "DeleteModelOperator")
         operation = hook.delete_model(
             model_id=self.model_id,
             location=self.location,
@@ -913,6 +1069,14 @@ class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
         self.log.info("Deletion is completed")
 
 
+@deprecated(
+    reason=(
+        "Class `AutoMLDeployModelOperator` has been deprecated and no longer available. Please use "
+        "`DeployModelOperator` instead"
+    ),
+    category=AirflowProviderDeprecationWarning,
+    action="error",
+)
 class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
     """
     Deploys a model; if a model is already deployed, deploying it with the same parameters has no effect.
@@ -923,6 +1087,10 @@ class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
     Only applicable for Text Classification, Image Object Detection and Tables; all other
     domains manage deployment automatically.
 
+    AutoMLDeployModelOperator has been deprecated and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.endpoint_service.DeployModelOperator`
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLDeployModelOperator`
@@ -989,6 +1157,16 @@ class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        model = hook.get_model(
+            model_id=self.model_id,
+            location=self.location,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        if not hasattr(model, "translation_model_metadata"):
+            _raise_exception_for_deprecated_operator(self.__class__.__name__, "DeployModelOperator")
         self.log.info("Deploying model_id %s", self.model_id)
 
         operation = hook.deploy_model(
@@ -1108,6 +1286,10 @@ class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
     """
     Lists AutoML Datasets in project.
 
+    AutoMLListDatasetOperator for tables, video intelligence, vision and natural language has been deprecated
+    and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.ListDatasetsOperator` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLListDatasetOperator`
@@ -1172,7 +1354,16 @@ class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        result = [Dataset.to_dict(dataset) for dataset in page_iterator]
+        result = []
+        for dataset in page_iterator:
+            if not hasattr(dataset, "translation_dataset_metadata"):
+                warnings.warn(
+                    "Class `AutoMLListDatasetOperator` has been deprecated and no longer available. "
+                    "Please use `ListDatasetsOperator` instead.",
+                    stacklevel=2,
+                )
+            else:
+                result.append(Dataset.to_dict(dataset))
         self.log.info("Datasets obtained.")
 
         self.xcom_push(
@@ -1190,6 +1381,10 @@ class AutoMLDeleteDatasetOperator(GoogleCloudBaseOperator):
     """
     Deletes a dataset and all of its contents.
 
+    AutoMLDeleteDatasetOperator for tables, video intelligence, vision and natural language has been
+    deprecated and no longer available. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.DeleteDatasetOperator` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AutoMLDeleteDatasetOperator`
@@ -1260,6 +1455,16 @@ class AutoMLDeleteDatasetOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        dataset: Dataset = hook.get_dataset(
+            dataset_id=self.dataset_id,
+            location=self.location,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        if not hasattr(dataset, "translation_dataset_metadata"):
+            _raise_exception_for_deprecated_operator(self.__class__.__name__, "DeleteDatasetOperator")
         dataset_id_list = self._parse_dataset_id(self.dataset_id)
         for dataset_id in dataset_id_list:
             self.log.info("Deleting dataset %s", dataset_id)