Skip to content

Commit 2f3baf1

Browse files
Lee-WLee-W
committed
feat(AIR301): apply AutoImport to cases that naes has been changed in Airflow 3
1 parent bcd4be6 commit 2f3baf1

7 files changed

+841
-275
lines changed

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

Lines changed: 22 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,51 +1,45 @@
11
from __future__ import annotations
22

3-
from airflow.auth.managers.base_auth_manager import is_authorized_dataset
3+
from airflow.api_connexion.security import requires_access_dataset
4+
from airflow.auth.managers.models.resource_details import (
5+
DatasetDetails,
6+
is_authorized_dataset,
7+
)
48
from airflow.datasets.manager import (
59
DatasetManager,
610
dataset_manager,
711
resolve_dataset_manager,
812
)
913
from airflow.lineage.hook import DatasetLineageInfo
10-
from airflow.listeners.spec.dataset import (
11-
on_dataset_changed,
12-
on_dataset_created,
13-
)
14-
from airflow.metrics.validators import (
15-
AllowListValidator,
16-
BlockListValidator,
17-
)
18-
from airflow.secrets.local_filesystem import load_connections
14+
from airflow.metrics.validators import AllowListValidator, BlockListValidator
15+
from airflow.secrets.local_filesystm import load_connections
1916
from airflow.security.permissions import RESOURCE_DATASET
20-
from airflow.timetables.simple import DatasetTriggeredTimetable
2117
from airflow.www.auth import has_access_dataset
2218

23-
# airflow.datasets.manager
24-
DatasetManager()
25-
dataset_manager
26-
resolve_dataset_manager
19+
requires_access_dataset()
2720

28-
# airflow.auth.managers
29-
is_authorized_dataset
3021
DatasetDetails()
22+
is_authorized_dataset()
23+
24+
DatasetManager()
25+
dataset_manager()
26+
resolve_dataset_manager()
3127

32-
# airflow.lineage.hook
3328
DatasetLineageInfo()
3429

35-
# airflow.listeners.spec.dataset
36-
on_dataset_changed
37-
on_dataset_created
38-
# airflow.metrics.validators
3930
AllowListValidator()
4031
BlockListValidator()
4132

42-
# airflow.secrets.local_filesystem
4333
load_connections()
4434

45-
# airflow.security.permissions
4635
RESOURCE_DATASET
47-
# airflow.timetables
48-
DatasetTriggeredTimetable()
4936

50-
# airflow.www.auth
51-
has_access_dataset
37+
has_access_dataset()
38+
39+
from airflow.listeners.spec.dataset import (
40+
on_dataset_changed,
41+
on_dataset_created,
42+
)
43+
44+
on_dataset_created()
45+
on_dataset_changed()

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
from __future__ import annotations
2+
3+
from airflow.providers.amazon.aws.auth_manager.avp.entities.AvpEntities import DATASET
4+
from airflow.providers.amazon.aws.datasets.s3 import (
5+
convert_dataset_to_openlineage as s3_convert_dataset_to_openlineage,
6+
)
7+
from airflow.providers.amazon.aws.datasets.s3 import create_dataset as s3_create_dataset
8+
from airflow.providers.common.io.dataset.file import (
9+
convert_dataset_to_openlineage as io_convert_dataset_to_openlineage,
10+
)
11+
from airflow.providers.common.io.dataset.file import create_dataset as io_create_dataset
12+
from airflow.providers.fab.auth_manager.fab_auth_manager import is_authorized_dataset as fab_is_authorized_dataset
13+
from airflow.providers.google.datasets.bigquery import (
14+
create_dataset as bigquery_create_dataset,
15+
)
16+
from airflow.providers.google.datasets.gcs import (
17+
convert_dataset_to_openlineage as gcs_convert_dataset_to_openlineage,
18+
)
19+
from airflow.providers.google.datasets.gcs import create_dataset as gcs_create_dataset
20+
from airflow.providers.openlineage.utils.utils import (
21+
DatasetInfo,
22+
translate_airflow_dataset,
23+
)
24+
25+
DATASET
26+
27+
s3_create_dataset()
28+
s3_convert_dataset_to_openlineage()
29+
30+
io_create_dataset()
31+
io_convert_dataset_to_openlineage()
32+
33+
fab_is_authorized_dataset()
34+
35+
# airflow.providers.google.datasets.bigquery
36+
bigquery_create_dataset()
37+
# airflow.providers.google.datasets.gcs
38+
gcs_create_dataset()
39+
gcs_convert_dataset_to_openlineage()
40+
# airflow.providers.openlineage.utils.utils
41+
DatasetInfo()
42+
translate_airflow_dataset()
43+
#
44+
# airflow.secrets.local_filesystem
45+
load_connections()
46+
#
47+
# airflow.security.permissions
48+
RESOURCE_DATASET
49+
50+
# airflow.timetables
51+
DatasetTriggeredTimetable()
52+
#
53+
# airflow.www.auth
54+
has_access_dataset

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ mod tests {
1818
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_args.py"))]
1919
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_names.py"))]
2020
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_names_fix.py"))]
21+
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_provider_names_fix.py"))]
2122
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_names_try.py"))]
2223
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_class_attribute.py"))]
2324
#[test_case(Rule::Airflow3Removal, Path::new("AIR301_airflow_plugin.py"))]

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

Lines changed: 49 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -756,7 +756,7 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
756756
}
757757
["airflow", "www", "auth", "has_access_dataset"] => Replacement::AutoImport {
758758
module: "airflow.www.auth",
759-
name: "has_access_dataset",
759+
name: "has_access_asset",
760760
},
761761
["airflow", "www", "utils", "get_sensitive_variables_fields"] => {
762762
Replacement::Name("airflow.utils.log.secrets_masker.get_sensitive_variables_fields")
@@ -766,31 +766,38 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
766766
}
767767

768768
// airflow.providers.amazon
769-
["airflow", "providers", "amazon", "aws", rest @ ..] => match &rest {
770-
["datasets", "s3", "create_dataset"] => {
771-
Replacement::Name("airflow.providers.amazon.aws.assets.s3.create_asset")
772-
}
773-
["datasets", "s3", "convert_dataset_to_openlineage"] => Replacement::Name(
774-
"airflow.providers.amazon.aws.assets.s3.convert_asset_to_openlineage",
775-
),
776-
["datasets", "s3", "sanitize_uri"] => {
769+
["airflow", "providers", "amazon", "aws", "datasets", "s3", rest] => match *rest {
770+
"create_dataset" => Replacement::AutoImport {
771+
module: "airflow.providers.amazon.aws.assets.s3",
772+
name: "create_asset",
773+
},
774+
"convert_dataset_to_openlineage" => Replacement::AutoImport {
775+
module: "airflow.providers.amazon.aws.assets.s3",
776+
name: "convert_asset_to_openlineage",
777+
},
778+
"sanitize_uri" => {
777779
Replacement::Name("airflow.providers.amazon.aws.assets.s3.sanitize_uri")
778780
}
779-
["auth_manager", "avp", "entities", "AvpEntities", "DATASET"] => Replacement::Name(
780-
"airflow.providers.amazon.aws.auth_manager.avp.entities.AvpEntities.ASSET",
781-
),
782781
_ => return,
783782
},
783+
["airflow", "providers", "amazon", "aws", "auth_manager", "avp", "entities", "AvpEntities", "DATASET"] => {
784+
Replacement::AutoImport {
785+
module: "airflow.providers.amazon.aws.auth_manager.avp.entities.AvpEntities",
786+
name: "ASSET",
787+
}
788+
}
784789

785790
// airflow.providers.common.io
786791
// airflow.providers.common.io.datasets.file
787792
["airflow", "providers", "common", "io", "datasets", "file", rest] => match *rest {
788-
"create_dataset" => {
789-
Replacement::Name("airflow.providers.common.io.assets.file.create_asset")
790-
}
791-
"convert_dataset_to_openlineage" => Replacement::Name(
792-
"airflow.providers.common.io.assets.file.convert_asset_to_openlineage",
793-
),
793+
"create_dataset" => Replacement::AutoImport {
794+
module: "airflow.providers.common.io.assets.file",
795+
name: "create_asset",
796+
},
797+
"convert_dataset_to_openlineage" => Replacement::AutoImport {
798+
module: "airflow.providers.common.io.assets.file",
799+
name: "convert_asset_to_openlineage",
800+
},
794801
"sanitize_uri" => {
795802
Replacement::Name("airflow.providers.common.io.assets.file.sanitize_uri")
796803
}
@@ -799,23 +806,27 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
799806

800807
// airflow.providers.fab
801808
["airflow", "providers", "fab", "auth_manager", "fab_auth_manager", "is_authorized_dataset"] => {
802-
Replacement::Name(
803-
"airflow.providers.fab.auth_manager.fab_auth_manager.is_authorized_asset",
804-
)
809+
Replacement::AutoImport {
810+
module: "airflow.providers.fab.auth_manager.fab_auth_manager",
811+
name: "is_authorized_asset",
812+
}
805813
}
806814

807815
// airflow.providers.google
808816
// airflow.providers.google.datasets
809817
["airflow", "providers", "google", "datasets", rest @ ..] => match &rest {
810-
["bigquery", "create_dataset"] => {
811-
Replacement::Name("airflow.providers.google.assets.bigquery.create_asset")
812-
}
813-
["gcs", "create_dataset"] => {
814-
Replacement::Name("airflow.providers.google.assets.gcs.create_asset")
815-
}
816-
["gcs", "convert_dataset_to_openlineage"] => Replacement::Name(
817-
"airflow.providers.google.assets.gcs.convert_asset_to_openlineage",
818-
),
818+
["bigquery", "create_dataset"] => Replacement::AutoImport {
819+
module: "airflow.providers.google.assets.bigquery",
820+
name: "create_asset",
821+
},
822+
["gcs", "create_dataset"] => Replacement::AutoImport {
823+
module: "airflow.providers.google.assets.gcs",
824+
name: "create_asset",
825+
},
826+
["gcs", "convert_dataset_to_openlineage"] => Replacement::AutoImport {
827+
module: "airflow.providers.google.assets.gcs",
828+
name: "convert_asset_to_openlineage",
829+
},
819830
["gcs", "sanitize_uri"] => {
820831
Replacement::Name("airflow.providers.google.assets.gcs.sanitize_uri")
821832
}
@@ -835,13 +846,15 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
835846
// airflow.providers.openlineage
836847
// airflow.providers.openlineage.utils.utils
837848
["airflow", "providers", "openlineage", "utils", "utils", rest] => match *rest {
838-
"DatasetInfo" => {
839-
Replacement::Name("airflow.providers.openlineage.utils.utils.AssetInfo")
840-
}
849+
"DatasetInfo" => Replacement::AutoImport {
850+
module: "airflow.providers.openlineage.utils.utils",
851+
name: "AssetInfo",
852+
},
841853

842-
"translate_airflow_dataset" => Replacement::Name(
843-
"airflow.providers.openlineage.utils.utils.translate_airflow_asset",
844-
),
854+
"translate_airflow_dataset" => Replacement::AutoImport {
855+
module: "airflow.providers.openlineage.utils.utils",
856+
name: "translate_airflow_asset",
857+
},
845858
_ => return,
846859
},
847860

0 commit comments

Comments
 (0)