Skip to content

Commit 6c37b86

Browse files
e-galane-galan
committed
Reroute AutoML operator links to Google Translation links
1 parent f509b0a commit 6c37b86

7 files changed

Lines changed: 445 additions & 36 deletions

File tree

‎airflow/providers/google/cloud/links/automl.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@
2121

2222
from typing import TYPE_CHECKING
2323

24+
from deprecated import deprecated
25+
26+
from airflow.exceptions import AirflowProviderDeprecationWarning
2427
from airflow.providers.google.cloud.links.base import BaseGoogleLink
2528

2629
if TYPE_CHECKING:
@@ -44,6 +47,13 @@
4447
)
4548

4649

50+
@deprecated(
51+
reason=(
52+
"Class `AutoMLDatasetLink` has been deprecated and will be removed after 31.12.2024. "
53+
"Please use `TranslationLegacyDatasetLink` from `airflow/providers/google/cloud/links/translate.py` instead."
54+
),
55+
category=AirflowProviderDeprecationWarning,
56+
)
4757
class AutoMLDatasetLink(BaseGoogleLink):
4858
"""Helper class for constructing AutoML Dataset link."""
4959

@@ -65,6 +75,13 @@ def persist(
6575
)
6676

6777

78+
@deprecated(
79+
reason=(
80+
"Class `AutoMLDatasetListLink` has been deprecated and will be removed after 31.12.2024. "
81+
"Please use `TranslationDatasetListLink` from `airflow/providers/google/cloud/links/translate.py` instead."
82+
),
83+
category=AirflowProviderDeprecationWarning,
84+
)
6885
class AutoMLDatasetListLink(BaseGoogleLink):
6986
"""Helper class for constructing AutoML Dataset List link."""
7087

@@ -87,6 +104,13 @@ def persist(
87104
)
88105

89106

107+
@deprecated(
108+
reason=(
109+
"Class `AutoMLModelLink` has been deprecated and will be removed after 31.12.2024. "
110+
"Please use `TranslationLegacyModelLink` from `airflow/providers/google/cloud/links/translate.py` instead."
111+
),
112+
category=AirflowProviderDeprecationWarning,
113+
)
90114
class AutoMLModelLink(BaseGoogleLink):
91115
"""Helper class for constructing AutoML Model link."""
92116

@@ -114,6 +138,13 @@ def persist(
114138
)
115139

116140

141+
@deprecated(
142+
reason=(
143+
"Class `AutoMLModelTrainLink` has been deprecated and will be removed after 31.12.2024. "
144+
"Please use `TranslationLegacyModelTrainLink` from `airflow/providers/google/cloud/links/translate.py` instead."
145+
),
146+
category=AirflowProviderDeprecationWarning,
147+
)
117148
class AutoMLModelTrainLink(BaseGoogleLink):
118149
"""Helper class for constructing AutoML Model Train link."""
119150

@@ -138,6 +169,13 @@ def persist(
138169
)
139170

140171

172+
@deprecated(
173+
reason=(
174+
"Class `AutoMLModelPredictLink` has been deprecated and will be removed after 31.12.2024. "
175+
"Please use `TranslationLegacyModelPredictLink` from `airflow/providers/google/cloud/links/translate.py` instead."
176+
),
177+
category=AirflowProviderDeprecationWarning,
178+
)
141179
class AutoMLModelPredictLink(BaseGoogleLink):
142180
"""Helper class for constructing AutoML Model Predict link."""
143181

‎airflow/providers/google/cloud/links/translate.py‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""This module contains Google Translate links."""
18+
19+
from __future__ import annotations
20+
21+
from typing import TYPE_CHECKING
22+
23+
from airflow.providers.google.cloud.links.base import BASE_LINK, BaseGoogleLink
24+
25+
if TYPE_CHECKING:
26+
from airflow.utils.context import Context
27+
28+
29+
TRANSLATION_BASE_LINK = BASE_LINK + "/translation"
30+
TRANSLATION_LEGACY_DATASET_LINK = (
31+
TRANSLATION_BASE_LINK + "/locations/{location}/datasets/{dataset_id}/sentences?project={project_id}"
32+
)
33+
TRANSLATION_DATASET_LIST_LINK = TRANSLATION_BASE_LINK + "/datasets?project={project_id}"
34+
TRANSLATION_LEGACY_MODEL_LINK = (
35+
TRANSLATION_BASE_LINK
36+
+ "/locations/{location}/datasets/{dataset_id}/evaluate;modelId={model_id}?project={project_id}"
37+
)
38+
TRANSLATION_LEGACY_MODEL_TRAIN_LINK = (
39+
TRANSLATION_BASE_LINK + "/locations/{location}/datasets/{dataset_id}/train?project={project_id}"
40+
)
41+
TRANSLATION_LEGACY_MODEL_PREDICT_LINK = (
42+
TRANSLATION_BASE_LINK
43+
+ "/locations/{location}/datasets/{dataset_id}/predict;modelId={model_id}?project={project_id}"
44+
)
45+
46+
47+
class TranslationLegacyDatasetLink(BaseGoogleLink):
48+
"""
49+
Helper class for constructing Legacy Translation Dataset link.
50+
51+
Legacy Datasets are created and managed by AutoML API.
52+
"""
53+
54+
name = "Translation Legacy Dataset"
55+
key = "translation_legacy_dataset"
56+
format_str = TRANSLATION_LEGACY_DATASET_LINK
57+
58+
@staticmethod
59+
def persist(
60+
context: Context,
61+
task_instance,
62+
dataset_id: str,
63+
project_id: str,
64+
):
65+
task_instance.xcom_push(
66+
context,
67+
key=TranslationLegacyDatasetLink.key,
68+
value={"location": task_instance.location, "dataset_id": dataset_id, "project_id": project_id},
69+
)
70+
71+
72+
class TranslationDatasetListLink(BaseGoogleLink):
73+
"""Helper class for constructing Translation Dataset List link."""
74+
75+
name = "Translation Dataset List"
76+
key = "translation_dataset_list"
77+
format_str = TRANSLATION_DATASET_LIST_LINK
78+
79+
@staticmethod
80+
def persist(
81+
context: Context,
82+
task_instance,
83+
project_id: str,
84+
):
85+
task_instance.xcom_push(
86+
context,
87+
key=TranslationDatasetListLink.key,
88+
value={
89+
"project_id": project_id,
90+
},
91+
)
92+
93+
94+
class TranslationLegacyModelLink(BaseGoogleLink):
95+
"""
96+
Helper class for constructing Translation Legacy Model link.
97+
98+
Legacy Models are created and managed by AutoML API.
99+
"""
100+
101+
name = "Translation Legacy Model"
102+
key = "translation_legacy_model"
103+
format_str = TRANSLATION_LEGACY_MODEL_LINK
104+
105+
@staticmethod
106+
def persist(
107+
context: Context,
108+
task_instance,
109+
dataset_id: str,
110+
model_id: str,
111+
project_id: str,
112+
):
113+
task_instance.xcom_push(
114+
context,
115+
key=TranslationLegacyModelLink.key,
116+
value={
117+
"location": task_instance.location,
118+
"dataset_id": dataset_id,
119+
"model_id": model_id,
120+
"project_id": project_id,
121+
},
122+
)
123+
124+
125+
class TranslationLegacyModelTrainLink(BaseGoogleLink):
126+
"""
127+
Helper class for constructing Translation Legacy Model Train link.
128+
129+
Legacy Models are created and managed by AutoML API.
130+
"""
131+
132+
name = "Translation Legacy Model Train"
133+
key = "translation_legacy_model_train"
134+
format_str = TRANSLATION_LEGACY_MODEL_TRAIN_LINK
135+
136+
@staticmethod
137+
def persist(
138+
context: Context,
139+
task_instance,
140+
project_id: str,
141+
):
142+
task_instance.xcom_push(
143+
context,
144+
key=TranslationLegacyModelTrainLink.key,
145+
value={
146+
"location": task_instance.location,
147+
"dataset_id": task_instance.model["dataset_id"],
148+
"project_id": project_id,
149+
},
150+
)
151+
152+
153+
class TranslationLegacyModelPredictLink(BaseGoogleLink):
154+
"""
155+
Helper class for constructing Translation Legacy Model Predict link.
156+
157+
Legacy Models are created and managed by AutoML API.
158+
"""
159+
160+
name = "Translation Legacy Model Predict"
161+
key = "translation_legacy_model_predict"
162+
format_str = TRANSLATION_LEGACY_MODEL_PREDICT_LINK
163+
164+
@staticmethod
165+
def persist(
166+
context: Context,
167+
task_instance,
168+
model_id: str,
169+
project_id: str,
170+
):
171+
task_instance.xcom_push(
172+
context,
173+
key=TranslationLegacyModelPredictLink.key,
174+
value={
175+
"location": task_instance.location,
176+
"dataset_id": task_instance.model.dataset_id,
177+
"model_id": model_id,
178+
"project_id": project_id,
179+
},
180+
)

0 commit comments

Comments
 (0)