[3/n][Selection syntax] Add stubs to allow Cloud to extend selection …

…syntax functionality. (dagster-io#27184) ## Summary & Motivation - Add `useAssetGraphSupplementaryData` so that Cloud can query for extra data to support the filtering - Feed that data all the way through to the visitor that does the filtering. Don't do anything with it in the OSS visitor, Cloud will use it in its visitor though to do that additional filtering. ## How I Tested These Changes Manual testing + relying on existing jest tests <img width="650" alt="Screenshot 2025-01-16 at 6 25 14 PM" src="https://github.com/user-attachments/assets/a5bbf8d7-c4f0-4d8b-bd53-0debdd3b794f" />
marijncv · Jan 21, 2025 · 9446dd8 · 9446dd8
1 parent d6520ef
commit 9446dd8
Show file tree

Hide file tree

Showing 15 changed files with 172 additions and 132 deletions.
diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/ComputeGraphData.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/ComputeGraphData.ts
@@ -4,7 +4,7 @@ import {ComputeGraphDataMessageType} from './ComputeGraphData.types';
 import {GraphData, buildGraphData, toGraphId} from './Utils';
 import {AssetNodeForGraphQueryFragment} from './types/useAssetGraphData.types';
 import {GraphDataState} from './useAssetGraphData';
-import {filterAssetSelectionByQuery} from '../asset-selection/AntlrAssetSelection';
+import {filterAssetSelectionByQuery} from '../asset-selection/filterAssetSelectionByQuery';
 import {doesFilterArrayMatchValueArray} from '../ui/Filters/doesFilterArrayMatchValueArray';
 
 export function computeGraphData({
@@ -13,6 +13,7 @@ export function computeGraphData({
   opsQuery,
   kinds: _kinds,
   hideEdgesToNodesOutsideQuery,
+  supplementaryData,
 }: Omit<ComputeGraphDataMessageType, 'id' | 'type'>): GraphDataState {
   if (repoFilteredNodes === undefined || graphQueryItems === undefined) {
     return {
@@ -26,7 +27,11 @@ export function computeGraphData({
   // In the future it might be ideal to move this server-side, but we currently
   // get to leverage the useQuery cache almost 100% of the time above, making this
   // super fast after the first load vs a network fetch on every page view.
-  const {all: allFilteredByOpQuery} = filterAssetSelectionByQuery(graphQueryItems, opsQuery);
+  const {all: allFilteredByOpQuery} = filterAssetSelectionByQuery(
+    graphQueryItems,
+    opsQuery,
+    supplementaryData,
+  );
   const kinds = _kinds?.map((c) => c.toLowerCase());
   const all = kinds?.length
     ? allFilteredByOpQuery.filter(

diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/ComputeGraphData.types.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/ComputeGraphData.types.ts
@@ -1,3 +1,4 @@
+import {AssetKey} from '../assets/types';
 import {AssetNodeForGraphQueryFragment} from './types/useAssetGraphData.types';
 import {AssetGraphFetchScope, AssetGraphQueryItem} from './useAssetGraphData';
 
@@ -13,6 +14,7 @@ export type ComputeGraphDataMessageType = BaseType & {
   opsQuery: string;
   kinds: AssetGraphFetchScope['kinds'];
   hideEdgesToNodesOutsideQuery?: boolean;
+  supplementaryData?: Record<string, AssetKey[]> | null;
 };
 
 export type BuildGraphDataMessageType = BaseType & {

diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx
@@ -3,6 +3,7 @@ import memoize from 'lodash/memoize';
 import reject from 'lodash/reject';
 import {useEffect, useMemo, useRef, useState} from 'react';
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
+import {useAssetGraphSupplementaryData} from 'shared/asset-graph/useAssetGraphSupplementaryData.oss';
 
 import {ASSET_NODE_FRAGMENT} from './AssetNode';
 import {GraphData, buildGraphData as buildGraphDataImpl, tokenForAssetKey} from './Utils';
@@ -160,8 +161,11 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop
   const lastProcessedRequestRef = useRef(0);
   const currentRequestRef = useRef(0);
 
+  const {loading: supplementaryDataLoading, data: supplementaryData} =
+    useAssetGraphSupplementaryData(opsQuery);
+
   useEffect(() => {
-    if (options.loading) {
+    if (options.loading || supplementaryDataLoading) {
       return;
     }
     const requestId = ++currentRequestRef.current;
@@ -173,6 +177,7 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop
       kinds,
       hideEdgesToNodesOutsideQuery,
       flagSelectionSyntax: featureEnabled(FeatureFlag.flagSelectionSyntax),
+      supplementaryData,
     })
       ?.then((data) => {
         if (lastProcessedRequestRef.current < requestId) {
@@ -197,6 +202,8 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop
     kinds,
     hideEdgesToNodesOutsideQuery,
     options.loading,
+    supplementaryData,
+    supplementaryDataLoading,
   ]);
 
   const loading = fetchResult.loading || graphDataLoading;

diff --git a/...odules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphSupplementaryData.oss.tsx b/...odules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphSupplementaryData.oss.tsx
@@ -0,0 +1,13 @@
+import {AssetKey} from '../assets/types';
+
+export type SupplementaryInformation = Record<string, AssetKey[]> | null | undefined;
+
+// Stub for cloud to supply extended selection syntax filtering capabilities
+export const useAssetGraphSupplementaryData = (
+  _: string,
+): {loading: boolean; data: SupplementaryInformation | null} => {
+  return {
+    loading: false,
+    data: null,
+  };
+};
diff --git a/...t-selection/AntlrAssetSelectionVisitor.ts → ...lection/AntlrAssetSelectionVisitor.oss.ts b/...t-selection/AntlrAssetSelectionVisitor.ts → ...lection/AntlrAssetSelectionVisitor.oss.ts
@@ -1,14 +1,13 @@
 import {AbstractParseTreeVisitor} from 'antlr4ts/tree/AbstractParseTreeVisitor';
-
+import {SupplementaryInformation} from 'shared/asset-graph/useAssetGraphSupplementaryData.oss';
 import {
   AllExpressionContext,
   AndExpressionContext,
+  AssetSelectionVisitor,
   AttributeExpressionContext,
   CodeLocationAttributeExprContext,
-  DownTraversalContext,
   DownTraversalExpressionContext,
   FunctionCallExpressionContext,
-  FunctionNameContext,
   GroupAttributeExprContext,
   KeyExprContext,
   KeySubstringExprContext,
@@ -21,43 +20,14 @@ import {
   TagAttributeExprContext,
   TraversalAllowedExpressionContext,
   UpAndDownTraversalExpressionContext,
-  UpTraversalContext,
   UpTraversalExpressionContext,
-  ValueContext,
-} from './generated/AssetSelectionParser';
-import {AssetSelectionVisitor} from './generated/AssetSelectionVisitor';
+} from 'shared/asset-selection/AssetSelectionAntlr.oss';
+
+import {getFunctionName, getTraversalDepth, getValue} from './util';
 import {GraphTraverser} from '../app/GraphQueryImpl';
 import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData';
 import {buildRepoPathForHuman} from '../workspace/buildRepoAddress';
 
-export function getTraversalDepth(ctx: UpTraversalContext | DownTraversalContext): number {
-  const digits = ctx.DIGITS();
-  if (digits) {
-    return parseInt(ctx.text);
-  }
-  return Number.MAX_SAFE_INTEGER;
-}
-
-export function getFunctionName(ctx: FunctionNameContext): string {
-  if (ctx.SINKS()) {
-    return 'sinks';
-  }
-  if (ctx.ROOTS()) {
-    return 'roots';
-  }
-  throw new Error('Invalid function name');
-}
-
-export function getValue(ctx: ValueContext): string {
-  if (ctx.QUOTED_STRING()) {
-    return ctx.text.slice(1, -1);
-  }
-  if (ctx.UNQUOTED_STRING()) {
-    return ctx.text;
-  }
-  throw new Error('Invalid value');
-}
-
 export class AntlrAssetSelectionVisitor
   extends AbstractParseTreeVisitor<Set<AssetGraphQueryItem>>
   implements AssetSelectionVisitor<Set<AssetGraphQueryItem>>
@@ -70,7 +40,8 @@ export class AntlrAssetSelectionVisitor
     return new Set<AssetGraphQueryItem>();
   }
 
-  constructor(all_assets: AssetGraphQueryItem[]) {
+  // Supplementary data is not used in oss
+  constructor(all_assets: AssetGraphQueryItem[], _supplementaryData?: SupplementaryInformation) {
     super();
     this.all_assets = new Set(all_assets);
     this.focus_assets = new Set();
@@ -192,15 +163,15 @@ export class AntlrAssetSelectionVisitor
 
   visitTagAttributeExpr(ctx: TagAttributeExprContext) {
     const key: string = getValue(ctx.value(0));
+    let value: string | undefined = undefined;
     if (ctx.EQUAL()) {
-      const value: string = getValue(ctx.value(1));
-      return new Set(
-        [...this.all_assets].filter((i) =>
-          i.node.tags.some((t) => t.key === key && t.value === value),
-        ),
-      );
+      value = getValue(ctx.value(1));
     }
-    return new Set([...this.all_assets].filter((i) => i.node.tags.some((t) => t.key === key)));
+    return new Set(
+      [...this.all_assets].filter((i) =>
+        i.node.tags.some((t) => t.key === key && (!value || t.value === value)),
+      ),
+    );
   }
 
   visitOwnerAttributeExpr(ctx: OwnerAttributeExprContext) {

diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AssetSelectionAntlr.oss.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AssetSelectionAntlr.oss.ts
@@ -0,0 +1,4 @@
+export * from './generated/AssetSelectionLexer';
+export * from './generated/AssetSelectionParser';
+export * from './generated/AssetSelectionVisitor';
+export * from './generated/AssetSelectionListener';
diff --git a/...les/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts b/...les/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts
@@ -7,7 +7,7 @@ import {
   buildRepositoryLocation,
   buildUserAssetOwner,
 } from '../../graphql/types';
-import {parseAssetSelectionQuery} from '../AntlrAssetSelection';
+import {parseAssetSelectionQuery} from '../filterAssetSelectionByQuery';
 
 const TEST_GRAPH: AssetGraphQueryItem[] = [
   // Top Layer
@@ -53,7 +53,6 @@ const TEST_GRAPH: AssetGraphQueryItem[] = [
 
 function assertQueryResult(query: string, expectedNames: string[]) {
   const result = parseAssetSelectionQuery(TEST_GRAPH, query);
-  expect(result).not.toBeInstanceOf(Error);
   if (result instanceof Error) {
     throw result;
   }

diff --git a/...rc/asset-selection/AntlrAssetSelection.ts → ...-selection/filterAssetSelectionByQuery.ts b/...rc/asset-selection/AntlrAssetSelection.ts → ...-selection/filterAssetSelectionByQuery.ts
@@ -6,14 +6,17 @@ import {
   Recognizer,
 } from 'antlr4ts';
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
+import {SupplementaryInformation} from 'shared/asset-graph/useAssetGraphSupplementaryData.oss';
+import {AntlrAssetSelectionVisitor} from 'shared/asset-selection/AntlrAssetSelectionVisitor.oss';
+import {
+  AssetSelectionLexer,
+  AssetSelectionParser,
+} from 'shared/asset-selection/AssetSelectionAntlr.oss';
 
-import {AntlrAssetSelectionVisitor} from './AntlrAssetSelectionVisitor';
-import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData';
-import {weakMapMemoize} from '../util/weakMapMemoize';
-import {AssetSelectionLexer} from './generated/AssetSelectionLexer';
-import {AssetSelectionParser} from './generated/AssetSelectionParser';
 import {featureEnabled} from '../app/Flags';
 import {filterByQuery} from '../app/GraphQueryImpl';
+import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData';
+import {weakMapMemoize} from '../util/weakMapMemoize';
 
 export class AntlrInputErrorListener implements ANTLRErrorListener<any> {
   syntaxError(
@@ -39,6 +42,7 @@ type AssetSelectionQueryResult = {
 export const parseAssetSelectionQuery = (
   all_assets: AssetGraphQueryItem[],
   query: string,
+  supplementaryData?: SupplementaryInformation,
 ): AssetSelectionQueryResult | Error => {
   try {
     const lexer = new AssetSelectionLexer(CharStreams.fromString(query));
@@ -53,7 +57,8 @@ export const parseAssetSelectionQuery = (
 
     const tree = parser.start();
 
-    const visitor = new AntlrAssetSelectionVisitor(all_assets);
+    const visitor = new AntlrAssetSelectionVisitor(all_assets, supplementaryData);
+
     const all_selection = visitor.visit(tree);
     const focus_selection = visitor.focus_assets;
 
@@ -67,9 +72,13 @@ export const parseAssetSelectionQuery = (
 };
 
 export const filterAssetSelectionByQuery = weakMapMemoize(
-  (all_assets: AssetGraphQueryItem[], query: string): AssetSelectionQueryResult => {
+  (
+    all_assets: AssetGraphQueryItem[],
+    query: string,
+    supplementaryData: SupplementaryInformation,
+  ): AssetSelectionQueryResult => {
     if (featureEnabled(FeatureFlag.flagSelectionSyntax)) {
-      const result = parseAssetSelectionQuery(all_assets, query);
+      const result = parseAssetSelectionQuery(all_assets, query, supplementaryData);
       if (result instanceof Error) {
         // fall back to old behavior
         return filterByQuery(all_assets, query);