Skip to content

Commit

Permalink
Merge pull request #579 from DataRecce/feature/drc-1026-ui-action-ban…
Browse files Browse the repository at this point in the history
…ner-when-row-count-diff-preset-check-is-not-run

[Feature] Recommend to run row count diff preset check
  • Loading branch information
wcchang1115 authored Jan 15, 2025
2 parents 694bacf + 35b5aa0 commit ec1eb61
Show file tree
Hide file tree
Showing 6 changed files with 308 additions and 16 deletions.
3 changes: 2 additions & 1 deletion js/src/components/check/CheckDetail.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,8 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {

const submittedRun = await submitRunFromCheck(checkId, { nowait: true });
setSubmittedRunId(submittedRun.run_id);
}, [check, checkId, setSubmittedRunId]);
queryClient.invalidateQueries({ queryKey: cacheKeys.check(checkId) });
}, [check, checkId, setSubmittedRunId, queryClient]);

const handleCancel = useCallback(async () => {
setAborting(true);
Expand Down
65 changes: 54 additions & 11 deletions js/src/components/lineage/GraphNode.tsx
Original file line number Diff line number Diff line change
@@ -1,4 +1,15 @@
import { Box, Flex, HStack, Icon, Spacer, Tooltip } from "@chakra-ui/react";
import {
Box,
Flex,
HStack,
Icon,
Spacer,
Tag,
TagLabel,
TagLeftIcon,
Text,
Tooltip,
} from "@chakra-ui/react";
import React, { useState } from "react";

import { Handle, NodeProps, Position, useStore } from "reactflow";
Expand All @@ -14,9 +25,41 @@ import { findByRunType } from "../run/registry";
import { isSchemaChanged } from "../schema/schemaDiff";
import { useLineageViewContext } from "./LineageViewContext";
import { FaCheckSquare, FaRegSquare, FaSquare } from "react-icons/fa";
import { RowCount } from "@/lib/api/models";
import { deltaPercentageString } from "../rowcount/delta";

interface GraphNodeProps extends NodeProps<LineageGraphNode> {}

function _RowCountDiffTag({ rowCount }: { rowCount: RowCount }) {
const base = rowCount.base;
const current = rowCount.curr;
const baseLabel = rowCount.base === null ? "N/A" : `${rowCount.base} Rows`;
const currentLabel = rowCount.curr === null ? "N/A" : `${rowCount.curr} Rows`;

let tagLabel;
let colorScheme;
if (base === null && current === null) {
tagLabel = "Failed to load";
colorScheme = "gray";
} else if (base === null || current === null) {
tagLabel = `${baseLabel} -> ${currentLabel}`;
colorScheme = base === null ? "green" : "red";
} else if (base === current) {
tagLabel = "=";
colorScheme = "gray";
} else if (base !== current) {
tagLabel = `${deltaPercentageString(base, current)} Rows`;
colorScheme = base < current ? "green" : "red";
}

return (
<Tag colorScheme={colorScheme}>
<TagLeftIcon as={findByRunType("row_count_diff")?.icon} />
<TagLabel>{tagLabel}</TagLabel>
</Tag>
);
}

const NodeRunsAggregated = ({
id,
inverted,
Expand Down Expand Up @@ -50,7 +93,7 @@ const NodeRunsAggregated = ({
const colorUnchanged = inverted ? "gray" : "lightgray";

return (
<Flex gap="5px">
<Flex flex="1">
{schemaChanged !== undefined && (
<Tooltip
label={`Schema (${schemaChanged ? "changed" : "no change"})`}
Expand All @@ -64,16 +107,14 @@ const NodeRunsAggregated = ({
</Box>
</Tooltip>
)}
{rowCountChanged !== undefined && (
<Spacer />
{runs && runs["row_count_diff"] && rowCountChanged !== undefined && (
<Tooltip
label={`Row count (${rowCountChanged ? "changed" : "no change"})`}
label={`Row count (${rowCountChanged ? "changed" : "="})`}
openDelay={500}
>
<Box height="16px">
<Icon
as={findByRunType("row_count_diff")?.icon}
color={rowCountChanged ? colorChanged : colorUnchanged}
/>
<Box>
<_RowCountDiffTag rowCount={runs["row_count_diff"].result} />
</Box>
</Tooltip>
)}
Expand Down Expand Up @@ -286,10 +327,12 @@ export function GraphNode({ data }: GraphNodeProps) {
})()}
/>
)}
<Spacer />
{data.isActionMode &&
(data.action ? (
<ActionTag node={data} action={data.action} />
<>
<Spacer />
<ActionTag node={data} action={data.action} />
</>
) : (
<></>
))}
Expand Down
8 changes: 7 additions & 1 deletion js/src/components/lineage/LineageView.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ import { useLocation } from "wouter";
import { Check } from "@/lib/api/checks";
import useValueDiffAlertDialog from "./useValueDiffAlertDialog";
import { trackMultiNodesAction } from "@/lib/api/track";
import { PresetCheckRecommendation } from "./PresetCheckRecommendation";

export interface LineageViewProps {
viewOptions?: LineageDiffViewOptions;
Expand Down Expand Up @@ -796,7 +797,12 @@ export function PrivateLineageView(
style={{ contain: "strict" }}
position="relative"
>
{interactive && <LineageViewTopBar />}
{interactive && (
<>
<LineageViewTopBar />
<PresetCheckRecommendation />
</>
)}
<ReactFlow
nodeTypes={nodeTypes}
edgeTypes={edgeTypes}
Expand Down
238 changes: 238 additions & 0 deletions js/src/components/lineage/PresetCheckRecommendation.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,238 @@
import {
HStack,
Button,
Spacer,
Text,
Modal,
ModalOverlay,
ModalContent,
ModalHeader,
ModalBody,
ModalFooter,
useDisclosure,
Stack,
Flex,
} from "@chakra-ui/react";
import { useCallback, useEffect, useMemo, useState } from "react";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import { cacheKeys } from "@/lib/api/cacheKeys";
import { getCheck, listChecks } from "@/lib/api/checks";
import { InfoOutlineIcon } from "@chakra-ui/icons";
import { select } from "@/lib/api/select";
import { useLineageGraphContext } from "@/lib/hooks/LineageGraphContext";
import { submitRunFromCheck } from "@/lib/api/runs";
import { useRecceActionContext } from "@/lib/hooks/RecceActionContext";
import { sessionStorageKeys } from "@/lib/api/sessionStorageKeys";
import { useRecceServerFlag } from "@/lib/hooks/useRecceServerFlag";

const usePresetCheckRecommendation = () => {
const queryChecks = useQuery({
queryKey: cacheKeys.checks(),
queryFn: listChecks,
});

const lastRecommendCheck = useMemo(() => {
// filter out the preset checks and it's the latest row count diff check
if (queryChecks.status === "success" && queryChecks.data.length > 0) {
const check = queryChecks.data
.filter((check) => check.is_preset)
.findLast((check) => check.type === "row_count_diff");
if (check) {
return check;
}
}
}, [queryChecks]);

const queryPresetCheck = useQuery({
queryKey: lastRecommendCheck?.check_id
? cacheKeys.check(lastRecommendCheck.check_id)
: [],
queryFn: async () => {
if (lastRecommendCheck?.check_id) {
return getCheck(lastRecommendCheck.check_id);
}
},
enabled: !!lastRecommendCheck?.check_id,
});

const querySelectedModels = useQuery({
queryKey: lastRecommendCheck?.check_id
? [...cacheKeys.check(lastRecommendCheck.check_id), "select"]
: [],
queryFn: async () =>
select({
select: lastRecommendCheck?.params?.select,
exclude: lastRecommendCheck?.params?.exclude,
}),
enabled: !!lastRecommendCheck?.params?.select,
});

const selectedNodes = useMemo(() => {
if (lastRecommendCheck) {
if (lastRecommendCheck.params?.node_names) {
return lastRecommendCheck.params.node_names;
}

if (lastRecommendCheck.params?.node_ids) {
return lastRecommendCheck.params.node_ids;
}
}

if (querySelectedModels.status === "success" && querySelectedModels.data) {
return querySelectedModels.data.nodes;
}
}, [lastRecommendCheck, querySelectedModels]);

return {
recommendedCheck: queryPresetCheck.data,
selectedNodes,
};
};

export const PresetCheckRecommendation = () => {
const { lineageGraph } = useLineageGraphContext();
const { showRunId } = useRecceActionContext();
const { data: flags } = useRecceServerFlag();
const queryClient = useQueryClient();
const { recommendedCheck, selectedNodes } = usePresetCheckRecommendation();
const [affectedModels, setAffectedModels] = useState<string>();
const [performedRecommend, setPerformedRecommend] = useState<boolean>(false);
const [ignoreRecommend, setIgnoreRecommend] = useState<boolean>(false);
const { isOpen, onOpen, onClose } = useDisclosure();
const recommendationKey = sessionStorageKeys.recommendationIgnored;

useEffect(() => {
const ignored = sessionStorage.getItem(recommendationKey);
if (ignored) {
setIgnoreRecommend(true);
}
}, [recommendationKey]);

useEffect(() => {
if (!recommendedCheck || !selectedNodes) {
return;
}

if (recommendedCheck.last_run?.run_id) {
setPerformedRecommend(true);
return;
}

const check = recommendedCheck;
if (selectedNodes.length > 0 && selectedNodes.length <= 3) {
if (check.params?.node_names) {
const nodeNames = check.params?.node_names.join(", ");
setAffectedModels(`'${nodeNames}'`);
} else if (check.params?.node_ids) {
const nodes = [];
for (const nodeId of check.params?.node_ids) {
const node = lineageGraph?.nodes[nodeId];
if (node) {
nodes.push(node.name);
}
}
const nodeNames = nodes.join(", ");
setAffectedModels(`'${nodeNames}'`);
} else if (selectedNodes) {
const nodes = [];
for (const nodeId of selectedNodes) {
const node = lineageGraph?.nodes[nodeId];
if (node) {
nodes.push(node.name);
}
}
const nodeNames = nodes.join(", ");
setAffectedModels(`'${nodeNames}'`);
}
} else if (lineageGraph?.modifiedSet?.length === selectedNodes.length) {
setAffectedModels("modified and potentially impacted models");
} else if (check.params?.select && !check.params?.exclude) {
setAffectedModels(`'${check.params?.select}'`);
} else {
setAffectedModels(`${selectedNodes.length} models`);
}
}, [recommendedCheck, selectedNodes, lineageGraph, recommendationKey, flags]);

const performPresetCheck = useCallback(async () => {
const check = recommendedCheck;
if (!check || check.last_run?.run_id) {
return;
}
const submittedRun = await submitRunFromCheck(check.check_id, {
nowait: true,
});
showRunId(submittedRun.run_id);
queryClient.invalidateQueries({
queryKey: cacheKeys.check(check.check_id),
});
}, [recommendedCheck, showRunId, queryClient]);

if (!recommendedCheck || !selectedNodes || flags?.single_env_onboarding) {
return <></>;
}
const numNodes = selectedNodes.length;

return (
!ignoreRecommend &&
!performedRecommend && (
<>
<HStack width="100%" padding="2pt 8pt" backgroundColor={"blue.50"}>
<HStack flex="1" fontSize={"10pt"} color="blue.600">
<InfoOutlineIcon />
<Text>
First Check: Perform a row count diff of {affectedModels} for
basic impact assessment
</Text>
<Spacer />
<Button
size="xs"
onClick={() => {
setIgnoreRecommend(true);
sessionStorage.setItem(recommendationKey, "true");
}}
>
Ignore
</Button>
<Button colorScheme="blue" size="xs" onClick={onOpen}>
Perform
</Button>
</HStack>
</HStack>
<Modal isOpen={isOpen} onClose={onClose} isCentered>
<ModalOverlay />
<ModalContent>
<ModalHeader>Row Count Check</ModalHeader>
<ModalBody>
<Stack spacing="4">
<Text>
Perform a row count check of the {numNodes} node(s) displayed
in the lineage diff DAG.
</Text>
<Flex bg="blue.100" color="blue.700">
<InfoOutlineIcon mt="10px" ml="5px" />
<Text margin="5px" paddingX="3px">
This is a recommended first check based on the preset checks
defined in your recce.yml file.
</Text>
</Flex>
</Stack>
</ModalBody>
<ModalFooter gap="5px">
<Button onClick={onClose}>Cancel</Button>
<Button
colorScheme="blue"
onClick={() => {
onClose();
performPresetCheck();
setPerformedRecommend(true);
}}
>
Execute on {numNodes} models
</Button>
</ModalFooter>
</ModalContent>
</Modal>
</>
)
);
};
5 changes: 2 additions & 3 deletions js/src/components/rowcount/RowCountDiffResultView.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -46,11 +46,10 @@ function _RowCountDiffResultView(
const result = runResult[key];
const base = isNumber(result?.base) ? result?.base : null;
const current = isNumber(result?.curr) ? result?.curr : null;
let delta = "No Change";
let delta = "=";

if (base !== null && current !== null) {
delta =
base !== current ? deltaPercentageString(base, current) : "No Change";
delta = base !== current ? deltaPercentageString(base, current) : "=";
} else {
if (base === current) {
delta = "N/A";
Expand Down
5 changes: 5 additions & 0 deletions js/src/lib/api/sessionStorageKeys.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
const prefix = "recce";

export const sessionStorageKeys = {
recommendationIgnored: `${prefix}-recommendation-ignored`,
};

0 comments on commit ec1eb61

Please sign in to comment.