package com.mogujie.trade.db; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.BeansException; import org.springframework.beans.MutablePropertyValues; import org.springframework.beans.factory.BeanCreationException; import org.springframework.beans.factory.BeanInitializationException; import org.springframework.beans.factory.FactoryBean; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; import org.springframework.beans.factory.config.ConstructorArgumentValues; import org.springframework.beans.factory.support.AutowireCandidateQualifier; import org.springframework.beans.factory.support.BeanDefinitionRegistry; import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor; import org.springframework.beans.factory.support.GenericBeanDefinition; import org.springframework.context.ApplicationContext; import org.springframework.context.ApplicationContextAware; import org.springframework.transaction.PlatformTransactionManager; import org.springframework.transaction.TransactionDefinition; import org.springframework.transaction.TransactionException; import org.springframework.transaction.TransactionStatus; import org.springframework.util.Assert; import javax.sql.DataSource; import java.io.IOException; import java.io.InputStream; import java.sql.SQLException; import java.util.EnumMap; import java.util.HashMap; import java.util.Map; import java.util.Properties; /** * @author by jiuru on 16/7/14. */ public class DataSourceScanner implements BeanDefinitionRegistryPostProcessor, ApplicationContextAware { private static final String PROPERTY_FILE_NAME = "jdbc.properties"; private static final String KEY_SEPARATOR = "."; private final Logger logger = LoggerFactory.getLogger(getClass()); private DataSourceFactory<? extends DataSource> dataSourceFactory; private ApplicationContext applicationContext; @Override public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException { // do nothing } @Override public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) throws BeansException { final Map<String, ReadWriteSplittingDataSource> dataSources = new HashMap<>(); InputStream in = this.getClass().getClassLoader().getResourceAsStream(PROPERTY_FILE_NAME); if (in != null) { Properties properties = new Properties(); try { properties.load(in); } catch (IOException e) { throw new BeanInitializationException("read property file error!", e); } try { Map<String, Map<DataSourceType, DataSource>> dataSourcesMapping = this.getDataSources(properties); this.registerDataSources(registry, dataSourcesMapping); int transcactionManagerCount = 0; String transactionManagerBeanName = null; for (Map.Entry<String, Map<DataSourceType, DataSource>> entry : dataSourcesMapping.entrySet()) { final String name = entry.getKey(); final DataSource masterDataSource = entry.getValue().get(DataSourceType.master); final ReadWriteSplittingDataSource readWriteSplittingDataSource = new ReadWriteSplittingDataSource( entry.getKey(), entry.getValue().get(DataSourceType.master), entry.getValue().get( DataSourceType.slave)); logger.info("init dataSource {}", readWriteSplittingDataSource); dataSources.put(name, readWriteSplittingDataSource); // 若无可写的数据源则跳过创建事务管理器 if (masterDataSource == null) { continue; } transactionManagerBeanName = name + "TransactionManager"; GenericBeanDefinition beanDefinition = new GenericBeanDefinition(); beanDefinition.setBeanClass(RoutingDataSourceTransactionManager.class); MutablePropertyValues propertyValues = new MutablePropertyValues(); propertyValues.add("dataSource", readWriteSplittingDataSource); propertyValues.add("name", name); beanDefinition.setPropertyValues(propertyValues); AutowireCandidateQualifier qualifier = new AutowireCandidateQualifier(Qualifier.class); qualifier.setAttribute(AutowireCandidateQualifier.VALUE_KEY, name); beanDefinition.addQualifier(qualifier); registry.registerBeanDefinition(transactionManagerBeanName, beanDefinition); PlatformTransactionManager transactionManager = this.applicationContext.getBean( transactionManagerBeanName, PlatformTransactionManager.class); Assert.notNull(transactionManager, "register BeanDefinition of " + transactionManagerBeanName + " error!"); transcactionManagerCount++; } // 兼容只有一个或无TransactionManager的情况 if (transcactionManagerCount == 1) {// 若只有一个则添加别名,兼容默认情况 registry.registerAlias(transactionManagerBeanName, "transcationManager"); } else if (transcactionManagerCount == 0) { // add an empty transcationManager GenericBeanDefinition beanDefinition = new GenericBeanDefinition(); beanDefinition.setBeanClass(EmptyTransactionManager.class); registry.registerBeanDefinition("transactionManager", beanDefinition); } } catch (SQLException e) { throw new BeanCreationException("initial dataSources error!", e); } } // register dataSourceLookup GenericBeanDefinition dataSourceLookupBeanDefinition = new GenericBeanDefinition(); dataSourceLookupBeanDefinition.setBeanClass(DataSourceLookup.class); ConstructorArgumentValues constructorArgumentValues = new ConstructorArgumentValues(); constructorArgumentValues.addIndexedArgumentValue(0, dataSources); dataSourceLookupBeanDefinition.setConstructorArgumentValues(constructorArgumentValues); registry.registerBeanDefinition("dataSourceLookup", dataSourceLookupBeanDefinition); } /** * 根据Properties配置解析得到数据源 * * @param properties * @return * @throws SQLException */ private Map<String, Map<DataSourceType, DataSource>> getDataSources(Properties properties) throws SQLException { Map<String, Map<DataSourceType, DataSource>> dataSourcesMapping = new HashMap<>(2); for (Map.Entry<Object, Object> entry : properties.entrySet()) { String[] parts = entry.getKey().toString().trim().split("\\" + KEY_SEPARATOR); if (parts.length == 3) { String name = parts[0]; if (!dataSourcesMapping.containsKey(name)) { for (DataSourceType dataSourceType : DataSourceType.values()) { DataSource ds = this.dataSourceFactory.getDataSource(this.parseDataSourceConfig(name, dataSourceType, properties)); Map<DataSourceType, DataSource> map = dataSourcesMapping.get(name); if (map == null) { map = new EnumMap<DataSourceType, DataSource>(DataSourceType.class); dataSourcesMapping.put(name, map); } DataSource preValue = map.put(dataSourceType, ds); if (preValue != null) { throw new IllegalArgumentException("dupilicated DataSource of" + name + " " + dataSourceType); } } } } else { // It's illegal, ignore. } } return dataSourcesMapping; } private DataSourceConfig parseDataSourceConfig(String name, DataSourceType dataSourceType, Properties properties) { String keyPrefix = name + KEY_SEPARATOR + dataSourceType + KEY_SEPARATOR; DataSourceConfig dataSourceConfig = new DataSourceConfig(); String url = properties.getProperty(keyPrefix + "url"); Assert.hasText(url, keyPrefix + "url is empty!"); dataSourceConfig.setUrl(url); String username = properties.getProperty(keyPrefix + "username"); Assert.hasText(username, keyPrefix + "username is empty!"); dataSourceConfig.setUsername(username); String password = properties.getProperty(keyPrefix + "password"); Assert.hasText(password, keyPrefix + "password is empty!"); dataSourceConfig.setPassword(password); String initialPoolSizeStr = properties.getProperty(keyPrefix + "initialPoolSize"); int initialPoolSize = initialPoolSizeStr == null ? DataSourceConfig.DEFAULT_INI_POOL_SIZE : Integer .parseInt(initialPoolSizeStr); dataSourceConfig.setInitialPoolSize(initialPoolSize); String minPoolSizeStr = properties.getProperty(keyPrefix + "minPoolSize"); int minPoolSize = minPoolSizeStr == null ? DataSourceConfig.DEFAULT_MIN_POOL_SIZE : Integer .parseInt(minPoolSizeStr); dataSourceConfig.setMinPoolSize(minPoolSize); String maxPoolSizeStr = properties.getProperty(keyPrefix + "maxPoolSize"); int maxPoolSize = maxPoolSizeStr == null ? DataSourceConfig.DEFAULT_MAX_POOL_SIZE : Integer .parseInt(maxPoolSizeStr); dataSourceConfig.setMaxPoolSize(maxPoolSize); return dataSourceConfig; } /** * 将数据源注入到Spring中 * * @param registry * @param dataSourcesMapping */ private void registerDataSources(BeanDefinitionRegistry registry, Map<String, Map<DataSourceType, DataSource>> dataSourcesMapping) { for (Map.Entry<String, Map<DataSourceType, DataSource>> entry : dataSourcesMapping.entrySet()) { final String name = entry.getKey(); for (Map.Entry<DataSourceType, DataSource> subEntry : entry.getValue().entrySet()) { GenericBeanDefinition dataSourceBeanDefinition = new GenericBeanDefinition(); dataSourceBeanDefinition.setBeanClass(DataSourceFactoryBean.class); ConstructorArgumentValues constructorArgumentValues = new ConstructorArgumentValues(); constructorArgumentValues.addIndexedArgumentValue(0, subEntry.getValue()); dataSourceBeanDefinition.setConstructorArgumentValues(constructorArgumentValues); String beanName = name + Character.toUpperCase(subEntry.getKey().name().charAt(0)) + subEntry.getKey().name().substring(1) + "DataSource"; registry.registerBeanDefinition(beanName, dataSourceBeanDefinition); } } } // --------------------Setters--------------- public void setDataSourceFactory(DataSourceFactory<? extends DataSource> dataSourceFactory) { this.dataSourceFactory = dataSourceFactory; } @Override public void setApplicationContext(ApplicationContext applicationContext) throws BeansException { this.applicationContext = applicationContext; } public static class EmptyTransactionManager implements PlatformTransactionManager { @Override public TransactionStatus getTransaction(TransactionDefinition definition) throws TransactionException { throw new UnsupportedOperationException(); } @Override public void commit(TransactionStatus status) throws TransactionException { throw new UnsupportedOperationException(); } @Override public void rollback(TransactionStatus status) throws TransactionException { throw new UnsupportedOperationException(); } } public static class DataSourceFactoryBean implements FactoryBean<DataSource> { private final DataSource dataSource; public DataSourceFactoryBean(DataSource dataSource) { Assert.notNull(dataSource); this.dataSource = dataSource; } @Override public DataSource getObject() throws Exception { return this.dataSource; } @Override public Class<?> getObjectType() { return this.dataSource.getClass(); } @Override public boolean isSingleton() { return true; } } }