Spring AOP实现多数据源切换

有时候我们会遇到这样的场景：一个应用系统中存在多个数据源，需要根据不同业务场景进行临时切换。比如读写分离（也可以考虑使用Mycat等数据库中间件）等。

Spring提供了动态数据源的功能，可以让我们实现在对数据库操作前进行切换。

下面我们演示怎么在项目中配置多数据源并根据不用业务场景进行切换（本文涉及到Spring Boot和Spring Data Jpa，相关内容及配置不做详解）。

1、在MySQL上面建立两个数据库（实际中可能在不同服务器），并建立相同的一张表student。

附上建表SQL：

DROP TABLE IF EXISTS `student`;
CREATE TABLE `student` (
  `id` bigint(11) NOT NULL AUTO_INCREMENT COMMENT '主键',
  `name` varchar(255) DEFAULT NULL COMMENT '姓名',
  `sex` tinyint(1) DEFAULT NULL COMMENT '性别：1男2女',
  `age` tinyint(3) DEFAULT NULL COMMENT '年龄',
  `birthday` datetime DEFAULT NULL COMMENT '生日',
  `address` varchar(255) DEFAULT NULL COMMENT '住址',
  PRIMARY KEY (`id`)
) ENGINE=InnoDB AUTO_INCREMENT=9 DEFAULT CHARSET=latin1;

2、建立对应的学生实体类。

/**
 * 学生实体类
 * @author z_hh  
 * @date 2018年11月30日
 */
@Getter
@Setter
@Builder
@ToString
@Entity
public class Student implements Serializable {

	private static final long serialVersionUID = -5636317592972887581L;

	/** 主键 */
	@Id
	@GeneratedValue(strategy = GenerationType.IDENTITY)
	private Long id;
	
	/** 姓名 */
	private String name;
	
	/** 性别 */
	private Integer sex;
	
	/** 年龄 */
	private Integer age;
	
	/** 生日 */
	private Date birthday;
	
	/** 地址 */
	private String address;
	
	public Student() {
	}

	public Student(Long id, String name, Integer sex, Integer age, Date birthday, String address) {
		super();
		this.id = id;
		this.name = name;
		this.sex = sex;
		this.age = age;
		this.birthday = birthday;
		this.address = address;
	}
	
}

3、配置两个不同数据源的连接信息。

4、自定义动态数据源（关键）。

这里定义一个动态数据源类，继承自AbstractRoutingDataSource并重写determineCurrentLookupKey方法，将数据源的切换粒度设置为线程。定义两个泛型（因为系统使用两个数据源）以表示数据源的类型，并提供线程可见变量CONTEX_THOLDER的获取和修改方法。

/**
 * 动态数据源
 * @author z_hh  
 * @date 2018年11月30日
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
	
	/**
	 * 本地线程可见变量，存储当前线程使用了哪一种数据源，初始设为Master
	 */
	private static final ThreadLocal<DatabaseType> CONTEX_THOLDER = new ThreadLocal<DatabaseType>() {
		protected DatabaseType initialValue() {
			return DatabaseType.MASTER;
		};
	};

	/**
	 * 重写AbstractRoutingDataSource方法，Spring从这里获取当前线程的数据源类型
	 */
	@Override
	protected Object determineCurrentLookupKey() {
		return CONTEX_THOLDER.get();
	}

	/**
	 * 数据源类型枚举类：Master主库，Slave从库
	 * @author z_hh  
	 * @date 2018年11月30日
	 */
	public static enum DatabaseType {
		MASTER, SLAVE
	}

	/**
	 * 将当前线程数据源类型设为Master
	 */
	public static void master() {
		CONTEX_THOLDER.set(DatabaseType.MASTER);
	}

	/**
	 * 将当前线程数据源类型设为Slave
	 */
	public static void slave() {
		CONTEX_THOLDER.set(DatabaseType.SLAVE);
	}

	/**
	 * 设置当前线程数据源类型
	 */
	public static void setDatabaseType(DatabaseType type) {
		CONTEX_THOLDER.set(type);
	}

	/**
	 * 获取当前线程数据源类型
	 */
	public static DatabaseType getType() {
		return CONTEX_THOLDER.get();
	}
}

5、数据源配置（关键）。

这里创建动态数据源对象，并创建两个类型的数据源作为其目标数据源，将其默认数据源设置为Master。最后交由Spring IOC容器管理。

/**
 * DataResource配置类
 * @author z_hh  
 * @date 2018年11月30日
 */
@Configuration
public class DataSourceConf {
	
	/**
	 * 将动态数据源注入Spring IOC容器
	 * @return
	 */
	@Bean
	public DataSource dynamicDataSource() {
		DataSource master = master();
		DataSource slave = slave();
		Map<Object, Object> targetDataSources = new HashMap<Object, Object>();
		targetDataSources.put(DynamicDataSource.DatabaseType.MASTER, master);
		targetDataSources.put(DynamicDataSource.DatabaseType.SLAVE, slave);
		DynamicDataSource dataSource = new DynamicDataSource();
		dataSource.setTargetDataSources(targetDataSources);// 该方法是AbstractRoutingDataSource的方法
		dataSource.setDefaultTargetDataSource(master);
		return dataSource;
	}

	/**
	 * 创建Master数据源对象
	 * @return
	 */
	public DataSource master() {
		HikariDataSource ds = new HikariDataSource();
        ds.setJdbcUrl(env.getProperty("master.datasource.url"));
        ds.setUsername(env.getProperty("master.datasource.username"));
        ds.setPassword(env.getProperty("master.datasource.password"));
        ds.setDriverClassName(env.getProperty("master.datasource.driver-class-name"));
		return ds;
	}

	/**
	 * 创建Slave数据源对象
	 * @return
	 */
	public DataSource slave() {
		HikariDataSource ds = new HikariDataSource();
        ds.setJdbcUrl(env.getProperty("slave.datasource.url"));
        ds.setUsername(env.getProperty("slave.datasource.username"));
        ds.setPassword(env.getProperty("slave.datasource.password"));
        ds.setDriverClassName(env.getProperty("slave.datasource.driver-class-name"));
		return ds;
	}

	@Autowired
	private Environment env;

}

6、定义两个注解，分别用于使用不同类型数据源的方法上面。

/**
 * 使用主库的注解
 * @author z_hh  
 * @date 2018年11月30日
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface Master {
}

/**
 * 使用从库的注解
 * @author z_hh  
 * @date 2018年11月30日
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface Slave {
}

7、定义Aspect切面（非常重要）。

关键点1：类上面使用了@Order注解，这是因为目标方法的@Transactional开启了事务也是一个AOP切面，需要通过Order属性去定义AOP切面的先后执行顺序。 order越小，在AOP的chain中越靠前，越先执行（chain模式）。而事务切面的优先级是默认最小的，所以我们需要将切换数据源的操作放在事务前面。

关键点2：切面使用@Around注解而不是@Before注解，这是因为我们在目标方法执行前进行了数据源切换，需要等目标方法执行完之后将其还原，恢复到之前的数据源（具体看业务场景吧）。

/**
 * 数据源的切入面
 * @author z_hh  
 * @date 2018年11月30日
 */
@Aspect
@Component
@Order(Ordered.LOWEST_PRECEDENCE - 1)
@Slf4j
public class DataSourceAOP {

	/**
	 * Master注解
	 * @param pjp
	 * @throws Throwable
	 */
	@Around("@annotation(cn.zhh.db.rw_separate.annotation.Master)")
	public Object setWriteDataSourceType(ProceedingJoinPoint pjp) throws Throwable {
		// 1、当前数据源类型
		DatabaseType currentType = DynamicDataSource.getType();
		try {
			// 2、如果当前是Slave，就切换到Master
			if (Objects.equals(currentType, DatabaseType.SLAVE)) {
				DynamicDataSource.master();
				log.info("dataSource切换到：master");
			}
			// 3、执行目标方法
			return pjp.proceed();
		} catch (Throwable t) {
			log.error("切换数据源异常！", t);
			throw t;
		} finally {
			// 4、需要将数据源还原
			DynamicDataSource.setDatabaseType(currentType);
		}
	}

	/**
	 * Slave注解
	 * @param pjp
	 * @throws Throwable
	 */
	@Around("@annotation(cn.zhh.db.rw_separate.annotation.Slave) && !@annotation(cn.zhh.db.rw_separate.annotation.Master)")
	public Object setReadDataSourceType(ProceedingJoinPoint pjp) throws Throwable {
		// 1、当前数据源类型
		DatabaseType currentType = DynamicDataSource.getType();
		try {
			// 2、如果当前是Master，就切换到Slave
			if (Objects.equals(currentType, DatabaseType.MASTER)) {
				DynamicDataSource.slave();
				log.info("dataSource切换到：slave");
			}
			// 3、执行目标方法
			return pjp.proceed();
		} catch (Throwable t) {
			log.error("切换数据源异常！", t);
			throw t;
		} finally {
			// 4、需要将数据源还原
			DynamicDataSource.setDatabaseType(currentType);
		}
	}

}

8、在业务代码里面使用自定义注解。

我们在写方法save使用了@Master注解，在读方法findById使用了@Slave方法。

/**
 * Student业务逻辑层
 * @author z_hh  
 * @date 2018年11月30日
 */
@Service
@Transactional
public class StudentServiceImpl implements StudentService {
	
	@Autowired
	private StudentDao dao;

	@Master
	@Override
	public Student save(Student student) {
		return dao.save(student);
	}

	@Override
	public void delete(Long id) {
		dao.deleteById(id);
	}

	@Slave
	@Override
	public Student findById(Long id) {
		return dao.findById(id).orElse(null);
	}

}

9、编写Junit测试代码。

测试前我们确认了两个数据库id为1的两条记录的name属性是不一样的，所以测试通过的话说明我们的通过AOP切换数据源功能生效了。

**
 * 测试类
 * @author z_hh  
 * @date 2018年11月30日
 */
@RunWith(SpringRunner.class)
@SpringBootTest
public class CommonTest {

	@Autowired
	private StudentService service;
	
	@Test
	public void test1() {
		// 保存student1到Master库，id=1，name=zhou
		Student student1 = Student.builder()
				.id(1L)
				.name("zhou")
				.sex(1)
				.age(24)
				.birthday(new Date())
				.build();
		assertNotNull(service.save(student1));
		
		// 从Slave库获取id=1的student2
		Student student2 = service.findById(1L);
		assertNotNull(student2);
		
		// 比较student1和student2的name，两个库里面是不相等的
		assertNotEquals(student1.getName(), student2.getName());
	}
}

10、我们也可以通过打印出来的日志证明。

Spring AOP实现多数据源切换

可以优化的地方：

（1）可以仅用一个自定义注解，定义一个属性用于设置需要使用的数据源，对应的切面逻辑也要修改。

（2）@Around注解可以改用@Before和@After注解。@Before里设置动态数据源本地线程共享对象里的值为想切换到的数据源的逻辑名称；@After里将动态数据源本地线程共享对象里的值移除掉，即可恢复到之前的数据源。

本文内容到此结束了，有什么问题或者建议，欢迎在评论区进行探讨！

猜你喜欢