From 5a5fe12735c2d4cdf55473714355b8e4a10ca4dc Mon Sep 17 00:00:00 2001 From: Sandy Ryza Date: Mon, 25 Mar 2024 09:11:56 -0700 Subject: [PATCH] Move `AssetSelection` from `BaseModel` to `DagsterModel` (#20638) ## Summary & Motivation Uses the `DagsterModel` introduced in https://github.com/dagster-io/dagster/pull/20637 for `AssetSelection`. ## How I Tested These Changes --- .../_core/definitions/asset_selection.py | 73 +++++-------------- .../asset_defs_tests/test_asset_selection.py | 2 +- .../test_external_sensor_data.py | 2 +- .../dbt_manifest_asset_selection.py | 6 +- 4 files changed, 22 insertions(+), 61 deletions(-) diff --git a/python_modules/dagster/dagster/_core/definitions/asset_selection.py b/python_modules/dagster/dagster/_core/definitions/asset_selection.py index 0527ac784992f..200f1192128d8 100644 --- a/python_modules/dagster/dagster/_core/definitions/asset_selection.py +++ b/python_modules/dagster/dagster/_core/definitions/asset_selection.py @@ -20,6 +20,7 @@ parse_clause, ) from dagster._core.storage.tags import TAG_NO_VALUE +from dagster._model import DagsterModel from dagster._serdes.serdes import whitelist_for_serdes from .asset_check_spec import AssetCheckKey @@ -42,7 +43,7 @@ ] -class AssetSelection(ABC, BaseModel, frozen=True): +class AssetSelection(ABC, DagsterModel): """An AssetSelection defines a query over a set of assets and asset checks, normally all that are defined in a code location. You can use the "|", "&", and "-" operators to create unions, intersections, and differences of selections, respectively. @@ -485,7 +486,7 @@ def operand__str__(self) -> str: @whitelist_for_serdes -class AllSelection(AssetSelection, frozen=True): +class AllSelection(AssetSelection): include_sources: Optional[bool] = None def resolve_inner( @@ -505,7 +506,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class AllAssetCheckSelection(AssetSelection, frozen=True): +class AllAssetCheckSelection(AssetSelection): def resolve_inner( self, asset_graph: BaseAssetGraph, allow_missing: bool ) -> AbstractSet[AssetKey]: @@ -524,7 +525,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class AssetChecksForAssetKeysSelection(AssetSelection, frozen=True): +class AssetChecksForAssetKeysSelection(AssetSelection): selected_asset_keys: Sequence[AssetKey] def resolve_inner( @@ -546,7 +547,7 @@ def to_serializable_asset_selection(self, asset_graph: BaseAssetGraph) -> "Asset @whitelist_for_serdes -class AssetCheckKeysSelection(AssetSelection, frozen=True): +class AssetCheckKeysSelection(AssetSelection): selected_asset_check_keys: Sequence[AssetCheckKey] def resolve_inner( @@ -574,11 +575,7 @@ def to_serializable_asset_selection(self, asset_graph: BaseAssetGraph) -> "Asset @whitelist_for_serdes -class AndAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class AndAssetSelection(AssetSelection): operands: Sequence[AssetSelection] def resolve_inner( @@ -618,11 +615,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class OrAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class OrAssetSelection(AssetSelection): operands: Sequence[AssetSelection] def resolve_inner( @@ -662,11 +655,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class SubtractAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class SubtractAssetSelection(AssetSelection): left: AssetSelection right: AssetSelection @@ -698,11 +687,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class SinksAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class SinksAssetSelection(AssetSelection): child: AssetSelection def resolve_inner( @@ -716,11 +701,7 @@ def to_serializable_asset_selection(self, asset_graph: BaseAssetGraph) -> "Asset @whitelist_for_serdes -class RequiredNeighborsAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class RequiredNeighborsAssetSelection(AssetSelection): child: AssetSelection def resolve_inner( @@ -737,11 +718,7 @@ def to_serializable_asset_selection(self, asset_graph: BaseAssetGraph) -> "Asset @whitelist_for_serdes -class RootsAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class RootsAssetSelection(AssetSelection): child: AssetSelection def resolve_inner( @@ -755,11 +732,7 @@ def to_serializable_asset_selection(self, asset_graph: BaseAssetGraph) -> "Asset @whitelist_for_serdes -class DownstreamAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class DownstreamAssetSelection(AssetSelection): child: AssetSelection depth: Optional[int] include_self: bool @@ -790,7 +763,7 @@ def to_serializable_asset_selection(self, asset_graph: BaseAssetGraph) -> "Asset @whitelist_for_serdes -class GroupsAssetSelection(AssetSelection, frozen=True): +class GroupsAssetSelection(AssetSelection): selected_groups: Sequence[str] include_sources: bool @@ -820,7 +793,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class TagAssetSelection(AssetSelection, frozen=True): +class TagAssetSelection(AssetSelection): key: str value: str include_sources: bool @@ -841,7 +814,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class KeysAssetSelection(AssetSelection, frozen=True): +class KeysAssetSelection(AssetSelection): selected_keys: Sequence[AssetKey] def resolve_inner( @@ -891,7 +864,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class KeyPrefixesAssetSelection(AssetSelection, frozen=True): +class KeyPrefixesAssetSelection(AssetSelection): selected_key_prefixes: Sequence[Sequence[str]] include_sources: bool @@ -946,11 +919,7 @@ def _fetch_all_upstream( @whitelist_for_serdes -class UpstreamAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class UpstreamAssetSelection(AssetSelection): child: AssetSelection depth: Optional[int] include_self: bool @@ -982,11 +951,7 @@ def __str__(self) -> str: @whitelist_for_serdes -class ParentSourcesAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class ParentSourcesAssetSelection(AssetSelection): child: AssetSelection def resolve_inner( diff --git a/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_selection.py b/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_selection.py index 2ed67b5945b08..7527b35bae6db 100644 --- a/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_selection.py +++ b/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_selection.py @@ -534,7 +534,7 @@ def test_all_asset_selection_subclasses_serializable(): def test_to_serializable_asset_selection(): - class UnserializableAssetSelection(AssetSelection, frozen=True): + class UnserializableAssetSelection(AssetSelection): def resolve_inner( self, asset_graph: BaseAssetGraph, allow_missing: bool ) -> AbstractSet[AssetKey]: diff --git a/python_modules/dagster/dagster_tests/core_tests/host_representation_tests/test_external_sensor_data.py b/python_modules/dagster/dagster_tests/core_tests/host_representation_tests/test_external_sensor_data.py index 513cb135ad688..7b243a251051d 100644 --- a/python_modules/dagster/dagster_tests/core_tests/host_representation_tests/test_external_sensor_data.py +++ b/python_modules/dagster/dagster_tests/core_tests/host_representation_tests/test_external_sensor_data.py @@ -24,7 +24,7 @@ def asset1(): ... @asset def asset2(): ... - class MySpecialAssetSelection(AssetSelection, frozen=True): + class MySpecialAssetSelection(AssetSelection): def resolve_inner( self, asset_graph: BaseAssetGraph, allow_missing: bool ) -> AbstractSet[AssetKey]: diff --git a/python_modules/libraries/dagster-dbt/dagster_dbt/dbt_manifest_asset_selection.py b/python_modules/libraries/dagster-dbt/dagster_dbt/dbt_manifest_asset_selection.py index 52d719b94be56..09f106b8c43fe 100644 --- a/python_modules/libraries/dagster-dbt/dagster_dbt/dbt_manifest_asset_selection.py +++ b/python_modules/libraries/dagster-dbt/dagster_dbt/dbt_manifest_asset_selection.py @@ -18,11 +18,7 @@ ) -class DbtManifestAssetSelection( - AssetSelection, - frozen=True, - arbitrary_types_allowed=True, -): +class DbtManifestAssetSelection(AssetSelection, arbitrary_types_allowed=True): """Defines a selection of assets from a dbt manifest wrapper and a dbt selection string. Args: