From 183932f995c80d376085d4f1b6e5a560db725ff5 Mon Sep 17 00:00:00 2001
From: wxz <wangxianzhang1992@163.com>
Date: Fri, 21 Aug 2020 14:05:00 +0800
Subject: [PATCH] =?UTF-8?q?=E5=8A=A8=E6=80=81=E6=95=B0=E6=8D=AE=E6=BA=90-?=
 =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=96=B9=E6=B3=95=E5=85=A5=E5=8F=82=E7=BA=A7?=
 =?UTF-8?q?=E5=88=AB=E7=9A=84=E8=AE=BE=E7=BD=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../datasource/annotation/DataSource.java     | 10 +++++
 .../datasource/aspect/DataSourceAspect.java   | 42 ++++++++++++++++++-
 .../datasource/bean/DataSourceParam.java      | 12 ++++++
 3 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/bean/DataSourceParam.java

diff --git a/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/annotation/DataSource.java b/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/annotation/DataSource.java
index 3a8ac1ae8f..e0675092c8 100644
--- a/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/annotation/DataSource.java
+++ b/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/annotation/DataSource.java
@@ -21,5 +21,15 @@ import java.lang.annotation.*;
 @Documented
 @Inherited
 public @interface DataSource {
+    /**
+     * 直接指定数据源名称
+     * @return
+     */
     String value() default "";
+
+    /**
+     * 是否从参数中获取数据源名称，优先级高于value
+     * @return
+     */
+    boolean datasourceNameFromArg() default false;
 }
diff --git a/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/aspect/DataSourceAspect.java b/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/aspect/DataSourceAspect.java
index d636568425..45113847c6 100644
--- a/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/aspect/DataSourceAspect.java
+++ b/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/aspect/DataSourceAspect.java
@@ -9,7 +9,9 @@
 package com.epmet.commons.dynamic.datasource.aspect;
 
 import com.epmet.commons.dynamic.datasource.annotation.DataSource;
+import com.epmet.commons.dynamic.datasource.bean.DataSourceParam;
 import com.epmet.commons.dynamic.datasource.config.DynamicContextHolder;
+import org.apache.commons.lang3.StringUtils;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
@@ -22,6 +24,7 @@ import org.springframework.core.annotation.Order;
 import org.springframework.stereotype.Component;
 
 import java.lang.reflect.Method;
+import java.lang.reflect.Parameter;
 
 /**
  * 多数据源，切面处理类
@@ -52,9 +55,9 @@ public class DataSourceAspect  {
         if(targetDataSource != null || methodDataSource != null){
             String value;
             if(methodDataSource != null){
-                value = methodDataSource.value();
+                value = getDatasourceName(methodDataSource, signature.getMethod().getParameters(), point.getArgs());
             }else {
-                value = targetDataSource.value();
+                value = getDatasourceName(targetDataSource, signature.getMethod().getParameters(), point.getArgs());
             }
 
             DynamicContextHolder.push(value);
@@ -68,4 +71,39 @@ public class DataSourceAspect  {
             logger.debug("clean datasource");
         }
     }
+
+    /**
+     * 获取要用到的数据源名称
+     * @param dataSource
+     * @return
+     */
+    public String getDatasourceName(DataSource dataSource, Parameter[] methodParameters, Object[] methodArgValues) {
+        if (dataSource.datasourceNameFromArg()) {
+            // 1.从参数中动态获取数据源名称
+            String datasourceNameFromParam = getDatasourceNameFromArg(methodParameters, methodArgValues);
+            if (StringUtils.isNotBlank(datasourceNameFromParam)) {
+                // 如果有DatasourceParam类型的参数并且设置了datasourceName值，那么返回这个值，否则使用硬编码的
+                return datasourceNameFromParam;
+            }
+        }
+        // 2.硬编码指定数据源名称
+        return dataSource.value();
+    }
+
+    /**
+     * 从参数中取数据源名称
+     * @param parameters
+     * @param argsObject
+     * @return
+     */
+    public String getDatasourceNameFromArg(Parameter[] parameters, Object[] argsObject) {
+        for (int i = 0; i < parameters.length; i++) {
+            if (parameters[i].getType() == DataSourceParam.class) {
+                DataSourceParam param = (DataSourceParam) argsObject[i];
+                return param.getDatasourceName();
+            }
+        }
+
+        return null;
+    }
 }
diff --git a/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/bean/DataSourceParam.java b/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/bean/DataSourceParam.java
new file mode 100644
index 0000000000..19b5905588
--- /dev/null
+++ b/epmet-commons/epmet-commons-dynamic-datasource/src/main/java/com/epmet/commons/dynamic/datasource/bean/DataSourceParam.java
@@ -0,0 +1,12 @@
+package com.epmet.commons.dynamic.datasource.bean;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+
+@Data
+@AllArgsConstructor
+public class DataSourceParam {
+
+    private String datasourceName;
+
+}