JPA/Hibernate 多租户架构：使用同一持久化单元动态访问不同数据源

import org.hibernate.context.spi.CurrentTenantIdentifierResolver;

public class TenantIdentifierResolver implements CurrentTenantIdentifierResolver {

    // 存储当前租户ID的ThreadLocal
    private static final ThreadLocal<String> currentTenant = new ThreadLocal<>();

    /**
     * 设置当前线程的租户ID。
     * 通常在请求进入时（例如通过Servlet Filter或JAX-RS Interceptor）设置。
     * @param tenantId 租户ID
     */
    public static void setTenantId(String tenantId) {
        currentTenant.set(tenantId);
    }

    /**
     * 清除当前线程的租户ID。
     * 通常在请求结束时调用，避免内存泄漏或租户ID混淆。
     */
    public static void clearTenantId() {
        currentTenant.remove();
    }

    @Override
    public String resolveCurrentTenantIdentifier() {
        String tenantId = currentTenant.get();
        if (tenantId != null) {
            return tenantId;
        }
        // 如果没有设置租户ID，可以返回一个默认值或抛出异常
        // 确保在生产环境中，租户ID始终被正确设置
        return "default_tenant"; // 例如，一个默认租户
    }

    @Override
    public boolean validateExistingCurrentSessions() {
        // 如果返回true，Hibernate会检查现有会话的租户ID是否与当前解析的租户ID匹配。
        // 如果不匹配，会抛出异常。这有助于防止租户ID混淆。
        return true;
    }
}

// 示例：在一个Servlet Filter中设置租户ID
public class TenantFilter implements Filter {

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        // 假设租户ID通过HTTP Header "X-Tenant-ID" 传递
        String tenantId = httpRequest.getHeader("X-Tenant-ID");
        if (tenantId == null || tenantId.isEmpty()) {
            // 或者从其他地方获取，例如子域名、URL路径等
            // tenantId = extractTenantIdFromUrl(httpRequest.getRequestURI());
            // 如果无法确定租户，可以返回错误或使用默认租户
            ((HttpServletResponse) response).sendError(HttpServletResponse.SC_BAD_REQUEST, "Tenant ID missing");
            return;
        }

        try {
            TenantIdentifierResolver.setTenantId(tenantId);
            chain.doFilter(request, response);
        } finally {
            TenantIdentifierResolver.clearTenantId(); // 确保清理，防止ThreadLocal内存泄漏或租户ID污染
        }
    }
    // ... 其他Filter方法
}

import org.hibernate.engine.jdbc.connections.spi.MultiTenantConnectionProvider;
import org.hibernate.service.spi.ServiceRegistryAwareService;
import org.hibernate.service.spi.ServiceRegistryImplementor;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;

public class CustomMultiTenantConnectionProvider implements MultiTenantConnectionProvider, ServiceRegistryAwareService {

    // 存储租户ID到DataSource的映射
    private final Map<String, DataSource> dataSources = new HashMap<>();

    // 默认数据源，当无法确定租户ID时使用
    private DataSource defaultDataSource;

    public CustomMultiTenantConnectionProvider() {
        // 实际应用中，这里应该从配置中加载数据源信息
        // 示例：初始化几个数据源
        // dataSources.put("tenant_a", createDataSource("jdbc:mysql://localhost:3306/db_tenant_a", "user_a", "pass_a"));
        // dataSources.put("tenant_b", createDataSource("jdbc:mysql://localhost:3306/db_tenant_b", "user_b", "pass_b"));
        // defaultDataSource = createDataSource("jdbc:mysql://localhost:3306/db_default", "user_default", "pass_default");
    }

    /**
     * 实际应用中，这里会通过JNDI查找或Spring配置注入数据源。
     * 假设我们通过JNDI查找Wildfly中的数据源。
     */
    @Override
    public void injectServices(ServiceRegistryImplementor serviceRegistry) {
        // 在Wildfly中，你可能需要通过JNDI查找数据源
        // 例如：
        // try {
        //     InitialContext context = new InitialContext();
        //     dataSources.put("tenant_a", (DataSource) context.lookup("java:/jdbc/TenantA_DS"));
        //     dataSources.put("tenant_b", (DataSource) context.lookup("java:/jdbc/TenantB_DS"));
        //     defaultDataSource = (DataSource) context.lookup("java:/jdbc/Default_DS");
        // } catch (NamingException e) {
        //     throw new RuntimeException("Error looking up datasources", e);
        // }
        // For demonstration, let's just create dummy datasources
        dataSources.put("tenant_a", new DummyDataSource("jdbc:h2:mem:tenant_a"));
        dataSources.put("tenant_b", new DummyDataSource("jdbc:h2:mem:tenant_b"));
        defaultDataSource = new DummyDataSource("jdbc:h2:mem:default");
    }

    @Override
    public Connection getConnection(String tenantIdentifier) throws SQLException {
        DataSource dataSource = dataSources.get(tenantIdentifier);
        if (dataSource == null) {
            // 如果找不到特定租户的数据源，可以使用默认数据源或抛出异常
            System.err.println("No specific datasource found for tenant: " + tenantIdentifier + ". Using default.");
            dataSource = defaultDataSource;
            if (dataSource == null) {
                throw new SQLException("No datasource found for tenant: " + tenantIdentifier + " and no default datasource configured.");
            }
        }
        return dataSource.getConnection();
    }

    @Override
    public Connection getAnyConnection() throws SQLException {
        // 获取任意一个连接，通常用于Schema管理（如Hibernate工具自动创建表）
        if (defaultDataSource != null) {
            return defaultDataSource.getConnection();
        }
        // 或者返回第一个可用的租户数据源连接
        if (!dataSources.isEmpty()) {
            return dataSources.values().iterator().next().getConnection();
        }
        throw new SQLException("No default or tenant datasources configured.");
    }

    @Override
    public void releaseConnection(String tenantIdentifier, Connection connection) throws SQLException {
        connection.close(); // 归还连接到连接池
    }

    @Override
    public void releaseAnyConnection(Connection connection) throws SQLException {
        connection.close(); // 归还连接到连接池
    }

    @Override
    public boolean supportsAggressiveRelease() {
        return false; // 通常为false，让连接池管理连接生命周期
    }

    @Override
    public boolean is  injected() {
        return true; // 表示服务已注入
    }

    // 内部类用于演示，实际使用JDBC DataSource
    private static class DummyDataSource implements DataSource {
        private String url;
        public DummyDataSource(String url) { this.url = url; }
        @Override public Connection getConnection() throws SQLException {
            System.out.println("Getting connection for URL: " + url);
            // 实际这里会返回一个真实的数据库连接
            return java.sql.DriverManager.getConnection(url);
        }
        @Override public Connection getConnection(String username, String password) throws SQLException { return getConnection(); }
        @Override public <T> T unwrap(Class<T> iface) throws SQLException { return null; }
        @Override public boolean isWrapperFor(Class<?> iface) throws SQLException { return false; }
        @Override public java.io.PrintWriter getLogWriter() throws SQLException { return null; }
        @Override public void setLogWriter(java.io.PrintWriter out) throws SQLException {}
        @Override public int getLoginTimeout() throws SQLException { return 0; }
        @Override public void setLoginTimeout(int seconds) throws SQLException {}
        @Override public java.util.logging.Logger getParentLogger() throws java.sql.SQLFeatureNotSupportedException { return null; }
    }
}

<persistence xmlns="http://xmlns.jcp.org/xml/ns/persistence"
             xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
             xsi:schemaLocation="http://xmlns.jcp.org/xml/ns/persistence http://xmlns.jcp.org/xml/ns/persistence/persistence_2_2.xsd"
             version="2.2">
    <persistence-unit name="YourPersistenceUnitName" transaction-type="JTA">
        <provider>org.hibernate.jpa.HibernatePersistenceProvider</provider>
        <!-- 你的实体类列表 -->
        <class>com.example.entity.YourEntity</class>
        <!-- ... 其他实体类 ... -->

        <properties>
            <!-- 启用Hibernate多租户模式，指定为DATABASE策略 -->
            <property name="hibernate.multiTenancy" value="DATABASE"/>

            <!-- 指定租户标识符解析器 -->
            <property name="hibernate.tenant_identifier_resolver"
                      value="com.example.multitenancy.TenantIdentifierResolver"/>

            <!-- 指定多租户连接提供者 -->
            <property name="hibernate.multi_tenant_connection_provider"
                      value="com.example.multitenancy.CustomMultiTenantConnectionProvider"/>

            <!-- 其他Hibernate和JPA属性 -->
            <property name="hibernate.dialect" value="org.hibernate.dialect.MySQL8Dialect"/>
            <property name="hibernate.hbm2ddl.auto" value="update"/>
            <property name="hibernate.show_sql" value="true"/>
            <property name="hibernate.format_sql" value="true"/>

            <!-- 注意：在这种模式下，不再直接配置JTA数据源，而是由CustomMultiTenantConnectionProvider管理 -->
            <!-- <property name="javax.persistence.jtaDataSource" value="java:/jdbc/YourDefaultDS"/> -->
        </properties>
    </persistence-unit>
</persistence>