feat(asset): make asset unique key an attrs

apache · Dec 10, 2024 · 5b27b1b · 5b27b1b
1 parent b42a42b
commit 5b27b1b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/task_sdk/src/airflow/sdk/definitions/asset/__init__.py b/task_sdk/src/airflow/sdk/definitions/asset/__init__.py
@@ -50,7 +50,14 @@
 log = logging.getLogger(__name__)
 
 
-class AssetUniqueKey(NamedTuple):
+@attrs.define(frozen=True)
+class AssetUniqueKey:
+    """
+    Columns to identify an unique asset.
+
+    :meta private:
+    """
+
     name: str
     uri: str
 
@@ -62,15 +69,15 @@ def to_asset(self) -> Asset:
         return Asset(name=self.name, uri=self.uri)
 
 
-@attrs.define()
-class AssetAliasUniqueKey:
 @attrs.define(frozen=True)
-
+class AssetAliasUniqueKey:
     """
     Columns to identify an unique asset alias.
+
+    :meta private:
     """
-    name: str
 
+    name: str
 
     @staticmethod
     def from_asset_alias(asset_alias: AssetAlias) -> AssetAliasUniqueKey:

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
@@ -77,7 +77,7 @@
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.providers.standard.sensors.python import PythonSensor
-from airflow.sdk.definitions.asset import AssetAlias
+from airflow.sdk.definitions.asset import Asset, AssetAlias
 from airflow.sensors.base import BaseSensorOperator
 from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG
 from airflow.stats import Stats