[Fix] Copy to clipboard should work on schema diff and lineage diff

Signed-off-by: Kent Huang <[email protected]>

js/src/components/check/CheckDetail.tsx

@@ -50,7 +50,7 @@ import { stripIndents } from "common-tags";
 import { useClipBoardToast } from "@/lib/hooks/useClipBoardToast";
 import { buildTitle, buildDescription, buildQuery } from "./check";
 import SqlEditor, { DualSqlEditor } from "../query/SqlEditor";
-import { useCallback, useState } from "react";
+import { useCallback, useRef, useState } from "react";
 import { cancelRun, submitRunFromCheck } from "@/lib/api/runs";
 import { Run } from "@/lib/api/types";
 import { RunView } from "../run/RunView";
@@ -62,6 +62,7 @@ import { VSplit } from "../split/Split";
 import { useCopyToClipboardButton } from "@/lib/hooks/ScreenShot";
 import { useRun } from "@/lib/hooks/useRun";
 import { useCheckToast } from "@/lib/hooks/useCheckToast";
+import { LineageViewRef } from "../lineage/LineageView";
 
 interface CheckDetailProps {
   checkId: string;
@@ -104,6 +105,8 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
   const runTypeEntry = check?.type ? findByRunType(check?.type) : undefined;
   const isPresetCheck = check?.is_preset || false;
 
+  const lineageViewRef = useRef<LineageViewRef>(null);
+
   const { mutate } = useMutation({
     mutationFn: (check: Partial<Check>) => updateCheck(checkId, check),
     onSuccess: () => {
@@ -179,6 +182,16 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
     mutate({ description });
   };
 
+  const shouldDisabledCopyButton = (
+    type: string,
+    run: Run | undefined
+  ): boolean => {
+    if (type === "schema_diff" || type === "lineage_diff") {
+      return false;
+    }
+    return !run?.result || !!run?.error;
+  };
+
   const { ref, onCopyToClipboard, onMouseEnter, onMouseLeave } =
     useCopyToClipboardButton();
 
@@ -321,11 +334,17 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
               <Button
                 leftIcon={<CopyIcon />}
                 variant="outline"
-                isDisabled={!run?.result || !!run?.error}
+                isDisabled={shouldDisabledCopyButton(check?.type ?? "", run)}
                 onMouseEnter={onMouseEnter}
                 onMouseLeave={onMouseLeave}
                 size="sm"
-                onClick={onCopyToClipboard}
+                onClick={() => {
+                  if (check?.type === "lineage_diff") {
+                    lineageViewRef.current?.copyToClipboard();
+                  } else {
+                    onCopyToClipboard();
+                  }
+                }}
               >
                 Copy to Clipboard
               </Button>
@@ -356,10 +375,10 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
                   </Center>
                 ))}
               {check && check.type === "schema_diff" && (
-                <SchemaDiffView check={check} />
+                <SchemaDiffView check={check} ref={ref} />
               )}
               {check && check.type === "lineage_diff" && (
-                <LineageDiffView check={check} />
+                <LineageDiffView check={check} ref={lineageViewRef} />
               )}
             </TabPanel>
             {(check?.type === "query" ||

js/src/components/check/LineageDiffView.tsx

@@ -1,20 +1,28 @@
 import { Check } from "@/lib/api/checks";
-import { LineageView } from "../lineage/LineageView";
+import { LineageView, LineageViewRef } from "../lineage/LineageView";
 import { Flex } from "@chakra-ui/react";
 import { ReactFlowProvider } from "reactflow";
+import { forwardRef, Ref } from "react";
 
 export interface LineageDiffViewProps {
   check: Check;
 }
 
-export function LineageDiffView({ check }: LineageDiffViewProps) {
+function _LineageDiffView(
+  { check }: LineageDiffViewProps,
+  ref: Ref<LineageViewRef>
+) {
   const viewOptions = { ...check.params, ...check.view_options };
 
   return (
     <Flex direction="column" height="100%">
       <ReactFlowProvider>
-        <LineageView viewOptions={viewOptions} interactive={false} />
+        <LineageView viewOptions={viewOptions} interactive={false} ref={ref} />
       </ReactFlowProvider>
     </Flex>
   );
 }
+
+export const LineageDiffView = forwardRef<LineageViewRef, LineageDiffViewProps>(
+  _LineageDiffView
+);

js/src/components/check/SchemaDiffView.tsx

@@ -7,7 +7,7 @@ import { select } from "@/lib/api/select";
 import { HSplit } from "../split/Split";
 import { Box, Center, Flex, Icon, List, ListItem } from "@chakra-ui/react";
 import { LineageGraphNode } from "../lineage/lineage";
-import { useMemo, useState } from "react";
+import { forwardRef, useMemo, useState } from "react";
 import {
   getIconForChangeStatus,
   getIconForResourceType,
@@ -83,7 +83,7 @@ const NodelistItem = ({
   );
 };
 
-export function SchemaDiffView({ check }: SchemaDiffViewProps) {
+export function _SchemaDiffView({ check }: SchemaDiffViewProps, ref: any) {
   const { lineageGraph } = useLineageGraphContext();
   const params = check.params as SchemaDiffParams;
 
@@ -192,8 +192,9 @@ export function SchemaDiffView({ check }: SchemaDiffViewProps) {
           base={node.data.base}
           current={node.data.current}
           enableScreenshot={true}
+          ref={ref}
         />
-        <List overflow="auto">
+        <List overflow="auto" backgroundColor="white">
           {nodes.map((node, i) => (
             <NodelistItem
               key={i}
@@ -214,3 +215,5 @@ export function SchemaDiffView({ check }: SchemaDiffViewProps) {
   // TODO: handle the edge case where the node is not found
   return <></>;
 }
+
+export const SchemaDiffView = forwardRef(_SchemaDiffView);

js/src/components/data-grid/ScreenshotDataGrid.tsx

@@ -1,17 +1,11 @@
 import "react-data-grid/lib/styles.css";
 import DataGrid, { DataGridProps } from "react-data-grid";
 import { Flex, forwardRef, Text } from "@chakra-ui/react";
-import {
-  useCopyToClipboardButton,
-  useImageDownloadModal,
-} from "@/lib/hooks/ScreenShot";
-import { useCallback } from "react";
 
 interface ScreenshotDataGridProps extends DataGridProps<any> {}
 
 export const ScreenshotDataGrid = forwardRef(
   ({ ...props }: ScreenshotDataGridProps, ref: any) => {
-    const { CopyToClipboardButton } = useCopyToClipboardButton();
     return (
       <>
         <DataGrid ref={ref} {...props} />

js/src/components/lineage/LineageView.tsx

@@ -34,9 +34,11 @@ import React, {
   RefObject,
   useCallback,
   useEffect,
+  useImperativeHandle,
   useLayoutEffect,
   useRef,
   useState,
+  forwardRef,
 } from "react";
 import ReactFlow, {
   Node,
@@ -92,6 +94,10 @@ export interface LineageViewProps {
   filterNodes?: (key: string, node: LineageGraphNode) => boolean;
 }
 
+export interface LineageViewRef {
+  copyToClipboard: () => void;
+}
+
 const nodeTypes = {
   customNode: GraphNode,
 };
@@ -187,11 +193,18 @@ const useResizeObserver = (
   }, [target, size, handler]);
 };
 
-export function LineageView({ ...props }: LineageViewProps) {
+export function _LineageView(
+  { ...props }: LineageViewProps,
+  ref: Ref<LineageViewRef>
+) {
   const reactFlow = useReactFlow();
   const refReactFlow = useRef<HTMLDivElement>(null);
   const { successToast, failToast } = useClipBoardToast();
-  const { copyToClipboard, ImageDownloadModal, ref } = useCopyToClipboard({
+  const {
+    copyToClipboard,
+    ImageDownloadModal,
+    ref: copyRef,
+  } = useCopyToClipboard({
     renderLibrary: "html-to-image",
     imageType: "png",
     shadowEffect: true,
@@ -231,6 +244,11 @@ export function LineageView({ ...props }: LineageViewProps) {
 
   const { showRunId, closeRunResult } = useRecceActionContext();
 
+  // Expose the function to the parent via the ref
+  useImperativeHandle(ref, () => ({
+    copyToClipboard,
+  }));
+
   /**
    * View mode
    * - all: show all nodes
@@ -601,7 +619,7 @@ export function LineageView({ ...props }: LineageViewProps) {
           minZoom={0.1}
           fitView={true}
           nodesDraggable={props.interactive}
-          ref={ref}
+          ref={copyRef}
         >
           <Background color="#ccc" />
           <Controls
@@ -704,3 +722,7 @@ export function LineageView({ ...props }: LineageViewProps) {
     </HSplit>
   );
 }
+
+export const LineageView = forwardRef<LineageViewRef, LineageViewProps>(
+  _LineageView
+);

js/src/components/schema/SchemaView.tsx

@@ -1,4 +1,4 @@
-import { useMemo } from "react";
+import { forwardRef, useMemo } from "react";
 
 import { mergeColumns, toDataGrid } from "./schema";
 import "react-data-grid/lib/styles.css";
@@ -16,11 +16,10 @@ interface SchemaViewProps {
   enableScreenshot?: boolean;
 }
 
-export function SchemaView({
-  base,
-  current,
-  enableScreenshot = false,
-}: SchemaViewProps) {
+export function _SchemaView(
+  { base, current, enableScreenshot = false }: SchemaViewProps,
+  ref: any
+) {
   const { columns, rows } = useMemo(() => {
     const schemaDiff = mergeColumns(base?.columns, current?.columns);
     const resourceType = current?.resource_type || base?.resource_type;
@@ -88,9 +87,12 @@ export function SchemaView({
             renderers={{ noRowsFallback: <EmptyRowsRenderer /> }}
             className="rdg-light"
             enableScreenshot={enableScreenshot}
+            ref={ref}
           />
         </>
       )}
     </Flex>
   );
 }
+
+export const SchemaView = forwardRef(_SchemaView);

js/src/components/valuediff/ValueDiffResultView.tsx

@@ -10,7 +10,6 @@ import {
   MenuGroup,
   MenuItem,
   MenuList,
-  Portal,
   Spacer,
 } from "@chakra-ui/react";
 
@@ -26,7 +25,6 @@ import {
   RecceActionOptions,
   useRecceActionContext,
 } from "@/lib/hooks/RecceActionContext";
-import { useRef } from "react";
 
 interface ValueDiffResultViewProp
   extends RunResultViewProps<ValueDiffParams, ValueDiffResult> {}
@@ -176,7 +174,6 @@ function _ValueDiffResultView({ run }: ValueDiffResultViewProp, ref: any) {
           maxHeight: "100%",
           overflow: "auto",
           borderBlock: "1px solid lightgray",
-          flex: "1",
         }}
         columns={columns}
         rows={result.data.data}
@@ -189,4 +186,4 @@ function _ValueDiffResultView({ run }: ValueDiffResultViewProp, ref: any) {
   );
 }
 
-export const ValueDiffResultView = forwardRef(_ValueDiffResultView);
+export const ValueDiffResultView = forwardRef(_ValueDiffResultView);