feat: 新增读写分离组件 DataSourceShardingStrategy.java

docs/.vitepress/config.ts

@@ -88,6 +88,7 @@ export default defineConfig({
                     {text: 'SQL 审计', link: '/zh/core/audit'},
                     {text: 'SQL 打印', link: '/zh/core/sql-print'},
                     {text: '多数据源', link: '/zh/core/multi-datasource'},
+                    {text: '读写分离 💥', link: '/zh/core/read-write-splitting'},
                     {text: '数据源加密', link: '/zh/core/datasource-encryption'},
                     {text: '动态表名', link: '/zh/core/dynamic-table'},
                     {text: '事务管理', link: '/zh/core/tx'},

docs/zh/core/read-write-splitting.md

@@ -0,0 +1,97 @@
+# 读写分离
+
+MyBatis-Flex 的读写分离功能是基于 【多数据源】 功能来实现的。
+
+读写分离的功能，要求当前的环境必须是多个数据库（也可理解为多个数据源），其原理是：
+让主数据库（master）处理事务性操作，比如：增、删、改（INSERT、DELETE、UPDATE），而从数据库（slave）处理 SELECT 查询操作。
+
+在 MyBatis 框架中，我们知道： 所有关于数据库的的操作都是通过 Mapper 来实现的，Mapper 里的一个方法，往往是和一个 SQL 一一对应。
+
+因此，在 MyBatis-Flex 中，提供了一种基于 Mapper 方法的读写分离策略。
+
+## 分片策略
+
+自定义 `DataSourceShardingStrategy` 例如：
+
+```java
+public class MyStrategy implements DataSourceShardingStrategy {
+
+    public String doSharding(String currentDataSourceKey
+        , Object mapper, Method mapperMethod, Object[] methodArgs){
+
+        //返回新的数据源 key
+        return "newDataSourceKey";
+    }
+}
+```
+
+doSharding 的参数分别为：
+
+- currentDataSourceKey：当前由用户端已配置的 key
+- mapper：当前的 mapper 对象
+- mapperMethod: 当前的 mapper 方法
+- methodArgs：当前的 mapper 方法的参数内容
+
+自定义好 数据源分片策略后，在项目启动时，需要通过 `DataSourceManager` 配置自己的自定义分片策略：
+
+```java
+DataSourceManager.setDataSourceShardingStrategy(new MyStrategy());
+```
+
+## 示例代码
+
+假设数据源配置如下：
+
+
+```yaml
+mybatis-flex:
+  datasource:
+    master:
+      type: druid
+      url: jdbc:mysql://127.0.0.1:3306/master-db
+      username: root
+      password: 123456
+    slave1:
+      type: com.your.datasource.type2
+      url: jdbc:mysql://127.0.0.1:3306/slave1
+      username: root
+      password: 123456
+    slave2:
+      type: com.your.datasource.type2
+      url: jdbc:mysql://127.0.0.1:3306/slave2
+      username: root
+      password: 123456
+    other:
+      type: com.your.datasource.type2
+      url: jdbc:mysql://127.0.0.1:3306/other
+      username: root
+      password: 123456
+```
+以上配置中，一共有 4 个数据源，分别为 `master`、`slave1`、`slave2`、`other`。
+假设我们的需求是：在 增删改 时，走 master，而在查询时，自动使用 `slave1`、`slave2` 进行负载均衡。
+
+
+那么，我们的分片策略代码如下：
+
+```java
+public class MyStrategy implements DataSourceShardingStrategy {
+
+    public String doSharding(String currentDataSourceKey
+        , Object mapper, Method mapperMethod, Object[] methodArgs){
+
+        // 不管 other 数据源的情况
+        if ("other".equals(currentDataSourceKey)){
+            return currentDataSourceKey;
+        }
+
+        // 如果 mapper 的方法属于 增删改，使用 master 数据源
+        if (StringUtil.startWithAny(mapperMethod.getName(),
+            "insert", "delete", "update")){
+            return "master";
+        }
+
+        //其他场景，使用 slave1 或者 slave2 进行负载均衡
+        return "slave*";
+    }
+}
+```

mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceKey.java

@@ -15,6 +15,7 @@
  */
 package com.mybatisflex.core.datasource;
 
+import java.lang.reflect.Method;
 import java.util.function.Supplier;
 
 /**
@@ -86,4 +87,9 @@ public class DataSourceKey {
     public static void setManualKeyThreadLocal(ThreadLocal<String> manualKeyThreadLocal) {
         DataSourceKey.manualKeyThreadLocal = manualKeyThreadLocal;
     }
+
+    public static String getByShardingStrategy(String dataSource, Object mapper, Method method, Object[] args) {
+        String shardingDsKey = DataSourceManager.getByShardingStrategy(dataSource, mapper, method, args);
+        return shardingDsKey != null ? shardingDsKey : dataSource;
+    }
 }

mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceManager.java

@@ -37,6 +37,15 @@ public class DataSourceManager {
         DataSourceManager.decipher = decipher;
     }
 
+    private static DataSourceShardingStrategy dataSourceShardingStrategy;
+
+    public static DataSourceShardingStrategy getDataSourceShardingStrategy() {
+        return dataSourceShardingStrategy;
+    }
+
+    public static void setDataSourceShardingStrategy(DataSourceShardingStrategy dataSourceShardingStrategy) {
+        DataSourceManager.dataSourceShardingStrategy = dataSourceShardingStrategy;
+    }
 
     public static void decryptDataSource(DataSource dataSource) {
         if (decipher == null) {
@@ -87,4 +96,7 @@ public class DataSourceManager {
     }
 
 
+    static String getByShardingStrategy(String dataSource, Object mapper, Method method, Object[] args) {
+        return dataSourceShardingStrategy != null ? dataSourceShardingStrategy.doSharding(dataSource, mapper, method, args) : null;
+    }
 }

mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceShardingStrategy.java

@@ -0,0 +1,22 @@
+/*
+ *  Copyright (c) 2022-2023, Mybatis-Flex (fuhai999@gmail.com).
+ *  <p>
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  <p>
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  <p>
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package com.mybatisflex.core.datasource;
+
+import java.lang.reflect.Method;
+
+public interface DataSourceShardingStrategy {
+    String doSharding(String currentDataSourceKey, Object mapper, Method mapperMethod, Object[] methodArgs);
+}

mybatis-flex-core/src/main/java/com/mybatisflex/core/mybatis/MapperInvocationHandler.java

@@ -67,6 +67,13 @@ public class MapperInvocationHandler implements InvocationHandler {
                 }
             }
 
+            //最终通过数据源 自定义分片 策略去获取
+            String shardingDataSourceKey = DataSourceKey.getByShardingStrategy(dataSourceKey, proxy, method, args);
+            if (shardingDataSourceKey != null && !shardingDataSourceKey.equals(dataSourceKey)) {
+                DataSourceKey.use(dataSourceKey);
+                needClearDsKey = true;
+            }
+
             //优先获取用户自己配置的 dbType
             DbType dbType = DialectFactory.getHintDbType();
             if (dbType == null) {