From bb394b06f16ebbec8c372daaf59631f23c1c18c2 Mon Sep 17 00:00:00 2001
From: mystarry-sky <zilan1986@163.com>
Date: Wed, 12 Mar 2025 19:36:23 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E9=87=8F=E6=A0=A1=E9=AA=8CopenGauss?=
 =?UTF-8?q?=20=E5=88=97=E8=BF=87=E6=BB=A4=E8=A7=84=E5=88=99=E4=B8=A2?=
 =?UTF-8?q?=E5=A4=B1=E4=B8=BB=E9=94=AE=E5=88=97=E5=AF=BC=E8=87=B4=E6=A0=A1?=
 =?UTF-8?q?=E9=AA=8C=E5=A4=B1=E8=B4=A5=E9=97=AE=E9=A2=98=E4=BF=AE=E5=A4=8D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../extract/data/BaseDataService.java          | 18 ++++++++++++++++++
 .../extract/service/RuleAdapterService.java    | 13 ++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/data/BaseDataService.java b/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/data/BaseDataService.java
index cc6ce38..469fec4 100644
--- a/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/data/BaseDataService.java
+++ b/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/data/BaseDataService.java
@@ -17,6 +17,8 @@ package org.opengauss.datachecker.extract.data;
 
 import com.alibaba.druid.pool.DruidDataSource;
 
+import cn.hutool.core.collection.CollectionUtil;
+
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.collections4.MapUtils;
 import org.apache.logging.log4j.Logger;
@@ -46,6 +48,7 @@ import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 /**
@@ -252,6 +255,21 @@ public class BaseDataService {
      */
     public List<ColumnsMetaData> queryTableColumnsMetaData(String tableName) {
         final List<ColumnsMetaData> columns = dataAccessService.queryTableColumnsMetaData(tableName);
+        List<PrimaryColumnBean> pkColList = dataAccessService.queryTablePrimaryColumns(tableName);
+        List<PrimaryColumnBean> ukColList = dataAccessService.queryTableUniqueColumns(tableName);
+        if (CollectionUtil.isNotEmpty(ukColList)) {
+            if (CollectionUtil.isNotEmpty(ukColList)) {
+                pkColList.addAll(ukColList);
+            }
+        }
+        Set<String> pkOrUkSet = pkColList.stream()
+            .map(PrimaryColumnBean::getColumnName)
+            .collect(Collectors.toUnmodifiableSet());
+        columns.forEach(meta -> {
+            if (pkOrUkSet.contains(meta.getColumnName())) {
+                meta.setColumnKey(ColumnKey.PRI);
+            }
+        });
         return ruleAdapterService.executeColumnRule(columns);
     }
 
diff --git a/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/service/RuleAdapterService.java b/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/service/RuleAdapterService.java
index 72b6c7e..a521200 100644
--- a/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/service/RuleAdapterService.java
+++ b/datachecker-extract/src/main/java/org/opengauss/datachecker/extract/service/RuleAdapterService.java
@@ -17,7 +17,9 @@ package org.opengauss.datachecker.extract.service;
 
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.logging.log4j.Logger;
+import org.opengauss.datachecker.common.config.ConfigCache;
 import org.opengauss.datachecker.common.entry.common.Rule;
+import org.opengauss.datachecker.common.entry.enums.CheckMode;
 import org.opengauss.datachecker.common.entry.enums.RuleType;
 import org.opengauss.datachecker.common.entry.extract.ColumnsMetaData;
 import org.opengauss.datachecker.common.entry.extract.TableMetadata;
@@ -25,9 +27,11 @@ import org.opengauss.datachecker.common.util.LogUtils;
 import org.springframework.stereotype.Service;
 
 import javax.annotation.Resource;
+
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * RuleAdapterService
@@ -55,7 +59,14 @@ public class RuleAdapterService {
      */
     public void init(Map<RuleType, List<Rule>> rules) {
         RULES.clear();
-        RULES.putAll(rules);
+        CheckMode checkMode = ConfigCache.getCheckMode();
+        if (Objects.equals(CheckMode.FULL, checkMode)) {
+            RULES.putAll(rules);
+        } else if (Objects.equals(CheckMode.INCREMENT, checkMode)) {
+            RULES.put(RuleType.COLUMN, rules.get(RuleType.COLUMN));
+        } else {
+            log.warn("check mode {} is not support rules", checkMode);
+        }
     }
 
     /**
-- 
Gitee