Asset Subselection (#7835)

dagster-io · May 18, 2022 · 422fa71 · 422fa71
1 parent fb025dc
commit 422fa71
Show file tree

Hide file tree

Showing 61 changed files with 10,123 additions and 7,494 deletions.
diff --git a/js_modules/dagit/packages/core/src/assets/LaunchAssetExecutionButton.tsx b/js_modules/dagit/packages/core/src/assets/LaunchAssetExecutionButton.tsx
@@ -111,11 +111,13 @@ export const LaunchAssetExecutionButton: React.FC<{
                 ],
               },
               runConfigData: {},
-              stepKeys: assets.map((o) => o.opNames).flat(),
               selector: {
                 repositoryLocationName: repoAddress.location,
                 repositoryName: repoAddress.name,
                 pipelineName: jobName,
+                assetSelection: assets.map((asset) => ({
+                  path: asset.assetKey.path,
+                })),
               },
             },
           })}

diff --git a/js_modules/dagit/packages/core/src/graphql/schema.graphql b/js_modules/dagit/packages/core/src/graphql/schema.graphql
diff --git a/js_modules/dagit/packages/core/src/instance/types/RunReExecutionQuery.ts b/js_modules/dagit/packages/core/src/instance/types/RunReExecutionQuery.ts
diff --git a/js_modules/dagit/packages/core/src/instigation/types/LaunchedRunListQuery.ts b/js_modules/dagit/packages/core/src/instigation/types/LaunchedRunListQuery.ts
diff --git a/js_modules/dagit/packages/core/src/partitions/types/PartitionRunListForStepQuery.ts b/js_modules/dagit/packages/core/src/partitions/types/PartitionRunListForStepQuery.ts
diff --git a/js_modules/dagit/packages/core/src/pipelines/types/PipelineRunsRootQuery.ts b/js_modules/dagit/packages/core/src/pipelines/types/PipelineRunsRootQuery.ts
diff --git a/js_modules/dagit/packages/core/src/runs/RunFragments.tsx b/js_modules/dagit/packages/core/src/runs/RunFragments.tsx
@@ -37,6 +37,11 @@ export const RunFragments = {
       parentRunId
       pipelineName
       solidSelection
+      assetSelection {
+        ... on AssetKey {
+          path
+        }
+      }
       pipelineSnapshotId
       executionPlan {
         artifactsPersisted

diff --git a/js_modules/dagit/packages/core/src/runs/RunTable.tsx b/js_modules/dagit/packages/core/src/runs/RunTable.tsx
@@ -164,6 +164,11 @@ export const RUN_TABLE_RUN_FRAGMENT = gql`
       repositoryLocationName
     }
     solidSelection
+    assetSelection {
+      ... on AssetKey {
+        path
+      }
+    }
     status
     tags {
       key

diff --git a/js_modules/dagit/packages/core/src/runs/RunUtils.tsx b/js_modules/dagit/packages/core/src/runs/RunUtils.tsx
@@ -164,6 +164,11 @@ export function getReexecutionVariables(input: {
       repositoryName,
       pipelineName: run.pipelineName,
       solidSelection: run.solidSelection,
+      assetSelection: run.assetSelection
+        ? run.assetSelection.map((asset_key) => ({
+            path: asset_key.path,
+          }))
+        : null,
     },
   };
 
@@ -180,7 +185,6 @@ export function getReexecutionVariables(input: {
       value: style.selection.query,
     });
   }
-
   return {executionParams};
 }
 

diff --git a/js_modules/dagit/packages/core/src/runs/types/RunActionButtonsTestQuery.ts b/js_modules/dagit/packages/core/src/runs/types/RunActionButtonsTestQuery.ts
diff --git a/js_modules/dagit/packages/core/src/runs/types/RunFragment.ts b/js_modules/dagit/packages/core/src/runs/types/RunFragment.ts
diff --git a/js_modules/dagit/packages/core/src/runs/types/RunRootQuery.ts b/js_modules/dagit/packages/core/src/runs/types/RunRootQuery.ts
diff --git a/js_modules/dagit/packages/core/src/runs/types/RunTableRunFragment.ts b/js_modules/dagit/packages/core/src/runs/types/RunTableRunFragment.ts
diff --git a/js_modules/dagit/packages/core/src/runs/types/RunsRootQuery.ts b/js_modules/dagit/packages/core/src/runs/types/RunsRootQuery.ts
diff --git a/js_modules/dagit/packages/core/src/schedules/types/PreviousRunsForScheduleQuery.ts b/js_modules/dagit/packages/core/src/schedules/types/PreviousRunsForScheduleQuery.ts
diff --git a/js_modules/dagit/packages/core/src/sensors/types/PreviousRunsForSensorQuery.ts b/js_modules/dagit/packages/core/src/sensors/types/PreviousRunsForSensorQuery.ts
diff --git a/js_modules/dagit/packages/core/src/types/globalTypes.ts b/js_modules/dagit/packages/core/src/types/globalTypes.ts
diff --git a/python_modules/dagster-graphql/dagster_graphql/implementation/execution/launch_execution.py b/python_modules/dagster-graphql/dagster_graphql/implementation/execution/launch_execution.py
@@ -33,9 +33,7 @@ def do_launch(graphene_info, execution_params, is_reexecuted=False):
         )
         check.str_param(execution_metadata.root_run_id, "root_run_id")
         check.str_param(execution_metadata.parent_run_id, "parent_run_id")
-
     external_pipeline = get_external_pipeline_or_raise(graphene_info, execution_params.selector)
-
     pipeline_run = create_valid_pipeline_run(graphene_info, external_pipeline, execution_params)
 
     return graphene_info.context.instance.submit_run(

diff --git a/python_modules/dagster-graphql/dagster_graphql/implementation/execution/run_lifecycle.py b/python_modules/dagster-graphql/dagster_graphql/implementation/execution/run_lifecycle.py
@@ -75,6 +75,9 @@ def create_valid_pipeline_run(graphene_info, external_pipeline, execution_params
         run_id=execution_params.execution_metadata.run_id
         if execution_params.execution_metadata.run_id
         else make_new_run_id(),
+        asset_selection=frozenset(execution_params.selector.asset_selection)
+        if execution_params.selector.asset_selection
+        else None,
         solid_selection=execution_params.selector.solid_selection,
         solids_to_execute=frozenset(execution_params.selector.solid_selection)
         if execution_params.selector.solid_selection

diff --git a/python_modules/dagster-graphql/dagster_graphql/implementation/external.py b/python_modules/dagster-graphql/dagster_graphql/implementation/external.py
@@ -29,7 +29,7 @@ def get_external_pipeline_or_raise(graphene_info, selector):
 
     full_pipeline = get_full_external_pipeline_or_raise(graphene_info, selector)
 
-    if selector.solid_selection is None:
+    if selector.solid_selection is None and selector.asset_selection is None:
         return full_pipeline
 
     return get_subset_external_pipeline(graphene_info.context, selector)

diff --git a/python_modules/dagster-graphql/dagster_graphql/implementation/utils.py b/python_modules/dagster-graphql/dagster_graphql/implementation/utils.py
@@ -5,6 +5,7 @@
 from graphql.execution.base import ResolveInfo
 
 import dagster._check as check
+from dagster.core.definitions.events import AssetKey
 from dagster.core.host_representation import GraphSelector, PipelineSelector
 from dagster.core.workspace.context import BaseWorkspaceRequestContext
 from dagster.utils.error import serializable_error_info_from_exc_info
@@ -65,11 +66,13 @@ def __init__(self, error):
 
 
 def pipeline_selector_from_graphql(data):
+    asset_selection = data.get("assetSelection", [])
     return PipelineSelector(
         location_name=data["repositoryLocationName"],
         repository_name=data["repositoryName"],
         pipeline_name=data.get("pipelineName") or data.get("jobName"),
         solid_selection=data.get("solidSelection"),
+        asset_selection=[AssetKey.from_graphql_input(asset_key) for asset_key in asset_selection],
     )
 
 

diff --git a/python_modules/dagster-graphql/dagster_graphql/schema/inputs.py b/python_modules/dagster-graphql/dagster_graphql/schema/inputs.py
@@ -105,6 +105,7 @@ class GrapheneJobOrPipelineSelector(graphene.InputObjectType):
     repositoryName = graphene.NonNull(graphene.String)
     repositoryLocationName = graphene.NonNull(graphene.String)
     solidSelection = graphene.List(graphene.NonNull(graphene.String))
+    assetSelection = graphene.List(graphene.NonNull(GrapheneAssetKeyInput))
 
     class Meta:
         description = """This type represents the fields necessary to identify a job or pipeline"""

diff --git a/python_modules/dagster-graphql/dagster_graphql/schema/pipelines/pipeline.py b/python_modules/dagster-graphql/dagster_graphql/schema/pipelines/pipeline.py
@@ -218,6 +218,7 @@ class GrapheneRun(graphene.ObjectType):
     pipelineName = graphene.NonNull(graphene.String)
     jobName = graphene.NonNull(graphene.String)
     solidSelection = graphene.List(graphene.NonNull(graphene.String))
+    assetSelection = graphene.List(graphene.NonNull(GrapheneAssetKey))
     resolvedOpSelection = graphene.List(graphene.NonNull(graphene.String))
     stats = graphene.NonNull(GrapheneRunStatsSnapshotOrError)
     stepStats = non_null_list(GrapheneRunStepStats)
@@ -287,6 +288,9 @@ def resolve_jobName(self, _graphene_info):
     def resolve_solidSelection(self, _graphene_info):
         return self._pipeline_run.solid_selection
 
+    def resolve_assetSelection(self, _graphene_info):
+        return self._pipeline_run.asset_selection
+
     def resolve_resolvedOpSelection(self, _graphene_info):
         return self._pipeline_run.solids_to_execute