Visitar URL original
[`airflow`] Refine and add rules to capture deprecated arguments and … · astral-sh/ruff@e347d86 · GitHub
Skip to content

Commit e347d86

Browse files
sjyangkevinsjyangkevin
authored
[airflow] Refine and add rules to capture deprecated arguments and a decorator (AIR301) (#23170)
## Summary Related: apache/airflow#41641 (comment) The PR add/refine rules for the following function signature change in Airflow 3.0. ### `Variable.get()` According to the following documents: 2.11.0: https://airflow.apache.org/docs/apache-airflow/2.11.0/core-concepts/variables.html 3.0.6: https://airflow.apache.org/docs/apache-airflow/3.0.6/core-concepts/variables.html More context: apache/airflow#41641 (comment) In Airflow 2.x, the class method `get` of `Variable` from `airflow.models` has a `keyword` argument named `default_var`. ```python from airflow.models import Variable # Returns the value of default_var (None) if the variable is not set baz = Variable.get("baz", default_var=None) ``` In AIR311, there is a rule to detect the migration from `airflow.models.Variable` to `airflow.sdk.Variable`. However, in Airflow 3, the change of the keyword argument from `default_var` to `default` is not yet captured. The syntax in Airflow 3 is as follow: ```python from airflow.sdk import Variable # Returns the value of default (None) if the variable is not set baz = Variable.get("baz", default=None) ``` Therefore, a new rule is added in AIR303 to detect the use of `default_var` when `Variable` is imported from `airflow.sdk`. The rule suggests a fix to use the `default` as argument name instead of `default_var`. ### `provide_context` is deprecated from python operators Second, the `provide_context` parameter is deprecated in Airflow 2 for python operators. It is still a valid syntax to set it in Airflow 2, but it is removed in Airflow 3. A rule is added to detect the presence of the keyword argument in `PythonOperator` and `PythonVirtualenvOperator`. The constructor call only accept keyword-only arguments. The rule will raise a warning message when this keyword argument is passed. ```python # This is ok in Airflow 2, but invalid in Airflow 3 PythonOperator(task_id="task", python_callable=lambda: None, provide_context=True) ``` ### Deprecated value for `trigger_rule` argument in TaskFlow API and Airflow operators The current rule implemented in `check_name` only catches `TriggerRule.NONE_FAILED_OR_SKIPPED` via qualified name resolution. However, the argument is usually passed in the following way. (Here, `PythonOperator` is an example, all Airflow operators can be configured with a `trigger_rule`, since they inherit from `BaseOperator`). Therefore, the existing rule is removed and re-implemented. There are two cases will be captured. A string value is passed into the `@task` decorator or the operator, the rule will fix `none_failed_or_skipped` to `none_failed_min_one_success` and `TriggerRule.NONE_FAILED_OR_SKIPPED` to `TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS` ```python @task(trigger_rule="none_failed_or_skipped") def invalid_trigger_rule_task(): pass @task(trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED) def invalid_trigger_rule_task(): pass @task(trigger_rule="all_success") def valid_trigger_rule_task(): pass PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule="none_failed_or_skipped") PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED) PythonOperator(task_id="valid_trigger_rule_task", python_callable=lambda: None, trigger_rule="all_success") ``` ### Deprecated `apply_defaults` decorator Similar to the above case, this decorator is also checked via qualified name resolution. The decorator is used in the following way. The existing rule is also removed, and re-implemented to capture the below use case, and suggest a fix to remove this decorator entirely. ```python class DecoratedOperator(BaseOperator): # this is deprecated in Airflow 3 @apply_defaults def __init__(self, message, **kwargs): super().__init__(**kwargs) self.message = message def execute(self, context): print(self.message) ``` ## Test Plan New test cases have been added to `AIR301_args.py` and `AIR301_decorator.py`. The test snapshots are updated. <!-- How was it tested? -->
1 parent 45bd045 commit e347d86

8 files changed

Lines changed: 979 additions & 193 deletions

‎crates/ruff_linter/resources/test/fixtures/airflow/AIR301_args.py‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,3 +106,244 @@ def decorator_deprecated_operator_args():
106106
GCSTaskHandler(filename_template="/tmp/test")
107107

108108
FabAuthManager(None)
109+
110+
111+
# airflow.sdk.Variable
112+
from airflow.sdk import Variable
113+
114+
Variable.get("key", default_var="deprecated")
115+
Variable.get(key="key", default_var="deprecated")
116+
117+
Variable.get("key", default="default")
118+
Variable.get(key="key", default="default")
119+
Variable.get("key")
120+
Variable.get("key", deserialize_json=True)
121+
122+
123+
# airflow.providers.standard.operators.python (provide_context parameter removed)
124+
def virtualenv_callable():
125+
pass
126+
127+
from airflow.providers.standard.operators.python import (
128+
PythonOperator,
129+
BranchPythonOperator,
130+
ShortCircuitOperator,
131+
PythonVirtualenvOperator,
132+
BranchPythonVirtualenvOperator,
133+
ExternalPythonOperator,
134+
BranchExternalPythonOperator,
135+
)
136+
137+
PythonOperator(task_id="task", python_callable=lambda: None, provide_context=True)
138+
BranchPythonOperator(task_id="task", python_callable=lambda: None, provide_context=True)
139+
ShortCircuitOperator(task_id="task", python_callable=lambda: None, provide_context=True)
140+
PythonVirtualenvOperator(
141+
task_id="virtualenv_python",
142+
python_callable=virtualenv_callable,
143+
requirements=["colorama==0.4.0"],
144+
system_site_packages=False,
145+
provide_context=True,
146+
)
147+
BranchPythonVirtualenvOperator(
148+
task_id="branching_venv",
149+
requirements=["numpy~=1.26.0"],
150+
venv_cache_path="some_path",
151+
python_callable=virtualenv_callable,
152+
op_args=["some_args"],
153+
provide_context=True,
154+
)
155+
ExternalPythonOperator(
156+
task_id="external_python",
157+
python_callable=virtualenv_callable,
158+
python="some_path",
159+
provide_context=True,
160+
)
161+
BranchExternalPythonOperator(
162+
task_id="branching_ext_python",
163+
python="some_path",
164+
python_callable=virtualenv_callable,
165+
op_args=["some_args"],
166+
provide_context=True,
167+
)
168+
169+
PythonOperator(task_id="task", python_callable=lambda: None)
170+
BranchPythonOperator(task_id="task", python_callable=lambda: None)
171+
ShortCircuitOperator(task_id="task", python_callable=lambda: None)
172+
PythonVirtualenvOperator(
173+
task_id="virtualenv_python",
174+
python_callable=virtualenv_callable,
175+
requirements=["colorama==0.4.0"],
176+
system_site_packages=False,
177+
)
178+
BranchPythonVirtualenvOperator(
179+
task_id="branching_venv",
180+
requirements=["numpy~=1.26.0"],
181+
venv_cache_path="some_path",
182+
python_callable=virtualenv_callable,
183+
op_args=["some_args"],
184+
)
185+
ExternalPythonOperator(
186+
task_id="external_python",
187+
python_callable=virtualenv_callable,
188+
python="some_path",
189+
)
190+
BranchExternalPythonOperator(
191+
task_id="branching_ext_python",
192+
python="some_path",
193+
python_callable=virtualenv_callable,
194+
op_args=["some_args"],
195+
)
196+
197+
# airflow.operators.python (provide_context parameter removed)
198+
from airflow.operators.python import (
199+
PythonOperator,
200+
BranchPythonOperator,
201+
ShortCircuitOperator,
202+
PythonVirtualenvOperator,
203+
BranchPythonVirtualenvOperator,
204+
ExternalPythonOperator,
205+
BranchExternalPythonOperator,
206+
)
207+
208+
PythonOperator(task_id="task", python_callable=lambda: None, provide_context=True)
209+
BranchPythonOperator(task_id="task", python_callable=lambda: None, provide_context=True)
210+
ShortCircuitOperator(task_id="task", python_callable=lambda: None, provide_context=True)
211+
PythonVirtualenvOperator(
212+
task_id="virtualenv_python",
213+
python_callable=virtualenv_callable,
214+
requirements=["colorama==0.4.0"],
215+
system_site_packages=False,
216+
provide_context=True,
217+
)
218+
BranchPythonVirtualenvOperator(
219+
task_id="branching_venv",
220+
requirements=["numpy~=1.26.0"],
221+
venv_cache_path="some_path",
222+
python_callable=virtualenv_callable,
223+
op_args=["some_args"],
224+
provide_context=True,
225+
)
226+
ExternalPythonOperator(
227+
task_id="external_python",
228+
python_callable=virtualenv_callable,
229+
python="some_path",
230+
provide_context=True,
231+
)
232+
BranchExternalPythonOperator(
233+
task_id="branching_ext_python",
234+
python="some_path",
235+
python_callable=virtualenv_callable,
236+
op_args=["some_args"],
237+
provide_context=True,
238+
)
239+
240+
PythonOperator(task_id="task", python_callable=lambda: None)
241+
BranchPythonOperator(task_id="task", python_callable=lambda: None)
242+
ShortCircuitOperator(task_id="task", python_callable=lambda: None)
243+
PythonVirtualenvOperator(
244+
task_id="virtualenv_python",
245+
python_callable=virtualenv_callable,
246+
requirements=["colorama==0.4.0"],
247+
system_site_packages=False,
248+
)
249+
BranchPythonVirtualenvOperator(
250+
task_id="branching_venv",
251+
requirements=["numpy~=1.26.0"],
252+
venv_cache_path="some_path",
253+
python_callable=virtualenv_callable,
254+
op_args=["some_args"],
255+
)
256+
ExternalPythonOperator(
257+
task_id="external_python",
258+
python_callable=virtualenv_callable,
259+
python="some_path",
260+
)
261+
BranchExternalPythonOperator(
262+
task_id="branching_ext_python",
263+
python="some_path",
264+
python_callable=virtualenv_callable,
265+
op_args=["some_args"],
266+
)
267+
268+
269+
# airflow.utils.trigger_rule.TriggerRule.NONE_FAILED_OR_SKIPPED is deprecated.
270+
from airflow.decorators import task
271+
from airflow.operators.python import PythonOperator
272+
from airflow.utils.trigger_rule import TriggerRule
273+
274+
@task(trigger_rule="none_failed_or_skipped")
275+
def invalid_trigger_rule_task():
276+
pass
277+
278+
@task(trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED)
279+
def invalid_trigger_rule_task():
280+
pass
281+
282+
@task(trigger_rule="all_success")
283+
def valid_trigger_rule_task():
284+
pass
285+
286+
PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule="none_failed_or_skipped")
287+
288+
PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED)
289+
290+
PythonOperator(task_id="valid_trigger_rule_task", python_callable=lambda: None, trigger_rule="all_success")
291+
292+
from airflow.sdk import task
293+
from airflow.providers.standard.operators.python import PythonOperator
294+
295+
@task(trigger_rule="none_failed_or_skipped")
296+
def invalid_trigger_rule_task():
297+
pass
298+
299+
@task(trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED)
300+
def invalid_trigger_rule_task():
301+
pass
302+
303+
@task(trigger_rule="all_success")
304+
def valid_trigger_rule_task():
305+
pass
306+
307+
PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule="none_failed_or_skipped")
308+
309+
PythonOperator(task_id="invalid_trigger_rule_task", python_callable=lambda: None, trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED)
310+
311+
PythonOperator(task_id="valid_trigger_rule_task", python_callable=lambda: None, trigger_rule="all_success")
312+
313+
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
314+
315+
execute_query = SQLExecuteQueryOperator(
316+
task_id="execute_query",
317+
sql="SELECT 1; SELECT * FROM AIRFLOW_DB_METADATA_TABLE LIMIT 1;",
318+
split_statements=True,
319+
return_last=False,
320+
trigger_rule="none_failed_or_skipped",
321+
)
322+
323+
execute_query = SQLExecuteQueryOperator(
324+
task_id="execute_query",
325+
sql="SELECT 1; SELECT * FROM AIRFLOW_DB_METADATA_TABLE LIMIT 1;",
326+
split_statements=True,
327+
return_last=False,
328+
trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED,
329+
)
330+
331+
from airflow.providers.amazon.aws.operators.s3 import S3FileTransformOperator
332+
333+
file_transform = S3FileTransformOperator(
334+
task_id="file_transform",
335+
source_s3_key="s3://bucket/key",
336+
dest_s3_key="s3://bucket_2/key_2",
337+
transform_script="cp",
338+
replace=True,
339+
trigger_rule="none_failed_or_skipped",
340+
)
341+
342+
file_transform = S3FileTransformOperator(
343+
task_id="file_transform",
344+
source_s3_key="s3://bucket/key",
345+
dest_s3_key="s3://bucket_2/key_2",
346+
transform_script="cp",
347+
replace=True,
348+
trigger_rule=TriggerRule.NONE_FAILED_OR_SKIPPED,
349+
)

‎crates/ruff_linter/resources/test/fixtures/airflow/AIR301_decorator.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
"""Test case for detecting @apply_defaults decorator usage in AIR301."""
2+
from airflow.models.baseoperator import BaseOperator
3+
from airflow.utils.decorators import apply_defaults
4+
5+
6+
class DecoratedOperator(BaseOperator):
7+
@apply_defaults
8+
def __init__(self, message, **kwargs):
9+
super().__init__(**kwargs)
10+
self.message = message
11+
12+
def execute(self, context):
13+
print(self.message)
14+
15+
class ExampleOperator(BaseOperator):
16+
def __init__(self, message, **kwargs):
17+
super().__init__(**kwargs)
18+
self.message = message
19+
20+
def execute(self, context):
21+
print(self.message)

‎crates/ruff_linter/resources/test/fixtures/airflow/AIR301_names.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
scale_time_units,
2929
)
3030
from airflow.utils.db import create_session
31-
from airflow.utils.decorators import apply_defaults
3231
from airflow.utils.file import mkdirs
3332
from airflow.utils.state import SHUTDOWN, terminating_states
3433
from airflow.utils.trigger_rule import TriggerRule
@@ -86,8 +85,6 @@
8685
# airflow.utils.db
8786
create_session
8887

89-
# airflow.utils.decorators
90-
apply_defaults
9188

9289
# airflow.utils.file
9390
mkdirs
@@ -99,7 +96,6 @@
9996

10097
# airflow.utils.trigger_rule
10198
TriggerRule.DUMMY
102-
TriggerRule.NONE_FAILED_OR_SKIPPED
10399

104100

105101
# airflow.www.auth

‎crates/ruff_linter/src/rules/airflow/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ mod tests {
2323
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_class_attribute.py"))]
2424
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_airflow_plugin.py"))]
2525
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_context.py"))]
26+
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_decorator.py"))]
2627
#[test_case(Rule::Airflow31Moved, Path::new("AIR321_names.py"))]
2728
#[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_amazon.py"))]
2829
#[test_case(Rule::Airflow3MovedToProvider, Path::new("AIR302_celery.py"))]

0 commit comments

Comments
 (0)