[Fix] refactor func isTaskShouldBeDisabled by isActionAvailable

Signed-off-by: Kent Huang <[email protected]>
DataRecce · Dec 11, 2024 · 39683c1 · 39683c1
1 parent a2d9cd0
commit 39683c1
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 39 deletions.
diff --git a/js/src/components/lineage/LineageViewTopBar.tsx b/js/src/components/lineage/LineageViewTopBar.tsx
@@ -302,7 +302,7 @@ const ControlItem = (props: {
 export const LineageViewTopBar = () => {
   const { nodes, deselect, selectMode, ...lineageViewContext } =
     useLineageViewContext();
-  const { isTaskShouldBeDisabled } = useLineageGraphContext();
+  const { isActionAvailable } = useLineageGraphContext();
   const selectNodes = useMemo(() => {
     return nodes.filter((node) => node.data.isSelected);
   }, [nodes]);
@@ -386,8 +386,7 @@ export const LineageViewTopBar = () => {
                   </MenuItem>
                   <Tooltip
                     label={
-                      isTaskShouldBeDisabled &&
-                      isTaskShouldBeDisabled("value_diff")
+                      !isActionAvailable("value_diff")
                         ? DisableTooltipMessages.audit_helper
                         : null
                     }
@@ -398,8 +397,7 @@ export const LineageViewTopBar = () => {
                       fontSize="10pt"
                       isDisabled={
                         !(isNoSelect || isSingleSelect || isMultiSelect) ||
-                        (isTaskShouldBeDisabled &&
-                          isTaskShouldBeDisabled("value_diff"))
+                        !isActionAvailable("value_diff")
                       }
                       icon={<Icon as={findByRunType("value_diff")?.icon} />}
                       onClick={() => {

diff --git a/js/src/components/lineage/NodeView.tsx b/js/src/components/lineage/NodeView.tsx
@@ -67,7 +67,7 @@ export function NodeView({ node, onCloseNode }: NodeViewProps) {
     onClose: onCodeDiffClose,
   } = useDisclosure();
   const { runAction } = useRecceActionContext();
-  const { envInfo, isTaskShouldBeDisabled } = useLineageGraphContext();
+  const { envInfo, isActionAvailable } = useLineageGraphContext();
   const { primaryKey } = useModelColumns(node.name);
   const refetchRowCount = () => {
     runAction(
@@ -87,7 +87,7 @@ export function NodeView({ node, onCloseNode }: NodeViewProps) {
     if (isAddedOrRemoved) {
       return DisableTooltipMessages.add_or_remove;
     }
-    if (isTaskShouldBeDisabled && isTaskShouldBeDisabled(runType)) {
+    if (!isActionAvailable(runType)) {
       if (runType === "value_diff") {
         return DisableTooltipMessages.audit_helper;
       } else if (runType === "profile_diff") {
@@ -126,9 +126,12 @@ export function NodeView({ node, onCloseNode }: NodeViewProps) {
                   } else if (envInfo?.adapterType === "sqlmesh") {
                     setSqlQuery(`select * from ${node.name}`);
                   }
-                  setPrimaryKeys(
-                    primaryKey !== undefined ? [primaryKey] : undefined
-                  );
+                  if (isActionAvailable("query_diff_with_primary_key")) {
+                    // Only set primary key if the action is available
+                    setPrimaryKeys(
+                      primaryKey !== undefined ? [primaryKey] : undefined
+                    );
+                  }
                   setLocation("/query");
                 }}
               >
@@ -160,9 +163,7 @@ export function NodeView({ node, onCloseNode }: NodeViewProps) {
                     icon={<Icon as={findByRunType("profile_diff")?.icon} />}
                     fontSize="14px"
                     isDisabled={
-                      isAddedOrRemoved ||
-                      (isTaskShouldBeDisabled &&
-                        isTaskShouldBeDisabled("profile_diff"))
+                      isAddedOrRemoved || !isActionAvailable("profile_diff")
                     }
                     onClick={() => {
                       runAction(
@@ -182,9 +183,7 @@ export function NodeView({ node, onCloseNode }: NodeViewProps) {
                     icon={<Icon as={findByRunType("value_diff")?.icon} />}
                     fontSize="14px"
                     isDisabled={
-                      isAddedOrRemoved ||
-                      (isTaskShouldBeDisabled &&
-                        isTaskShouldBeDisabled("value_diff"))
+                      isAddedOrRemoved || !isActionAvailable("value_diff")
                     }
                     onClick={() => {
                       runAction(

diff --git a/js/src/components/query/QueryForm.tsx b/js/src/components/query/QueryForm.tsx
@@ -20,7 +20,8 @@ export const QueryForm = ({
   onPrimaryKeysChange,
   ...prob
 }: QueryFormProps) => {
-  const { lineageGraph, isTaskShouldBeDisabled } = useLineageGraphContext();
+  const { lineageGraph, isActionAvailable } = useLineageGraphContext();
+
   const labelInfo =
     "Provide a primary key to perform query diff in data warehouse and only return changed rows.";
 
@@ -62,11 +63,7 @@ export const QueryForm = ({
           size="xs"
           width={"240px"}
           placeholder="Start by typing key name..."
-          isDisabled={
-            (isTaskShouldBeDisabled &&
-              isTaskShouldBeDisabled("query_diff_with_primary_key")) ||
-            false
-          }
+          isDisabled={!isActionAvailable("query_diff_with_primary_key")}
         />
       </FormControl>
     </Flex>

diff --git a/js/src/components/query/QueryPage.tsx b/js/src/components/query/QueryPage.tsx
@@ -181,12 +181,5 @@ export const QueryPage = () => {
         )}
       </Box>
     </Flex>
-    /* <QueryForm
-      p="5px"
-      border="1px"
-      borderColor="gray.300"
-      defaultPrimaryKeys={primaryKeys}
-      onPrimaryKeysChange={setPrimaryKeys}
-    /> */
   );
 };
diff --git a/js/src/lib/hooks/LineageGraphContext.tsx b/js/src/lib/hooks/LineageGraphContext.tsx
@@ -53,14 +53,16 @@ export interface LineageGraphContextType {
   isLoading?: boolean;
   error?: string;
   supportTasks?: { [key: string]: boolean };
-  isTaskShouldBeDisabled?: (taskName: string) => boolean;
   retchLineageGraph?: () => void;
+  isActionAvailable: (actionName: string) => boolean;
 
   runsAggregated?: RunsAggregated;
   refetchRunsAggregated?: () => void;
 }
 
-const defaultLineageGraphsContext: LineageGraphContextType = {};
+const defaultLineageGraphsContext: LineageGraphContextType = {
+  isActionAvailable: () => true,
+};
 
 const LineageGraphContext = createContext(defaultLineageGraphsContext);
 
@@ -208,15 +210,13 @@ export function LineageGraphContextProvider({ children }: LineageGraphProps) {
 
   const { connectionStatus, connect } = useLineageWatcher();
 
-  const isTaskShouldBeDisabled = useCallback(
-    (taskName: string) => {
+  const isActionAvailable = useCallback(
+    (name: string) => {
       if (supportTasks) {
-        const isSupport = supportTasks[taskName];
-        if (isSupport !== undefined) {
-          return !isSupport;
-        }
+        return supportTasks[name] ?? true; // default to true if action not found in supportTasks
       }
-      return false;
+      // If the supportTasks does not be provided, all actions are available
+      return true;
     },
     [supportTasks]
   );
@@ -238,7 +238,7 @@ export function LineageGraphContextProvider({ children }: LineageGraphProps) {
           isLoading,
           runsAggregated,
           supportTasks,
-          isTaskShouldBeDisabled,
+          isActionAvailable,
           refetchRunsAggregated: () => {
             refetchRunsAggregated();
           },