基于请求头在Hibernate会话工厂运行时切换数据源

嘿，我来帮你搞定基于请求头动态切换Hibernate数据源的需求！咱们一步步来实现这个功能：

1. 先实现动态路由数据源 #

首先需要一个能根据上下文切换数据源的类，继承Spring的AbstractRoutingDataSource，它会帮我们自动路由到对应的数据源：

public class DynamicDataSource extends AbstractRoutingDataSource {
    // 用ThreadLocal存储当前线程的数据源标识，保证请求线程隔离
    private static final ThreadLocal<String> CURRENT_DATA_SOURCE = new ThreadLocal<>();

    // 设置当前线程对应的数据源
    public static void setCurrentDataSource(String dataSourceKey) {
        CURRENT_DATA_SOURCE.set(dataSourceKey);
    }

    // 清理当前线程的数据源标识，避免线程复用导致的问题
    public static void clearCurrentDataSource() {
        CURRENT_DATA_SOURCE.remove();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        // 返回当前线程绑定的数据源标识，Spring会根据这个key路由到对应数据源
        return CURRENT_DATA_SOURCE.get();
    }
}

2. 配置两个实际的数据源实例 #

在你的HibernateConfiguration里，添加两个分别对应不同客户端的数据源配置：

@Autowired
private Environment environment;

// 客户端1对应的数据源
@Bean
public DataSource dataSourceClient1() {
    DriverManagerDataSource dataSource = new DriverManagerDataSource();
    dataSource.setDriverClassName(environment.getProperty("spring.datasource.client1.driver-class-name"));
    dataSource.setUrl(environment.getProperty("spring.datasource.client1.url"));
    dataSource.setUsername(environment.getProperty("spring.datasource.client1.username"));
    dataSource.setPassword(environment.getProperty("spring.datasource.client1.password"));
    return dataSource;
}

// 客户端2对应的数据源
@Bean
public DataSource dataSourceClient2() {
    DriverManagerDataSource dataSource = new DriverManagerDataSource();
    dataSource.setDriverClassName(environment.getProperty("spring.datasource.client2.driver-class-name"));
    dataSource.setUrl(environment.getProperty("spring.datasource.client2.url"));
    dataSource.setUsername(environment.getProperty("spring.datasource.client2.username"));
    dataSource.setPassword(environment.getProperty("spring.datasource.client2.password"));
    return dataSource;
}

3. 配置动态数据源Bean #

把上面两个数据源注入到动态数据源中，同时设置默认的 fallback 数据源（比如默认用客户端1的数据源）：

@Bean
public DataSource dynamicDataSource() {
    DynamicDataSource dynamicDataSource = new DynamicDataSource();
    // 配置目标数据源映射，key要和请求头里的客户端标识对应
    Map<Object, Object> targetDataSources = new HashMap<>();
    targetDataSources.put("client1", dataSourceClient1());
    targetDataSources.put("client2", dataSourceClient2());
    dynamicDataSource.setTargetDataSources(targetDataSources);
    // 设置默认数据源，当请求头没有标识时使用
    dynamicDataSource.setDefaultTargetDataSource(dataSourceClient1());
    return dynamicDataSource;
}

4. 添加请求拦截器处理请求头 #

创建一个拦截器，在请求进入时读取请求头里的客户端标识，绑定到当前线程；请求结束后清理绑定，避免内存泄漏：

@Component
public class DataSourceSwitchInterceptor implements HandlerInterceptor {
    // 假设用来区分客户端的请求头key是"X-Client-Id"
    private static final String CLIENT_HEADER_KEY = "X-Client-Id";

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String clientId = request.getHeader(CLIENT_HEADER_KEY);
        if (clientId != null) {
            // 根据请求头设置当前线程对应的数据源
            DynamicDataSource.setCurrentDataSource(clientId);
        }
        // 如果没有请求头，就使用默认数据源
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        // 必须清理ThreadLocal，否则线程池复用线程时会携带旧的数据源标识
        DynamicDataSource.clearCurrentDataSource();
    }
}

别忘了把这个拦截器注册到Spring MVC的配置里，让它生效：

@Configuration
public class WebConfig implements WebMvcConfigurer {
    @Autowired
    private DataSourceSwitchInterceptor dataSourceSwitchInterceptor;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(dataSourceSwitchInterceptor)
                // 拦截需要切换数据源的请求路径，这里拦截所有请求
                .addPathPatterns("/**");
    }
}

5. 修改Hibernate会话工厂配置 #

最后把原来sessionFactory()里的数据源替换成动态数据源：

@Bean
public LocalSessionFactoryBean sessionFactory() {
    LocalSessionFactoryBean sessionFactory = new LocalSessionFactoryBean();
    // 这里替换为我们的动态数据源
    sessionFactory.setDataSource(dynamicDataSource());
    sessionFactory.setPackagesToScan(new String[]{"com.mhycomoant.test.model"});
    sessionFactory.setHibernateProperties(hibernateProperties());
    return sessionFactory;
}

// 原有Hibernate属性配置保持不变
private Properties hibernateProperties() {
    Properties properties = new Properties();
    properties.put("hibernate.dialect", environment.getProperty("hibernate.dialect"));
    properties.put("hibernate.show_sql", environment.getProperty("hibernate.show_sql"));
    // 其他Hibernate配置项...
    return properties;
}

一些关键注意事项 #

• 配置文件里要补充两个客户端的数据源信息，比如application.properties：

  spring.datasource.client1.url=jdbc:mysql://localhost:3306/db_client1
  spring.datasource.client1.driver-class-name=com.mysql.cj.jdbc.Driver
  spring.datasource.client1.username=root
  spring.datasource.client1.password=123456
  
  spring.datasource.client2.url=jdbc:mysql://localhost:3306/db_client2
  spring.datasource.client2.driver-class-name=com.mysql.cj.jdbc.Driver
  spring.datasource.client2.username=root
  spring.datasource.client2.password=123456
  

• 确保请求头X-Client-Id的值是client1或client2，和动态数据源里的key完全对应
• 一定要在afterCompletion方法中清理ThreadLocal，否则线程池复用线程会导致数据源混乱

内容的提问来源于stack exchange，提问作者William Buttlicker