[airflow] extend and fix AIR302 rules (#15525)

This commit is contained in:
Wei Lee 2025-01-16 17:40:00 +08:00 committed by GitHub
parent 79e52c7fdf
commit fc9dd63d64
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 592 additions and 569 deletions

View File

@ -66,7 +66,7 @@ from airflow.providers.openlineage.utils.utils import (
) )
from airflow.providers.postgres.datasets import postgres from airflow.providers.postgres.datasets import postgres
from airflow.providers.trino.datasets import trino from airflow.providers.trino.datasets import trino
from airflow.secrets.local_filesystem import get_connection, load_connections from airflow.secrets.local_filesystem import LocalFilesystemBackend, load_connections
from airflow.security.permissions import RESOURCE_DATASET from airflow.security.permissions import RESOURCE_DATASET
from airflow.sensors.base_sensor_operator import BaseSensorOperator from airflow.sensors.base_sensor_operator import BaseSensorOperator
from airflow.sensors.date_time_sensor import DateTimeSensor from airflow.sensors.date_time_sensor import DateTimeSensor
@ -84,6 +84,7 @@ from airflow.timetables.simple import DatasetTriggeredTimetable
from airflow.triggers.external_task import TaskStateTrigger from airflow.triggers.external_task import TaskStateTrigger
from airflow.utils import dates from airflow.utils import dates
from airflow.utils.dag_cycle_tester import test_cycle from airflow.utils.dag_cycle_tester import test_cycle
from airflow.utils.dag_parsing_context import get_parsing_context
from airflow.utils.dates import ( from airflow.utils.dates import (
date_range, date_range,
datetime_to_nano, datetime_to_nano,
@ -225,7 +226,10 @@ postgres.sanitize_uri
trino.sanitize_uri trino.sanitize_uri
# airflow.secrets # airflow.secrets
get_connection, load_connections # get_connection
lfb = LocalFilesystemBackend()
lfb.get_connections()
load_connections
# airflow.security.permissions # airflow.security.permissions
RESOURCE_DATASET RESOURCE_DATASET
@ -272,6 +276,9 @@ dates.datetime_to_nano
# airflow.utils.dag_cycle_tester # airflow.utils.dag_cycle_tester
test_cycle test_cycle
# airflow.utils.dag_parsing_context
get_parsing_context
# airflow.utils.decorators # airflow.utils.decorators
apply_defaults apply_defaults

View File

@ -294,6 +294,11 @@ fn check_method(checker: &mut Checker, call_expr: &ExprCall) {
)), )),
_ => None, _ => None,
}, },
["airflow", "secrets", "local_filesystem", "LocalFilesystemBackend"] => match attr.as_str()
{
"get_connections" => Some(Replacement::Name("get_connection")),
_ => None,
},
["airflow", "datasets", ..] | ["airflow", "Dataset"] => match attr.as_str() { ["airflow", "datasets", ..] | ["airflow", "Dataset"] => match attr.as_str() {
"iter_datasets" => Some(Replacement::Name("iter_assets")), "iter_datasets" => Some(Replacement::Name("iter_assets")),
"iter_dataset_aliases" => Some(Replacement::Name("iter_asset_aliases")), "iter_dataset_aliases" => Some(Replacement::Name("iter_asset_aliases")),
@ -543,12 +548,13 @@ fn check_name(checker: &mut Checker, expr: &Expr, range: TextRange) {
["airflow", "secrets", "local_filesystem", "load_connections"] => { ["airflow", "secrets", "local_filesystem", "load_connections"] => {
Replacement::Name("airflow.secrets.local_filesystem.load_connections_dict") Replacement::Name("airflow.secrets.local_filesystem.load_connections_dict")
} }
["airflow", "secrets", "local_filesystem", "get_connection"] => {
Replacement::Name("airflow.secrets.local_filesystem.load_connections_dict")
}
// airflow.utils.dag_parsing_context
["airflow", "utils", "dag_parsing_context", "get_parsing_context"] => {
Replacement::Name("airflow.sdk.get_parsing_context")
}
// airflow.utils.dates // airflow.utils.dates
["airflow", "utils", "dates", "date_range"] => Replacement::Name("airflow.timetables."), ["airflow", "utils", "dates", "date_range"] => Replacement::None,
["airflow", "utils", "dates", "days_ago"] => { ["airflow", "utils", "dates", "days_ago"] => {
Replacement::Name("pendulum.today('UTC').add(days=-N, ...)") Replacement::Name("pendulum.today('UTC').add(days=-N, ...)")
} }