Dumb Programmer

So, how to use the same connection for all the partitions by switching the catalog? And before that why should one go for that? To answer the later, there are lots of advantages in going for that,

• You can reduce the number of connection pools
• The number of session factories can be limited as the pools have been reduced.

The first thing we have to do is to make the DataSourceIdentifier return the catalog,

public abstract class DataSourceIdentifier {

	protected String dataSourceId;
	protected String dataSourceCatalog;
	
	public String getDataSourceId() {
		if(dataSourceId == null) {
			resolve();
		}
		return dataSourceId;
	}

	public String getDataSourceCatalog() {
		if(dataSourceCatalog == null) {
			resolve();
		}
		return dataSourceCatalog;
	}

	// Add the partitioning logic here
	public abstract void resolve();
}

Then we have create one session factory (assuming that all partitions reside in the same db instance otherwise create one per db instance). We have to create a custom session factory bean to make sure the correct catalog is set before the connection is used by the session.

public class CustomLocalSessionFactoryBean extends LocalSessionFactoryBean {
	
	@Override
	protected void postProcessConfiguration(Configuration config) throws HibernateException {
		String connectionProvider = config.getProperty(Environment.CONNECTION_PROVIDER);
		if (connectionProvider != null) {
			if (connectionProvider.equals(TransactionAwareDataSourceConnectionProvider.class.getName())) {
				connectionProvider = CustomTransactionAwareDataSourceConnectionProvider.class.getName();
			} else if (connectionProvider.equals(LocalJtaDataSourceConnectionProvider.class.getName())) {
				connectionProvider = CustomLocalJtaDataSourceConnectionProvider.class.getName();
			} else if (connectionProvider.equals(LocalDataSourceConnectionProvider.class.getName())) {
				connectionProvider = CustomLocalDataSourceConnectionProvider.class.getName();
			}
			// Set Spring-provided DataSource as Hibernate ConnectionProvider.
			config.setProperty(Environment.CONNECTION_PROVIDER, connectionProvider);
		}
		super.postProcessConfiguration(config);
	}
	
	@Override
	protected void afterSessionFactoryCreation() throws Exception {
		SessionFactoryImplementor implementor = (SessionFactoryImplementor) getSessionFactory();
		ConnectionProvider provider = implementor.getConnectionProvider();
		if(provider instanceof CustomConnectionProvider) {
			((CustomConnectionProvider)provider).setFactory(implementor);
		}
		super.afterSessionFactoryCreation();
	}



	public interface CustomConnectionProvider {
		
		public void setFactory(SessionFactory factory);
	}

	public static class CustomLocalDataSourceConnectionProvider extends LocalDataSourceConnectionProvider implements CustomConnectionProvider {
		
		private SessionFactoryImplementor factory;
		
		@Override
		public void setFactory(SessionFactory factory) {
			this.factory = (SessionFactoryImplementor) factory;
		}

		@Override
		public Connection getConnection() throws SQLException {
			Connection connection = super.getConnection();
			DataSourceIdentifier dataSourceId = ContextHolder.getDataSourceId(this.factory);
			if(dataSourceId != null && dataSourceId.getDataSourceCatalog() != null) {
				connection.setCatalog(dataSourceId.getDataSourceCatalog());
			}
			connection.setAutoCommit(false);
			return connection;
		}
		
	}

	// Repeat the same for CustomTransactionAwareDataSourceConnectionProvider and CustomLocalJtaDataSourceConnectionProvider as well
	
}

The session factory configuration will look like this,

<bean id="db0" class="CustomLocalSessionFactoryBean">
	<property name="hibernateProperties" ref="hibernateProperties" />
	<property name="mappingResources" ref="dbPartitionMappings"/>
	<property name="dataSource" ref="db0DataSource" />
</bean>

And that is it. We are ready to test this out. It has been working pretty good in my laptop I would love to get back from you guys on optimizing the design or the issues that you find with this logic.