Merge pull request #28 from arenadata/pxf_6.x/ADBDEV-3104

pxf-6.x/ADBDEV-3104: Add close connection if something wrong in openForRead() method
arenadata · Oct 17, 2022 · 8031b6f · 8031b6f
2 parents d649732 + 2f3d064
commit 8031b6f
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/server/pxf-jdbc/src/main/java/org/greenplum/pxf/plugins/jdbc/JdbcAccessor.java b/server/pxf-jdbc/src/main/java/org/greenplum/pxf/plugins/jdbc/JdbcAccessor.java
@@ -22,7 +22,6 @@
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang.StringUtils;
 import org.greenplum.pxf.api.OneRow;
-import org.greenplum.pxf.api.error.PxfRuntimeException;
 import org.greenplum.pxf.api.model.Accessor;
 import org.greenplum.pxf.api.model.ConfigurationFactory;
 import org.greenplum.pxf.api.security.SecureLogin;
@@ -36,12 +35,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.charset.Charset;
-import java.sql.Connection;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.SQLTimeoutException;
-import java.sql.Statement;
+import java.sql.*;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
@@ -101,8 +95,18 @@ public boolean openForRead() throws SQLException, SQLTimeoutException {
         if (statementRead != null && !statementRead.isClosed()) {
             return true;
         }
-
         Connection connection = super.getConnection();
+        try {
+            return openForReadInner(connection);
+        } catch (Throwable e) {
+            if (statementRead == null) {
+                JdbcBasePlugin.closeConnection(connection);
+            }
+            throw new RuntimeException(e.getMessage(), e);
+        }
+    }
+
+    private boolean openForReadInner(Connection connection) throws SQLException {
         SQLQueryBuilder sqlQueryBuilder = new SQLQueryBuilder(context, connection.getMetaData(), getQueryText());
 
         // Build SELECT query
@@ -230,7 +234,7 @@ public boolean openForWrite() throws SQLException, SQLTimeoutException {
         return true;
     }
 
-     /**
+    /**
      * writeNextObject() implementation
      * <p>
      * If batchSize is not 0 or 1, add a tuple to the batch of statementWrite

diff --git a/server/pxf-jdbc/src/main/java/org/greenplum/pxf/plugins/jdbc/JdbcBasePlugin.java b/server/pxf-jdbc/src/main/java/org/greenplum/pxf/plugins/jdbc/JdbcBasePlugin.java
@@ -488,7 +488,7 @@ private Connection getConnectionInternal() throws Exception {
      * @param connection connection to close
      * @throws SQLException throws when a SQLException occurs
      */
-    private static void closeConnection(Connection connection) throws SQLException {
+    static void closeConnection(Connection connection) throws SQLException {
         if (connection == null) {
             LOG.warn("Call to close connection is ignored as connection provided was null");
             return;

diff --git a/server/pxf-jdbc/src/test/java/org/greenplum/pxf/plugins/jdbc/JdbcAccessorTest.java b/server/pxf-jdbc/src/test/java/org/greenplum/pxf/plugins/jdbc/JdbcAccessorTest.java
@@ -82,7 +82,7 @@ public void testReadFromQueryFailsWhenServerDirectoryIsNotSpecified() throws SQL
         context.setDataSource("query:foo");
         accessor.setRequestContext(context);
         accessor.afterPropertiesSet();
-        Exception e = assertThrows(IllegalStateException.class,
+        Exception e = assertThrows(RuntimeException.class,
                 () -> accessor.openForRead());
         assertEquals("No server configuration directory found for server unknown", e.getMessage());
     }