/*
    * 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.commons.dbcp2.managed;
  
  import java.sql.Connection;
  import java.sql.SQLException;
  import java.util.concurrent.locks.Lock;
  import java.util.concurrent.locks.ReentrantLock;
  
  import org.apache.commons.dbcp2.DelegatingConnection;
  import org.apache.commons.pool2.ObjectPool;
  
  /**
   * ManagedConnection is responsible for managing a database connection in a transactional environment (typically called
   * "Container Managed"). A managed connection operates like any other connection when no global transaction (a.k.a. XA
   * transaction or JTA Transaction) is in progress. When a global transaction is active a single physical connection to
   * the database is used by all ManagedConnections accessed in the scope of the transaction. Connection sharing means
   * that all data access during a transaction has a consistent view of the database. When the global transaction is
   * committed or rolled back the enlisted connections are committed or rolled back. Typically, upon transaction
   * completion, a connection returns to the auto commit setting in effect before being enlisted in the transaction, but
   * some vendors do not properly implement this.
   * <p>
   * When enlisted in a transaction the setAutoCommit(), commit(), rollback(), and setReadOnly() methods throw a
   * SQLException. This is necessary to assure that the transaction completes as a single unit.
   * </p>
   *
   * @param <C>
   *            the Connection type
   *
   * @since 2.0
   */
  public class ManagedConnection<C extends Connection> extends DelegatingConnection<C> {
  
      /**
       * Delegates to {@link ManagedConnection#transactionComplete()} for transaction completion events.
       *
       * @since 2.0
       */
      protected class CompletionListener implements TransactionContextListener {
          @Override
          public void afterCompletion(final TransactionContext completedContext, final boolean committed) {
              if (completedContext == transactionContext) {
                  transactionComplete();
              }
          }
      }
  
      private final ObjectPool<C> pool;
      private final TransactionRegistry transactionRegistry;
      private final boolean accessToUnderlyingConnectionAllowed;
      private TransactionContext transactionContext;
      private boolean isSharedConnection;
      private final Lock lock;
  
      /**
       * Constructs a new instance responsible for managing a database connection in a transactional environment.
       *
       * @param pool
       *            The connection pool.
       * @param transactionRegistry
       *            The transaction registry.
       * @param accessToUnderlyingConnectionAllowed
       *            Whether or not to allow access to the underlying Connection.
       * @throws SQLException
       *             Thrown when there is problem managing transactions.
       */
      public ManagedConnection(final ObjectPool<C> pool, final TransactionRegistry transactionRegistry,
              final boolean accessToUnderlyingConnectionAllowed) throws SQLException {
          super(null);
          this.pool = pool;
          this.transactionRegistry = transactionRegistry;
          this.accessToUnderlyingConnectionAllowed = accessToUnderlyingConnectionAllowed;
          this.lock = new ReentrantLock();
          updateTransactionStatus();
      }
  
      @Override
      protected void checkOpen() throws SQLException {
          super.checkOpen();
          updateTransactionStatus();
      }
  
      @Override
      public void close() throws SQLException {
          if (!isClosedInternal()) {
             // Don't actually close the connection if in a transaction. The
             // connection will be closed by the transactionComplete method.
             //
             // DBCP-484 we need to make sure setClosedInternal(true) being
             // invoked if transactionContext is not null as this value will
             // be modified by the transactionComplete method which could run
             // in the different thread with the transaction calling back.
             lock.lock();
             try {
                 if (transactionContext == null || transactionContext.isTransactionComplete()) {
                     super.close();
                 }
             } finally {
                 try {
                     setClosedInternal(true);
                 } finally {
                     lock.unlock();
                 }
             }
         }
     }
 
     @Override
     public void commit() throws SQLException {
         if (transactionContext != null) {
             throw new SQLException("Commit can not be set while enrolled in a transaction");
         }
         super.commit();
     }
 
     @Override
     public C getDelegate() {
         if (isAccessToUnderlyingConnectionAllowed()) {
             return getDelegateInternal();
         }
         return null;
     }
 
     //
     // The following methods can't be used while enlisted in a transaction
     //
 
     @Override
     public Connection getInnermostDelegate() {
         if (isAccessToUnderlyingConnectionAllowed()) {
             return super.getInnermostDelegateInternal();
         }
         return null;
     }
 
     /**
      * @return The transaction context.
      * @since 2.6.0
      */
     public TransactionContext getTransactionContext() {
         return transactionContext;
     }
 
     /**
      * @return The transaction registry.
      * @since 2.6.0
      */
     public TransactionRegistry getTransactionRegistry() {
         return transactionRegistry;
     }
 
     /**
      * If false, getDelegate() and getInnermostDelegate() will return null.
      *
      * @return if false, getDelegate() and getInnermostDelegate() will return null
      */
     public boolean isAccessToUnderlyingConnectionAllowed() {
         return accessToUnderlyingConnectionAllowed;
     }
 
     @Override
     public void rollback() throws SQLException {
         if (transactionContext != null) {
             throw new SQLException("Commit can not be set while enrolled in a transaction");
         }
         super.rollback();
     }
 
     @Override
     public void setAutoCommit(final boolean autoCommit) throws SQLException {
         if (transactionContext != null) {
             throw new SQLException("Auto-commit can not be set while enrolled in a transaction");
         }
         super.setAutoCommit(autoCommit);
     }
 
     @Override
     public void setReadOnly(final boolean readOnly) throws SQLException {
         if (transactionContext != null) {
             throw new SQLException("Read-only can not be set while enrolled in a transaction");
         }
         super.setReadOnly(readOnly);
     }
 
     /**
      * Completes the transaction.
      */
     protected void transactionComplete() {
         lock.lock();
         try {
             transactionContext.completeTransaction();
         } finally {
             lock.unlock();
         }
 
         // If we were using a shared connection, clear the reference now that
         // the transaction has completed
         if (isSharedConnection) {
             setDelegate(null);
             isSharedConnection = false;
         }
 
         // autoCommit may have been changed directly on the underlying connection
         clearCachedState();
 
         // If this connection was closed during the transaction and there is
         // still a delegate present close it
         final Connection delegate = getDelegateInternal();
         if (isClosedInternal() && delegate != null) {
             try {
                 setDelegate(null);
 
                 if (!delegate.isClosed()) {
                     delegate.close();
                 }
             } catch (final SQLException ignored) {
                 // Not a whole lot we can do here as connection is closed
                 // and this is a transaction callback so there is no
                 // way to report the error.
             }
         }
     }
 
     private void updateTransactionStatus() throws SQLException {
         // if there is an active transaction context, assure the transaction context hasn't changed
         if (transactionContext != null && !transactionContext.isTransactionComplete()) {
             if (transactionContext.isActive()) {
                 if (transactionContext != transactionRegistry.getActiveTransactionContext()) {
                     throw new SQLException("Connection can not be used while enlisted in another transaction");
                 }
                 return;
             }
             // transaction should have been cleared up by TransactionContextListener, but in
             // rare cases another lister could have registered which uses the connection before
             // our listener is called. In that rare case, trigger the transaction complete call now
             transactionComplete();
         }
 
         // the existing transaction context ended (or we didn't have one), get the active transaction context
         transactionContext = transactionRegistry.getActiveTransactionContext();
 
         // if there is an active transaction context, and it already has a shared connection, use it
         if (transactionContext != null && transactionContext.getSharedConnection() != null) {
             // A connection for the connection factory has already been enrolled
             // in the transaction, replace our delegate with the enrolled connection
 
             // return current connection to the pool
             @SuppressWarnings("resource")
             final C connection = getDelegateInternal();
             setDelegate(null);
             if (connection != null && transactionContext.getSharedConnection() != connection) {
                 try {
                     pool.returnObject(connection);
                 } catch (final Exception e) {
                     // whatever... try to invalidate the connection
                     try {
                         pool.invalidateObject(connection);
                     } catch (final Exception ignored) {
                         // no big deal
                     }
                 }
             }
 
             // add a listener to the transaction context
             transactionContext.addTransactionContextListener(new CompletionListener());
 
             // Set our delegate to the shared connection. Note that this will
             // always be of type C since it has been shared by another
             // connection from the same pool.
             @SuppressWarnings("unchecked")
             final C shared = (C) transactionContext.getSharedConnection();
             setDelegate(shared);
 
             // remember that we are using a shared connection, so it can be cleared after the
             // transaction completes
             isSharedConnection = true;
         } else {
             C connection = getDelegateInternal();
             // if our delegate is null, create one
             if (connection == null) {
                 try {
                     // borrow a new connection from the pool
                     connection = pool.borrowObject();
                     setDelegate(connection);
                 } catch (final Exception e) {
                     throw new SQLException("Unable to acquire a new connection from the pool", e);
                 }
             }
 
             // if we have a transaction, out delegate becomes the shared delegate
             if (transactionContext != null) {
                 // add a listener to the transaction context
                 transactionContext.addTransactionContextListener(new CompletionListener());
 
                 // register our connection as the shared connection
                 try {
                     transactionContext.setSharedConnection(connection);
                 } catch (final SQLException e) {
                     // transaction is hosed
                     transactionContext = null;
                     try {
                         pool.invalidateObject(connection);
                     } catch (final Exception ignored) {
                         // we are try but no luck
                     }
                     throw e;
                 }
             }
         }
         // autoCommit may have been changed directly on the underlying
         // connection
         clearCachedState();
     }
 }