Sample class jdbc.wrapper.XADataSourceProxy

/*
 * ============================================================================
 * (C) Copyright IBM Corporation 2006, 2007.  All Rights Reserved.
 * 
 * You may only copy, modify, and distribute these samples internally.
 * These samples have not been tested under all conditions and are provided
 * to you by IBM without obligation of support of any kind.
 *
 * IBM PROVIDES THESE SAMPLES "AS IS" SUBJECT TO ANY STATUTORY WARRANTIES THAT
 * CANNOT BE EXCLUDED. IBM MAKES NO WARRANTIES OR CONDITIONS, EITHER EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OR
 * CONDITIONS OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND
 * NON-INFRINGEMENT REGARDING THESE SAMPLES OR TECHNICAL SUPPORT, IF ANY.
 * ============================================================================
 */
package jdbc.wrapper;

import java.io.PrintWriter;
import java.io.Serializable;
import java.sql.CallableStatement;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Savepoint;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import javax.sql.DataSource;
import javax.sql.PooledConnection;
import javax.sql.XAConnection;
import javax.sql.XADataSource;
import javax.transaction.Synchronization;
import javax.transaction.Transaction;
import javax.transaction.xa.XAResource;
import jdbc.wrapper.XAResourceInfoImpl;
import com.ibm.tx.jta.DestroyXAResourceException;
import com.ibm.tx.jta.ExtendedTransactionManager;
import com.ibm.tx.jta.TransactionManagerFactory;
import com.ibm.tx.jta.XAResourceFactory;
import com.ibm.tx.jta.XAResourceNotAvailableException;

public class XADataSourceProxy implements DataSource {
	private static final ExtendedTransactionManager TRANSACTION_MANAGER = TransactionManagerFactory
			.getTransactionManager();

	private XAResourceInfoImpl _xaResourceInfo;
	private XADataSource _xaDataSource;

	private ArrayList<PooledConnection> _availableConnections;
	private Map<Transaction, PooledConnection> _inUseConnectionsByTransaction;

	private static final TraceComponent _tc = Tr.register(
			XADataSourceProxy.class, null, null);

	private boolean _autoCommit;
	private boolean _autoCommitOverride = false;
	public void setAutoCommit(boolean c) {
		_autoCommit = c;
		_autoCommitOverride = true;
	}
	public void close(){}

	public XADataSourceProxy(String xaDataSourceClassName,
			Properties configurationProperties) {
		_xaResourceInfo = new XAResourceInfoImpl(xaDataSourceClassName, configurationProperties);
		_xaDataSource = _xaResourceInfo.createDataSource();
		_availableConnections = new ArrayList<PooledConnection>();
		_inUseConnectionsByTransaction = new HashMap<Transaction, PooledConnection>();
	}

	public Connection getConnection() throws SQLException {
		if (_tc.isEntryEnabled())
			Tr.entry(_tc, "getConnection", this);
		XAConnection xaConnection;
		if (_availableConnections.isEmpty()) {
			xaConnection = null;
		} else {
			xaConnection = (XAConnection) _availableConnections.remove(0);
		}

		// A connection has been requested by the application. If there is
		// a transaction on the thread, enlist the connection in the
		// transaction.
		try {
			final Transaction transaction = TRANSACTION_MANAGER
					.getTransaction();
			final Connection connection;
			if (transaction != null) {
				if (xaConnection == null) {
					xaConnection = (XAConnection) _inUseConnectionsByTransaction.get(transaction);
					if (xaConnection == null) {
						xaConnection = _xaDataSource.getXAConnection();
						_inUseConnectionsByTransaction.put(transaction,	xaConnection);
						final XAResource xaResource = xaConnection.getXAResource();
						if (xaResource instanceof Serializable) {
							transaction.enlistResource(xaResource);
						} else {
							// Enlisting NON-serializable resource
							final int recoveryId = TRANSACTION_MANAGER
									.registerResourceInfo(XADataSourceProxyFactory.class.getName(), 
									                      _xaResourceInfo);
							TRANSACTION_MANAGER.enlist(xaResource, recoveryId);
						}

						// Register a synchronization with the transaction. Once
						// that transaction has completed we will be notified 
						// via afterCompletion at which point we can free-up the 
						// connection and make it available for re-use.
						transaction.registerSynchronization(new SynchronizationImpl(
															xaConnection, transaction));
					} else {
						// Re-using connection within same transaction.
					}
				}
				connection = new ConnectionWrapper(xaConnection.getConnection());
				if (_autoCommitOverride) {
					connection.setAutoCommit(_autoCommit);
				}
			} else {
				if (xaConnection == null) {
					xaConnection = _xaDataSource.getXAConnection();
				}
				connection = xaConnection.getConnection();
			}
			return connection;
		} catch (Exception e) {
			e.printStackTrace();
			throw new SQLException();
		}
	}

	public Connection getConnection(String username, String password)
			throws SQLException {
		throw new UnsupportedOperationException();
	}

	public PrintWriter getLogWriter() throws SQLException {
		return _xaDataSource.getLogWriter();
	}

	public int getLoginTimeout() throws SQLException {
		return _xaDataSource.getLoginTimeout();
	}

	public void setLogWriter(PrintWriter out) throws SQLException {
		_xaDataSource.setLogWriter(out);
	}

	public void setLoginTimeout(int seconds) throws SQLException {
		_xaDataSource.setLoginTimeout(seconds);
	}

	public static class XADataSourceProxyFactory implements XAResourceFactory {
		public void destroyXAResource(XAResource xaRes)
				throws DestroyXAResourceException {
		}
		public XAResource getXAResource(Serializable xaresinfo)
				throws XAResourceNotAvailableException {
			try {
				final XAResourceInfoImpl xaResourceInfo = (XAResourceInfoImpl) xaresinfo;
				return xaResourceInfo.createDataSource().getXAConnection().getXAResource();
			} catch (Exception e) {
				throw new XAResourceNotAvailableException(e);
			}
		}
	}

	private class ConnectionWrapper implements Connection {
		private Connection _connection;

		public ConnectionWrapper(Connection connection) {
			_connection = connection;
		}

		public void clearWarnings() throws SQLException {
			_connection.clearWarnings();
		}

		public void close() throws SQLException {

		}

        // Don't call commit() on the ConnectionWrapper - it will be done via the transaction manager
		public void commit() throws SQLException {
			throw new SQLException();
		}

		public Statement createStatement() throws SQLException {
			return _connection.createStatement();
		}

		public Statement createStatement(int resultSetType,
				int resultSetConcurrency) throws SQLException {
			return _connection.createStatement(resultSetType,
					resultSetConcurrency);
		}

		public Statement createStatement(int resultSetType,
				int resultSetConcurrency, int resultSetHoldability)
				throws SQLException {
			return _connection.createStatement(resultSetType,
					resultSetConcurrency, resultSetHoldability);
		}

		public boolean getAutoCommit() throws SQLException {
			return _connection.getAutoCommit();
		}

		public String getCatalog() throws SQLException {
			return _connection.getCatalog();
		}

		public int getHoldability() throws SQLException {
			return _connection.getHoldability();
		}

		public DatabaseMetaData getMetaData() throws SQLException {
			return _connection.getMetaData();
		}

		public int getTransactionIsolation() throws SQLException {
			return _connection.getTransactionIsolation();
		}

		public Map<String, Class<?>> getTypeMap() throws SQLException {
			return _connection.getTypeMap();
		}

		public SQLWarning getWarnings() throws SQLException {
			return _connection.getWarnings();
		}

		public boolean isClosed() throws SQLException {
			return _connection.isClosed();
		}

		public boolean isReadOnly() throws SQLException {
			return _connection.isReadOnly();
		}

		public String nativeSQL(String sql) throws SQLException {

			return _connection.nativeSQL(sql);
		}

		public CallableStatement prepareCall(String sql) throws SQLException {
			return _connection.prepareCall(sql);
		}

		public CallableStatement prepareCall(String sql, int resultSetType,
				int resultSetConcurrency) throws SQLException {
			return _connection.prepareCall(sql, resultSetType,
					resultSetConcurrency);
		}

		public CallableStatement prepareCall(String sql, int resultSetType,
				int resultSetConcurrency, int resultSetHoldability)
				throws SQLException {
			return _connection.prepareCall(sql, resultSetType,
					resultSetConcurrency);
		}

		public PreparedStatement prepareStatement(String sql)
				throws SQLException {
			return _connection.prepareStatement(sql);
		}

		public PreparedStatement prepareStatement(String sql,
				int autoGeneratedKeys) throws SQLException {
			return _connection.prepareStatement(sql, autoGeneratedKeys);
		}

		public PreparedStatement prepareStatement(String sql,
				int[] columnIndexes) throws SQLException {
			return _connection.prepareStatement(sql, columnIndexes);
		}

		public PreparedStatement prepareStatement(String sql,
				String[] columnNames) throws SQLException {
			return _connection.prepareStatement(sql, columnNames);
		}

		public PreparedStatement prepareStatement(String sql,
				int resultSetType, int resultSetConcurrency)
				throws SQLException {
			return _connection.prepareStatement(sql, resultSetType,
					resultSetConcurrency);
		}

		public PreparedStatement prepareStatement(String sql,
				int resultSetType, int resultSetConcurrency,
				int resultSetHoldability) throws SQLException {
			return _connection.prepareStatement(sql, resultSetType,
					resultSetConcurrency, resultSetHoldability);
		}

		public void releaseSavepoint(Savepoint savepoint) throws SQLException {
			_connection.releaseSavepoint(savepoint);
		}

		public void rollback() throws SQLException {
			throw new SQLException();
		}

		public void rollback(Savepoint savepoint) throws SQLException {
			throw new SQLException();
		}

		public void setAutoCommit(boolean autoCommit) throws SQLException {
			_connection.setAutoCommit(autoCommit);
		}

		public void setCatalog(String catalog) throws SQLException {
			_connection.setCatalog(catalog);
		}

		public void setHoldability(int holdability) throws SQLException {
			_connection.setHoldability(holdability);
		}

		public void setReadOnly(boolean readOnly) throws SQLException {
			_connection.setReadOnly(readOnly);
		}

		public Savepoint setSavepoint() throws SQLException {
			return _connection.setSavepoint();
		}

		public Savepoint setSavepoint(String name) throws SQLException {
			return _connection.setSavepoint(name);
		}

		public void setTransactionIsolation(int level) throws SQLException {
			_connection.setTransactionIsolation(level);
		}

		public void setTypeMap(Map<String, Class<?>> map) throws SQLException {
			_connection.setTypeMap(map);
		}

	}

	private class SynchronizationImpl implements Synchronization {
		private PooledConnection _pooledConnection;
		private Transaction _transaction;

		public SynchronizationImpl(PooledConnection pooledConnection, Transaction transaction) {
			_pooledConnection = pooledConnection;
			_transaction = transaction;
		}

		public void afterCompletion(int status) {
			_availableConnections.add(0, _pooledConnection);
			_inUseConnectionsByTransaction.remove(_transaction);
		}

		public void beforeCompletion() {
		}
	}
}