Working!

dagster-io · Mar 16, 2024 · 74f58c8 · 74f58c8
1 parent 99fe064
commit 74f58c8
Show file tree

Hide file tree

Showing 11 changed files with 444 additions and 173 deletions.
diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetColumnsNode.tsx b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetColumnsNode.tsx
@@ -8,37 +8,43 @@ import {
 } from './AssetNode';
 import {AssetNodeFragment} from './types/AssetNode.types';
 import {AssetColumnLineageLocal} from '../assets/lineage/useColumnLineageDataForAssets';
-import {AssetComputeKindTag} from '../graph/OpTags';
 import {TypeTag} from '../metadata/TableSchema';
 
-export const AssetColumnsNode = ({
+export const AssetColumnsGroupNode = ({
   selected,
   definition,
-  columns,
+  height,
 }: {
   selected: boolean;
   definition: AssetNodeFragment;
-  columns: AssetColumnLineageLocal[''][];
+  height: number;
 }) => {
   return (
-    <AssetInsetForHoverEffect style={{opacity: columns.length > 0 ? 1 : 0.4}}>
+    <AssetInsetForHoverEffect>
       <AssetNodeContainer $selected={selected}>
         <div style={{minHeight: 24}} />
         <AssetNodeBox $selected={selected} $isSource={definition.isSource}>
           <AssetNameRow definition={definition} />
-          <Box style={{padding: '6px 8px'}} flex={{direction: 'column'}}>
-            {columns.map((c) => (
-              <Tooltip key={c.name} content={c.description || 'No description provided'}>
-                <Box flex={{gap: 4, justifyContent: 'space-between', alignItems: 'center'}}>
-                  <Caption>{c.name}</Caption>
-                  <TypeTag type={c.type || ''} />
-                </Box>
-              </Tooltip>
-            ))}
-          </Box>
+          <Box style={{height: height - 60}} flex={{direction: 'column'}}></Box>
         </AssetNodeBox>
-        <AssetComputeKindTag definition={definition} style={{right: -2, paddingTop: 7}} />
       </AssetNodeContainer>
     </AssetInsetForHoverEffect>
   );
 };
+
+export const AssetColumnNode = ({column}: {column: AssetColumnLineageLocal['']}) => {
+  return (
+    <Box style={{width: '100%', height: 32}} flex={{direction: 'column'}}>
+      <Tooltip key={column.name} content={column.description || 'No description provided'}>
+        <Box
+          style={{height: 32}}
+          padding={{left: 8, right: 4}}
+          flex={{gap: 4, justifyContent: 'space-between', alignItems: 'center'}}
+        >
+          <Caption>{column.name}</Caption>
+          <TypeTag type={column.type || ''} />
+        </Box>
+      </Tooltip>
+    </Box>
+  );
+};
diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetGraphExplorer.tsx b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetGraphExplorer.tsx
@@ -16,6 +16,7 @@ import pickBy from 'lodash/pickBy';
 import uniq from 'lodash/uniq';
 import without from 'lodash/without';
 import * as React from 'react';
+import {useMemo} from 'react';
 import styled from 'styled-components';
 
 import {AssetEdges} from './AssetEdges';
@@ -241,7 +242,11 @@ const AssetGraphExplorerWithData = ({
   });
   const focusGroupIdAfterLayoutRef = React.useRef('');
 
-  const {layout, loading, async} = useAssetLayout(assetGraphData, expandedGroups, direction);
+  const {layout, loading, async} = useAssetLayout(
+    assetGraphData,
+    expandedGroups,
+    useMemo(() => ({direction}), [direction]),
+  );
 
   React.useEffect(() => {
     if (!loading) {

diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/Utils.tsx b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/Utils.tsx
@@ -44,7 +44,7 @@ export function isHiddenAssetGroupJob(jobName: string) {
 //
 export type GraphId = string;
 export const toGraphId = (key: {path: string[]}): GraphId => JSON.stringify(key.path);
-export const fromGraphID = (graphId: GraphId): AssetNodeKeyFragment => ({
+export const fromGraphId = (graphId: GraphId): AssetNodeKeyFragment => ({
   path: JSON.parse(graphId),
   __typename: 'AssetKey',
 });
@@ -269,7 +269,7 @@ export const itemWithAssetKey = (key: {path: string[]}) => {
   return (asset: {assetKey: {path: string[]}}) => tokenForAssetKey(asset.assetKey) === token;
 };
 
-export const isGroupId = (str: string) => /^[^@:]+@[^@:]+:[^@:]+$/.test(str);
+export const isGroupId = (str: string) => /^[^@:]+@[^@:]+:.+$/.test(str);
 
 export const groupIdForNode = (node: GraphNode) =>
   [
@@ -284,7 +284,7 @@ export const groupIdForNode = (node: GraphNode) =>
 export const getUpstreamNodes = memoize(
   (assetKey: AssetNodeKeyFragment, graphData: GraphData): AssetNodeKeyFragment[] => {
     const upstream = Object.keys(graphData.upstream[toGraphId(assetKey)] || {});
-    const currentUpstream = upstream.map((graphId) => fromGraphID(graphId));
+    const currentUpstream = upstream.map((graphId) => fromGraphId(graphId));
     return [
       assetKey,
       ...currentUpstream,

diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/layout.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/layout.ts
@@ -34,39 +34,61 @@ export type AssetGraphLayout = {
 };
 const MARGIN = 100;
 
+export type LayoutAssetGraphConfig = dagre.GraphLabel & {
+  direction: AssetLayoutDirection;
+  /** Pass `auto` to use getAssetNodeDimensions, or a value to give nodes a fixed height */
+  nodeHeight: number | 'auto';
+  /** Our asset groups have "title bars" - use these numbers to adjust the bounding boxes.
+   * Note that these adjustments are applied post-dagre layout. For padding > nodesep, you
+   * may need to set "clusterpaddingtop", "clusterpaddingbottom" so Dagre lays out the boxes
+   * with more spacing.
+   */
+  groupPaddingTop: number;
+  groupPaddingBottom: number;
+};
+
 export type LayoutAssetGraphOptions = {
   direction: AssetLayoutDirection;
+  overrides?: Partial<LayoutAssetGraphConfig>;
+};
+
+export const Config = {
+  horizontal: {
+    ranker: 'tight-tree',
+    direction: 'horizontal',
+    marginx: MARGIN,
+    marginy: MARGIN,
+    ranksep: 60,
+    rankdir: 'LR',
+    edgesep: 90,
+    nodesep: -10,
+    nodeHeight: 'auto',
+    groupPaddingTop: 65,
+    groupPaddingBottom: -15,
+  },
+  vertical: {
+    ranker: 'tight-tree',
+    direction: 'horizontal',
+    marginx: MARGIN,
+    marginy: MARGIN,
+    ranksep: 20,
+    rankdir: 'TB',
+    nodesep: 40,
+    edgesep: 10,
+    nodeHeight: 'auto',
+    groupPaddingTop: 40,
+    groupPaddingBottom: -20,
+  },
 };
 
 export const layoutAssetGraph = (
   graphData: GraphData,
   opts: LayoutAssetGraphOptions,
 ): AssetGraphLayout => {
   const g = new dagre.graphlib.Graph({compound: true});
+  const config = Object.assign({}, Config[opts.direction], opts.overrides || {});
 
-  const ranker = 'tight-tree';
-
-  g.setGraph(
-    opts.direction === 'horizontal'
-      ? {
-          rankdir: 'LR',
-          marginx: MARGIN,
-          marginy: MARGIN,
-          nodesep: -10,
-          edgesep: 90,
-          ranksep: 60,
-          ranker,
-        }
-      : {
-          rankdir: 'TB',
-          marginx: MARGIN,
-          marginy: MARGIN,
-          nodesep: 40,
-          edgesep: 10,
-          ranksep: 20,
-          ranker,
-        },
-  );
+  g.setGraph(config);
   g.setDefaultEdgeLabel(() => ({}));
 
   // const shouldRender = (node?: GraphNode) => node && node.definition.opNames.length > 0;
@@ -95,17 +117,24 @@ export const layoutAssetGraph = (
   if (groupsPresent) {
     Object.keys(groups).forEach((groupId) => {
       if (expandedGroups.includes(groupId)) {
-        g.setNode(groupId, {}); // sized based on it's children
+        // sized based on it's children, but "border" tells Dagre we want cluster-level
+        // spacing between the node and others. Necessary because our groups have title bars.
+        g.setNode(groupId, {borderType: 'borderRight'});
       } else {
-        g.setNode(groupId, {width: 320, height: 110});
+        g.setNode(groupId, {width: ASSET_NODE_WIDTH, height: 110});
       }
     });
   }
 
   // Add all the nodes inside expanded groups to the graph
   renderedNodes.forEach((node) => {
     if (!groupsPresent || expandedGroups.includes(groupIdForNode(node))) {
-      g.setNode(node.id, getAssetNodeDimensions(node.definition));
+      const label =
+        config.nodeHeight === 'auto'
+          ? getAssetNodeDimensions(node.definition)
+          : {width: ASSET_NODE_WIDTH, height: config.nodeHeight};
+
+      g.setNode(node.id, label);
       if (groupsPresent && node.definition.groupName) {
         g.setParent(node.id, groupIdForNode(node));
       }
@@ -204,10 +233,11 @@ export const layoutAssetGraph = (
     }
     for (const group of Object.values(groups)) {
       if (group.expanded) {
-        group.bounds =
-          opts.direction === 'horizontal'
-            ? padBounds(group.bounds, {x: 15, top: 65, bottom: -15})
-            : padBounds(group.bounds, {x: 15, top: 40, bottom: -20});
+        group.bounds = padBounds(group.bounds, {
+          x: 15,
+          top: config.groupPaddingTop,
+          bottom: config.groupPaddingBottom,
+        });
       }
     }
   }
@@ -275,6 +305,7 @@ export const extendBounds = (a: IBounds, b: IBounds) => {
   return {x: xmin, y: ymin, width: xmax - xmin, height: ymax - ymin};
 };
 
+export const ASSET_NODE_WIDTH = 320;
 export const ASSET_NODE_NAME_MAX_LENGTH = 38;
 
 export const getAssetNodeDimensions = (def: {
@@ -287,7 +318,7 @@ export const getAssetNodeDimensions = (def: {
   description?: string | null;
   computeKind: string | null;
 }) => {
-  const width = 320;
+  const width = ASSET_NODE_WIDTH;
 
   let height = 100; // top tags area + name + description