How to implement dynamic switching of data sources in JPA?

One way to achieve dynamic switching of data sources in JPA is by using multiple EntityManagerFactory instances.

1. Configure connection information for multiple data sources: Define connection information for multiple data sources in the configuration file, such as database URL, username, password, etc. You can use either application.properties or application.yml file for configuration.
2. Create multiple EntityManagerFactories: Multiple EntityManagerFactories can be created based on the data source connection information defined in the configuration file. EntityManagerFactories can be created as beans using the @Configuration annotation, specifying the corresponding data source connection information.

@Configuration
public class DataSourceConfig {
    
    @Bean
    public DataSource dataSource1() {
        // 数据源1的连接信息
        // ...
    }

    @Bean
    public DataSource dataSource2() {
        // 数据源2的连接信息
        // ...
    }

    @Bean
    @Primary
    public EntityManagerFactory entityManagerFactory1() {
        LocalContainerEntityManagerFactoryBean factory = new LocalContainerEntityManagerFactoryBean();
        factory.setDataSource(dataSource1());
        // 设置其他配置信息
        // ...
        factory.afterPropertiesSet();
        return factory.getObject();
    }

    @Bean
    public EntityManagerFactory entityManagerFactory2() {
        LocalContainerEntityManagerFactoryBean factory = new LocalContainerEntityManagerFactoryBean();
        factory.setDataSource(dataSource2());
        // 设置其他配置信息
        // ...
        factory.afterPropertiesSet();
        return factory.getObject();
    }
}

1. A DataSource that dynamically routes connections to multiple databases based on a certain criteria.
2. AbstractRoutingDataSource is a class in Java that allows for dynamic switching of data sources at runtime.
3. Find the current lookup key.

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

public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        // 根据自定义的规则来返回当前要使用的数据源的名称
        // ...
    }
}

1. Configure dynamic data source: Configure dynamic data source in the configuration file, and associate multiple EntityManagerFactory with the dynamic data source.

@Configuration
public class DataSourceConfig {
    
    // ...

    @Bean
    public DataSource dynamicDataSource() {
        DynamicDataSource dataSource = new DynamicDataSource();
        dataSource.setDefaultTargetDataSource(entityManagerFactory1().createEntityManager().getEntityManagerFactory().getDataSource());
        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put("dataSource1", entityManagerFactory1().createEntityManager().getEntityManagerFactory().getDataSource());
        targetDataSources.put("dataSource2", entityManagerFactory2().createEntityManager().getEntityManagerFactory().getDataSource());
        dataSource.setTargetDataSources(targetDataSources);
        return dataSource;
    }

    @Bean
    public JpaTransactionManager transactionManager() {
        JpaTransactionManager transactionManager = new JpaTransactionManager();
        transactionManager.setEntityManagerFactory(entityManagerFactory1());
        return transactionManager;
    }
}

1. a special variable that is local to each thread
2. Find the current lookup key.

public class DataSourceContextHolder {
    private static final ThreadLocal<String> contextHolder = new ThreadLocal<>();

    public static void setDataSource(String dataSourceName) {
        contextHolder.set(dataSourceName);
    }

    public static String getDataSource() {
        return contextHolder.get();
    }

    public static void clearDataSource() {
        contextHolder.remove();
    }
}

// 使用动态数据源
@Service
public class UserServiceImpl implements UserService {

    @Autowired
    private UserRepository userRepository;

    @Override
    public User findById(Long id) {
        String dataSourceName = determineDataSourceName(id);
        DataSourceContextHolder.setDataSource(dataSourceName);
        User user = userRepository.findById(id);
        DataSourceContextHolder.clearDataSource();
        return user;
    }

    private String determineDataSourceName(Long id) {
        // 根据id判断要使用的数据源的名称
        // ...
    }
}

By following the above steps, you can dynamically switch data sources in JPA. It is important to make sure to promptly clear the data source name in ThreadLocal after each operation to avoid affecting subsequent operations.