聊聊sharding-jdbc的SingleXADataSource

public final class SingleXADataSource extends AbstractUnsupportedSingleXADataSource {
    
    @Getter
    private final String resourceName;
    
    @Getter
    private final XADataSource xaDataSource;
    
    private final DatabaseType databaseType;
    
    private final DataSource originalDataSource;
    
    private final boolean isOriginalXADataSource;
    
    public SingleXADataSource(final DatabaseType databaseType, final String resourceName, final DataSource dataSource) {
        this.databaseType = databaseType;
        this.resourceName = resourceName;
        originalDataSource = dataSource;
        if (dataSource instanceof XADataSource) {
            xaDataSource = (XADataSource) dataSource;
            isOriginalXADataSource = true;
        } else {
            xaDataSource = XADataSourceFactory.build(databaseType, dataSource);
            isOriginalXADataSource = false;
        }
    }
    
    @Override
    public SingleXAConnection getXAConnection() throws SQLException {
        return isOriginalXADataSource ? getXAConnectionFromXADataSource() : getXAConnectionFromNoneXADataSource();
    }
    
    private SingleXAConnection getXAConnectionFromXADataSource() throws SQLException {
        XAConnection xaConnection = xaDataSource.getXAConnection();
        return new SingleXAConnection(resourceName, xaConnection.getConnection(), xaConnection);
    }
    
    private SingleXAConnection getXAConnectionFromNoneXADataSource() throws SQLException {
        Connection originalConnection = originalDataSource.getConnection();
        XAConnection xaConnection = XAConnectionFactory.createXAConnection(databaseType, xaDataSource, originalConnection);
        return new SingleXAConnection(resourceName, originalConnection, xaConnection);
    }
}

public abstract class AbstractUnsupportedSingleXADataSource implements XADataSource {
    
    @Override
    public final XAConnection getXAConnection(final String user, final String password) throws SQLException {
        throw new SQLFeatureNotSupportedException("getXAConnection by user and password");
    }
    
    @Override
    public final PrintWriter getLogWriter() throws SQLException {
        throw new SQLFeatureNotSupportedException("getLogWriter");
    }
    
    @Override
    public final void setLogWriter(final PrintWriter out) throws SQLException {
        throw new SQLFeatureNotSupportedException("setLogWriter");
    }
    
    @Override
    public final void setLoginTimeout(final int seconds) throws SQLException {
        throw new SQLFeatureNotSupportedException("setLoginTimeout");
    }
    
    @Override
    public final int getLoginTimeout() throws SQLException {
        throw new SQLFeatureNotSupportedException("getLoginTimeout");
    }
    
    @Override
    public final Logger getParentLogger() throws SQLFeatureNotSupportedException {
        throw new SQLFeatureNotSupportedException("getParentLogger");
    }
}

public final class SingleXADataSourceTest {
    
    @Test
    public void assertBuildSingleXADataSourceOfXA() {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.MySQL, "ds1");
        SingleXADataSource actual = new SingleXADataSource(DatabaseType.MySQL, "ds1", dataSource);
        assertThat(actual.getResourceName(), is("ds1"));
        assertThat(actual.getXaDataSource(), is((XADataSource) dataSource));
    }
    
    @Test
    public void assertBuildSingleXADataSourceOfNoneXA() {
        DataSource dataSource = DataSourceUtils.build(HikariDataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource actual = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        assertThat(actual.getResourceName(), is("ds1"));
        assertThat(actual.getXaDataSource(), instanceOf(JdbcDataSource.class));
        JdbcDataSource jdbcDataSource = (JdbcDataSource) actual.getXaDataSource();
        assertThat(jdbcDataSource.getUser(), is("root"));
        assertThat(jdbcDataSource.getPassword(), is("root"));
    }
    
    @Test
    @SneakyThrows
    public void assertGetXAConnectionOfXA() {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        SingleXAConnection actual = shardingXADataSource.getXAConnection();
        assertThat(actual.getConnection(), instanceOf(Connection.class));
    }
    
    @Test
    @SneakyThrows
    public void assertGetXAConnectionOfNoneXA() {
        DataSource dataSource = DataSourceUtils.build(HikariDataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        SingleXAConnection actual = shardingXADataSource.getXAConnection();
        assertThat(actual.getConnection(), instanceOf(Connection.class));
    }
    
    @Test(expected = SQLFeatureNotSupportedException.class)
    public void assertGetLoginTimeout() throws SQLException {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        shardingXADataSource.getLoginTimeout();
    }
    
    @Test(expected = SQLFeatureNotSupportedException.class)
    public void assertSetLogWriter() throws SQLException {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        shardingXADataSource.setLogWriter(null);
    }
    
    @Test(expected = SQLFeatureNotSupportedException.class)
    public void assertSetLoginTimeout() throws SQLException {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        shardingXADataSource.setLoginTimeout(10);
    }
    
    @Test(expected = SQLFeatureNotSupportedException.class)
    public void assertGetParentLogger() throws SQLException {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        shardingXADataSource.getParentLogger();
    }
    
    @Test(expected = SQLFeatureNotSupportedException.class)
    public void assertGetLogWriter() throws SQLException {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        shardingXADataSource.getLogWriter();
    }
    
    @Test(expected = SQLFeatureNotSupportedException.class)
    public void assertGetXAConnectionByUserAndPassword() throws SQLException {
        DataSource dataSource = DataSourceUtils.build(DruidXADataSource.class, DatabaseType.H2, "ds1");
        SingleXADataSource shardingXADataSource = new SingleXADataSource(DatabaseType.H2, "ds1", dataSource);
        shardingXADataSource.getXAConnection("root", "root");
    }
}