From 144a6cc6bbf1b75e1020e584204b92c75d45eb86 Mon Sep 17 00:00:00 2001
From: gibsondan <daniel@elementl.com>
Date: Tue, 13 Feb 2024 11:25:40 -0600
Subject: [PATCH] Add "assets" resolver to GrapheneAssetSelection (#19771)

Summary:
Will allow us to break out the list of asset keys in the sensor view by
whether or not they have an AMP policy.

Test Plan: BK

## Summary & Motivation

## How I Tested These Changes
---
 .../packages/ui-core/src/graphql/schema.graphql  |  1 +
 .../packages/ui-core/src/graphql/types.ts        |  2 ++
 .../dagster_graphql/schema/asset_selections.py   | 13 +++++++++++++
 .../graphql/test_sensors.py                      | 16 ++++++++++++++++
 4 files changed, 32 insertions(+)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql b/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql
index 66c5a753da5db..11a0d68face88 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql
+++ b/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql
@@ -2906,6 +2906,7 @@ enum SensorType {
 type AssetSelection {
   assetSelectionString: String
   assetKeys: [AssetKey!]!
+  assets: [Asset!]!
 }
 
 union SensorOrError = Sensor | SensorNotFoundError | UnauthorizedError | PythonError
diff --git a/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts b/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts
index 3731c85817785..2f082a83a3ea2 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts
@@ -568,6 +568,7 @@ export type AssetSelection = {
   __typename: 'AssetSelection';
   assetKeys: Array<AssetKey>;
   assetSelectionString: Maybe<Scalars['String']>;
+  assets: Array<Asset>;
 };
 
 export type AssetSubset = {
@@ -6373,6 +6374,7 @@ export const buildAssetSelection = (
       overrides && overrides.hasOwnProperty('assetSelectionString')
         ? overrides.assetSelectionString!
         : 'dolores',
+    assets: overrides && overrides.hasOwnProperty('assets') ? overrides.assets! : [],
   };
 };
 
diff --git a/python_modules/dagster-graphql/dagster_graphql/schema/asset_selections.py b/python_modules/dagster-graphql/dagster_graphql/schema/asset_selections.py
index 9094b5b2c9553..12676c126399d 100644
--- a/python_modules/dagster-graphql/dagster_graphql/schema/asset_selections.py
+++ b/python_modules/dagster-graphql/dagster_graphql/schema/asset_selections.py
@@ -3,6 +3,7 @@
 from dagster._core.definitions.external_asset_graph import ExternalAssetGraph
 from dagster._core.host_representation.external import ExternalRepository
 
+from ..implementation.fetch_assets import get_asset_nodes_by_asset_key
 from .asset_key import GrapheneAssetKey
 from .util import non_null_list
 
@@ -10,6 +11,7 @@
 class GrapheneAssetSelection(graphene.ObjectType):
     assetSelectionString = graphene.String()
     assetKeys = non_null_list(GrapheneAssetKey)
+    assets = non_null_list("dagster_graphql.schema.pipelines.pipeline.GrapheneAsset")
 
     def __init__(self, asset_selection: AssetSelection, external_repository: ExternalRepository):
         self._asset_selection = asset_selection
@@ -25,5 +27,16 @@ def resolve_assetKeys(self, _graphene_info):
             for asset_key in self._asset_selection.resolve(asset_graph)
         ]
 
+    def resolve_assets(self, graphene_info):
+        from dagster_graphql.schema.pipelines.pipeline import GrapheneAsset
+
+        asset_graph = ExternalAssetGraph.from_external_repository(self._external_repository)
+        asset_nodes_by_asset_key = get_asset_nodes_by_asset_key(graphene_info)
+
+        return [
+            GrapheneAsset(key=asset_key, definition=asset_nodes_by_asset_key.get(asset_key))
+            for asset_key in self._asset_selection.resolve(asset_graph)
+        ]
+
     class Meta:
         name = "AssetSelection"
diff --git a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_sensors.py b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_sensors.py
index f4f6c28042ecd..36c3d8c92ab7b 100644
--- a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_sensors.py
+++ b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_sensors.py
@@ -155,6 +155,16 @@
         assetKeys {
           path
         }
+        assets {
+          key {
+            path
+          }
+          definition {
+            assetKey {
+              path
+            }
+          }
+        }
       }
     }
   }
@@ -1420,3 +1430,9 @@ def test_asset_selection(graphql_context):
     assert result.data["sensorOrError"]["assetSelection"]["assetKeys"] == [
         {"path": ["fresh_diamond_bottom"]}
     ]
+    assert result.data["sensorOrError"]["assetSelection"]["assets"] == [
+        {
+            "key": {"path": ["fresh_diamond_bottom"]},
+            "definition": {"assetKey": {"path": ["fresh_diamond_bottom"]}},
+        }
+    ]