From 9fc4665c1a2ed4c9ceda43ea5dc3b8648c9e29a5 Mon Sep 17 00:00:00 2001
From: RuoYi <yzz_ivy@163.com>
Date: Fri, 26 Jul 2019 20:19:47 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A4=9A=E6=95=B0=E6=8D=AE=E6=BA=90=E6=94=AF?=
 =?UTF-8?q?=E6=8C=81=E7=B1=BB=E6=B3=A8=E8=A7=A3=EF=BC=88=E5=85=81=E8=AE=B8?=
 =?UTF-8?q?=E7=BB=A7=E6=89=BF=E7=88=B6=E7=B1=BB=E7=9A=84=E6=B3=A8=E8=A7=A3?=
 =?UTF-8?q?=EF=BC=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../ruoyi/common/annotation/DataSource.java   |  6 +++-
 .../framework/aspectj/DataSourceAspect.java   | 29 +++++++++++++++----
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/ruoyi-common/src/main/java/com/ruoyi/common/annotation/DataSource.java b/ruoyi-common/src/main/java/com/ruoyi/common/annotation/DataSource.java
index 5a41a2c6..a272a5b7 100644
--- a/ruoyi-common/src/main/java/com/ruoyi/common/annotation/DataSource.java
+++ b/ruoyi-common/src/main/java/com/ruoyi/common/annotation/DataSource.java
@@ -1,6 +1,8 @@
 package com.ruoyi.common.annotation;
 
+import java.lang.annotation.Documented;
 import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
@@ -11,8 +13,10 @@ import com.ruoyi.common.enums.DataSourceType;
  * 
  * @author ruoyi
  */
-@Target(ElementType.METHOD)
+@Target({ ElementType.METHOD, ElementType.TYPE })
 @Retention(RetentionPolicy.RUNTIME)
+@Documented
+@Inherited
 public @interface DataSource
 {
     /**
diff --git a/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/DataSourceAspect.java b/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/DataSourceAspect.java
index 89d28fac..4f13ddb3 100644
--- a/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/DataSourceAspect.java
+++ b/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/DataSourceAspect.java
@@ -26,7 +26,8 @@ public class DataSourceAspect
 {
     protected Logger logger = LoggerFactory.getLogger(getClass());
 
-    @Pointcut("@annotation(com.ruoyi.common.annotation.DataSource)")
+    @Pointcut("@annotation(com.ruoyi.common.annotation.DataSource)"
+            + "|| @within(com.ruoyi.common.annotation.DataSource)")
     public void dsPointCut()
     {
 
@@ -35,11 +36,7 @@ public class DataSourceAspect
     @Around("dsPointCut()")
     public Object around(ProceedingJoinPoint point) throws Throwable
     {
-        MethodSignature signature = (MethodSignature) point.getSignature();
-
-        Method method = signature.getMethod();
-
-        DataSource dataSource = method.getAnnotation(DataSource.class);
+        DataSource dataSource = getDataSource(point);
 
         if (StringUtils.isNotNull(dataSource))
         {
@@ -56,4 +53,24 @@ public class DataSourceAspect
             DynamicDataSourceContextHolder.clearDataSourceType();
         }
     }
+
+    /**
+     * 获取需要切换的数据源
+     */
+    public DataSource getDataSource(ProceedingJoinPoint point)
+    {
+        MethodSignature signature = (MethodSignature) point.getSignature();
+        Class<? extends Object> targetClass = point.getTarget().getClass();
+        DataSource targetDataSource = targetClass.getAnnotation(DataSource.class);
+        if (StringUtils.isNotNull(targetDataSource))
+        {
+            return targetDataSource;
+        }
+        else
+        {
+            Method method = signature.getMethod();
+            DataSource dataSource = method.getAnnotation(DataSource.class);
+            return dataSource;
+        }
+    }
 }