move RemoteAssetGraph.from_workspace to cached property on IWorkspace (…

…#20469)

same logic as #20468

## How I Tested These Changes

existing coverage

python_modules/dagster-graphql/dagster_graphql/implementation/asset_checks_loader.py

@@ -6,7 +6,6 @@
 from dagster._core.definitions.asset_check_evaluation import AssetCheckEvaluation
 from dagster._core.definitions.asset_check_spec import AssetCheckKey
 from dagster._core.definitions.events import AssetKey
-from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
 from dagster._core.definitions.selector import RepositorySelector
 from dagster._core.instance import DagsterInstance
 from dagster._core.remote_representation.code_location import CodeLocation
@@ -125,7 +124,7 @@ def _fetch_checks(
             self._context.instance, check_keys=all_check_keys
         )
 
-        asset_graph = RemoteAssetGraph.from_workspace(self._context)
+        asset_graph = self._context.asset_graph
         graphene_checks: Mapping[AssetKey, AssetChecksOrErrorUnion] = {}
         for asset_key in self._asset_keys:
             if asset_key in errors:

python_modules/dagster-graphql/dagster_graphql/implementation/execution/backfill.py

@@ -2,7 +2,6 @@
 
 import dagster._check as check
 import pendulum
-from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
 from dagster._core.definitions.selector import PartitionsByAssetSelector, RepositorySelector
 from dagster._core.errors import (
     DagsterError,
@@ -46,7 +45,7 @@ def get_asset_backfill_preview(
 ) -> Sequence["GrapheneAssetPartitions"]:
     from ...schema.backfill import GrapheneAssetPartitions
 
-    asset_graph = RemoteAssetGraph.from_workspace(graphene_info.context)
+    asset_graph = graphene_info.context.asset_graph
 
     check.invariant(backfill_preview_params.get("assetSelection") is not None)
     check.invariant(backfill_preview_params.get("partitionNames") is not None)
@@ -198,7 +197,7 @@ def create_and_launch_partition_backfill(
         if backfill_params.get("fromFailure"):
             raise DagsterError("fromFailure is not supported for pure asset backfills")
 
-        asset_graph = RemoteAssetGraph.from_workspace(graphene_info.context)
+        asset_graph = graphene_info.context.asset_graph
 
         assert_permission_for_asset_graph(
             graphene_info, asset_graph, asset_selection, Permissions.LAUNCH_PARTITION_BACKFILL
@@ -227,7 +226,7 @@ def create_and_launch_partition_backfill(
         if backfill_params.get("fromFailure"):
             raise DagsterError("fromFailure is not supported for pure asset backfills")
 
-        asset_graph = RemoteAssetGraph.from_workspace(graphene_info.context)
+        asset_graph = graphene_info.context.asset_graph
         assert_permission_for_asset_graph(
             graphene_info, asset_graph, asset_selection, Permissions.LAUNCH_PARTITION_BACKFILL
         )
@@ -265,7 +264,7 @@ def cancel_partition_backfill(
         check.failed(f"No backfill found for id: {backfill_id}")
 
     if backfill.is_asset_backfill:
-        asset_graph = RemoteAssetGraph.from_workspace(graphene_info.context)
+        asset_graph = graphene_info.context.asset_graph
         assert_permission_for_asset_graph(
             graphene_info,
             asset_graph,

python_modules/dagster-graphql/dagster_graphql/implementation/fetch_assets.py

@@ -28,7 +28,6 @@
     PartitionsDefinition,
     PartitionsSubset,
 )
-from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
 from dagster._core.definitions.time_window_partitions import (
     BaseTimeWindowPartitionsSubset,
     PartitionRangeStatus,
@@ -211,7 +210,7 @@ def get_asset_nodes_by_asset_key(
 
     stale_status_loader = StaleStatusLoader(
         instance=graphene_info.context.instance,
-        asset_graph=lambda: RemoteAssetGraph.from_workspace(graphene_info.context),
+        asset_graph=lambda: graphene_info.context.asset_graph,
     )
 
     dynamic_partitions_loader = CachingDynamicPartitionsLoader(graphene_info.context.instance)

python_modules/dagster-graphql/dagster_graphql/schema/roots/mutation.py

@@ -3,7 +3,6 @@
 import dagster._check as check
 import graphene
 from dagster._core.definitions.events import AssetKey
-from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
 from dagster._core.errors import DagsterInvariantViolationError
 from dagster._core.nux import get_has_seen_nux, set_nux_seen
 from dagster._core.workspace.permissions import Permissions
@@ -727,7 +726,7 @@ def mutate(
 
         reporting_user_tags = {**graphene_info.context.get_reporting_user_tags()}
 
-        asset_graph = RemoteAssetGraph.from_workspace(graphene_info.context)
+        asset_graph = graphene_info.context.asset_graph
 
         assert_permission_for_asset_graph(
             graphene_info, asset_graph, [asset_key], Permissions.REPORT_RUNLESS_ASSET_EVENTS

python_modules/dagster-graphql/dagster_graphql/schema/roots/query.py

@@ -954,7 +954,7 @@ def load_asset_graph() -> RemoteAssetGraph:
             if repo is not None:
                 return repo.asset_graph
             else:
-                return RemoteAssetGraph.from_workspace(graphene_info.context)
+                return graphene_info.context.asset_graph
 
         stale_status_loader = StaleStatusLoader(
             instance=graphene_info.context.instance,

python_modules/dagster/dagster/_core/definitions/remote_asset_graph.py

@@ -24,7 +24,6 @@
 from dagster._core.definitions.metadata import ArbitraryMetadataMapping
 from dagster._core.remote_representation.external import ExternalRepository
 from dagster._core.remote_representation.handle import RepositoryHandle
-from dagster._core.workspace.workspace import IWorkspace
 
 from .backfill_policy import BackfillPolicy
 from .base_asset_graph import AssetKeyOrCheckKey, BaseAssetGraph, BaseAssetNode
@@ -208,34 +207,6 @@ def __init__(
         self._asset_checks_by_key = asset_checks_by_key
         self._asset_check_execution_sets_by_key = asset_check_execution_sets_by_key
 
-    @classmethod
-    def from_workspace(cls, context: IWorkspace) -> "RemoteAssetGraph":
-        code_locations = (
-            location_entry.code_location
-            for location_entry in context.get_workspace_snapshot().values()
-            if location_entry.code_location
-        )
-        repos = (
-            repo
-            for code_location in code_locations
-            for repo in code_location.get_repositories().values()
-        )
-        repo_handle_external_asset_nodes: Sequence[
-            Tuple[RepositoryHandle, "ExternalAssetNode"]
-        ] = []
-        asset_checks: Sequence["ExternalAssetCheck"] = []
-
-        for repo in repos:
-            for external_asset_node in repo.get_external_asset_nodes():
-                repo_handle_external_asset_nodes.append((repo.handle, external_asset_node))
-
-            asset_checks.extend(repo.get_external_asset_checks())
-
-        return cls.from_repository_handles_and_external_asset_nodes(
-            repo_handle_external_asset_nodes=repo_handle_external_asset_nodes,
-            external_asset_checks=asset_checks,
-        )
-
     @classmethod
     def from_repository_handles_and_external_asset_nodes(
         cls,

python_modules/dagster/dagster/_core/execution/asset_backfill.py

@@ -883,8 +883,7 @@ def execute_asset_backfill_iteration(
     logger.info(f"Evaluating asset backfill {backfill.backfill_id}")
 
     workspace_context = workspace_process_context.create_request_context()
-
-    asset_graph = RemoteAssetGraph.from_workspace(workspace_context)
+    asset_graph = workspace_context.asset_graph
 
     if not backfill.is_asset_backfill:
         check.failed("Backfill must be an asset backfill")

python_modules/dagster/dagster/_core/execution/backfill.py

@@ -7,7 +7,6 @@
 from dagster._core.definitions import AssetKey
 from dagster._core.definitions.base_asset_graph import BaseAssetGraph
 from dagster._core.definitions.partition import PartitionsSubset
-from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
 from dagster._core.errors import DagsterDefinitionChangedDeserializationError
 from dagster._core.execution.bulk_actions import BulkActionType
 from dagster._core.instance import DynamicPartitionsStore
@@ -154,7 +153,7 @@ def is_valid_serialization(self, workspace: IWorkspace) -> bool:
             if self.serialized_asset_backfill_data:
                 return AssetBackfillData.is_valid_serialization(
                     self.serialized_asset_backfill_data,
-                    RemoteAssetGraph.from_workspace(workspace),
+                    workspace.asset_graph,
                 )
             else:
                 return True
@@ -171,7 +170,7 @@ def get_backfill_status_per_asset_key(
             return []
 
         if self.is_asset_backfill:
-            asset_graph = RemoteAssetGraph.from_workspace(workspace)
+            asset_graph = workspace.asset_graph
             try:
                 asset_backfill_data = self.get_asset_backfill_data(asset_graph)
             except DagsterDefinitionChangedDeserializationError:
@@ -188,7 +187,7 @@ def get_target_partitions_subset(
             return None
 
         if self.is_asset_backfill:
-            asset_graph = RemoteAssetGraph.from_workspace(workspace)
+            asset_graph = workspace.asset_graph
             try:
                 asset_backfill_data = self.get_asset_backfill_data(asset_graph)
             except DagsterDefinitionChangedDeserializationError:
@@ -205,7 +204,7 @@ def get_target_root_partitions_subset(
             return None
 
         if self.is_asset_backfill:
-            asset_graph = RemoteAssetGraph.from_workspace(workspace)
+            asset_graph = workspace.asset_graph
             try:
                 asset_backfill_data = self.get_asset_backfill_data(asset_graph)
             except DagsterDefinitionChangedDeserializationError:
@@ -220,7 +219,7 @@ def get_num_partitions(self, workspace: IWorkspace) -> Optional[int]:
             return 0
 
         if self.is_asset_backfill:
-            asset_graph = RemoteAssetGraph.from_workspace(workspace)
+            asset_graph = workspace.asset_graph
             try:
                 asset_backfill_data = self.get_asset_backfill_data(asset_graph)
             except DagsterDefinitionChangedDeserializationError:
@@ -238,7 +237,7 @@ def get_partition_names(self, workspace: IWorkspace) -> Optional[Sequence[str]]:
             return []
 
         if self.is_asset_backfill:
-            asset_graph = RemoteAssetGraph.from_workspace(workspace)
+            asset_graph = workspace.asset_graph
             try:
                 asset_backfill_data = self.get_asset_backfill_data(asset_graph)
             except DagsterDefinitionChangedDeserializationError:

python_modules/dagster/dagster/_core/execution/submit_asset_runs.py

@@ -188,7 +188,7 @@ def _create_asset_run(
         # likely is outdated and targeting the wrong job, refetch the asset
         # graph from the workspace
         workspace = workspace_process_context.create_request_context()
-        asset_graph = RemoteAssetGraph.from_workspace(workspace)
+        asset_graph = workspace.asset_graph
 
     check.failed(
         f"Failed to target asset selection {run_request.asset_selection} in run after retrying."

python_modules/dagster/dagster/_core/workspace/workspace.py

@@ -1,11 +1,18 @@
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Mapping, NamedTuple, Optional, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Mapping, NamedTuple, Optional, Sequence, Tuple
 
 from dagster._utils.error import SerializableErrorInfo
 
 if TYPE_CHECKING:
+    from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
     from dagster._core.remote_representation import CodeLocation, CodeLocationOrigin
+    from dagster._core.remote_representation.external_data import (
+        ExternalAssetCheck,
+        ExternalAssetNode,
+    )
+    from dagster._core.remote_representation.handle import RepositoryHandle
 
 
 # For locations that are loaded asynchronously
@@ -48,6 +55,37 @@ def get_workspace_snapshot(self) -> Mapping[str, CodeLocationEntry]:
     def get_code_location_statuses(self) -> Sequence[CodeLocationStatusEntry]:
         pass
 
+    @cached_property
+    def asset_graph(self) -> "RemoteAssetGraph":
+        """Returns a workspace scoped RemoteAssetGraph."""
+        from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
+
+        code_locations = (
+            location_entry.code_location
+            for location_entry in self.get_workspace_snapshot().values()
+            if location_entry.code_location
+        )
+        repos = (
+            repo
+            for code_location in code_locations
+            for repo in code_location.get_repositories().values()
+        )
+        repo_handle_external_asset_nodes: Sequence[
+            Tuple["RepositoryHandle", "ExternalAssetNode"]
+        ] = []
+        asset_checks: Sequence["ExternalAssetCheck"] = []
+
+        for repo in repos:
+            for external_asset_node in repo.get_external_asset_nodes():
+                repo_handle_external_asset_nodes.append((repo.handle, external_asset_node))
+
+            asset_checks.extend(repo.get_external_asset_checks())
+
+        return RemoteAssetGraph.from_repository_handles_and_external_asset_nodes(
+            repo_handle_external_asset_nodes=repo_handle_external_asset_nodes,
+            external_asset_checks=asset_checks,
+        )
+
 
 def location_status_from_location_entry(
     entry: CodeLocationEntry,

python_modules/dagster/dagster/_daemon/asset_daemon.py

@@ -486,7 +486,7 @@ def _run_iteration_impl(
                 if not get_has_migrated_to_sensors(instance):
                     # Do a one-time migration to create the cursors for each sensor, based on the
                     # existing cursor for the legacy AMP tick
-                    asset_graph = RemoteAssetGraph.from_workspace(workspace)
+                    asset_graph = workspace.asset_graph
                     pre_sensor_cursor = _get_pre_sensor_auto_materialize_cursor(
                         instance, asset_graph
                     )
@@ -648,7 +648,7 @@ def _process_auto_materialize_tick_generator(
 
         workspace = workspace_process_context.create_request_context()
 
-        asset_graph = RemoteAssetGraph.from_workspace(workspace)
+        asset_graph = workspace.asset_graph
 
         instance: DagsterInstance = workspace_process_context.instance
         error_info = None

python_modules/dagster/dagster/_scheduler/stale.py

@@ -6,7 +6,6 @@
     StaleStatus,
 )
 from dagster._core.definitions.events import AssetKey
-from dagster._core.definitions.remote_asset_graph import RemoteAssetGraph
 from dagster._core.definitions.run_request import RunRequest
 from dagster._core.remote_representation.external import (
     ExternalSchedule,
@@ -20,7 +19,7 @@ def resolve_stale_or_missing_assets(
     run_request: RunRequest,
     instigator: Union[ExternalSensor, ExternalSchedule],
 ) -> Sequence[AssetKey]:
-    asset_graph = RemoteAssetGraph.from_workspace(context.create_request_context())
+    asset_graph = context.create_request_context().asset_graph
     asset_selection = (
         run_request.asset_selection
         if run_request.asset_selection is not None