AOP 之 动态数据源

package com.github.mengweijin.mybatisplus.demo.DynamicDataSource;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * @description 自定义多数据源切换注解
 *
 * @author Meng Wei Jin
 **/
@Target({ElementType.TYPE, ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface DS {
    /**
     * 切换数据源名称
     */
    String value() default DynamicDataSourceHolder.MASTER;
}

package com.github.mengweijin.mybatisplus.demo.DynamicDataSource;

import lombok.Getter;

public class DynamicDataSourceHolder {

    public static final String MASTER = "master";

    public static final String SLAVE = "slave";

    @Getter
    public static final ThreadLocal<String> dataSourceThreadLocal = new ThreadLocal<>();
}

package com.github.mengweijin.mybatisplus.demo.DynamicDataSource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

/**
 * @description 动态数据源
 *
 * @author Meng Wei Jin
 **/
public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceHolder.getDataSourceThreadLocal().get();
    }
}

package com.github.mengweijin.mybatisplus.demo.DynamicDataSource;

import com.zaxxer.hikari.HikariDataSource;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

@Configuration
public class DynamicDataSourceConfiguration {

    @Bean("masterDataSource")
    @ConfigurationProperties("spring.datasource.master")
    public DataSource masterDataSource() {
        return new HikariDataSource();
    }

    @Bean("slaveDataSource")
    @ConfigurationProperties("spring.datasource.slave")
    @ConditionalOnProperty(prefix = "spring.datasource.slave")
    public DataSource slaveDataSource() {
        return new HikariDataSource();
    }

    @Bean
    public DynamicDataSource dataSource() {
        Map<Object, Object> targetDataSources = new HashMap<>(2);
        targetDataSources.put(DynamicDataSourceHolder.MASTER, masterDataSource());
        targetDataSources.put(DynamicDataSourceHolder.SLAVE, slaveDataSource());

        DynamicDataSource dataSource = new DynamicDataSource();
        dataSource.setDefaultTargetDataSource(masterDataSource());
        dataSource.setTargetDataSources(targetDataSources);
        return dataSource;
    }
}

package com.github.mengweijin.mybatisplus.demo.DynamicDataSource;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import java.lang.reflect.Method;

@Aspect
@Order(1)
@Component
public class DynamicDataSourceAspect {
    @Pointcut("@annotation(com.github.mengweijin.mybatisplus.demo.DynamicDataSource.DS)")
    public void dsPointCut() {
    }

    @Around("dsPointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        Signature signature = joinPoint.getSignature();
        MethodSignature methodSignature = (MethodSignature) signature;
        Method method = methodSignature.getMethod();
        DS ds = method.getAnnotation(DS.class);

        if (ds != null) {
            DynamicDataSourceHolder.getDataSourceThreadLocal().set(ds.value());
        }

        try {
            return joinPoint.proceed();
        } finally {
            // 在执行方法之后，从 threadLocal 中移除
            DynamicDataSourceHolder.getDataSourceThreadLocal().remove();
        }
    }
}

使用

• 可以添加在类上面，也可以在方法上面。
• 可以指定使用哪个数据源。如果只添加注解而不具体指定，默认使用 DynamicDataSourceHolder.MASTER

@DS(DynamicDataSourceHolder.SLAVE)
@Service
public class UserService {

    @DS(DynamicDataSourceHolder.SLAVE)
    public void add(User user) {
        // add user to database
    }
}