From 38d8a5d0d2bbe318419d459bee00e438a7ec9ff5 Mon Sep 17 00:00:00 2001
From: Rebecca Schlussel <rschlussel@fb.com>
Date: Thu, 15 Aug 2024 13:49:22 -0400
Subject: [PATCH] Fix failure with enabling JDBC autocommit

There was a bug where enabling autocommit when it had previously been
false would cause failure like:
java.sql.SQLException: Connection is in auto-commit mode

This was due to calling commit() from setAutoCommit() after autocommit
had already been changed to true.

Fixes #15916
---
 .../presto/jdbc/PrestoConnection.java         |  3 +-
 .../presto/jdbc/TestJdbcConnection.java       | 82 +++++++++++++++++--
 2 files changed, 75 insertions(+), 10 deletions(-)

diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java
index 46d48b5dc928..1f3f76063185 100644
--- a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java
+++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoConnection.java
@@ -187,10 +187,11 @@ public void setAutoCommit(boolean autoCommit)
             throws SQLException
     {
         checkOpen();
-        boolean wasAutoCommit = this.autoCommit.getAndSet(autoCommit);
+        boolean wasAutoCommit = this.autoCommit.get();
         if (autoCommit && !wasAutoCommit) {
             commit();
         }
+        this.autoCommit.set(autoCommit);
     }
 
     @Override
diff --git a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcConnection.java b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcConnection.java
index ed77488e1949..7e80977af634 100644
--- a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcConnection.java
+++ b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcConnection.java
@@ -105,21 +105,85 @@ public void teardownServer()
     public void testCommit()
             throws SQLException
     {
-        try (Connection connection = createConnection()) {
-            connection.setAutoCommit(false);
-            try (Statement statement = connection.createStatement()) {
-                statement.execute("CREATE TABLE test_commit (x bigint)");
+        try {
+            try (Connection connection = createConnection()) {
+                connection.setAutoCommit(false);
+                try (Statement statement = connection.createStatement()) {
+                    statement.execute("CREATE TABLE test_commit (x bigint)");
+                }
+
+                try (Connection otherConnection = createConnection()) {
+                    assertThat(listTables(otherConnection)).doesNotContain("test_commit");
+                }
+
+                connection.commit();
             }
 
-            try (Connection otherConnection = createConnection()) {
-                assertThat(listTables(otherConnection)).doesNotContain("test_commit");
+            try (Connection connection = createConnection()) {
+                assertThat(listTables(connection)).contains("test_commit");
+            }
+        }
+        finally {
+            try (Connection connection = createConnection()) {
+                try (Statement statement = connection.createStatement()) {
+                    statement.execute("DROP TABLE test_commit");
+                }
+            }
+        }
+    }
+
+    @Test
+    public void testAutoCommit()
+            throws SQLException
+    {
+        try {
+            try (Connection connection = createConnection()) {
+                connection.setAutoCommit(true);
+                try (Statement statement = connection.createStatement()) {
+                    statement.execute("CREATE TABLE test_autocommit (x bigint)");
+                }
             }
 
-            connection.commit();
+            try (Connection connection = createConnection()) {
+                assertThat(listTables(connection)).contains("test_autocommit");
+            }
+        }
+        finally {
+            try (Connection connection = createConnection()) {
+                try (Statement statement = connection.createStatement()) {
+                    statement.execute("DROP TABLE test_autocommit");
+                }
+            }
         }
+    }
 
-        try (Connection connection = createConnection()) {
-            assertThat(listTables(connection)).contains("test_commit");
+    @Test
+    public void testResetAutoCommit()
+            throws SQLException
+    {
+        try {
+            try (Connection connection = createConnection()) {
+                connection.setAutoCommit(false);
+                try (Statement statement = connection.createStatement()) {
+                    statement.execute("CREATE TABLE test_reset_autocommit (x bigint)");
+                }
+
+                try (Connection otherConnection = createConnection()) {
+                    assertThat(listTables(otherConnection)).doesNotContain("test_reset_autocommit");
+                }
+                connection.setAutoCommit(true);
+            }
+
+            try (Connection connection = createConnection()) {
+                assertThat(listTables(connection)).contains("test_reset_autocommit");
+            }
+        }
+        finally {
+            try (Connection connection = createConnection()) {
+                try (Statement statement = connection.createStatement()) {
+                    statement.execute("DROP TABLE test_reset_autocommit");
+                }
+            }
         }
     }