ManagedConnection.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.openejb.resource.jdbc.managed.local;

import org.apache.openejb.util.LogCategory;
import org.apache.openejb.util.Logger;

import javax.sql.CommonDataSource;
import javax.sql.DataSource;
import javax.sql.XAConnection;
import javax.sql.XADataSource;
import javax.transaction.RollbackException;
import javax.transaction.Status;
import javax.transaction.Synchronization;
import javax.transaction.SystemException;
import javax.transaction.Transaction;
import javax.transaction.TransactionManager;
import javax.transaction.TransactionSynchronizationRegistry;
import javax.transaction.xa.XAResource;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Wrapper;

public class ManagedConnection implements InvocationHandler {
    private final TransactionManager transactionManager;
    private final Key key;
    private final TransactionSynchronizationRegistry registry;
    protected XAResource xaResource;
    protected Connection delegate;
    protected XAConnection xaConnection;
    private Transaction currentTransaction;
    private boolean closed;

    public ManagedConnection(final CommonDataSource ds,
                             final TransactionManager txMgr,
                             final TransactionSynchronizationRegistry txRegistry,
                             final String user, final String password) {
        transactionManager = txMgr;
        registry = txRegistry;
        closed = false;
        key = new Key(ds, user, password);
    }

    public XAResource getXAResource() throws SQLException {
        if (xaResource == null) {
            newConnection();
        }
        return xaResource;
    }

    @Override
    public Object invoke(final Object proxy, final Method method, final Object[] args) throws Throwable {
        // first some Object method management
        final String mtdName = method.getName();
        if ("toString".equals(mtdName)) {
            return "ManagedConnection{" + delegate + "}";
        }
        if ("hashCode".equals(mtdName)) {
            return hashCode();
        }
        if ("equals".equals(mtdName)) {
            InvocationHandler handler;
            return args[0] == this || ((handler = unwrapHandler(args[0])) == this) || (delegate != null && delegate.equals(unwrapDelegate(args[0], handler)));
        }

        // allow to get delegate if needed by the underlying program
        if (Wrapper.class == method.getDeclaringClass() && args.length == 1 && Connection.class == args[0]) {
            if ("isWrapperFor".equals(mtdName)) {
                return true;
            }
            if ("unwrap".equals(mtdName)) {
                return delegate;
            }
        }

        // here the real logic starts
        try {
            final Transaction transaction = transactionManager.getTransaction();

            // shouldn't be used without a transaction but if so just delegate to the actual connection
            if (transaction == null) {
                if ("close".equals(mtdName)) {
                    if (delegate == null) { // no need to get a connection
                        return close();
                    }

                    closeConnection(true);
                    return null;
                }
                if ("isClosed".equals(mtdName) && closed) {
                    return true;
                }
                if (delegate == null) {
                    newConnection();
                }
                return invoke(method, delegate, args);
            }

            // if we have a tx check it is the same this connection is linked to
            if (currentTransaction != null && isUnderTransaction(currentTransaction.getStatus())) {
                if (!currentTransaction.equals(transaction)) {
                    throw new SQLException("Connection can not be used while enlisted in another transaction");
                }
                return invokeUnderTransaction(method, args);
            }

            // get the already bound connection to the current transaction or enlist this one in the tx
            final int transactionStatus = transaction.getStatus();
            if (isUnderTransaction(transactionStatus)) {
                Connection connection = Connection.class.cast(registry.getResource(key));
                if (connection == null && delegate == null) {
                    newConnection();

                    currentTransaction = transaction;
                    try {
                        if (!transaction.enlistResource(getXAResource())) {
                            closeConnection(true);
                            throw new SQLException("Unable to enlist connection in transaction: enlistResource returns 'false'.");
                        }
                    } catch (final RollbackException ignored) {
                        // no-op
                    } catch (final SystemException e) {
                        throw new SQLException("Unable to enlist connection the transaction", e);
                    }

                    registry.putResource(key, delegate);
                    transaction.registerSynchronization(new ClosingSynchronization());

                    if (xaConnection == null) {
                        try {
                            setAutoCommit(false);
                        } catch (final SQLException xae) { // we are alreay in a transaction so this can't be called from a user perspective - some XA DataSource prevents it in their code
                            final String message = "Can't set auto commit to false cause the XA datasource doesn't support it, this is likely an issue";
                            final Logger logger = Logger.getInstance(LogCategory.OPENEJB_RESOURCE_JDBC, ManagedConnection.class);
                            if (logger.isDebugEnabled()) { // we don't want to print the exception by default
                                logger.warning(message, xae);
                            } else {
                                logger.warning(message);
                            }
                        }
                    }
                } else if (delegate == null) { // shouldn't happen
                    delegate = connection;
                }

                return invokeUnderTransaction(method, args);
            }

            if ("isClosed".equals(mtdName) && closed) {
                return true;
            }
            if ("close".equals(mtdName)) { // let it be handled by the ClosingSynchronisation since we have a tx there
                return close();
            }

            // we shouldn't come here, tempted to just throw an exception
            if (delegate == null) {
                newConnection();
            }
            return invoke(method, delegate, args);
        } catch (final InvocationTargetException ite) {
            throw ite.getTargetException();
        }
    }

    private InvocationHandler unwrapHandler(final Object arg) {
        if (arg == null || !Proxy.isProxyClass(arg.getClass())) {
            return null;
        }
        return Proxy.getInvocationHandler(arg);
    }

    private Object unwrapDelegate(final Object arg, final InvocationHandler handler) {
        return handler != null && ManagedConnection.class.isInstance(handler) ? ManagedConnection.class.cast(handler).delegate : arg;
    }

    protected Object newConnection() throws SQLException {
        final Object connection = DataSource.class.isInstance(key.ds) ?
                (key.user == null ? DataSource.class.cast(key.ds).getConnection() : DataSource.class.cast(key.ds).getConnection(key.user, key.pwd)) :
                (key.user == null ? XADataSource.class.cast(key.ds).getXAConnection() : XADataSource.class.cast(key.ds).getXAConnection(key.user, key.pwd));
        if (XAConnection.class.isInstance(connection)) {
            xaConnection = XAConnection.class.cast(connection);
            xaResource = xaConnection.getXAResource();
            delegate = xaConnection.getConnection();
        } else {
            delegate = Connection.class.cast(connection);
            xaResource = new LocalXAResource(delegate);
        }
        return connection;
    }

    protected void setAutoCommit(final boolean value) throws SQLException {
        if (delegate == null) {
            newConnection();
        }
        delegate.setAutoCommit(value);
    }

    private static Object invoke(final Method method, final Connection delegate, final Object[] args) throws Throwable {
        try {
            return method.invoke(delegate, args);
        } catch (final InvocationTargetException ite) {
            throw ite.getCause();
        }
    }

    private Object invokeUnderTransaction(final Method method, final Object[] args) throws Exception {
        final String mtdName = method.getName();
        if ("setAutoCommit".equals(mtdName)
                || "commit".equals(mtdName)
                || "rollback".equals(mtdName)
                || "setSavepoint".equals(mtdName)
                || "setReadOnly".equals(mtdName)) {
            throw forbiddenCall(mtdName);
        }
        if ("close".equals(mtdName)) {
            return close();
        }
        if ("isClosed".equals(mtdName) && closed) {
            return true; // if !closed let's delegate to the underlying connection
        }
        return method.invoke(delegate, args);
    }

    // will be done later
    // we need to delay it in case of rollback
    private Object close() {
        closed = true;
        return null;
    }

    private static boolean isUnderTransaction(final int status) {
        return status == Status.STATUS_ACTIVE || status == Status.STATUS_MARKED_ROLLBACK;
    }

    private static SQLException forbiddenCall(final String mtdName) {
        return new SQLException("can't call " + mtdName + " when the connection is JtaManaged");
    }

    private class ClosingSynchronization implements Synchronization {
        @Override
        public void beforeCompletion() {
            // no-op
        }

        @Override
        public void afterCompletion(final int status) {
            closeConnection(true);
        }
    }

    private void closeConnection(final boolean force) {
        if (!force && closed) {
            return;
        }
        try {
            if (xaConnection != null) { // handles the underlying connection
                xaConnection.close();
            } else if (delegate != null && !delegate.isClosed()) {
                delegate.close();
            }
        } catch (final SQLException e) {
            // no-op
        } finally {
            close(); // set the flag
        }
    }

    private static final class Key {
        private final CommonDataSource ds;
        private final String user;
        private final String pwd;
        private final int hash;

        private Key(final CommonDataSource ds, final String user, final String pwd) {
            this.ds = ds;
            this.user = user;
            this.pwd = pwd;

            int result = ds.hashCode();
            result = 31 * result + (user != null ? user.hashCode() : 0);
            result = 31 * result + (pwd != null ? pwd.hashCode() : 0);
            hash = result;
        }

        @Override
        public boolean equals(final Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }

            Key key = Key.class.cast(o);
            return (ds == key.ds || ds.equals(key.ds)) &&
                    !(user != null ? !user.equals(key.user) : key.user != null) &&
                    !(pwd != null ? !pwd.equals(key.pwd) : key.pwd != null);
        }

        @Override
        public int hashCode() {
            return hash;
        }
    }
}