基于AbstractRoutingDataSource的多库切换TODO实现方案咨询

Hey, I’ve dealt with this exact problem when building dynamic data source routing for large-scale databases—your issue boils down to unclosed thread-bound connections sticking around, which prevents the routing data source from picking up the new key you set. Let’s walk through how to resolve that TODO and get your cross-database queries working.

Why the Switch Isn’t Working #

When you use AbstractRoutingDataSource, Spring (and most ORM frameworks like MyBatis) binds a database connection to the current thread for efficiency. If you switch the data source key in your ThreadLocalStorage but don’t release the old connection, the framework will keep reusing that existing connection instead of fetching a new one from the target data source. That’s why your cross-database queries are failing to switch.

How to Implement the Connection Cleanup #

The key is to release the current thread’s bound connection before setting the new data source key. We’ll use Spring’s DataSourceUtils class—it’s designed to manage thread-bound connections correctly, ensuring they’re returned to the pool instead of being destroyed outright.

Step 1: Add Connection Cleanup Logic #

Update your setPrimaryDB and setSecondaryDB methods (or create a helper method) to clear the thread’s connection first:

import org.springframework.jdbc.datasource.DataSourceUtils;
import java.sql.Connection;
import java.sql.SQLException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

// Inside your HistoryAwareRoutingSource or the class managing data source switching
private static final Logger log = LoggerFactory.getLogger(HistoryAwareRoutingSource.class);

public void setPrimaryDB() {
    // Clean up existing connection before switching
    clearThreadBoundConnection();
    ThreadLocalStorage.setDataSourceKey("short_term");
}

public void setSecondaryDB() {
    clearThreadBoundConnection();
    ThreadLocalStorage.setDataSourceKey("long_term");
}

private void clearThreadBoundConnection() {
    try {
        // Get the connection bound to the current thread (if any)
        Connection currentConn = DataSourceUtils.getConnection(this);
        if (currentConn != null) {
            // Release the connection back to the pool (DataSourceUtils handles this safely)
            DataSourceUtils.releaseConnection(currentConn, this);
        }
    } catch (SQLException e) {
        log.error("Failed to release thread-bound database connection", e);
    }
}

Step 2: Handle ORM Framework-Specific Cleanup (If Applicable) #

If you’re using MyBatis or Hibernate, you’ll also need to clean up their thread-bound resources:

• MyBatis: Close the current SqlSession to ensure it doesn’t hold onto the old connection:

  import org.mybatis.spring.SqlSessionUtils;
  import org.apache.ibatis.session.SqlSessionFactory;
  
  // Inject SqlSessionFactory via constructor or @Autowired
  private final SqlSessionFactory sqlSessionFactory;
  
  private void clearMyBatisResources() {
      SqlSessionUtils.closeSqlSession(SqlSessionUtils.getSqlSession(sqlSessionFactory), sqlSessionFactory);
  }
  

• Hibernate: Flush and close the current EntityManager if it’s bound to the thread.

Step 3: Clean Up ThreadLocal After Requests #

To avoid issues with thread pool reuse (since threads in the pool retain their ThreadLocal values), add cleanup logic at the end of each request:

• Use a Spring HandlerInterceptor or Filter to call ThreadLocalStorage.clearDataSourceKey() in the afterCompletion phase.
• If using transactions, ensure the ThreadLocal is cleared after the transaction completes.

Critical Notes #

• Transactions: You cannot switch data sources mid-transaction—transactions are bound to a single connection/data source. If your cross-database query needs to be atomic, you’ll need to use distributed transactions or split the operation into separate non-transactional steps.
• Connection Pool Behavior: Releasing the connection via DataSourceUtils returns it to the pool, which is far more efficient than destroying it. Connection pools are designed to reuse connections, so we want to leverage that.

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