VladaZakharova commented on code in PR #38635: URL: https://github.com/apache/airflow/pull/38635#discussion_r1547493116
########## airflow/providers/google/cloud/hooks/automl.py: ########## @@ -15,166 +15,511 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from __future__ import annotations +""" +This module contains a Google AutoML hook. -import warnings +.. spelling:word-list:: Review Comment: There is a separate file that should be used to list all the words ########## airflow/providers/google/cloud/hooks/automl.py: ########## @@ -15,166 +15,511 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from __future__ import annotations +""" +This module contains a Google AutoML hook. -import warnings +.. spelling:word-list:: -from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning + PredictResponse +""" +from __future__ import annotations -class CloudAutoMLHook: - """ - Former Google Cloud AutoML hook. +from functools import cached_property +from typing import TYPE_CHECKING, Sequence + +from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault +from google.cloud.automl_v1beta1 import ( + AutoMlClient, + BatchPredictInputConfig, + BatchPredictOutputConfig, + Dataset, + ExamplePayload, + InputConfig, + Model, + PredictionServiceClient, + PredictResponse, +) + +from airflow.exceptions import AirflowException +from airflow.providers.google.common.consts import CLIENT_INFO +from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook + +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.automl_v1beta1.services.auto_ml.pagers import ( + ListDatasetsPager, + ) + from google.protobuf.field_mask_pb2 import FieldMask - Deprecated as AutoML API becomes unusable starting March 31, 2024: - https://cloud.google.com/automl/docs - """ - deprecation_warning = ( - "CloudAutoMLHook has been deprecated, as AutoML API becomes unusable starting " - "March 31, 2024, and will be removed in future release. Please use an equivalent " - " Vertex AI hook available in" - "airflow.providers.google.cloud.hooks.vertex_ai instead." - ) +class CloudAutoMLHook(GoogleBaseHook): + """ + Google Cloud AutoML hook. - method_exception = "This method cannot be used as AutoML API becomes unusable." + All the methods in the hook where project_id is used must be called with + keyword arguments rather than positional. + """ - def __init__(self, **_) -> None: - warnings.warn(self.deprecation_warning, AirflowProviderDeprecationWarning) + def __init__( + self, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + if kwargs.get("delegate_to") is not None: + raise RuntimeError( + "The `delegate_to` parameter has been deprecated before and finally removed in this version" + " of Google Provider. You MUST convert it to `impersonate_chain`" + ) + super().__init__( + gcp_conn_id=gcp_conn_id, + impersonation_chain=impersonation_chain, + ) + self._client: AutoMlClient | None = None @staticmethod def extract_object_id(obj: dict) -> str: """Return unique id of the object.""" - warnings.warn( - "'extract_object_id' method is deprecated and will be removed in future release.", - AirflowProviderDeprecationWarning, - ) return obj["name"].rpartition("/")[-1] - def get_conn(self): - """ - Retrieve connection to AutoML (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def wait_for_operation(self, **_): - """ - Wait for long-lasting operation to complete (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def prediction_client(self, **_): - """ - Create a PredictionServiceClient (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def create_model(self, **_): - """ - Create a model_id and returns a Model in the `response` field when it completes (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def batch_predict(self, **_): - """ - Perform a batch prediction (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def predict(self, **_): - """ - Perform an online prediction (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def create_dataset(self, **_): - """ - Create a dataset (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def import_data(self, **_): - """ - Import data (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def list_column_specs(self, **_): - """ - List column specs (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def get_model(self, **_): - """ - Get a model (deprecated). - - :raises: AirflowException - """ - raise AirflowException(self.method_exception) - - def delete_model(self, **_): - """ - Delete a model (deprecated). + def get_conn(self) -> AutoMlClient: + """ + Retrieve connection to AutoML. + + :return: Google Cloud AutoML client object. + """ + if self._client is None: + self._client = AutoMlClient(credentials=self.get_credentials(), client_info=CLIENT_INFO) + return self._client + + def wait_for_operation(self, operation: Operation, timeout: float | None = None): + """Wait for long-lasting operation to complete.""" + try: + return operation.result(timeout=timeout) + except Exception: + error = operation.exception(timeout=timeout) + raise AirflowException(error) + + @cached_property + def prediction_client(self) -> PredictionServiceClient: + """ + Creates PredictionServiceClient. + + :return: Google Cloud AutoML PredictionServiceClient client object. + """ + return PredictionServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO) + + @GoogleBaseHook.fallback_to_default_project_id + def create_model( + self, + model: dict | Model, + location: str, + project_id: str = PROVIDE_PROJECT_ID, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + retry: Retry | _MethodDefault = DEFAULT, + ) -> Operation: + """ + Create a model_id and returns a Model in the `response` field when it completes. + + When you create a model, several model evaluations are created for it: + a global evaluation, and one evaluation for each annotation spec. + + :param model: The model_id to create. If a dict is provided, it must be of the same form + as the protobuf message `google.cloud.automl_v1beta1.types.Model` + :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 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. + + :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance + """ + client = self.get_conn() + parent = f"projects/{project_id}/locations/{location}" + return client.create_model( + request={"parent": parent, "model": model}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - :raises: AirflowException - """ - raise AirflowException(self.method_exception) + @GoogleBaseHook.fallback_to_default_project_id + def batch_predict( + self, + model_id: str, + input_config: dict | BatchPredictInputConfig, + output_config: dict | BatchPredictOutputConfig, + location: str, + project_id: str = PROVIDE_PROJECT_ID, + params: dict[str, str] | None = None, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Operation: + """ + Perform a batch prediction and returns a long-running operation object. + + Unlike the online `Predict`, batch prediction result won't be immediately + available in the response. Instead, a long-running operation object is returned. + + :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 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. + + :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance + """ + client = self.prediction_client + name = f"projects/{project_id}/locations/{location}/models/{model_id}" + result = client.batch_predict( + request={ + "name": name, + "input_config": input_config, + "output_config": output_config, + "params": params, + }, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def predict( + self, + model_id: str, + payload: dict | ExamplePayload, + location: str, + project_id: str = PROVIDE_PROJECT_ID, + params: dict[str, str] | None = None, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> PredictResponse: + """ + Perform an online prediction and returns the prediction result in the response. + + :param model_id: Name of the model_id requested to serve the prediction. + :param payload: Required. Payload to perform a prediction on. The payload must match the problem type + that the model_id was trained to solve. If a dict is provided, it must be of + the same form as the protobuf message `google.cloud.automl_v1beta1.types.ExamplePayload` + :param params: Additional domain-specific parameters, 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. + + :return: `google.cloud.automl_v1beta1.types.PredictResponse` instance + """ + client = self.prediction_client + name = f"projects/{project_id}/locations/{location}/models/{model_id}" + result = client.predict( + request={"name": name, "payload": payload, "params": params}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def create_dataset( + self, + dataset: dict | Dataset, + location: str, + project_id: str = PROVIDE_PROJECT_ID, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Dataset: + """ + Create a dataset. + + :param dataset: The dataset to create. If a dict is provided, it must be of the + same form as the protobuf message Dataset. + :param project_id: ID of the Google Cloud project where dataset 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. + + :return: `google.cloud.automl_v1beta1.types.Dataset` instance. + """ + client = self.get_conn() + parent = f"projects/{project_id}/locations/{location}" + result = client.create_dataset( + request={"parent": parent, "dataset": dataset}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result + + @GoogleBaseHook.fallback_to_default_project_id + def import_data( + self, + dataset_id: str, + location: str, + input_config: dict | InputConfig, + project_id: str = PROVIDE_PROJECT_ID, + retry: Retry | _MethodDefault = DEFAULT, + timeout: float | None = None, + metadata: Sequence[tuple[str, str]] = (), + ) -> Operation: + """ + Import data into a dataset. For Tables this method can only be called on an empty Dataset. + + :param dataset_id: Name of the AutoML dataset. + :param input_config: The desired input location and its domain specific semantics, if any. + If a dict is provided, it must be of the same form as the protobuf message InputConfig. + :param project_id: ID of the Google Cloud project where dataset 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. + + :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance + """ + client = self.get_conn() + name = f"projects/{project_id}/locations/{location}/datasets/{dataset_id}" + result = client.import_data( + request={"name": name, "input_config": input_config}, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + return result - def update_dataset(self, **_): + def list_column_specs(self, **kwargs) -> None: """ - Update a model (deprecated). + List column specs in a table spec (Deprecated). :raises: AirflowException """ - raise AirflowException(self.method_exception) - - def deploy_model(self, **_): - """ - Deploy a model (deprecated). + raise AirflowException( Review Comment: For deprecations we need to use decorator @deprecated(reason="some reason here", category=AirflowProviderDeprecationWarning, action="error") that will raise error as well as outputting beautiful message -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
