Merge pull request #424 from DataRecce/feature/drc-628-make-actions-t…

…ogether Make the node actions together
DataRecce · Sep 20, 2024 · 8d1c704 · 8d1c704
2 parents df07719 + bb69cde
commit 8d1c704
Show file tree

Hide file tree

Showing 5 changed files with 187 additions and 162 deletions.
diff --git a/js/src/components/lineage/ActionTag.tsx b/js/src/components/lineage/ActionTag.tsx
@@ -100,13 +100,7 @@ export const ActionTag = ({ node, action }: ActionTagProps) => {
 
   if (run.type === "row_count_diff") {
     const result = run.result as RowCountDiffResult;
-    return (
-      <RowCountTag
-        rowCount={result[node.name]}
-        node={node}
-        isInteractive={false}
-      />
-    );
+    return <RowCountTag rowCount={result[node.name]} node={node} />;
   }
 
   return <>{run_id}</>;

diff --git a/js/src/components/lineage/NodeFilter.tsx b/js/src/components/lineage/NodeFilter.tsx
@@ -318,7 +318,7 @@ const MoreActionMenu = (props: NodeFilterProps) => {
           </MenuItem>
         </MenuGroup>
         <MenuDivider />
-        <MenuGroup title="Add check" m="0" px="12px">
+        <MenuGroup title="Add to Checklist" m="0" px="12px">
           <MenuItem
             as={Text}
             size="sm"

diff --git a/js/src/components/lineage/NodeTag.tsx b/js/src/components/lineage/NodeTag.tsx
@@ -106,33 +106,25 @@ export function ModelRowCount({ rowCount }: ModelRowCountProps) {
 export interface RowCountTagProps {
   node: LineageGraphNode;
   rowCount?: RowCount;
-  isInteractive?: boolean; // if true, allows user to manually fetch row count
+  onRefresh?: () => Promise<any>;
+  isFetching?: boolean;
+  error?: Error | null;
 }
 
 export function RowCountTag({
-  rowCount: defaultRowCount,
+  rowCount: fetchedRowCount,
   node,
-  isInteractive,
+  onRefresh,
+  isFetching,
 }: RowCountTagProps) {
   const { runsAggregated, refetchRunsAggregated } = useLineageGraphContext();
   const lastRowCount: RowCount | undefined =
     runsAggregated?.[node.id]?.["row_count_diff"]?.result;
 
-  const {
-    data: fetchedRowCount,
-    refetch: invokeRowCountQuery,
-    isFetching,
-  } = useQuery({
-    queryKey: cacheKeys.rowCount(node.name),
-    queryFn: () => queryModelRowCount(node.name),
-    enabled: false,
-    initialData: defaultRowCount,
-  });
-
-  const rowCount = fetchedRowCount || defaultRowCount || lastRowCount;
   const icon = findByRunType("row_count_diff")?.icon;
 
   let label;
+  const rowCount = fetchedRowCount || lastRowCount;
   if (rowCount) {
     const base = rowCount?.base === null ? "N/A" : rowCount?.base;
     const current = rowCount?.curr === null ? "N/A" : rowCount?.curr;
@@ -158,15 +150,15 @@ export function RowCountTag({
             <>row count</>
           )}
         </TagLabel>
-        {isInteractive && (
+        {onRefresh && (
           <TagRightIcon
             as={IconButton}
             isLoading={isFetching}
             aria-label="Query Row Count"
             icon={<RepeatIcon />}
             size="xs"
             onClick={async () => {
-              await invokeRowCountQuery();
+              await onRefresh();
               refetchRunsAggregated?.();
             }}
           />