From 1ce598284e293a79a082e4a06291ec3184e8dffb Mon Sep 17 00:00:00 2001
From: rusher <diego.dupin@gmail.com>
Date: Wed, 16 Oct 2024 15:56:07 +0200
Subject: [PATCH] [misc] XA Pool datasource ensuring close when using
 XAConnection.close()

---
 .../org/mariadb/jdbc/MariaDbDataSource.java   | 49 ++++++++++++-------
 .../mariadb/jdbc/MariaDbPoolConnection.java   |  3 ++
 2 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/src/main/java/org/mariadb/jdbc/MariaDbDataSource.java b/src/main/java/org/mariadb/jdbc/MariaDbDataSource.java
index c9dbe941a..3ed7659b7 100644
--- a/src/main/java/org/mariadb/jdbc/MariaDbDataSource.java
+++ b/src/main/java/org/mariadb/jdbc/MariaDbDataSource.java
@@ -208,9 +208,12 @@ public Logger getParentLogger() {
   public PooledConnection getPooledConnection() throws SQLException {
     if (conf == null) config();
     org.mariadb.jdbc.Connection conn = Driver.connect(conf);
-    return conf.pinGlobalTxToPhysicalConnection()
-        ? new MariaDbPoolPinnedConnection(conn)
-        : new MariaDbPoolConnection(conn);
+    MariaDbPoolConnection poolConnection =
+        conf.pinGlobalTxToPhysicalConnection()
+            ? new MariaDbPoolPinnedConnection(conn)
+            : new MariaDbPoolConnection(conn);
+    setConnectionCloseListener(poolConnection);
+    return poolConnection;
   }
 
   @Override
@@ -219,28 +222,40 @@ public PooledConnection getPooledConnection(String username, String password)
     if (conf == null) config();
     Configuration conf = this.conf.clone(username, password);
     org.mariadb.jdbc.Connection conn = Driver.connect(conf);
-    return conf.pinGlobalTxToPhysicalConnection()
-        ? new MariaDbPoolPinnedConnection(conn)
-        : new MariaDbPoolConnection(conn);
+    MariaDbPoolConnection poolConnection =
+        conf.pinGlobalTxToPhysicalConnection()
+            ? new MariaDbPoolPinnedConnection(conn)
+            : new MariaDbPoolConnection(conn);
+    setConnectionCloseListener(poolConnection);
+    return poolConnection;
+  }
+
+  private void setConnectionCloseListener(MariaDbPoolConnection mariaDbPoolConnection) {
+    mariaDbPoolConnection.addConnectionEventListener(
+        new ConnectionEventListener() {
+
+          @Override
+          public void connectionClosed(ConnectionEvent event) {
+            try {
+              mariaDbPoolConnection.realClose();
+            } catch (SQLException e) {
+              throw new RuntimeException(e);
+            }
+          }
+
+          @Override
+          public void connectionErrorOccurred(ConnectionEvent event) {}
+        });
   }
 
   @Override
   public XAConnection getXAConnection() throws SQLException {
-    if (conf == null) config();
-    org.mariadb.jdbc.Connection conn = Driver.connect(conf);
-    return conf.pinGlobalTxToPhysicalConnection()
-        ? new MariaDbPoolPinnedConnection(conn)
-        : new MariaDbPoolConnection(conn);
+    return (MariaDbPoolConnection) getPooledConnection();
   }
 
   @Override
   public XAConnection getXAConnection(String username, String password) throws SQLException {
-    if (conf == null) config();
-    Configuration conf = this.conf.clone(username, password);
-    org.mariadb.jdbc.Connection conn = Driver.connect(conf);
-    return conf.pinGlobalTxToPhysicalConnection()
-        ? new MariaDbPoolPinnedConnection(conn)
-        : new MariaDbPoolConnection(conn);
+    return (MariaDbPoolConnection) getPooledConnection(username, password);
   }
 
   /**
diff --git a/src/main/java/org/mariadb/jdbc/MariaDbPoolConnection.java b/src/main/java/org/mariadb/jdbc/MariaDbPoolConnection.java
index 0a42fe1fa..9302373a4 100644
--- a/src/main/java/org/mariadb/jdbc/MariaDbPoolConnection.java
+++ b/src/main/java/org/mariadb/jdbc/MariaDbPoolConnection.java
@@ -131,6 +131,9 @@ public void fireConnectionErrorOccurred(SQLException returnEx) {
   @Override
   public void close() throws SQLException {
     fireConnectionClosed(new ConnectionEvent(this));
+  }
+
+  protected void realClose() throws SQLException {
     connection.setPoolConnection(null);
     connection.close();
   }