apache-airflow-providers-google 12.0.0rc2__py3-none-any.whl → 14.0.0__py3-none-any.whl

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

@@ -20,21 +20,10 @@
 from __future__ import annotations
 
 import ast
-import warnings
 from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, cast
 
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.automl_v1beta1 import (
-    BatchPredictResult,
-    ColumnSpec,
-    Dataset,
-    Model,
-    PredictResponse,
-    TableSpec,
-)
-
 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
@@ -48,35 +37,31 @@ from airflow.providers.google.cloud.links.translate import (
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.automl_v1beta1 import (
+    ColumnSpec,
+    Dataset,
+    Model,
+    PredictResponse,
+    TableSpec,
+)
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 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."
-    )
-
-
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTabularTrainingJobOperator, "
+    "airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLVideoTrainingJobOperator, "
+    "airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLImageTrainingJobOperator, "
+    "airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator, "
+    "airflow.providers.google.cloud.operators.translate.TranslateCreateModelOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
     """
     Creates Google Cloud AutoML model.
@@ -88,6 +73,7 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
         :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.generative_model.SupervisedFineTuningTrainOperator`,
+        :class:`airflow.providers.google.cloud.operators.translate.TranslateCreateModelOperator`.
         instead.
 
     .. seealso::
@@ -149,17 +135,6 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
 
     def execute(self, context: Context):
-        # Raise exception if running not AutoML Translation prediction job
-        if "translation_model_metadata" not in self.model:
-            _raise_exception_for_deprecated_operator(
-                self.__class__.__name__,
-                [
-                    "CreateAutoMLTabularTrainingJobOperator",
-                    "CreateAutoMLVideoTrainingJobOperator",
-                    "CreateAutoMLImageTrainingJobOperator",
-                    "SupervisedFineTuningTrainOperator",
-                ],
-            )
         hook = CloudAutoMLHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
@@ -195,6 +170,11 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
         return result
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.translate.TranslateTextOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLPredictOperator(GoogleCloudBaseOperator):
     """
     Runs prediction operation on Google Cloud AutoML.
@@ -297,20 +277,10 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
             )
         return None
 
-    def _check_model_type(self):
-        if not hasattr(self.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(
@@ -350,161 +320,19 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
 
 
 @deprecated(
-    planned_removal_date="January 01, 2025",
-    use_instead="airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job",
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.CreateDatasetOperator, "
+    "airflow.providers.google.cloud.operators.translate.TranslateCreateDatasetOperator",
     category=AirflowProviderDeprecationWarning,
 )
-class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
-    """
-    Perform a batch prediction on Google Cloud AutoML.
-
-    .. warning::
-        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`
-
-    :param project_id: ID of the Google Cloud project where model will be created if None then
-        default project_id is used.
-    :param location: The location of the project.
-    :param model_id: Name of the model_id requested to serve the batch prediction.
-    :param input_config: Required. The input configuration for batch prediction.
-        If a dict is provided, it must be of the same form as the protobuf message
-        `google.cloud.automl_v1beta1.types.BatchPredictInputConfig`
-    :param output_config: Required. The Configuration specifying where output predictions should be
-        written. If a dict is provided, it must be of the same form as the protobuf message
-        `google.cloud.automl_v1beta1.types.BatchPredictOutputConfig`
-    :param prediction_params: Additional domain-specific parameters for the predictions,
-        any string must be up to 25000 characters long.
-    :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.
-    :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
-        retried.
-    :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
-        `retry` is specified, the timeout applies to each individual attempt.
-    :param metadata: Additional metadata that is provided to the method.
-    :param gcp_conn_id: The connection ID to use to connect 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
-        of the last account in the list, which will be impersonated in the request.
-        If set as a string, the account must grant the originating account
-        the Service Account Token Creator IAM role.
-        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).
-    """
-
-    template_fields: Sequence[str] = (
-        "model_id",
-        "input_config",
-        "output_config",
-        "location",
-        "project_id",
-        "impersonation_chain",
-    )
-    operator_extra_links = (TranslationLegacyModelPredictLink(),)
-
-    def __init__(
-        self,
-        *,
-        model_id: str,
-        input_config: dict,
-        output_config: dict,
-        location: str,
-        project_id: str = PROVIDE_PROJECT_ID,
-        prediction_params: dict[str, str] | None = None,
-        metadata: MetaData = (),
-        timeout: float | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-
-        self.model_id = model_id
-        self.location = location
-        self.project_id = project_id
-        self.prediction_params = prediction_params
-        self.metadata = metadata
-        self.timeout = timeout
-        self.retry = retry
-        self.gcp_conn_id = gcp_conn_id
-        self.impersonation_chain = impersonation_chain
-        self.input_config = input_config
-        self.output_config = output_config
-
-    @cached_property
-    def hook(self) -> CloudAutoMLHook:
-        return CloudAutoMLHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-
-    @cached_property
-    def model(self) -> Model:
-        return self.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,
-        )
-
-    def execute(self, context: Context):
-        if not hasattr(self.model, "translation_model_metadata"):
-            _raise_exception_for_deprecated_operator(
-                self.__class__.__name__,
-                [
-                    "CreateBatchPredictionJobOperator",
-                    "GetBatchPredictionJobOperator",
-                    "ListBatchPredictionJobsOperator",
-                    "DeleteBatchPredictionJobOperator",
-                ],
-            )
-        self.log.info("Fetch batch prediction.")
-        operation = self.hook.batch_predict(
-            model_id=self.model_id,
-            input_config=self.input_config,
-            output_config=self.output_config,
-            project_id=self.project_id,
-            location=self.location,
-            params=self.prediction_params,
-            retry=self.retry,
-            timeout=self.timeout,
-            metadata=self.metadata,
-        )
-        operation_result = self.hook.wait_for_operation(timeout=self.timeout, operation=operation)
-        result = BatchPredictResult.to_dict(operation_result)
-        self.log.info("Batch prediction is ready.")
-        project_id = self.project_id or self.hook.project_id
-        if project_id:
-            TranslationLegacyModelPredictLink.persist(
-                context=context,
-                task_instance=self,
-                model_id=self.model_id,
-                project_id=project_id,
-                dataset_id=self.model.dataset_id,
-            )
-        return result
-
-
 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.
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.CreateDatasetOperator`,
+    :class:`airflow.providers.google.cloud.operators.translate.TranslateCreateDatasetOperator` instead.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -565,8 +393,6 @@ 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,
@@ -596,6 +422,12 @@ class AutoMLCreateDatasetOperator(GoogleCloudBaseOperator):
         return result
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.ImportDataOperator, "
+    "airflow.providers.google.cloud.operators.translate.TranslateImportDataOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLImportDataOperator(GoogleCloudBaseOperator):
     """
     Imports data to a Google Cloud AutoML dataset.
@@ -672,7 +504,7 @@ class AutoMLImportDataOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-        dataset: Dataset = hook.get_dataset(
+        hook.get_dataset(
             dataset_id=self.dataset_id,
             location=self.location,
             project_id=self.project_id,
@@ -680,8 +512,6 @@ class AutoMLImportDataOperator(GoogleCloudBaseOperator):
             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,
@@ -704,6 +534,11 @@ class AutoMLImportDataOperator(GoogleCloudBaseOperator):
             )
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    category=AirflowProviderDeprecationWarning,
+    reason="Shutdown of legacy version of AutoML Tables on March 31, 2024.",
+)
 class AutoMLTablesListColumnSpecsOperator(GoogleCloudBaseOperator):
     """
     Lists column specs in a table.
@@ -787,11 +622,6 @@ class AutoMLTablesListColumnSpecsOperator(GoogleCloudBaseOperator):
         self.retry = retry
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
-        raise AirflowException(
-            "Operator AutoMLTablesListColumnSpecsOperator has been deprecated due to shutdown of "
-            "a legacy version of AutoML Tables on March 31, 2024. "
-            "For additional information see: https://cloud.google.com/automl-tables/docs/deprecations."
-        )
 
     def execute(self, context: Context):
         hook = CloudAutoMLHook(
@@ -824,6 +654,12 @@ class AutoMLTablesListColumnSpecsOperator(GoogleCloudBaseOperator):
         return result
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.UpdateDatasetOperator",
+    category=AirflowProviderDeprecationWarning,
+    reason="Shutdown of legacy version of AutoML Tables on March 31, 2024.",
+)
 class AutoMLTablesUpdateDatasetOperator(GoogleCloudBaseOperator):
     """
     Updates a dataset.
@@ -892,12 +728,6 @@ class AutoMLTablesUpdateDatasetOperator(GoogleCloudBaseOperator):
         self.retry = retry
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
-        raise AirflowException(
-            "Operator AutoMLTablesUpdateDatasetOperator has been deprecated due to shutdown of "
-            "a legacy version of AutoML Tables on March 31, 2024. "
-            "For additional information see: https://cloud.google.com/automl-tables/docs/deprecations. "
-            "Please use UpdateDatasetOperator from Vertex AI instead."
-        )
 
     def execute(self, context: Context):
         hook = CloudAutoMLHook(
@@ -924,6 +754,11 @@ class AutoMLTablesUpdateDatasetOperator(GoogleCloudBaseOperator):
         return Dataset.to_dict(result)
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLGetModelOperator(GoogleCloudBaseOperator):
     """
     Get Google Cloud AutoML model.
@@ -1003,8 +838,6 @@ 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:
@@ -1018,6 +851,12 @@ class AutoMLGetModelOperator(GoogleCloudBaseOperator):
         return model
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator, "
+    "airflow.providers.google.cloud.operators.translate.TranslateDeleteModelOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
     """
     Delete Google Cloud AutoML model.
@@ -1088,7 +927,7 @@ class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-        model: Model = hook.get_model(
+        hook.get_model(
             model_id=self.model_id,
             location=self.location,
             project_id=self.project_id,
@@ -1096,8 +935,6 @@ class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
             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,
@@ -1110,6 +947,11 @@ class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
         self.log.info("Deletion is completed")
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.endpoint_service.DeployModelOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
     """
     Deploys a model; if a model is already deployed, deploying it with the same parameters has no effect.
@@ -1187,13 +1029,6 @@ class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
         self.retry = retry
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
-        raise AirflowException(
-            "Operator AutoMLDeployModelOperator has been deprecated due to shutdown of "
-            "a legacy version of AutoML AutoML Natural Language, Vision, Video Intelligence "
-            "on March 31, 2024. "
-            "For additional information see: https://cloud.google.com/vision/automl/docs/deprecations. "
-            "Please use DeployModelOperator from Vertex AI instead."
-        )
 
     def execute(self, context: Context):
         hook = CloudAutoMLHook(
@@ -1214,6 +1049,11 @@ class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
         self.log.info("Model was deployed successfully.")
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    category=AirflowProviderDeprecationWarning,
+    reason="Shutdown of legacy version of AutoML Tables on March 31, 2024.",
+)
 class AutoMLTablesListTableSpecsOperator(GoogleCloudBaseOperator):
     """
     Lists table specs in a dataset.
@@ -1288,11 +1128,6 @@ class AutoMLTablesListTableSpecsOperator(GoogleCloudBaseOperator):
         self.retry = retry
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
-        raise AirflowException(
-            "Operator AutoMLTablesListTableSpecsOperator has been deprecated due to shutdown of "
-            "a legacy version of AutoML Tables on March 31, 2024. "
-            "For additional information see: https://cloud.google.com/automl-tables/docs/deprecations. "
-        )
 
     def execute(self, context: Context):
         hook = CloudAutoMLHook(
@@ -1324,6 +1159,12 @@ class AutoMLTablesListTableSpecsOperator(GoogleCloudBaseOperator):
         return result
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.ListDatasetsOperator, "
+    "airflow.providers.google.cloud.operators.translate.TranslateDatasetsListOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
     """
     Lists AutoML Datasets in project.
@@ -1399,14 +1240,7 @@ class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
         )
         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))
+            result.append(Dataset.to_dict(dataset))
         self.log.info("Datasets obtained.")
 
         self.xcom_push(
@@ -1420,6 +1254,12 @@ class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
         return result
 
 
+@deprecated(
+    planned_removal_date="September 30, 2025",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.ListDatasetsOperator, "
+    "airflow.providers.google.cloud.operators.translate.TranslateDatasetsListOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class AutoMLDeleteDatasetOperator(GoogleCloudBaseOperator):
     """
     Deletes a dataset and all of its contents.
@@ -1498,7 +1338,7 @@ class AutoMLDeleteDatasetOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-        dataset: Dataset = hook.get_dataset(
+        hook.get_dataset(
             dataset_id=self.dataset_id,
             location=self.location,
             project_id=self.project_id,
@@ -1506,8 +1346,6 @@ class AutoMLDeleteDatasetOperator(GoogleCloudBaseOperator):
             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)