Sharing a transaction while using AbstractRoutingDataSource to switch datasources

Question

How to share a transaction between datasources when using AbstractRoutingDataSource to switch the active data source?

So far, without the transaction the queries get executed on both databases correctly, but when I start a transaction, everything executes on the same database (i.e. I cannot switch to the second database anymore).

Any ideas?

@Transactional
public void crossDbTransactionTest() {
    // Selects a datasource from my pool of AbstractRoutingDataSources
    DbConnectionContextHolder.setDbConnectionByYear(2012);

    // execute something in the first database
    this.executeSomeJpaQuery("xyz"); 

    // switch to the second database
    DbConnectionContextHolder.setDbConnectionByYear(2011);

    // execute something in the second database
    this.executeSomeJpaQuery("xyz"); // on any errors rollback changes in both databases
}

EDIT1 (added configuration files):

persistence.xml:

<persistence-unit name="primarnaKonekcija" transaction-type="RESOURCE_LOCAL">
    <provider>org.hibernate.ejb.HibernatePersistence</provider>
    <properties>
        <property name="hibernate.dialect" value="org.hibernate.dialect.SQLServerDialect" />
        <property name="hibernate.max_fetch_depth" value="1" />
        <property name="hibernate.transaction.manager_lookup_class"
                              value="org.hibernate.transaction.JBossTransactionManagerLookup" />
    </properties>
</persistence-unit>

spring-jpa.xml:

<!-- Shared DB credentials -->
<context:property-placeholder location="classpath:config.properties" />

<!-- DB connections by year -->
<bean id="parentDataSource" class="org.springframework.jdbc.datasource.DriverManagerDataSource" abstract="true">
    <property name="driverClassName" value="${db.driver}" />
    <property name="username" value="${db.user}" />
    <property name="password" value="${db.password}" />
</bean>
<bean id="dataSource" class="myPackage.DbConnectionRoutingDataSource">
    <!-- Placeholder that is replaced in BeanFactoryPostProcessor -->
    <property name="targetDataSources">
        <map key-type="int">
            <entry key="0" value-ref="placeholderDs" />
        </map>
    </property>
    <property name="defaultTargetDataSource" ref="placeholderDs" />
</bean>

<!-- EntityManager configuration -->
<bean id="entityManagerFactory" class="org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean">
    <property name="persistenceUnitName" value="primarnaKonekcija" />
    <property name="dataSource" ref="dataSource" />
    <property name="jpaVendorAdapter">
        <bean class="org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter">
            <property name="databasePlatform" value="org.hibernate.dialect.SQLServerDialect" />
            <property name="showSql" value="true" />
        </bean>
    </property>
</bean>

<tx:annotation-driven />
<tx:jta-transaction-manager />

EDIT 2:

Tried switching everything to JTA and JNDI provided datasources.

Changing transaction-type="RESOURCE_LOCAL" to transaction-type="JTA" didn't work either - JtaStatusHelper throws an NullPointerException, saying that the transactionManager is null.

EDIT 3:

Added JBossTransactionManagerLookup to persistence.xml, now I get the "Adding multiple last resources is disallowed" when switching to second datasource within the transaction.

EDIT 4:

Tried setting JBOSS so I get past that error - database switching works now with the expected warning: "Multiple last resources have been added to the current transaction. This is transactionally unsafe and should not be relied upon.". Gonna try to configure MSSQL XA driver in JBOSS next.

EDIT 5:

After configuring MSSQL XA, everything works as intended, will post an answer with steps needed to set this up.

Answer 1

This is a solution that I wouldn't recommend unless you're having no other choice. Level 2 cache just can't possibly work with a solution such as this, but it's a (stable) solution that I was forced to use to buy time until the underlying legacy databases are merged into one.

First in JBoss standalone.xml set up your database connections as XA datasources. If using MS SQL server, follow the instructions how to properly set up XA at http://msdn.microsoft.com/en-us/library/aa342335.aspx

standalone.xml

<datasources>
    <datasource jndi-name="java:jboss/datasources/ExampleDS" pool-name="ExampleDS" enabled="true" use-java-context="true">
        <connection-url>jdbc:h2:mem:test;DB_CLOSE_DELAY=-1</connection-url>
        <driver>h2</driver>
        <security>
            <user-name>sa</user-name>
            <password>sa</password>
        </security>
    </datasource>
    <xa-datasource jta="true" jndi-name="java:jboss/datasources/MYDB_ONE" pool-name="MYDB_ONE" enabled="true" use-java-context="true" use-ccm="true">
        <xa-datasource-property name="ServerName">
            localhost
        </xa-datasource-property>
        <xa-datasource-property name="DatabaseName">
            MYDB_ONE
        </xa-datasource-property>
        <xa-datasource-property name="SelectMethod">
            cursor
        </xa-datasource-property>
        <xa-datasource-class>com.microsoft.sqlserver.jdbc.SQLServerXADataSource</xa-datasource-class>
        <driver>sqljdbc</driver>
        <xa-pool>
            <is-same-rm-override>false</is-same-rm-override>
        </xa-pool>
        <security>
            <user-name>some_user</user-name>
            <password>some_password</password>
        </security>
        <validation>
            <valid-connection-checker class-name="org.jboss.jca.adapters.jdbc.extensions.mssql.MSSQLValidConnectionChecker"/>
        </validation>
    </xa-datasource>
    <xa-datasource jta="true" jndi-name="java:jboss/datasources/MYDB_TWO" pool-name="MYDB_TWO" enabled="true" use-java-context="true" use-ccm="true">
        <xa-datasource-property name="ServerName">
            localhost
        </xa-datasource-property>
        <xa-datasource-property name="DatabaseName">
            MYDB_TWO
        </xa-datasource-property>
        <xa-datasource-property name="SelectMethod">
            cursor
        </xa-datasource-property>
        <xa-datasource-class>com.microsoft.sqlserver.jdbc.SQLServerXADataSource</xa-datasource-class>
        <driver>sqljdbc</driver>
        <xa-pool>
            <is-same-rm-override>false</is-same-rm-override>
        </xa-pool>
        <security>
            <user-name>some_user</user-name>
            <password>some_password</password>
        </security>
        <validation>
            <valid-connection-checker class-name="org.jboss.jca.adapters.jdbc.extensions.mssql.MSSQLValidConnectionChecker"/>
        </validation>
    </xa-datasource>
    <drivers>
        <driver name="h2" module="com.h2database.h2">
            <xa-datasource-class>org.h2.jdbcx.JdbcDataSource</xa-datasource-class>
        </driver>
        <driver name="sqljdbc" module="com.microsoft.sqlserver.jdbc">
            <driver-class>com.microsoft.sqlserver.jdbc.SQLServerDriver</driver-class>
        </driver>
        <driver name="postgresql" module="org.postgresql">
            <xa-datasource-class>org.postgresql.xa.PGXADataSource</xa-datasource-class>
        </driver>
    </drivers>
</datasources>

Then I set up my entityManager bean that uses my implementation of the AbstractRoutingDataSource as its dataSource. This is the Spring powered JPA setup without using the persistence.xml file; as far as I know, this is the only the way to get automatic package scanning of entities when using JBoss 7.

springJpaConfig.xml

<!-- Use @PersistenceContext annotations for injecting entity managers -->
<bean class="org.springframework.orm.jpa.support.PersistenceAnnotationBeanPostProcessor" />

<!-- Set up JTA transaction manager -->
<tx:jta-transaction-manager />

<bean id="entityManagerFactoryMyDB" class="org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean">
    <property name="persistenceUnitName" value="MyDB" />
    <property name="dataSource" ref="dataSourceMyDB" />
    <property name="packagesToScan" value="my.package.with.jpa.entities" />
    <property name="jpaVendorAdapter">
        <bean class="org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter">
            <property name="showSql" value="true" />
        </bean>
    </property>
    <property name="jpaPropertyMap">
        <map>
            <entry key="javax.persistence.transactionType" value="jta" />

            <entry key="hibernate.transaction.jta.platform" value="org.hibernate.service.jta.platform.internal.JBossAppServerJtaPlatform" />
            <entry key="jboss.entity.manager.factory.jndi.name" value="java:app/MyDBEntityManagerFactory" />

            <entry key="hibernate.dialect" value="org.hibernate.dialect.SQLServer2008Dialect" />
        </map>
    </property>
</bean>

<bean id="dataSourceMyDB" class="some.package.AbstractRoutingDataSourceMyDB">
    <property name="lenientFallback" value="false" />
    <property name="defaultTargetDataSource" value="java:jboss/datasources/ExampleDS" />
    <property name="targetDataSources">
        <map key-type="String">
            <!-- This is a placeholder that will be filled in by BeanFactoryPostProcessor -->
        </map>
    </property>
</bean>

<!-- This allows us to modify Spring configuration load the list of datasources -->
<bean class="some.package.DatasourceRegisteringBeanFactoryPostProcessor" />

I use ExampleDS as the default in AbstractRoutingDataSourceMyDB because you have to provide a defaultTargetDataSource, but I always want to select a valid DB manually, hence if anyone tries to access the DB without first manually selecting the connection they will try to execute their query on the non existant ExampleDS database which will throw an exception (very hacky, but it gets the job done).

In the BeanFactoryPostProcessor I now need to fill in the list of my datasources:

DatasourceRegisteringBeanFactoryPostProcessor.java

package some.package
class DatasourceRegisteringBeanFactoryPostProcessor implements BeanFactoryPostProcessor {
    public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) {
        HashMap<String, String> connectionsListMyDB = new HashMap<>();

        // Load your connection list from wherever you need to, you can
        // enumerate them directly from JNDI or some configuration location
        connectionsListMyDB.put("db1", "java:jboss/datasources/MYDB_ONE");
        connectionsListMyDB.put("db2", "java:jboss/datasources/MYDB_TWO");

        if (connectionsList.isEmpty())
            throw new RuntimeException("No JPA connections defined");

        // Configure the dataSource bean properties
        BeanDefinitionRegistry factory = (BeanDefinitionRegistry) beanFactory;
        MutablePropertyValues mpv = factory.getBeanDefinition("dataSourceMyDB").getPropertyValues();

        ManagedMap<String, String> mm = (ManagedMap<String, String>) mpv.getPropertyValue(
                "targetDataSources").getValue();
        mm.clear();
        for (Entry<String, String> e : connectionsListMyDB.entrySet()) {
            mm.put(e.getKey(), e.getValue());
        }
    }
}

This is my implementation of AbstractRoutingDataSource which allows me to switch connections at runtime:

AbstractRoutingDataSourceMyDB.java

public class AbstractRoutingDataSourceMyDB extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        return getDbConnectionMyDB();
    }

    // ThreadLocal variable so that the connection gets set for the current thread
    // using spring's request scope on the class instead of ThreadLocal would also work here.
    private final ThreadLocal<String> contextHolder = new ThreadLocal<String>();

    public void setDbConnectionMyDB(String myKey) {
        Assert.notNull(myKey, "myKey cannot be null");

        contextHolder.set(myKey);

        String k = contextHolder.get();
    }

    public String getDbConnectionMyDB() {
        return (String) contextHolder.get();
    }

    public void clearDbConnectionMyDB() {
        contextHolder.remove();
    }
}

Beware that you have to call entitymanager.flush() and clear() before you change the current connection from within your DAO classes or all pending operations in that transaction's scope will get executed on the new connection on transaction commit. This is because Hibernate session is oblivious that the connection ever changed, as far as it knows - it's always the same database.

---

So in your DAO you can do this now:

SomeTableDAO.java

@PersistenceContext(unitName = "MyDB")
private EntityManager em;

@Autowired
private AbstractRoutingDataSourceMyDB routingSource;

public void someMethod(int id) {
    em.flush();
    em.clear();
    routingSource.setDbConnectionMyDB("db1");
    em.remove(em.getReference(Something.class, id)); // delete something in db1

    em.flush();
    em.clear();
    routingSource.setDbConnectionMyDB("db2");
    em.remove(em.getReference(Something.class, id)); // delete something else with the same id in db2
}

So there you go, while it's not pretty - it can be done :)