Cassandra数据源动态切换

一、背景

        项目以cassandra作为非关系型数据库，并根据租户分库，程序查询时需要根据租户信息查询对应的数据库，因此需要程序可以切换数据源。

二、SpringBoot集成cassandra

1、导入pom依赖

        org.springframework.data
        spring-data-cassandra

2、application.yml添加配置

spring:
  data:
    cassandra:
      keyspace-name: system
      contact-points:
      - 10.5.11.131
      port: 9043
      username: cassandra
      password: cassandra

3、使用cassandraTemplate存储数据

String sql = "insert into table (xx, xx, xx) values (?, ?, ?)";
cassandraTemplate.getCqlOperations().execute(sql, XX, XX, XX);

三、SpringBoot集成CassandraTemplate源码研究

1、cassandraTemplate如何获取数据源

cassandraTemplate的execute方法如下

@Override
public boolean execute(String cql) throws DataAccessException {

    Assert.hasText(cql, "CQL must not be empty");

	return queryForResultSet(cql).wasApplied();
}

继续追踪，query方法如下，在这里getCurrentSession方法即是获取数据源 

@Override
@Nullable
public  T query(String cql, ResultSetExtractor resultSetExtractor) throws     DataAccessException {

		Assert.hasText(cql, "CQL must not be empty");
		Assert.notNull(resultSetExtractor, "ResultSetExtractor must not be null");

		try {
			if (logger.isDebugEnabled()) {
				logger.debug("Executing CQL Statement [{}]", cql);
			}

			SimpleStatement statement = applyStatementSettings(new SimpleStatement(cql));
            
            //这边会获取数据源，然后调用execute方法存储
			ResultSet results = getCurrentSession().execute(statement);

			return resultSetExtractor.extractData(results);
		} catch (DriverException e) {
			throw translateException("Query", cql, e);
		}
	}

进入getCurrentSession方法，可以发现是先获取SessionFactory，然后再得到session对象。SessionFactory是一个接口，它有2个实现类，AbstractRoutingSessionFactory和DefaultSessionFactory，默认实例化的是DefaultSessionFactory

private Session getCurrentSession() {

		SessionFactory sessionFactory = getSessionFactory();

		Assert.state(sessionFactory != null, "SessionFactory is null");

		return sessionFactory.getSession();
}

DefaultSessionFactory的getSession方法很简单，直接返回session对象。而AbstractRoutingSessionFactory是个抽象类，它获取sessionFactory的方法是protected，所以我们可以创建类继承AbstractRoutingSessionFactory，然后重写determineTargetSessionFactory或者直接重写getSession方法

@Override
public Session getSession() {
        return determineTargetSessionFactory().getSession();
}

protected SessionFactory determineTargetSessionFactory() {

	Assert.notNull(this.resolvedSessionFactories, "SessionFactory router not initialized");

	Object lookupKey = determineCurrentLookupKey();
	SessionFactory sessionFactory = this.resolvedSessionFactories.get(lookupKey);

	if (sessionFactory == null && (this.lenientFallback || lookupKey == null)) {
		sessionFactory = this.resolvedDefaultSessionFactory;
	}

	if (sessionFactory == null) {
		throw new IllegalStateException(
				String.format("Cannot determine target SessionFactory for lookup key [%s]", lookupKey));
	}

	return sessionFactory;
}

2、SessionFactory实例化源码

那么这些SessionFactory是如何实例化的呢？是在CassandraDataAutoConfiguration中，public CassandraTemplate cassandraTemplate这个方法，在CassandraTemplate的构造函数中，是默认使用DefaultSessionFactory的，所以我们可以自己创建CassandraTemplate，指定自定义的SessionFactory，从而获取自己想要的数据源

@Configuration
@ConditionalOnClass({ Cluster.class, CassandraAdminOperations.class })
@ConditionalOnBean(Cluster.class)
@EnableConfigurationProperties(CassandraProperties.class)
@AutoConfigureAfter(CassandraAutoConfiguration.class)
public class CassandraDataAutoConfiguration {

	private final BeanFactory beanFactory;

	private final CassandraProperties properties;

	private final Cluster cluster;

	private final Environment environment;

	public CassandraDataAutoConfiguration(BeanFactory beanFactory, CassandraProperties properties, Cluster cluster,
			Environment environment) {
		this.beanFactory = beanFactory;
		this.properties = properties;
		this.cluster = cluster;
		this.environment = environment;
	}

	@Bean
	@ConditionalOnMissingBean
	public CassandraMappingContext cassandraMapping(CassandraCustomConversions conversions)
			throws ClassNotFoundException {
		CassandraMappingContext context = new CassandraMappingContext();
		List packages = EntityScanPackages.get(this.beanFactory).getPackageNames();
		if (packages.isEmpty() && AutoConfigurationPackages.has(this.beanFactory)) {
			packages = AutoConfigurationPackages.get(this.beanFactory);
		}
		if (!packages.isEmpty()) {
			context.setInitialEntitySet(CassandraEntityClassScanner.scan(packages));
		}
		PropertyMapper.get().from(this.properties::getKeyspaceName).whenHasText().as(this::createSimpleUserTypeResolver)
				.to(context::setUserTypeResolver);
		context.setCustomConversions(conversions);
		return context;
	}

	private SimpleUserTypeResolver createSimpleUserTypeResolver(String keyspaceName) {
		return new SimpleUserTypeResolver(this.cluster, keyspaceName);
	}

	@Bean
	@ConditionalOnMissingBean
	public CassandraConverter cassandraConverter(CassandraMappingContext mapping,
			CassandraCustomConversions conversions) {
		MappingCassandraConverter converter = new MappingCassandraConverter(mapping);
		converter.setCustomConversions(conversions);
		return converter;
	}

	@Bean
	@ConditionalOnMissingBean(Session.class)
	public CassandraSessionFactoryBean cassandraSession(CassandraConverter converter) throws Exception {
		CassandraSessionFactoryBean session = new CassandraSessionFactoryBean();
		session.setCluster(this.cluster);
		session.setConverter(converter);
		session.setKeyspaceName(this.properties.getKeyspaceName());
		Binder binder = Binder.get(this.environment);
		binder.bind("spring.data.cassandra.schema-action", SchemaAction.class).ifBound(session::setSchemaAction);
		return session;
	}

	@Bean
	@ConditionalOnMissingBean
	public CassandraTemplate cassandraTemplate(Session session, CassandraConverter converter) throws Exception {
		return new CassandraTemplate(session, converter);
	}

	@Bean
	@ConditionalOnMissingBean
	public CassandraCustomConversions cassandraCustomConversions() {
		return new CassandraCustomConversions(Collections.emptyList());
	}

}

3、创建数据源源码

数据源是封装到SessionManager对象中的，它是Session的子类。创建它，就需要Cluster对象，如下所示

@Configuration
@ConditionalOnClass({ Cluster.class })
@EnableConfigurationProperties(CassandraProperties.class)
public class CassandraAutoConfiguration {

	private final CassandraProperties properties;

	private final ObjectProvider builderCustomizers;

	public CassandraAutoConfiguration(CassandraProperties properties,
			ObjectProvider builderCustomizers) {
		this.properties = properties;
		this.builderCustomizers = builderCustomizers;
	}

	@Bean
	@ConditionalOnMissingBean
	@SuppressWarnings("deprecation")
	public Cluster cassandraCluster() {
		PropertyMapper map = PropertyMapper.get();
		CassandraProperties properties = this.properties;
        //Cluster名，端口号
		Cluster.Builder builder = Cluster.builder().withClusterName(properties.getClusterName())
				.withPort(properties.getPort());
        //用户名
		map.from(properties::getUsername).whenNonNull()
				.to((username) -> builder.withCredentials(username, properties.getPassword()));
		map.from(properties::getCompression).whenNonNull().to(builder::withCompression);
		map.from(properties::getLoadBalancingPolicy).whenNonNull().as(BeanUtils::instantiateClass)
				.to(builder::withLoadBalancingPolicy);
		map.from(this::getQueryOptions).to(builder::withQueryOptions);
		map.from(properties::getReconnectionPolicy).whenNonNull().as(BeanUtils::instantiateClass)
				.to(builder::withReconnectionPolicy);
		map.from(properties::getRetryPolicy).whenNonNull().as(BeanUtils::instantiateClass).to(builder::withRetryPolicy);
		map.from(this::getSocketOptions).to(builder::withSocketOptions);
		map.from(properties::isSsl).whenTrue().toCall(builder::withSSL);
		map.from(this::getPoolingOptions).to(builder::withPoolingOptions);
		map.from(properties::getContactPoints).as(StringUtils::toStringArray).to(builder::addContactPoints);
		map.from(properties::isJmxEnabled).whenFalse().toCall(builder::withoutJMXReporting);
		customize(builder);
		return builder.build();
	}

	private void customize(Cluster.Builder builder) {
		this.builderCustomizers.orderedStream().forEach((customizer) -> customizer.customize(builder));
	}

	private QueryOptions getQueryOptions() {
		PropertyMapper map = PropertyMapper.get();
		QueryOptions options = new QueryOptions();
		CassandraProperties properties = this.properties;
		map.from(properties::getConsistencyLevel).whenNonNull().to(options::setConsistencyLevel);
		map.from(properties::getSerialConsistencyLevel).whenNonNull().to(options::setSerialConsistencyLevel);
		map.from(properties::getFetchSize).to(options::setFetchSize);
		return options;
	}

	private SocketOptions getSocketOptions() {
		PropertyMapper map = PropertyMapper.get();
		SocketOptions options = new SocketOptions();
		map.from(this.properties::getConnectTimeout).whenNonNull().asInt(Duration::toMillis)
				.to(options::setConnectTimeoutMillis);
		map.from(this.properties::getReadTimeout).whenNonNull().asInt(Duration::toMillis)
				.to(options::setReadTimeoutMillis);
		return options;
	}

	private PoolingOptions getPoolingOptions() {
		PropertyMapper map = PropertyMapper.get();
		CassandraProperties.Pool properties = this.properties.getPool();
		PoolingOptions options = new PoolingOptions();
		map.from(properties::getIdleTimeout).whenNonNull().asInt(Duration::getSeconds)
				.to(options::setIdleTimeoutSeconds);
		map.from(properties::getPoolTimeout).whenNonNull().asInt(Duration::toMillis).to(options::setPoolTimeoutMillis);
		map.from(properties::getHeartbeatInterval).whenNonNull().asInt(Duration::getSeconds)
				.to(options::setHeartbeatIntervalSeconds);
		map.from(properties::getMaxQueueSize).to(options::setMaxQueueSize);
		return options;
	}

}

创建完Cluster对象，在CassandraDataAutoConfiguration类创建CassandraSessionFactoryBean时，因为CassandraSessionFactoryBean继承CassandraCqlSessionFactoryBean，而CassandraCqlSessionFactoryBean的afterPropertiesSet会调用Cluster的connect方法，从而创建SessionManager对象，这个对象就是数据源

public class CassandraCqlSessionFactoryBean
		implements FactoryBean, InitializingBean, DisposableBean, PersistenceExceptionTranslator {

	protected final Logger logger = LoggerFactory.getLogger(getClass());

	private final PersistenceExceptionTranslator exceptionTranslator = new CassandraExceptionTranslator();

	private @Nullable Cluster cluster;

	private List startupScripts = Collections.emptyList();

	private List shutdownScripts = Collections.emptyList();

	private @Nullable Session session;

	private @Nullable String keyspaceName;

	
	@Override
	public void afterPropertiesSet() throws Exception {

		this.session = connect(getKeyspaceName());

		executeScripts(getStartupScripts());
	}

	Session connect(@Nullable String keyspaceName) {
		return (StringUtils.hasText(keyspaceName) ? getCluster().connect(keyspaceName) : getCluster().connect());
	}
}

四、Cassandra数据源切换方案

1、方案一，创建AbstractRoutingSessionFactory子类

1.1、手动创建数据源

//ip，可多个，用逗号分隔开
String cassandraContactPoints = tenantCassandraDO.getCassandraContactPoints();
//端口
Integer cassandraPort =tenantCassandraDO.getCassandraPort();
//用户名
String cassandraUserName = tenantCassandraDO.getCassandraUserName();
//密码
String cassandraPassword = tenantCassandraDO.getCassandraPassword();
//这边一定要加withoutJMXReporting，否则会去报缺少com/codahale/metrics/JmxReporter这个类
Cluster.Builder builder = Cluster.builder().withPort(cassandraPort).withoutJMXReporting();
//将每个ip地址添加进去
for (String cassandraContactPoint : cassandraContactPoints.split(",")) {
        builder.addContactPoint(cassandraContactPoint);
}
//验证用户名密码
builder.withCredentials(cassandraUserName, cassandraPassword);
//创建Cluster
Cluster cluster = builder.build();
//创建session
Session session = cluster.connect();

1.2、创建自己的AbstractRoutingSessionFactory实现类

public class CassandraDataSessionFactory extends AbstractRoutingSessionFactory {

	@Override
	public void afterPropertiesSet() {
	}
	
	@Override
	public Session getSession() {
		return 之前创建的session;
	}

	@Override
	protected Object determineCurrentLookupKey() {
		return null;
	}

}

1.3、去除CassandraDataAutoConfiguration类的自动配置。因为CassandraDataAutoConfiguration会自动连接cassandra，我们手动创建数据源后，SpringBoot内部的Cluster是没有ip、端口的，所以连接的时候会报错。因子要去除它的自动配置

@SpringBootApplication(exclude = {CassandraDataAutoConfiguration.class})

1.4、创建CassandraTemplate，并将自己的AbstractRoutingSessionFactory实现类注册进去。我是将CassandraAutoConfiguration的源码考出来，然后做了一些修改。

@Configuration
public class DasCassandraDataAutoConfiguration {
	
	private final BeanFactory beanFactory;

	private final CassandraProperties properties;

	private final Cluster cluster;

	public DasCassandraDataAutoConfiguration(BeanFactory beanFactory, CassandraProperties properties, Cluster cluster,
			Environment environment) {
		this.beanFactory = beanFactory;
		this.properties = properties;
		this.cluster = cluster;
	}

	@Bean
	@ConditionalOnMissingBean
	public CassandraMappingContext cassandraMapping(CassandraCustomConversions conversions)
			throws ClassNotFoundException {
		CassandraMappingContext context = new CassandraMappingContext();
		List packages = EntityScanPackages.get(this.beanFactory).getPackageNames();
		if (packages.isEmpty() && AutoConfigurationPackages.has(this.beanFactory)) {
			packages = AutoConfigurationPackages.get(this.beanFactory);
		}
		if (!packages.isEmpty()) {
			context.setInitialEntitySet(CassandraEntityClassScanner.scan(packages));
		}
		PropertyMapper.get().from(this.properties::getKeyspaceName).whenHasText().as(this::createSimpleUserTypeResolver)
				.to(context::setUserTypeResolver);
		context.setCustomConversions(conversions);
		return context;
	}

	private SimpleUserTypeResolver createSimpleUserTypeResolver(String keyspaceName) {
		return new SimpleUserTypeResolver(this.cluster, keyspaceName);
	}

	@Bean
	@ConditionalOnMissingBean
	public CassandraConverter cassandraConverter(CassandraMappingContext mapping,
			CassandraCustomConversions conversions) {
		MappingCassandraConverter converter = new MappingCassandraConverter(mapping);
		converter.setCustomConversions(conversions);
		return converter;
	}

	@Bean
	public CassandraTemplate cassandraTemplate(CassandraConverter converter) {
        //将自己创建的AbstractRoutingSessionFactory实现类传进去
		return new CassandraTemplate(new CassandraDataSessionFactory(), converter);
	}

	@Bean
	@ConditionalOnMissingBean
	public CassandraCustomConversions cassandraCustomConversions() {
		return new CassandraCustomConversions(Collections.emptyList());
	}
}

2、方案二，创建ArgumentPreparedStatementBinder子类

方案一理论上可以，但会有个致命的问题。比如有10个站点的数据要insert，每个站点对应一个数据源，它们最终都会调用execute方法，只会把SQL语句和参数传递过去，然后到达AbstractRoutingSessionFactory的getSession方法时，是没有传参过来的，所以很难判断这条SQL应该调用哪个数据源。我唯一想到的方法就是用ThreadLocal，每个进程在获取数据时就得到对应的数据源然后存入ThreadLocal，在getSession方法里再取出。这样就显的很笨重。

所以就研究出了方案二方法，可以根据SQL的参数得到数据源Session。