From 186518ad4c8702d92ef29b2631ad80c58bc41e92 Mon Sep 17 00:00:00 2001
From: Tom Jenkinson <tom.jenkinson@redhat.com>
Date: Mon, 23 Jul 2018 10:48:47 +0100
Subject: [PATCH] DBCP-515 Added support for obtaining a reference to the
 connection

---
 .../dbcp2/managed/ManagedConnection.java      |  2 +-
 .../dbcp2/managed/TransactionContext.java     | 36 ++++++++++---------
 .../managed/TransactionContextListener.java   |  4 +--
 .../managed/TestManagedDataSourceInTx.java    | 31 ++++++++++++++++
 4 files changed, 53 insertions(+), 20 deletions(-)

diff --git a/src/main/java/org/apache/commons/dbcp2/managed/ManagedConnection.java b/src/main/java/org/apache/commons/dbcp2/managed/ManagedConnection.java
index 1b1f9dee5..b28500fe2 100644
--- a/src/main/java/org/apache/commons/dbcp2/managed/ManagedConnection.java
+++ b/src/main/java/org/apache/commons/dbcp2/managed/ManagedConnection.java
@@ -203,7 +203,7 @@ public void close() throws SQLException {
      */
     protected class CompletionListener implements TransactionContextListener {
         @Override
-        public void afterCompletion(final TransactionContext completedContext, final boolean commited) {
+        public void afterCompletion(final TransactionContext completedContext) {
             if (completedContext == transactionContext) {
                 transactionComplete();
             }
diff --git a/src/main/java/org/apache/commons/dbcp2/managed/TransactionContext.java b/src/main/java/org/apache/commons/dbcp2/managed/TransactionContext.java
index f4065d6c0..dbc1321d1 100644
--- a/src/main/java/org/apache/commons/dbcp2/managed/TransactionContext.java
+++ b/src/main/java/org/apache/commons/dbcp2/managed/TransactionContext.java
@@ -111,23 +111,27 @@ public void setSharedConnection(final Connection sharedConnection) throws SQLExc
      *             if a problem occurs adding the listener to the transaction
      */
     public void addTransactionContextListener(final TransactionContextListener listener) throws SQLException {
-        try {
-            getTransaction().registerSynchronization(new Synchronization() {
-                @Override
-                public void beforeCompletion() {
-                    // empty
-                }
+        if (isActive()) {
+            try {
+                getTransaction().registerSynchronization(new Synchronization() {
+                    @Override
+                    public void beforeCompletion() {
+                        // empty
+                    }
 
-                @Override
-                public void afterCompletion(final int status) {
-                    listener.afterCompletion(TransactionContext.this, status == Status.STATUS_COMMITTED);
-                }
-            });
-        } catch (final RollbackException e) {
-            // JTA spec doesn't let us register with a transaction marked rollback only
-            // just ignore this and the tx state will be cleared another way.
-        } catch (final Exception e) {
-            throw new SQLException("Unable to register transaction context listener", e);
+                    @Override
+                    public void afterCompletion(final int status) {
+                        listener.afterCompletion(TransactionContext.this);
+                    }
+                });
+            } catch (final RollbackException e) {
+                // JTA spec doesn't let us register with a transaction marked rollback only
+                // just ignore this and the tx state will be cleared another way.
+            } catch (final Exception e) {
+                throw new SQLException("Unable to register transaction context listener", e);
+            }
+        } else {
+            listener.afterCompletion(this);
         }
     }
 
diff --git a/src/main/java/org/apache/commons/dbcp2/managed/TransactionContextListener.java b/src/main/java/org/apache/commons/dbcp2/managed/TransactionContextListener.java
index cd47079bf..84f1c9e78 100644
--- a/src/main/java/org/apache/commons/dbcp2/managed/TransactionContextListener.java
+++ b/src/main/java/org/apache/commons/dbcp2/managed/TransactionContextListener.java
@@ -28,8 +28,6 @@ public interface TransactionContextListener {
      *
      * @param transactionContext
      *            the transaction context that completed
-     * @param commited
-     *            true if the transaction committed; false otherwise
      */
-    void afterCompletion(TransactionContext transactionContext, boolean commited);
+    void afterCompletion(TransactionContext transactionContext);
 }
diff --git a/src/test/java/org/apache/commons/dbcp2/managed/TestManagedDataSourceInTx.java b/src/test/java/org/apache/commons/dbcp2/managed/TestManagedDataSourceInTx.java
index a8d45a918..8eb38f03c 100644
--- a/src/test/java/org/apache/commons/dbcp2/managed/TestManagedDataSourceInTx.java
+++ b/src/test/java/org/apache/commons/dbcp2/managed/TestManagedDataSourceInTx.java
@@ -36,6 +36,7 @@
 import java.sql.ResultSet;
 import java.sql.PreparedStatement;
 
+import javax.transaction.Synchronization;
 import javax.transaction.Transaction;
 
 /**
@@ -59,6 +60,36 @@ public void tearDown() throws Exception {
         super.tearDown();
     }
 
+    @Test
+    public void testGetConnectionInAfterCompletion() throws Exception {
+
+        final DelegatingConnection<?> connection = (DelegatingConnection<?>) newConnection();
+        // Don't close so we can check it for warnings in afterCompletion
+        transactionManager.getTransaction().registerSynchronization(new Synchronization() {
+            @Override
+            public void beforeCompletion() {
+
+            }
+
+            @Override
+            public void afterCompletion(int i) {
+                try {
+                    Connection connection1 = ds.getConnection();
+                    try {
+                        connection1.getWarnings();
+                        fail("Could operate on closed connection");
+                    } catch (SQLException e) {
+                        // This is expected
+                    }
+                } catch (SQLException e) {
+                    fail("Should have been able to get connection");
+                }
+            }
+        });
+        connection.close();
+        transactionManager.commit();
+    }
+
     /**
      * @see #testSharedConnection()
      */