[ui] Use location-based feature flag on partitions dialog (#17017)

## Summary & Motivation

Add a feature flag utility to read flags from location entries, use it
to gate the `SHOW_SINGLE_RUN_BACKFILL_TOGGLE` UI in the materialization
partition dialog.

<img width="758" alt="Screenshot 2023-10-04 at 5 41 44 PM"
src="https://github.com/dagster-io/dagster/assets/2823852/4f51346d-ff95-4af4-be10-95d07a4918d9">

## How I Tested These Changes

Jest to verify flag checking.

View
http://localhost:3000/locations/partitioned_assets_repository@dagster_test.toys.repo/asset-groups/default/lineage/hourly_asset1,
materialize. Verify that the radio container is not rendered.

Force the flag utility hook to return true, verify that the radio
container is rendered.

js_modules/dagster-ui/packages/ui-core/src/assets/LaunchAssetChoosePartitionsDialog.tsx

@@ -56,6 +56,7 @@ import {DimensionRangeWizard} from '../partitions/DimensionRangeWizard';
 import {assembleIntoSpans, stringForSpan} from '../partitions/SpanRepresentation';
 import {DagsterTag} from '../runs/RunTag';
 import {testId} from '../testing/testId';
+import {useFeatureFlagForCodeLocation} from '../workspace/WorkspaceContext';
 import {RepoAddress} from '../workspace/types';
 
 import {partitionCountString} from './AssetNodePartitionCounts';
@@ -149,6 +150,11 @@ const LaunchAssetChoosePartitionsDialogBody: React.FC<Props> = ({
   const [previewCount, setPreviewCount] = React.useState(0);
   const morePreviewsCount = partitionedAssets.length - previewCount;
 
+  const showSingleRunBackfillToggle = useFeatureFlagForCodeLocation(
+    repoAddress.location,
+    'SHOW_SINGLE_RUN_BACKFILL_TOGGLE',
+  );
+
   const [lastRefresh, setLastRefresh] = React.useState(Date.now());
 
   const refetch = async () => {
@@ -586,40 +592,42 @@ const LaunchAssetChoosePartitionsDialogBody: React.FC<Props> = ({
                 disabled={launchWithRangesAsTags}
                 onChange={() => setMissingFailedOnly(!missingFailedOnly)}
               />
-              <RadioContainer>
-                <Subheading>Launch as...</Subheading>
-                <Radio
-                  data-testid={testId('ranges-as-tags-true-radio')}
-                  checked={canLaunchWithRangesAsTags && launchWithRangesAsTags}
-                  disabled={!canLaunchWithRangesAsTags}
-                  onChange={() => setLaunchWithRangesAsTags(!launchWithRangesAsTags)}
-                >
-                  <Box flex={{direction: 'row', alignItems: 'center', gap: 8}}>
-                    <span>Single run</span>
-                    <Tooltip
-                      targetTagName="div"
-                      position="top-left"
-                      content={
-                        <div style={{maxWidth: 300}}>
-                          This option requires that your assets are written to operate on a
-                          partition key range via context.asset_partition_key_range_for_output or
-                          context.asset_partitions_time_window_for_output.
-                        </div>
-                      }
-                    >
-                      <Icon name="info" color={Colors.Gray500} />
-                    </Tooltip>
-                  </Box>
-                </Radio>
-                <Radio
-                  data-testid={testId('ranges-as-tags-false-radio')}
-                  checked={!canLaunchWithRangesAsTags || !launchWithRangesAsTags}
-                  disabled={!canLaunchWithRangesAsTags}
-                  onChange={() => setLaunchWithRangesAsTags(!launchWithRangesAsTags)}
-                >
-                  Multiple runs (One per selected partition)
-                </Radio>
-              </RadioContainer>
+              {showSingleRunBackfillToggle ? (
+                <RadioContainer>
+                  <Subheading>Launch as...</Subheading>
+                  <Radio
+                    data-testid={testId('ranges-as-tags-true-radio')}
+                    checked={canLaunchWithRangesAsTags && launchWithRangesAsTags}
+                    disabled={!canLaunchWithRangesAsTags}
+                    onChange={() => setLaunchWithRangesAsTags(!launchWithRangesAsTags)}
+                  >
+                    <Box flex={{direction: 'row', alignItems: 'center', gap: 8}}>
+                      <span>Single run</span>
+                      <Tooltip
+                        targetTagName="div"
+                        position="top-left"
+                        content={
+                          <div style={{maxWidth: 300}}>
+                            This option requires that your assets are written to operate on a
+                            partition key range via context.asset_partition_key_range_for_output or
+                            context.asset_partitions_time_window_for_output.
+                          </div>
+                        }
+                      >
+                        <Icon name="info" color={Colors.Gray500} />
+                      </Tooltip>
+                    </Box>
+                  </Radio>
+                  <Radio
+                    data-testid={testId('ranges-as-tags-false-radio')}
+                    checked={!canLaunchWithRangesAsTags || !launchWithRangesAsTags}
+                    disabled={!canLaunchWithRangesAsTags}
+                    onChange={() => setLaunchWithRangesAsTags(!launchWithRangesAsTags)}
+                  >
+                    Multiple runs (One per selected partition)
+                  </Radio>
+                </RadioContainer>
+              ) : null}
             </Box>
           )}
         </ToggleableSection>

js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/LaunchAssetExecutionButton.test.tsx

@@ -8,6 +8,7 @@ import {CustomAlertProvider} from '../../app/CustomAlertProvider';
 import {LaunchPartitionBackfillMutation} from '../../instance/backfill/types/BackfillUtils.types';
 import {LaunchPipelineExecutionMutation} from '../../runs/types/RunUtils.types';
 import {TestProvider} from '../../testing/TestProvider';
+import * as WorkspaceContext from '../../workspace/WorkspaceContext';
 import {
   AssetsInScope,
   ERROR_INVALID_ASSET_SELECTION,
@@ -38,6 +39,8 @@ import {
 // This file must be mocked because Jest can't handle `import.meta.url`.
 jest.mock('../../graph/asyncGraphLayout', () => ({}));
 
+const flagSpy = jest.spyOn(WorkspaceContext, 'useFeatureFlagForCodeLocation');
+
 describe('LaunchAssetExecutionButton', () => {
   describe('labeling', () => {
     it('should say "Materialize all" for an `all` scope', async () => {
@@ -306,61 +309,100 @@ describe('LaunchAssetExecutionButton', () => {
       await expectLaunchExecutesMutationAndCloses('Launch 1 run', launchMock);
     });
 
-    it('should launch backfills as pure-asset backfills if no job is in context', async () => {
-      const launchMock = buildExpectedLaunchBackfillMutation({
-        selector: undefined,
-        assetSelection: [{path: ['asset_daily']}],
-        partitionNames: ASSET_DAILY_PARTITION_KEYS,
-        fromFailure: false,
-        tags: [],
+    describe('Single run backfill toggle', () => {
+      afterEach(() => {
+        flagSpy.mockClear();
       });
-      renderButton({
-        scope: {all: [ASSET_DAILY]},
-        preferredJobName: undefined,
-        launchMock,
-      });
-      await clickMaterializeButton();
-      await screen.findByTestId('choose-partitions-dialog');
 
-      // missing-and-failed only option is available
-      expect(screen.getByTestId('missing-only-checkbox')).toBeEnabled();
+      it('should launch backfills as pure-asset backfills if no job is in context', async () => {
+        flagSpy.mockReturnValue(true);
+        const launchMock = buildExpectedLaunchBackfillMutation({
+          selector: undefined,
+          assetSelection: [{path: ['asset_daily']}],
+          partitionNames: ASSET_DAILY_PARTITION_KEYS,
+          fromFailure: false,
+          tags: [],
+        });
+        renderButton({
+          scope: {all: [ASSET_DAILY]},
+          preferredJobName: undefined,
+          launchMock,
+        });
+        await clickMaterializeButton();
+        await screen.findByTestId('choose-partitions-dialog');
 
-      // ranges-as-tags option is available
-      const rangesAsTags = screen.getByTestId('ranges-as-tags-true-radio');
-      await waitFor(async () => expect(rangesAsTags).toBeEnabled());
+        // missing-and-failed only option is available
+        expect(screen.getByTestId('missing-only-checkbox')).toBeEnabled();
 
-      await expectLaunchExecutesMutationAndCloses('Launch 1148-run backfill', launchMock);
-    });
+        // ranges-as-tags option is available
+        const rangesAsTags = screen.getByTestId('ranges-as-tags-true-radio');
+        await waitFor(async () => expect(rangesAsTags).toBeEnabled());
 
-    it('should launch a single run if you choose to pass the partition range using tags', async () => {
-      const launchMock = buildExpectedLaunchSingleRunMutation({
-        mode: 'default',
-        executionMetadata: {
-          tags: [
-            {key: 'dagster/asset_partition_range_start', value: '2020-01-02'},
-            {key: 'dagster/asset_partition_range_end', value: '2023-02-22'},
-          ],
-        },
-        runConfigData: '{}\n',
-        selector: {
-          repositoryLocationName: 'test.py',
-          repositoryName: 'repo',
-          pipelineName: 'my_asset_job',
-          assetSelection: [{path: ['asset_daily']}],
-        },
+        await expectLaunchExecutesMutationAndCloses('Launch 1148-run backfill', launchMock);
       });
-      renderButton({
-        scope: {all: [ASSET_DAILY]},
-        preferredJobName: 'my_asset_job',
-        launchMock,
+
+      it('should launch a single run if you choose to pass the partition range using tags', async () => {
+        flagSpy.mockReturnValue(true);
+        const launchMock = buildExpectedLaunchSingleRunMutation({
+          mode: 'default',
+          executionMetadata: {
+            tags: [
+              {key: 'dagster/asset_partition_range_start', value: '2020-01-02'},
+              {key: 'dagster/asset_partition_range_end', value: '2023-02-22'},
+            ],
+          },
+          runConfigData: '{}\n',
+          selector: {
+            repositoryLocationName: 'test.py',
+            repositoryName: 'repo',
+            pipelineName: 'my_asset_job',
+            assetSelection: [{path: ['asset_daily']}],
+          },
+        });
+        renderButton({
+          scope: {all: [ASSET_DAILY]},
+          preferredJobName: 'my_asset_job',
+          launchMock,
+        });
+        await clickMaterializeButton();
+        await screen.findByTestId('choose-partitions-dialog');
+
+        const rangesAsTags = screen.getByTestId('ranges-as-tags-true-radio');
+        await waitFor(async () => expect(rangesAsTags).toBeEnabled());
+        await userEvent.click(rangesAsTags);
+        await expectLaunchExecutesMutationAndCloses('Launch 1 run', launchMock);
       });
-      await clickMaterializeButton();
-      await screen.findByTestId('choose-partitions-dialog');
 
-      const rangesAsTags = screen.getByTestId('ranges-as-tags-true-radio');
-      await waitFor(async () => expect(rangesAsTags).toBeEnabled());
-      await userEvent.click(rangesAsTags);
-      await expectLaunchExecutesMutationAndCloses('Launch 1 run', launchMock);
+      it('should not show the backfill toggle if the flag is false', async () => {
+        flagSpy.mockReturnValue(false);
+
+        const launchMock = buildExpectedLaunchSingleRunMutation({
+          mode: 'default',
+          executionMetadata: {
+            tags: [
+              {key: 'dagster/asset_partition_range_start', value: '2020-01-02'},
+              {key: 'dagster/asset_partition_range_end', value: '2023-02-22'},
+            ],
+          },
+          runConfigData: '{}\n',
+          selector: {
+            repositoryLocationName: 'test.py',
+            repositoryName: 'repo',
+            pipelineName: 'my_asset_job',
+            assetSelection: [{path: ['asset_daily']}],
+          },
+        });
+        renderButton({
+          scope: {all: [ASSET_DAILY]},
+          preferredJobName: 'my_asset_job',
+          launchMock,
+        });
+        await clickMaterializeButton();
+        await screen.findByTestId('choose-partitions-dialog');
+
+        const rangesAsTags = screen.queryByTestId('ranges-as-tags-true-radio');
+        expect(rangesAsTags).toBeNull();
+      });
     });
   });
 

js_modules/dagster-ui/packages/ui-core/src/workspace/WorkspaceContext.tsx

@@ -77,6 +77,10 @@ export const ROOT_WORKSPACE_QUERY = gql`
       ...WorkspaceDisplayMetadata
     }
     updatedTimestamp
+    featureFlags {
+      name
+      enabled
+    }
     locationOrLoadError {
       ... on RepositoryLocation {
         id
@@ -369,6 +373,27 @@ export const useActivePipelineForName = (pipelineName: string, snapshotId?: stri
   return null;
 };
 
+export const getFeatureFlagForCodeLocation = (
+  locationEntries: WorkspaceLocationNodeFragment[],
+  locationName: string,
+  flagName: string,
+) => {
+  const matchingLocation = locationEntries.find(({id}) => id === locationName);
+  if (matchingLocation) {
+    const {featureFlags} = matchingLocation;
+    const matchingFlag = featureFlags.find(({name}) => name === flagName);
+    if (matchingFlag) {
+      return matchingFlag.enabled;
+    }
+  }
+  return false;
+};
+
+export const useFeatureFlagForCodeLocation = (locationName: string, flagName: string) => {
+  const {locationEntries} = useWorkspaceState();
+  return getFeatureFlagForCodeLocation(locationEntries, locationName, flagName);
+};
+
 export const isThisThingAJob = (repo: DagsterRepoOption | null, pipelineOrJobName: string) => {
   const pipelineOrJob = repo?.repository.pipelines.find(
     (pipelineOrJob) => pipelineOrJob.name === pipelineOrJobName,

js_modules/dagster-ui/packages/ui-core/src/workspace/__tests__/getFeatureFlagForCodeLocation.test.tsx

@@ -0,0 +1,43 @@
+import {buildFeatureFlag, buildWorkspaceLocationEntry} from '../../graphql/types';
+import {getFeatureFlagForCodeLocation} from '../WorkspaceContext';
+import {WorkspaceLocationNodeFragment} from '../types/WorkspaceContext.types';
+
+describe('getFeatureFlagForCodeLocation', () => {
+  const locationEntries: WorkspaceLocationNodeFragment[] = [
+    buildWorkspaceLocationEntry({
+      id: 'foo',
+      featureFlags: [
+        buildFeatureFlag({
+          name: 'LOREM',
+          enabled: false,
+        }),
+        buildFeatureFlag({
+          name: 'IPSUM',
+          enabled: true,
+        }),
+      ],
+    }),
+  ];
+
+  it('returns true for an enabled flag in a matching code location', () => {
+    expect(getFeatureFlagForCodeLocation(locationEntries, 'foo', 'IPSUM')).toBe(true);
+  });
+
+  it('returns false for a disabled flag in a matching code location', () => {
+    expect(getFeatureFlagForCodeLocation(locationEntries, 'foo', 'LOREM')).toBe(false);
+  });
+
+  it('returns false for a flag that cannot be found in a matching code location', () => {
+    expect(getFeatureFlagForCodeLocation(locationEntries, 'foo', 'FAKE_FLAG')).toBe(false);
+  });
+
+  it('returns false for an existing flag that cannot be found in an unknown code location', () => {
+    expect(getFeatureFlagForCodeLocation(locationEntries, 'FAKE_LOCATION', 'IPSUM')).toBe(false);
+  });
+
+  it('returns false for an unknown flag in an unknown code location', () => {
+    expect(getFeatureFlagForCodeLocation(locationEntries, 'FAKE_LOCATION', 'FAKE_FLAG')).toBe(
+      false,
+    );
+  });
+});