From 58bd547824097acaaa667c1e068fbbb841a6bdee Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Sat, 7 Dec 2024 15:43:50 +0800
Subject: [PATCH] feat(datasets): add backward compat for DatasetAll,
 DatasetAny, expand_alias_to_datasets and DatasetAliasEvent (#44635)

---
 airflow/datasets/__init__.py   | 39 ++++++++++++++++++++++++++--------
 tests/datasets/test_dataset.py | 20 +++++++++++++++--
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/airflow/datasets/__init__.py b/airflow/datasets/__init__.py
index d7c51a3082377..d082d6fbed854 100644
--- a/airflow/datasets/__init__.py
+++ b/airflow/datasets/__init__.py
@@ -24,17 +24,38 @@
 # lib.)  This is required by some IDEs to resolve the import paths.
 from __future__ import annotations
 
+import importlib
 import warnings
 
-from airflow.sdk.definitions.asset import AssetAlias as DatasetAlias, Dataset
-
 # TODO: Remove this module in Airflow 3.2
 
-warnings.warn(
-    "Import from the airflow.dataset module is deprecated and "
-    "will be removed in the Airflow 3.2. Please import it from 'airflow.sdk.definitions.asset'.",
-    DeprecationWarning,
-    stacklevel=2,
-)
+_names_moved = {
+    "DatasetAlias": ("airflow.sdk.definitions.asset", "AssetAlias"),
+    "DatasetAll": ("airflow.sdk.definitions.asset", "AssetAll"),
+    "DatasetAny": ("airflow.sdk.definitions.asset", "DatasetAny"),
+    "Dataset": ("airflow.sdk.definitions.asset", "Asset"),
+    "expand_alias_to_datasets": ("airflow.models.asset", "expand_alias_to_assets"),
+}
+
+
+def __getattr__(name: str):
+    # PEP-562: Lazy loaded attributes on python modules
+    if name not in _names_moved:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    module_path, new_name = _names_moved[name]
+    warnings.warn(
+        f"Import 'airflow.dataset.{name}' is deprecated and "
+        f"will be removed in the Airflow 3.2. Please import it from '{module_path}.{new_name}'.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    mod = importlib.import_module(module_path, __name__)
+    val = getattr(mod, new_name)
+
+    # Store for next time
+    globals()[name] = val
+    return val
+
 
-__all__ = ["Dataset", "DatasetAlias"]
+__all__ = ["Dataset", "DatasetAlias", "DatasetAll", "DatasetAny", "expand_alias_to_datasets"]
diff --git a/tests/datasets/test_dataset.py b/tests/datasets/test_dataset.py
index 4898bd5fd47e6..44e62839fa9fb 100644
--- a/tests/datasets/test_dataset.py
+++ b/tests/datasets/test_dataset.py
@@ -36,8 +36,24 @@
             "airflow.datasets",
             "Dataset",
             (
-                "Import from the airflow.dataset module is deprecated and "
-                "will be removed in the Airflow 3.2. Please import it from 'airflow.sdk.definitions.asset'."
+                "Import 'airflow.dataset.Dataset' is deprecated and "
+                "will be removed in the Airflow 3.2. Please import it from 'airflow.sdk.definitions.asset.Asset'."
+            ),
+        ),
+        (
+            "airflow.datasets",
+            "DatasetAlias",
+            (
+                "Import 'airflow.dataset.DatasetAlias' is deprecated and "
+                "will be removed in the Airflow 3.2. Please import it from 'airflow.sdk.definitions.asset.AssetAlias'."
+            ),
+        ),
+        (
+            "airflow.datasets",
+            "expand_alias_to_datasets",
+            (
+                "Import 'airflow.dataset.expand_alias_to_datasets' is deprecated and "
+                "will be removed in the Airflow 3.2. Please import it from 'airflow.models.asset.expand_alias_to_assets'."
             ),
         ),
         (