Close pending statements on connection close

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 #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: #101

Author: staticlibs

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);
+        }
+    }
 }

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

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 {

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));
     }