Skip to content

Commit 5333e05

Browse files
Lee-WLee-W
committed
fix(AIR301): fix airflow.assets.manager.AssetManager and introduce SourceModuleMoved for refactoring
1 parent 92027c9 commit 5333e05

File tree

4 files changed

+54
-45
lines changed

4 files changed

+54
-45
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ pub(crate) enum Replacement {
1313
path: &'static str,
1414
name: &'static str,
1515
},
16+
SourceModuleMoved {
17+
name: String,
18+
module: &'static str,
19+
},
1620
}
1721

1822
pub(crate) fn is_guarded_by_try_except(

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

Lines changed: 48 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,9 @@ impl Violation for Airflow3Removal {
6464
Replacement::AutoImport { path: _, name: _ } => {
6565
format!("`{deprecated}` is removed in Airflow 3.0")
6666
}
67+
Replacement::SourceModuleMoved { name: _, module: _ } => {
68+
format!("`{deprecated}` is removed in Airflow 3.0")
69+
}
6770
}
6871
}
6972

@@ -72,6 +75,9 @@ impl Violation for Airflow3Removal {
7275
match replacement {
7376
Replacement::Name(name) => Some(format!("Use `{name}` instead")),
7477
Replacement::AutoImport { path, name } => Some(format!("Use `{path}.{name}` instead")),
78+
Replacement::SourceModuleMoved { name, module } => {
79+
Some(format!("Use `{module}.{name}` instead"))
80+
}
7581
_ => None,
7682
}
7783
}
@@ -593,43 +599,38 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
593599
}
594600

595601
// airflow.configuration
596-
["airflow", "configuration", rest] => match *rest {
597-
"get" => Replacement::Name("airflow.configuration.conf.get"),
598-
"getboolean" => Replacement::Name("airflow.configuration.conf.getboolean"),
599-
"getfloat" => Replacement::Name("airflow.configuration.conf.getfloat"),
600-
"getint" => Replacement::Name("airflow.configuration.conf.getint"),
601-
"has_option" => Replacement::Name("airflow.configuration.conf.has_option"),
602-
"remove_option" => Replacement::Name("airflow.configuration.conf.remove_option"),
603-
"as_dict" => Replacement::Name("airflow.configuration.conf.as_dict"),
604-
"set" => Replacement::Name("airflow.configuration.conf.set"),
605-
_ => return,
602+
["airflow", "configuration", rest @ ("as_dict" | "get" | "getboolean" | "getfloat" | "getint" | "has_option"
603+
| "remove_option" | "set")] => Replacement::SourceModuleMoved {
604+
name: rest.to_string(),
605+
module: "airflow.configuration.conf",
606606
},
607607

608608
// airflow.contrib.*
609609
["airflow", "contrib", ..] => {
610610
Replacement::Message("The whole `airflow.contrib` module has been removed.")
611611
}
612612

613+
// airflow.datasets.manager
614+
["airflow", "datasets", "manager", rest] => match *rest {
615+
"DatasetManager" => Replacement::Name("airflow.assets.manager.AssetManager"),
616+
"dataset_manager" => Replacement::Name("airflow.assets.manager.asset_manager"),
617+
"resolve_dataset_manager" => Replacement::Name("airflow.assets.resolve_asset_manager"),
618+
_ => return,
619+
},
620+
["airflow", "datasets", "metadata", "Metadata"] => {
621+
Replacement::Name("airflow.sdk.Metadata")
622+
}
613623
// airflow.datasets
614624
["airflow", "Dataset"] | ["airflow", "datasets", "Dataset"] => Replacement::AutoImport {
615625
path: "airflow.sdk",
616626
name: "Asset",
617627
},
618-
["airflow", "datasets", rest @ ..] => match &rest {
619-
["DatasetAliasEvent"] => Replacement::None,
620-
["DatasetAlias"] => Replacement::Name("airflow.sdk.AssetAlias"),
621-
["DatasetAll"] => Replacement::Name("airflow.sdk.AssetAll"),
622-
["DatasetAny"] => Replacement::Name("airflow.sdk.AssetAny"),
623-
["expand_alias_to_datasets"] => Replacement::Name("airflow.sdk.expand_alias_to_assets"),
624-
["metadata", "Metadata"] => Replacement::Name("airflow.sdk.Metadata"),
625-
// airflow.datasets.manager
626-
["manager", "DatasetManager"] => Replacement::Name("airflow.assets.AssetManager"),
627-
["manager", "dataset_manager"] => {
628-
Replacement::Name("airflow.assets.manager.asset_manager")
629-
}
630-
["manager", "resolve_dataset_manager"] => {
631-
Replacement::Name("airflow.assets.resolve_asset_manager")
632-
}
628+
["airflow", "datasets", rest] => match *rest {
629+
"DatasetAliasEvent" => Replacement::None,
630+
"DatasetAlias" => Replacement::Name("airflow.sdk.AssetAlias"),
631+
"DatasetAll" => Replacement::Name("airflow.sdk.AssetAll"),
632+
"DatasetAny" => Replacement::Name("airflow.sdk.AssetAny"),
633+
"expand_alias_to_datasets" => Replacement::Name("airflow.sdk.expand_alias_to_assets"),
633634
_ => return,
634635
},
635636

@@ -645,11 +646,11 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
645646

646647
// airflow.listeners.spec
647648
// TODO: this is removed
648-
["airflow", "listeners", "spec", "dataset", rest @ ..] => match &rest {
649-
["on_dataset_created"] => {
649+
["airflow", "listeners", "spec", "dataset", rest] => match *rest {
650+
"on_dataset_created" => {
650651
Replacement::Name("airflow.listeners.spec.asset.on_asset_created")
651652
}
652-
["on_dataset_changed"] => {
653+
"on_dataset_changed" => {
653654
Replacement::Name("airflow.listeners.spec.asset.on_asset_changed")
654655
}
655656
_ => return,
@@ -668,9 +669,10 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
668669

669670
// airflow.models.baseoperator
670671
["airflow", "models", "baseoperator", rest] => match *rest {
671-
"chain" => Replacement::Name("airflow.sdk.chain"),
672-
"chain_linear" => Replacement::Name("airflow.sdk.chain_linear"),
673-
"cross_downstream" => Replacement::Name("airflow.sdk.cross_downstream"),
672+
"chain" | "chain_linear" | "cross_downstream" => Replacement::SourceModuleMoved {
673+
name: (*rest).to_string(),
674+
module: "airflow.sdk",
675+
},
674676
"BaseOperatorLink" => {
675677
Replacement::Name("airflow.sdk.definitions.baseoperatorlink.BaseOperatorLink")
676678
}
@@ -797,14 +799,15 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
797799
},
798800

799801
// airflow.providers.common.io
800-
["airflow", "providers", "common", "io", rest @ ..] => match &rest {
801-
["datasets", "file", "create_dataset"] => {
802+
// airflow.providers.common.io.datasets.file
803+
["airflow", "providers", "common", "io", "datasets", "file", rest] => match *rest {
804+
"create_dataset" => {
802805
Replacement::Name("airflow.providers.common.io.assets.file.create_asset")
803806
}
804-
["datasets", "file", "convert_dataset_to_openlineage"] => Replacement::Name(
807+
"convert_dataset_to_openlineage" => Replacement::Name(
805808
"airflow.providers.common.io.assets.file.convert_asset_to_openlineage",
806809
),
807-
["datasets", "file", "sanitize_uri"] => {
810+
"sanitize_uri" => {
808811
Replacement::Name("airflow.providers.common.io.assets.file.sanitize_uri")
809812
}
810813
_ => return,
@@ -818,17 +821,18 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
818821
}
819822

820823
// airflow.providers.google
821-
["airflow", "providers", "google", rest @ ..] => match &rest {
822-
["datasets", "bigquery", "create_dataset"] => {
824+
// airflow.providers.google.datasets
825+
["airflow", "providers", "google", "datasets", rest @ ..] => match &rest {
826+
["bigquery", "create_dataset"] => {
823827
Replacement::Name("airflow.providers.google.assets.bigquery.create_asset")
824828
}
825-
["datasets", "gcs", "create_dataset"] => {
829+
["gcs", "create_dataset"] => {
826830
Replacement::Name("airflow.providers.google.assets.gcs.create_asset")
827831
}
828-
["datasets", "gcs", "convert_dataset_to_openlineage"] => Replacement::Name(
832+
["gcs", "convert_dataset_to_openlineage"] => Replacement::Name(
829833
"airflow.providers.google.assets.gcs.convert_asset_to_openlineage",
830834
),
831-
["datasets", "gcs", "sanitize_uri"] => {
835+
["gcs", "sanitize_uri"] => {
832836
Replacement::Name("airflow.providers.google.assets.gcs.sanitize_uri")
833837
}
834838
_ => return,
@@ -845,12 +849,13 @@ fn check_name(checker: &Checker, expr: &Expr, range: TextRange) {
845849
}
846850

847851
// airflow.providers.openlineage
848-
["airflow", "providers", "openlineage", rest @ ..] => match &rest {
849-
["utils", "utils", "DatasetInfo"] => {
852+
// airflow.providers.openlineage.utils.utils
853+
["airflow", "providers", "openlineage", "utils", "utils", rest] => match *rest {
854+
"DatasetInfo" => {
850855
Replacement::Name("airflow.providers.openlineage.utils.utils.AssetInfo")
851856
}
852857

853-
["utils", "utils", "translate_airflow_dataset"] => Replacement::Name(
858+
"translate_airflow_dataset" => Replacement::Name(
854859
"airflow.providers.openlineage.utils.utils.translate_airflow_asset",
855860
),
856861
_ => return,

‎crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR301_AIR301_class_attribute.py.snap‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ AIR301_class_attribute.py:42:6: AIR301 `airflow.datasets.manager.DatasetManager`
244244
43 | dm.register_dataset_change()
245245
44 | dm.create_datasets()
246246
|
247-
= help: Use `airflow.assets.AssetManager` instead
247+
= help: Use `airflow.assets.manager.AssetManager` instead
248248

249249
AIR301_class_attribute.py:43:4: AIR301 [*] `register_dataset_change` is removed in Airflow 3.0
250250
|

‎crates/ruff_linter/src/rules/airflow/snapshots/ruff_linter__rules__airflow__tests__AIR301_AIR301_names.py.snap‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -321,7 +321,7 @@ AIR301_names.py:118:1: AIR301 `airflow.datasets.manager.DatasetManager` is remov
321321
119 | dataset_manager
322322
120 | resolve_dataset_manager
323323
|
324-
= help: Use `airflow.assets.AssetManager` instead
324+
= help: Use `airflow.assets.manager.AssetManager` instead
325325

326326
AIR301_names.py:119:1: AIR301 `airflow.datasets.manager.dataset_manager` is removed in Airflow 3.0
327327
|

0 commit comments

Comments
 (0)