From bc1aefc43b94e0d800336aadf40f662a61e4dd7e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=A8=E5=B7=9D?= <chuan.yang@aliyun.com>
Date: Tue, 7 May 2024 13:42:42 +0800
Subject: [PATCH] =?UTF-8?q?fix=EF=BC=9A=E5=A4=9A=E6=AC=A1=E6=90=9C?=
 =?UTF-8?q?=E7=B4=A2=E5=AF=BC=E8=87=B4=E5=B7=B2=E9=80=89=E5=BA=93=E8=A1=A8?=
 =?UTF-8?q?=E8=A2=AB=E6=B8=85=E7=A9=BA=E9=97=AE=E9=A2=98=20refs=20#209612?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../Task/component/DatabaseSelecter/index.tsx | 21 ++++++++++++----
 .../Task/component/TableSelecter/index.tsx    | 24 +++++++++++--------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/src/component/Task/component/DatabaseSelecter/index.tsx b/src/component/Task/component/DatabaseSelecter/index.tsx
index 2965823d3..fa5dca137 100644
--- a/src/component/Task/component/DatabaseSelecter/index.tsx
+++ b/src/component/Task/component/DatabaseSelecter/index.tsx
@@ -20,8 +20,8 @@ import ExportCard from '@/component/ExportCard';
 import DataBaseStatusIcon from '@/component/StatusIcon/DatabaseIcon';
 import { DeleteOutlined } from '@ant-design/icons';
 import { Empty, Popconfirm, Space, Spin, Tree, Typography, Checkbox } from 'antd';
-import React, { useEffect, useState } from 'react';
-import { DataNode } from 'antd/lib/tree';
+import React, { useCallback, useEffect, useState } from 'react';
+import { DataNode, TreeProps } from 'antd/lib/tree';
 import classnames from 'classnames';
 import styles from './index.less';
 
@@ -150,6 +150,19 @@ const DatabaseSelecter: React.FC<IProps> = function ({
   const handleSearch = (value) => {
     setSourceSearchValue(value);
   };
+  /**
+   * 选中一个库后
+   */
+  const handleChosenDataBase: TreeProps['onCheck'] = useCallback(
+    (_checkedKeys, { checked, node: { key: curNodeKey } }) => {
+      if (checked) {
+        onChange([...checkedKeys, curNodeKey]);
+      } else {
+        onChange(checkedKeys.filter((key) => key !== curNodeKey));
+      }
+    },
+    [checkedKeys, onChange],
+  );
 
   const allTreeDataKeys = getAllTreeDataKeys();
   const checkAll = allTreeDataKeys?.length && allTreeDataKeys.length === checkedKeys.length;
@@ -192,9 +205,7 @@ const DatabaseSelecter: React.FC<IProps> = function ({
               className={styles.allTree}
               treeData={allTreeData}
               checkedKeys={checkedKeys}
-              onCheck={(_checkedKeys) => {
-                onChange(_checkedKeys as string[]);
-              }}
+              onCheck={handleChosenDataBase}
             />
           </ExportCard>
         </Spin>
diff --git a/src/component/Task/component/TableSelecter/index.tsx b/src/component/Task/component/TableSelecter/index.tsx
index b439aaaee..04c7601b0 100644
--- a/src/component/Task/component/TableSelecter/index.tsx
+++ b/src/component/Task/component/TableSelecter/index.tsx
@@ -146,7 +146,7 @@ const getTreeData = (validTableList: IDataBaseWithTable[]) => {
         </Space>
       ),
       key: id,
-      icon: getDataSourceStyleByConnectType(dataSource.type).dbIcon.component,
+      icon: <Icon component={getDataSourceStyleByConnectType(dataSource.type).dbIcon.component} />,
       checkable: false,
       expandable: true,
       children,
@@ -269,7 +269,10 @@ const TableSelecter: React.ForwardRefRenderFunction<TableSelecterRef, IProps> =
         const nodeKey = key as string;
         remainKeys = checkedKeys.filter((key) => key !== nodeKey);
       }
-      onChange(remainKeys.map(parseDataBaseIdAndTableNamebByKey));
+      const newValue = remainKeys.map(parseDataBaseIdAndTableNamebByKey);
+      onChange(newValue);
+      const willExpandKeys: number[] = newValue.map(({ databaseId }) => databaseId);
+      setSelectedExpandKeys(Array.from(new Set(willExpandKeys)));
     },
     [checkedKeys, onChange],
   );
@@ -309,18 +312,19 @@ const TableSelecter: React.ForwardRefRenderFunction<TableSelecterRef, IProps> =
    * 选中一张表后
    */
   const handleChosenTable: TreeProps['onCheck'] = useCallback(
-    (_checkedKeys: string[]) => {
-      const newValue: TableItem[] = [];
-      const willExpandKeys: number[] = [];
-      _checkedKeys.forEach((key) => {
-        const tableItem = parseDataBaseIdAndTableNamebByKey(key);
+    (_checkedKeys: string[], { checked, node: { key: curNodeKey } }) => {
+      const preCheckKeys = checked ? checkedKeys : checkedKeys.filter((key) => key !== curNodeKey);
+      const newValue: TableItem[] = preCheckKeys.map(parseDataBaseIdAndTableNamebByKey);
+      const willExpandKeys: number[] = newValue.map(({ databaseId }) => databaseId);
+      if (checked) {
+        const tableItem = parseDataBaseIdAndTableNamebByKey(curNodeKey as string);
         newValue.push(tableItem);
         willExpandKeys.push(tableItem.databaseId);
-      });
+      }
       onChange(newValue);
-      setSelectedExpandKeys(willExpandKeys);
+      setSelectedExpandKeys(Array.from(new Set(willExpandKeys)));
     },
-    [onChange],
+    [checkedKeys, onChange],
   );
 
   useImperativeHandle(