Skip to content

Commit

Permalink
[dagit] Move data fetching for the asset graph into hooks (#7363)
Browse files Browse the repository at this point in the history
  • Loading branch information
bengotow committed Apr 12, 2022
1 parent fac5a5f commit 6d5eae5
Show file tree
Hide file tree
Showing 3 changed files with 207 additions and 144 deletions.
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
import {gql, useQuery} from '@apollo/client';
import {Box, Checkbox, NonIdealState, SplitPanelContainer} from '@dagster-io/ui';
import _, {flatMap, uniq, uniqBy, without} from 'lodash';
import React from 'react';
import styled from 'styled-components/macro';

import {filterByQuery, GraphQueryItem} from '../../app/GraphQueryImpl';
import {GraphQueryItem} from '../../app/GraphQueryImpl';
import {
FIFTEEN_SECONDS,
QueryRefreshCountdown,
Expand Down Expand Up @@ -36,29 +35,16 @@ import {GraphQueryInput} from '../../ui/GraphQueryInput';
import {Loading} from '../../ui/Loading';

import {AssetLinks} from './AssetLinks';
import {AssetNode, ASSET_NODE_FRAGMENT, ASSET_NODE_LIVE_FRAGMENT} from './AssetNode';
import {AssetNode} from './AssetNode';
import {ForeignNode} from './ForeignNode';
import {LaunchAssetExecutionButton} from './LaunchAssetExecutionButton';
import {OmittedAssetsNotice} from './OmittedAssetsNotice';
import {SidebarAssetInfo} from './SidebarAssetInfo';
import {
buildGraphData,
buildLiveData,
GraphData,
graphHasCycles,
REPOSITORY_LIVE_FRAGMENT,
layoutGraph,
LiveData,
Node,
isSourceAsset,
} from './Utils';
import {AssetGraphLiveQuery, AssetGraphLiveQueryVariables} from './types/AssetGraphLiveQuery';
import {
AssetGraphQuery,
AssetGraphQueryVariables,
AssetGraphQuery_assetNodes,
} from './types/AssetGraphQuery';
import {GraphData, graphHasCycles, layoutGraph, LiveData, Node, isSourceAsset} from './Utils';
import {AssetGraphQuery_assetNodes} from './types/AssetGraphQuery';
import {useAssetGraphData} from './useAssetGraphData';
import {useFindAssetInWorkspace} from './useFindAssetInWorkspace';
import {useLiveDataForAssetKeys} from './useLiveDataForAssetKeys';

type AssetNode = AssetGraphQuery_assetNodes;

Expand All @@ -79,77 +65,22 @@ interface Props {
}

export const AssetGraphExplorer: React.FC<Props> = (props) => {
const {pipelineSelector, explorerPath} = props;

const fetchResult = useQuery<AssetGraphQuery, AssetGraphQueryVariables>(ASSETS_GRAPH_QUERY, {
variables: {pipelineSelector},
notifyOnNetworkStatusChange: true,
});

const fetchResultFilteredNodes = React.useMemo(() => {
const nodes = fetchResult.data?.assetNodes;
if (!nodes) {
return undefined;
}
return props.filterNodes ? nodes.filter(props.filterNodes) : nodes;
}, [fetchResult.data, props.filterNodes]);

const {
fetchResult,
assetGraphData,
graphQueryItems,
graphAssetKeys,
applyingEmptyDefault,
} = React.useMemo(() => {
if (fetchResultFilteredNodes === undefined) {
return {
graphAssetKeys: [],
graphQueryItems: [],
assetGraphData: null,
applyingEmptyDefault: false,
};
}
const graphQueryItems = buildGraphQueryItems(fetchResultFilteredNodes);
const {all, applyingEmptyDefault} = filterByQuery(graphQueryItems, explorerPath.opsQuery);

return {
graphAssetKeys: all.map((n) => ({path: n.node.assetKey.path})),
assetGraphData: buildGraphData(all.map((n) => n.node)),
graphQueryItems,
applyingEmptyDefault,
};
}, [fetchResultFilteredNodes, explorerPath.opsQuery]);

const liveResult = useQuery<AssetGraphLiveQuery, AssetGraphLiveQueryVariables>(
ASSETS_GRAPH_LIVE_QUERY,
{
skip: graphAssetKeys.length === 0,
variables: {
assetKeys: graphAssetKeys,
repositorySelector: pipelineSelector
? {
repositoryLocationName: pipelineSelector.repositoryLocationName,
repositoryName: pipelineSelector.repositoryName,
}
: undefined,
},
notifyOnNetworkStatusChange: true,
},
} = useAssetGraphData(props.pipelineSelector, props.explorerPath.opsQuery, props.filterNodes);

const {liveResult, liveDataByNode} = useLiveDataForAssetKeys(
props.pipelineSelector,
assetGraphData,
graphAssetKeys,
);

const liveDataRefreshState = useQueryRefreshAtInterval(liveResult, FIFTEEN_SECONDS);

const liveDataByNode = React.useMemo(() => {
if (!liveResult.data || !assetGraphData) {
return {};
}

const {repositoriesOrError, assetNodes: liveAssetNodes} = liveResult.data;
const repos =
repositoriesOrError.__typename === 'RepositoryConnection' ? repositoriesOrError.nodes : [];

return buildLiveData(assetGraphData, liveAssetNodes, repos);
}, [assetGraphData, liveResult]);

useDocumentTitle('Assets');
useDidLaunchEvent(liveResult.refetch);

Expand All @@ -175,7 +106,7 @@ export const AssetGraphExplorer: React.FC<Props> = (props) => {

return (
<AssetGraphExplorerWithData
key={explorerPath.pipelineName}
key={props.explorerPath.pipelineName}
assetKeys={assetKeys}
assetGraphData={assetGraphData}
graphQueryItems={graphQueryItems}
Expand Down Expand Up @@ -438,44 +369,6 @@ const AssetGraphExplorerWithData: React.FC<
);
};

const ASSETS_GRAPH_LIVE_QUERY = gql`
query AssetGraphLiveQuery($repositorySelector: RepositorySelector, $assetKeys: [AssetKeyInput!]) {
repositoriesOrError(repositorySelector: $repositorySelector) {
__typename
... on RepositoryConnection {
nodes {
__typename
id
...RepositoryLiveFragment
}
}
}
assetNodes(assetKeys: $assetKeys, loadMaterializations: true) {
id
...AssetNodeLiveFragment
}
}
${REPOSITORY_LIVE_FRAGMENT}
${ASSET_NODE_LIVE_FRAGMENT}
`;

const ASSETS_GRAPH_QUERY = gql`
query AssetGraphQuery($pipelineSelector: PipelineSelector) {
assetNodes(pipeline: $pipelineSelector) {
id
...AssetNodeFragment
jobNames
dependencyKeys {
path
}
dependedByKeys {
path
}
}
}
${ASSET_NODE_FRAGMENT}
`;

const SVGContainer = styled.svg`
overflow: visible;
border-radius: 0;
Expand Down Expand Up @@ -544,29 +437,6 @@ const opsInRange = (
return uniq(ledToTarget);
};

const buildGraphQueryItems = (nodes: AssetNode[]) => {
const items: {
[name: string]: GraphQueryItem & {
node: AssetNode;
};
} = {};

for (const node of nodes) {
const name = tokenForAssetKey(node.assetKey);
items[name] = {
node: node,
name: name,
inputs: node.dependencyKeys.map((key) => ({
dependsOn: [{solid: {name: tokenForAssetKey(key)}}],
})),
outputs: node.dependedByKeys.map((key) => ({
dependedBy: [{solid: {name: tokenForAssetKey(key)}}],
})),
};
}
return Object.values(items);
};

const titleForLaunch = (nodes: Node[], liveDataByNode: LiveData) => {
const isRematerializeForAll = (nodes.length
? nodes.map((n) => liveDataByNode[n.id])
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
import {gql, useQuery} from '@apollo/client';
import React from 'react';

import {filterByQuery, GraphQueryItem} from '../../app/GraphQueryImpl';
import {tokenForAssetKey} from '../../app/Util';
import {PipelineSelector} from '../../types/globalTypes';

import {ASSET_NODE_FRAGMENT} from './AssetNode';
import {buildGraphData} from './Utils';
import {
AssetGraphQuery,
AssetGraphQueryVariables,
AssetGraphQuery_assetNodes,
} from './types/AssetGraphQuery';

/** Fetches data for rendering an asset graph:
*
* @param pipelineSelector: Optionally scope to an asset job, or pass null for the global graph
*
* @param opsQuery: filter the returned graph using selector syntax string (eg: asset_name++)
*
* @param filterNodes: filter the returned graph using the provided function. The global graph
* uses this option to implement the "3 of 4 repositories" picker.
*/
export function useAssetGraphData(
pipelineSelector: PipelineSelector | null | undefined,
opsQuery: string,
filterNodes?: (assetNode: AssetGraphQuery_assetNodes) => boolean,
) {
const fetchResult = useQuery<AssetGraphQuery, AssetGraphQueryVariables>(ASSET_GRAPH_QUERY, {
variables: {pipelineSelector},
notifyOnNetworkStatusChange: true,
});

const fetchResultFilteredNodes = React.useMemo(() => {
const nodes = fetchResult.data?.assetNodes;
if (!nodes) {
return undefined;
}
return filterNodes ? nodes.filter(filterNodes) : nodes;
}, [fetchResult.data, filterNodes]);

const {
assetGraphData,
graphQueryItems,
graphAssetKeys,
applyingEmptyDefault,
} = React.useMemo(() => {
if (fetchResultFilteredNodes === undefined) {
return {
graphAssetKeys: [],
graphQueryItems: [],
assetGraphData: null,
applyingEmptyDefault: false,
};
}
const graphQueryItems = buildGraphQueryItems(fetchResultFilteredNodes);
const {all, applyingEmptyDefault} = filterByQuery(graphQueryItems, opsQuery);

return {
graphAssetKeys: all.map((n) => ({path: n.node.assetKey.path})),
assetGraphData: buildGraphData(all.map((n) => n.node)),
graphQueryItems,
applyingEmptyDefault,
};
}, [fetchResultFilteredNodes, opsQuery]);

return {
fetchResult,
assetGraphData,
graphQueryItems,
graphAssetKeys,
applyingEmptyDefault,
};
}

type AssetNode = AssetGraphQuery_assetNodes;

const buildGraphQueryItems = (nodes: AssetNode[]) => {
const items: {
[name: string]: GraphQueryItem & {
node: AssetNode;
};
} = {};

for (const node of nodes) {
const name = tokenForAssetKey(node.assetKey);
items[name] = {
node: node,
name: name,
inputs: node.dependencyKeys.map((key) => ({
dependsOn: [{solid: {name: tokenForAssetKey(key)}}],
})),
outputs: node.dependedByKeys.map((key) => ({
dependedBy: [{solid: {name: tokenForAssetKey(key)}}],
})),
};
}
return Object.values(items);
};

const ASSET_GRAPH_QUERY = gql`
query AssetGraphQuery($pipelineSelector: PipelineSelector) {
assetNodes(pipeline: $pipelineSelector) {
id
...AssetNodeFragment
jobNames
dependencyKeys {
path
}
dependedByKeys {
path
}
}
}
${ASSET_NODE_FRAGMENT}
`;

0 comments on commit 6d5eae5

Please sign in to comment.