Skip to content

Commit

Permalink
Add deprecation warnings and raise exception for already deprecated ones
Browse files Browse the repository at this point in the history
  • Loading branch information
@molcay
molcay committed Apr 30, 2024
1 parent 546199a commit c7748ae
Show file tree
Hide file tree
Showing 2 changed files with 95 additions and 50 deletions.
56 changes: 28 additions & 28 deletions airflow/providers/google/cloud/operators/automl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,8 @@ def _raise_exception_for_deprecated_operator(
alternative_class_name_str += f" or `{alternative_class_names[-1]}`"

raise AirflowException(
f"{deprecated_class_name} for text, image and video prediction is "
f"deprecated and will be removed after 31.03.2024. All the functionality of "
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."
Expand All @@ -80,8 +80,8 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
"""
Creates Google Cloud AutoML model.
AutoMLTrainModelOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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`,
Expand Down Expand Up @@ -276,8 +276,8 @@ def hook(self) -> CloudAutoMLHook | PredictionServiceHook:
)

def _check_model_type(self):
hook: CloudAutoMLHook = self.hook
model: Model = hook.get_model(
hook = self.hook
model = hook.get_model(
model_id=self.model_id,
location=self.location,
project_id=self.project_id,
Expand All @@ -287,8 +287,8 @@ def _check_model_type(self):
)
if not hasattr(model, "translation_model_metadata"):
raise AirflowException(
"AutoMLPredictOperator for text, image and video prediction is "
"deprecated. Please use endpoint_id param instead of model_id param."
"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):
Expand Down Expand Up @@ -338,8 +338,8 @@ class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
"""
Perform a batch prediction on Google Cloud AutoML.
AutoMLBatchPredictOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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`,
Expand Down Expand Up @@ -434,7 +434,7 @@ def execute(self, context: Context):
timeout=self.timeout,
metadata=self.metadata,
)
print(model, hasattr(model, "translation_model_metadata"), hasattr(model, "tables_model_metadata"))

if not hasattr(model, "translation_model_metadata"):
_raise_exception_for_deprecated_operator(
self.__class__.__name__,
Expand Down Expand Up @@ -475,8 +475,8 @@ class AutoMLCreateDatasetOperator(GoogleCloudBaseOperator):
"""
Creates a Google Cloud AutoML dataset.
AutoMLCreateDatasetOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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::
Expand Down Expand Up @@ -573,8 +573,8 @@ class AutoMLImportDataOperator(GoogleCloudBaseOperator):
"""
Imports data to a Google Cloud AutoML dataset.
AutoMLImportDataOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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::
Expand Down Expand Up @@ -788,7 +788,7 @@ def execute(self, context: Context):

@deprecated(
reason=(
"Class `AutoMLTablesUpdateDatasetOperator` is deprecated and will be removed after 31.03.2024. "
"Class `AutoMLTablesUpdateDatasetOperator` has been deprecated and no longer available. "
"Please use `UpdateDatasetOperator` instead"
),
category=AirflowProviderDeprecationWarning,
Expand All @@ -798,7 +798,7 @@ class AutoMLTablesUpdateDatasetOperator(GoogleCloudBaseOperator):
"""
Updates a dataset.
AutoMLTablesUpdateDatasetOperator is deprecated and will be removed after 31.03.2024. Please use
AutoMLTablesUpdateDatasetOperator has been deprecated and no longer available. Please use
:class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.UpdateDatasetOperator`
instead.
Expand Down Expand Up @@ -889,8 +889,8 @@ class AutoMLGetModelOperator(GoogleCloudBaseOperator):
"""
Get Google Cloud AutoML model.
AutoMLGetModelOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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::
Expand Down Expand Up @@ -982,8 +982,8 @@ class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
"""
Delete Google Cloud AutoML model.
AutoMLDeleteModelOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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::
Expand Down Expand Up @@ -1071,7 +1071,7 @@ def execute(self, context: Context):

@deprecated(
reason=(
"Class `AutoMLDeployModelOperator` is deprecated and will be removed after 31.03.2024. Please use "
"Class `AutoMLDeployModelOperator` has been deprecated and no longer available. Please use "
"`DeployModelOperator` instead"
),
category=AirflowProviderDeprecationWarning,
Expand All @@ -1087,7 +1087,7 @@ class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
Only applicable for Text Classification, Image Object Detection and Tables; all other
domains manage deployment automatically.
AutoMLDeployModelOperator is deprecated and will be removed after 31.03.2024. Please use
AutoMLDeployModelOperator has been deprecated and no longer available. Please use
:class:`airflow.providers.google.cloud.operators.vertex_ai.endpoint_service.DeployModelOperator`
instead.
Expand Down Expand Up @@ -1286,8 +1286,8 @@ class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
"""
Lists AutoML Datasets in project.
AutoMLListDatasetOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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::
Expand Down Expand Up @@ -1358,7 +1358,7 @@ def execute(self, context: Context):
for dataset in page_iterator:
if not hasattr(dataset, "translation_dataset_metadata"):
warnings.warn(
"Class `AutoMLListDatasetOperator` is deprecated and will be removed after 31.03.2024. "
"Class `AutoMLListDatasetOperator` has been deprecated and no longer available. "
"Please use `ListDatasetsOperator` instead.",
stacklevel=2,
)
Expand All @@ -1381,8 +1381,8 @@ class AutoMLDeleteDatasetOperator(GoogleCloudBaseOperator):
"""
Deletes a dataset and all of its contents.
AutoMLDeleteDatasetOperator for tables, video intelligence, vision and natural language is deprecated
and will be removed after 31.03.2024. Please use
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::
Expand Down
89 changes: 67 additions & 22 deletions tests/providers/google/cloud/operators/test_automl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

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.operators.automl import (
AutoMLBatchPredictOperator,
AutoMLCreateDatasetOperator,
Expand Down Expand Up @@ -110,8 +111,8 @@ def test_execute_deprecated(self, mock_hook):
task_id=TASK_ID,
)
expected_exception_str = (
"AutoMLTrainModelOperator for text, image and video prediction is "
"deprecated and will be removed after"
"AutoMLTrainModelOperator for text, image, and video prediction has been "
"deprecated and no longer available"
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
Expand Down Expand Up @@ -183,8 +184,8 @@ def test_execute_deprecated(self, mock_hook):
prediction_params={},
)
expected_exception_str = (
"AutoMLBatchPredictOperator for text, image and video prediction is "
"deprecated and will be removed after"
"AutoMLBatchPredictOperator for text, image, and video prediction has been "
"deprecated and no longer available"
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
Expand Down Expand Up @@ -271,6 +272,51 @@ def test_templating(self, create_task_instance_of_operator):
assert task.location == "location"
assert task.impersonation_chain == "impersonation-chain"

@mock.patch("airflow.providers.google.cloud.operators.automl.CloudAutoMLHook")
def test_execute_deprecation(self, mock_hook):
returned_model = mock.MagicMock(**MODEL_DEPRECATED)
del returned_model.translation_model_metadata
mock_hook.return_value.get_model.return_value = returned_model

mock_hook.return_value.predict.return_value = PredictResponse()

op = AutoMLPredictOperator(
model_id=MODEL_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
payload=PAYLOAD,
task_id=TASK_ID,
operation_params={"TEST_KEY": "TEST_VALUE"},
)
expected_exception_str = (
"AutoMLPredictOperator for text, image, and video prediction has been "
"deprecated. Please use endpoint_id param instead of model_id param."
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
mock_hook.return_value.predict.assert_not_called()

def test_hook_type(self):
op = AutoMLPredictOperator(
model_id=MODEL_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
payload=PAYLOAD,
task_id=TASK_ID,
operation_params={"TEST_KEY": "TEST_VALUE"},
)
assert isinstance(op.hook, CloudAutoMLHook)

op = AutoMLPredictOperator(
endpoint_id="endpoint_id",
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
payload=PAYLOAD,
task_id=TASK_ID,
operation_params={"TEST_KEY": "TEST_VALUE"},
)
assert isinstance(op.hook, PredictionServiceHook)


class TestAutoMLCreateImportOperator:
@mock.patch("airflow.providers.google.cloud.operators.automl.CloudAutoMLHook")
Expand Down Expand Up @@ -303,8 +349,8 @@ def test_execute_deprecated(self, mock_hook):
task_id=TASK_ID,
)
expected_exception_str = (
"AutoMLCreateDatasetOperator for text, image and video prediction is "
"deprecated and will be removed after"
"AutoMLCreateDatasetOperator for text, image, and video prediction has been "
"deprecated and no longer available"
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
Expand Down Expand Up @@ -401,8 +447,7 @@ def test_execute(self, mock_hook):

expected_exception_str = (
r"Call to deprecated class AutoMLTablesUpdateDatasetOperator. \(Class "
r"`AutoMLTablesUpdateDatasetOperator` is deprecated and will be removed "
r"after"
r"`AutoMLTablesUpdateDatasetOperator` has been deprecated and no longer available"
)
with pytest.raises(AirflowProviderDeprecationWarning, match=expected_exception_str):
AutoMLTablesUpdateDatasetOperator(
Expand Down Expand Up @@ -430,7 +475,7 @@ def test_templating(self, create_task_instance_of_operator):
)
assert str(err.value).startswith(
"Call to deprecated class AutoMLTablesUpdateDatasetOperator. "
"(Class `AutoMLTablesUpdateDatasetOperator` is deprecated and will be removed after"
"(Class `AutoMLTablesUpdateDatasetOperator` has been deprecated and no longer available"
)


Expand Down Expand Up @@ -469,8 +514,8 @@ def test_execute_deprecated(self, mock_hook):
task_id=TASK_ID,
)
expected_exception_str = (
"AutoMLGetModelOperator for text, image and video prediction is "
"deprecated and will be removed after"
"AutoMLGetModelOperator for text, image, and video prediction has been "
"deprecated and no longer available"
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
Expand Down Expand Up @@ -537,8 +582,8 @@ def test_execute_deprecated(self, mock_hook):
task_id=TASK_ID,
)
expected_exception_str = (
"AutoMLDeleteModelOperator for text, image and video prediction is "
"deprecated and will be removed after"
"AutoMLDeleteModelOperator for text, image, and video prediction has been "
"deprecated and no longer available"
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
Expand Down Expand Up @@ -580,8 +625,8 @@ def test_execute(self, mock_hook):
image_detection_metadata = {}

expected_exception_str = (
r"Call to deprecated class AutoMLDeployModelOperator. \(Class `AutoMLDeployModelOperator` is "
"deprecated and will be removed after"
r"Call to deprecated class AutoMLDeployModelOperator. \(Class `AutoMLDeployModelOperator` has "
r"been deprecated and no longer available"
)
with pytest.raises(AirflowProviderDeprecationWarning, match=expected_exception_str):
AutoMLDeployModelOperator(
Expand Down Expand Up @@ -612,7 +657,7 @@ def test_templating(self, create_task_instance_of_operator):

assert str(err.value).startswith(
"Call to deprecated class AutoMLDeployModelOperator. "
"(Class `AutoMLDeployModelOperator` is deprecated and will be removed after"
"(Class `AutoMLDeployModelOperator` has been deprecated and no longer available"
)


Expand Down Expand Up @@ -651,8 +696,8 @@ def test_execute_deprecated(self, mock_hook):
task_id=TASK_ID,
)
expected_exception_str = (
"AutoMLImportDataOperator for text, image and video prediction is "
"deprecated and will be removed after"
"AutoMLImportDataOperator for text, image, and video prediction has been "
"deprecated and no longer available"
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
Expand Down Expand Up @@ -760,8 +805,8 @@ def test_execute_deprecated(self, mock_hook):
mock_hook.return_value.list_datasets.return_value = [DATASET, not_valid_dataset]
op = AutoMLListDatasetOperator(location=GCP_LOCATION, project_id=GCP_PROJECT_ID, task_id=TASK_ID)
expected_warning_str = (
"Class `AutoMLListDatasetOperator` is deprecated and will be removed after "
"31.03.2024. Please use `ListDatasetsOperator` instead"
"Class `AutoMLListDatasetOperator` has been deprecated and no longer available. "
"Please use `ListDatasetsOperator` instead"
)
with pytest.warns(UserWarning, match=expected_warning_str):
op.execute(context=mock.MagicMock())
Expand Down Expand Up @@ -826,8 +871,8 @@ def test_execute_deprecated(self, mock_hook):
task_id=TASK_ID,
)
expected_exception_str = (
"AutoMLDeleteDatasetOperator for text, image and video prediction is "
"deprecated and will be removed after"
"AutoMLDeleteDatasetOperator for text, image, and video prediction has been "
"deprecated and no longer available"
)
with pytest.raises(AirflowException, match=expected_exception_str):
op.execute(context=mock.MagicMock())
Expand Down

0 comments on commit c7748ae

Please sign in to comment.