From 61a29043ac4e64b40c8f575e202ae7985c6ec96a Mon Sep 17 00:00:00 2001
From: guoaomen <www.guoao@foxmail.com>
Date: Mon, 21 Aug 2023 15:22:07 +0800
Subject: [PATCH] [inner-2320] fix: shardingTable can direct route only when
 the rules and nodes are the same globalTable nodes are the same, they can be
 direct route

---
 .../parser/druid/impl/DruidSelectParser.java  | 95 +++++++++++--------
 1 file changed, 54 insertions(+), 41 deletions(-)

diff --git a/src/main/java/com/actiontech/dble/route/parser/druid/impl/DruidSelectParser.java b/src/main/java/com/actiontech/dble/route/parser/druid/impl/DruidSelectParser.java
index 52762bc47..d6eeb3230 100644
--- a/src/main/java/com/actiontech/dble/route/parser/druid/impl/DruidSelectParser.java
+++ b/src/main/java/com/actiontech/dble/route/parser/druid/impl/DruidSelectParser.java
@@ -12,6 +12,7 @@ import com.actiontech.dble.config.ErrorCode;
 import com.actiontech.dble.config.model.SystemConfig;
 import com.actiontech.dble.config.model.sharding.SchemaConfig;
 import com.actiontech.dble.config.model.sharding.table.BaseTableConfig;
+import com.actiontech.dble.config.model.sharding.table.ChildTableConfig;
 import com.actiontech.dble.config.model.sharding.table.GlobalTableConfig;
 import com.actiontech.dble.config.model.sharding.table.ShardingTableConfig;
 import com.actiontech.dble.config.privileges.ShardingPrivileges;
@@ -240,56 +241,68 @@ public class DruidSelectParser extends DefaultDruidParser {
         }
 
         List<Pair<String, String>> tables = ctx.getTables();
-        int index = 0;
-        AbstractPartitionAlgorithm firstRule = null;
-        boolean directRoute = true;
-        Set<String> firstDataNodes = new HashSet<>();
         Map<String, BaseTableConfig> tableConfigMap = schemaConfig.getTables() == null ? null : schemaConfig.getTables();
 
-        if (tableConfigMap != null) {
-            for (Pair<String, String> table : tables) {
-                String tableName = table.getValue();
-                BaseTableConfig tc = tableConfigMap.get(tableName);
-                if (tc == null) {
-                    Map<String, String> tableAliasMap = ctx.getTableAliasMap();
-                    if (tableAliasMap != null && tableAliasMap.get(tableName) != null) {
-                        tc = tableConfigMap.get(tableAliasMap.get(tableName));
-                    }
-                }
-
-                if (index == 0) {
-                    if (tc != null) {
-                        if (!(tc instanceof ShardingTableConfig)) {
-                            continue;
-                        }
-                        firstRule = ((ShardingTableConfig) tc).getFunction();
-                        firstDataNodes.addAll(tc.getShardingNodes());
-                    }
-                } else {
-                    if (tc != null) {
-                        if (!(tc instanceof ShardingTableConfig)) {
-                            continue;
-                        }
-                        AbstractPartitionAlgorithm ruleCfg = ((ShardingTableConfig) tc).getFunction();
-                        Set<String> dataNodes = new HashSet<>(tc.getShardingNodes());
-                        if (firstRule != null && ((!ruleCfg.equals(firstRule)) || !dataNodes.equals(firstDataNodes))) {
-                            directRoute = false;
-                            break;
-                        }
-                    }
-                }
-                index++;
-            }
-        }
-
         RouteResultset rrsResult = rrs;
-        if (directRoute) {
+        if (canDirectRoute(tables, tableConfigMap)) {
             rrs.setStatement(RouterUtil.removeSchema(rrs.getStatement(), schemaConfig.getName()));
             rrsResult = tryRoute(schemaConfig, rrs);
         }
         return rrsResult;
     }
 
+    private boolean canDirectRoute(List<Pair<String, String>> tables, Map<String, BaseTableConfig> tableConfigMap) {
+        if (CollectionUtil.isEmpty(tableConfigMap) || CollectionUtil.isEmpty(tables)) {
+            return true;
+        }
+        boolean directRoute = true;
+        AbstractPartitionAlgorithm firstRule = null;
+        Set<String> firstShardingDataNodes = null;
+        for (Pair<String, String> table : tables) {
+            String tableName = table.getValue();
+            BaseTableConfig tc = tableConfigMap.get(tableName);
+            if (tc == null) {
+                Map<String, String> tableAliasMap = ctx.getTableAliasMap();
+                if (tableAliasMap != null && tableAliasMap.get(tableName) != null) {
+                    tc = tableConfigMap.get(tableAliasMap.get(tableName));
+                }
+            }
+            if (tc != null) {
+                AbstractPartitionAlgorithm ruleCfg = null;
+                Set<String> dataNodes;
+                if (tc instanceof ShardingTableConfig) {
+                    ruleCfg = ((ShardingTableConfig) tc).getFunction();
+                    dataNodes = new HashSet<>(tc.getShardingNodes());
+                } else if (tc instanceof ChildTableConfig) {
+                    ruleCfg = ((ChildTableConfig) tc).getDirectRouteTC().getFunction();
+                    dataNodes = new HashSet<>(tc.getShardingNodes());
+                } else if (tc instanceof GlobalTableConfig) {
+                    dataNodes = new HashSet<>(tc.getShardingNodes());
+                } else {
+                    directRoute = false;
+                    break;
+                }
+                if (firstRule == null && ruleCfg != null) {
+                    firstRule = ruleCfg;
+                }
+                if (firstShardingDataNodes == null) {
+                    firstShardingDataNodes = dataNodes;
+                }
+
+                if (firstRule != null && ruleCfg != null && ((!ruleCfg.equals(firstRule)) || !dataNodes.equals(firstShardingDataNodes))) {
+                    //sharding/child
+                    directRoute = false;
+                    break;
+                } else if ((firstRule == null || ruleCfg == null) && !dataNodes.equals(firstShardingDataNodes)) {
+                    //global
+                    directRoute = false;
+                    break;
+                }
+            }
+        }
+        return directRoute;
+    }
+
     private RouteResultset tryRoute(SchemaConfig schema, RouteResultset rrs) throws SQLException {
         if ((ctx.getTables() == null || ctx.getTables().size() == 0) && (ctx.getTableAliasMap() == null || ctx.getTableAliasMap().isEmpty())) {
             rrs = RouterUtil.routeToSingleNode(rrs, schema.getRandomShardingNode(), null);