Close pending statements on connection close

staticlibs · staticlibs · commit bebd33754fe5 · 2025-03-20T15:54:03.000Z
When connection is closed without closing its statements, the JDBC spec does NOT require to close such statements automatically. But pending statements are causing problems (see duckdb#101) and having them auto-closed seems to be a reasonable expectation from a client point of view. This change adds tracking and auto-closing of statements on connection close. It also adds/improves synchronizaton on closing connections/ statements/result sets. Straighforward sync on current object is used deliberately instead of other concurrency methods to keep the synchronization impl as simple as possible. Testing: new tests are added to check auto-closing of result sets and statements. Fixes: duckdb#101
diff --git a/src/main/java/org/duckdb/DuckDBConnection.java b/src/main/java/org/duckdb/DuckDBConnection.java
@@ -20,8 +20,7 @@
 import java.sql.Savepoint;
 import java.sql.Statement;
 import java.sql.Struct;
-import java.util.Map;
-import java.util.Properties;
+import java.util.*;
 import java.util.concurrent.Executor;
 import org.duckdb.user.DuckDBMap;
 import org.duckdb.user.DuckDBUserArray;
@@ -37,6 +36,9 @@ public final class DuckDBConnection implements java.sql.Connection {
     boolean transactionRunning;
     final String url;
     private final boolean readOnly;
+    // Tracking set for statements created from this connection,
+    // must only be used from outside with track/untrack accessors.
+    private final Set<DuckDBPreparedStatement> statements = new LinkedHashSet<>();
 
     public static DuckDBConnection newConnection(String url, boolean readOnly, Properties properties)
         throws SQLException {
@@ -110,15 +112,32 @@ protected void finalize() throws Throwable {
         close();
     }
 
-    public synchronized void close() throws SQLException {
-        if (conn_ref != null) {
+    public void close() throws SQLException {
+        if (conn_ref == null) {
+            return;
+        }
+        synchronized (this) {
+            if (conn_ref == null) {
+                return;
+            }
+            // Closing remaining statements is not required by JDBC spec,
+            // but it is a reasonable expectation from clients point of view.
+            List<DuckDBPreparedStatement> stmtList = new ArrayList<>(statements);
+            for (DuckDBPreparedStatement stmt : stmtList) {
+                stmt.close();
+            }
             DuckDBNative.duckdb_jdbc_disconnect(conn_ref);
             conn_ref = null;
         }
     }
 
     public boolean isClosed() throws SQLException {
-        return conn_ref == null;
+        if (conn_ref == null) {
+            return true;
+        }
+        synchronized (this) {
+            return conn_ref == null;
+        }
     }
 
     public boolean isValid(int timeout) throws SQLException {
@@ -371,4 +390,19 @@ public void registerArrowStream(String name, Object arrow_array_stream) {
         long array_stream_address = getArrowStreamAddress(arrow_array_stream);
         DuckDBNative.duckdb_jdbc_arrow_register(conn_ref, array_stream_address, name.getBytes(StandardCharsets.UTF_8));
     }
+
+    void trackStatement(DuckDBPreparedStatement stmt) throws SQLException {
+        synchronized (this) {
+            if (isClosed()) {
+                throw new SQLException("Connection was closed");
+            }
+            this.statements.add(stmt);
+        }
+    }
+
+    void untrackStatement(DuckDBPreparedStatement stmt) throws SQLException {
+        synchronized (this) {
+            this.statements.remove(stmt);
+        }
+    }
 }
diff --git a/src/main/java/org/duckdb/DuckDBPreparedStatement.java b/src/main/java/org/duckdb/DuckDBPreparedStatement.java
@@ -1,6 +1,5 @@
 package org.duckdb;
 
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.Reader;
@@ -31,7 +30,6 @@
 import java.time.LocalDateTime;
 import java.time.OffsetDateTime;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Calendar;
 import java.util.List;
 import java.util.logging.Level;
@@ -61,16 +59,14 @@ public DuckDBPreparedStatement(DuckDBConnection conn) throws SQLException {
             throw new SQLException("connection parameter cannot be null");
         }
         this.conn = conn;
+        this.conn.trackStatement(this);
     }
 
     public DuckDBPreparedStatement(DuckDBConnection conn, String sql) throws SQLException {
-        if (conn == null) {
-            throw new SQLException("connection parameter cannot be null");
-        }
+        this(conn);
         if (sql == null) {
             throw new SQLException("sql query parameter cannot be null");
         }
-        this.conn = conn;
         this.isPreparedStatement = true;
         prepare(sql);
     }
@@ -299,15 +295,24 @@ public void clearParameters() throws SQLException {
 
     @Override
     public void close() throws SQLException {
-        if (select_result != null) {
-            select_result.close();
-            select_result = null;
+        if (conn == null) {
+            return;
         }
-        if (stmt_ref != null) {
-            DuckDBNative.duckdb_jdbc_release(stmt_ref);
-            stmt_ref = null;
+        synchronized (this) {
+            if (conn == null) {
+                return;
+            }
+            if (select_result != null) {
+                select_result.close();
+                select_result = null;
+            }
+            if (stmt_ref != null) {
+                DuckDBNative.duckdb_jdbc_release(stmt_ref);
+                stmt_ref = null;
+            }
+            conn.untrackStatement(this);
+            conn = null; // we use this as a check for closed-ness
         }
-        conn = null; // we use this as a check for closed-ness
     }
 
     protected void finalize() throws Throwable {
@@ -352,9 +357,14 @@ public void setQueryTimeout(int seconds) throws SQLException {
      * It is not safe to call this function when the connection is already closed.
      */
     @Override
-    public synchronized void cancel() throws SQLException {
-        if (conn.conn_ref != null) {
-            DuckDBNative.duckdb_jdbc_interrupt(conn.conn_ref);
+    public void cancel() throws SQLException {
+        synchronized (this) {
+            if (isClosed()) {
+                throw new SQLException("Statement was closed");
+            }
+            if (conn.conn_ref != null) {
+                DuckDBNative.duckdb_jdbc_interrupt(conn.conn_ref);
+            }
         }
     }
 
@@ -559,7 +569,12 @@ public int getResultSetHoldability() throws SQLException {
 
     @Override
     public boolean isClosed() throws SQLException {
-        return conn == null;
+        if (conn == null) {
+            return true;
+        }
+        synchronized (this) {
+            return conn == null;
+        }
     }
 
     @Override
@@ -576,7 +591,9 @@ public boolean isPoolable() throws SQLException {
     public void closeOnCompletion() throws SQLException {
         if (isClosed())
             throw new SQLException("Statement is closed");
-        closeOnCompletion = true;
+        synchronized (this) {
+            closeOnCompletion = true;
+        }
     }
 
     @Override
diff --git a/src/main/java/org/duckdb/DuckDBResultSet.java b/src/main/java/org/duckdb/DuckDBResultSet.java
@@ -31,7 +31,6 @@
 import java.time.LocalDateTime;
 import java.time.OffsetDateTime;
 import java.time.OffsetTime;
-import java.util.Arrays;
 import java.util.Calendar;
 import java.util.Map;
 import java.util.Objects;
@@ -92,12 +91,19 @@ public synchronized boolean next() throws SQLException {
         return true;
     }
 
-    public synchronized void close() throws SQLException {
-        if (result_ref != null) {
+    public void close() throws SQLException {
+        if (result_ref == null) {
+            return;
+        }
+        synchronized (this) {
+            if (result_ref == null) {
+                return;
+            }
             DuckDBNative.duckdb_jdbc_free_result(result_ref);
             // Nullness is used to determine whether we're closed
             result_ref = null;
-
+        }
+        synchronized (stmt) {
             // isCloseOnCompletion() throws if already closed, and we can't check for isClosed() because it could change
             // between when we check and call isCloseOnCompletion, so access the field directly.
             if (stmt.closeOnCompletion) {
@@ -110,8 +116,13 @@ protected void finalize() throws Throwable {
         close();
     }
 
-    public synchronized boolean isClosed() throws SQLException {
-        return result_ref == null;
+    public boolean isClosed() throws SQLException {
+        if (result_ref == null) {
+            return true;
+        }
+        synchronized (this) {
+            return result_ref == null;
+        }
     }
 
     private void check(int columnIndex) throws SQLException {
diff --git a/src/test/java/org/duckdb/TestDuckDBJDBC.java b/src/test/java/org/duckdb/TestDuckDBJDBC.java
@@ -20,6 +20,7 @@
 import static org.duckdb.test.Assertions.fail;
 import static org.duckdb.test.Runner.runTests;
 
+import java.io.File;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
@@ -4740,6 +4741,58 @@ public static void test_typed_connection_properties() throws Exception {
         }
     }
 
+    // https://github.com/duckdb/duckdb-java/issues/101
+    public static void test_unclosed_statement_does_not_hang() throws Exception {
+        String dbName = "test_issue_101.db";
+        String url = JDBC_URL + dbName;
+        Connection conn = DriverManager.getConnection(url);
+        Statement stmt = conn.createStatement();
+        stmt.execute("select 42");
+        // statement not closed explicitly
+        conn.close();
+        assertTrue(stmt.isClosed());
+        Connection connOther = DriverManager.getConnection(url);
+        connOther.close();
+        assertTrue(new File(dbName).delete());
+    }
+
+    public static void test_result_set_auto_closed() throws Exception {
+        try (Connection conn = DriverManager.getConnection(JDBC_URL)) {
+            Statement stmt = conn.createStatement();
+            ResultSet rs1 = stmt.executeQuery("select 42");
+            ResultSet rs2 = stmt.executeQuery("select 43");
+            assertTrue(rs1.isClosed());
+            stmt.close();
+            assertTrue(rs2.isClosed());
+        }
+    }
+
+    public static void test_statements_auto_closed_on_conn_close() throws Exception {
+        Connection conn = DriverManager.getConnection(JDBC_URL);
+        Statement stmt1 = conn.createStatement();
+        stmt1.execute("select 42");
+        PreparedStatement stmt2 = conn.prepareStatement("select 43");
+        stmt2.execute();
+        Statement stmt3 = conn.createStatement();
+        stmt3.execute("select 44");
+        stmt3.close();
+        conn.close();
+        assertTrue(stmt1.isClosed());
+        assertTrue(stmt2.isClosed());
+    }
+
+    public static void test_statement_auto_closed_on_completion() throws Exception {
+        try (Connection conn = DriverManager.getConnection(JDBC_URL)) {
+            Statement stmt = conn.createStatement();
+            stmt.closeOnCompletion();
+            assertTrue(stmt.isCloseOnCompletion());
+            try (ResultSet rs = stmt.executeQuery("select 42")) {
+                rs.next();
+            }
+            assertTrue(stmt.isClosed());
+        }
+    }
+
     public static void main(String[] args) throws Exception {
         System.exit(runTests(args, TestDuckDBJDBC.class, TestExtensionTypes.class));
     }
