From f1e91cbe13b906671f8bfc2a7c64bca03976705b Mon Sep 17 00:00:00 2001
From: Hishida Masato <hishida@nautilus-technologies.com>
Date: Tue, 21 Jan 2025 09:44:45 +0900
Subject: [PATCH] fix(tgsql): Delay close of PreparedStatement

---
 .../tgsql/core/executor/sql/BasicSqlProcessor.java     |  6 ++++--
 .../core/executor/sql/PreparedStatementResult.java     |  8 ++++++--
 .../tgsql/core/executor/engine/BasicEngineTest.java    | 10 +++++-----
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/BasicSqlProcessor.java b/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/BasicSqlProcessor.java
index 84bf808..ebb4d7f 100644
--- a/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/BasicSqlProcessor.java
+++ b/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/BasicSqlProcessor.java
@@ -38,6 +38,7 @@
 import com.tsurugidb.tsubakuro.sql.StatementMetadata;
 import com.tsurugidb.tsubakuro.sql.TableMetadata;
 import com.tsurugidb.tsubakuro.sql.exception.TargetNotFoundException;
+import com.tsurugidb.tsubakuro.util.Owner;
 
 /**
  * A basic implementation of {@link SqlProcessor}.
@@ -190,12 +191,13 @@ public void rollbackTransaction() throws ServerException, IOException, Interrupt
         LOG.debug("start prepare: '{}'", statement);
         desireActive();
         var client = getSqlClient();
-        try (var prepared = client.prepare(statement).await()) {
+        try (var preparedOwner = Owner.of(client.prepare(statement).await())) {
+            var prepared = preparedOwner.get();
             var t = transaction.getTransaction();
             if (prepared.hasResultRecords()) {
                 LOG.debug("start query: '{}'", statement);
                 var result = t.executeQuery(prepared).await();
-                return new PreparedStatementResult(result);
+                return new PreparedStatementResult(result, preparedOwner.release());
             }
             LOG.debug("start execute: '{}'", statement);
             var result = t.executeStatement(prepared).await();
diff --git a/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/PreparedStatementResult.java b/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/PreparedStatementResult.java
index 95a5b7c..fdb014b 100755
--- a/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/PreparedStatementResult.java
+++ b/modules/tgsql/core/src/main/java/com/tsurugidb/tgsql/core/executor/sql/PreparedStatementResult.java
@@ -19,21 +19,25 @@
 
 import com.tsurugidb.tsubakuro.exception.ServerException;
 import com.tsurugidb.tsubakuro.sql.ExecuteResult;
+import com.tsurugidb.tsubakuro.sql.PreparedStatement;
 import com.tsurugidb.tsubakuro.sql.ResultSet;
 
 public class PreparedStatementResult implements AutoCloseable {
 
     private final ResultSet resultSet;
     private final ExecuteResult executeResult;
+    private final PreparedStatement preparedStatement;
 
-    public PreparedStatementResult(ResultSet resultSet) {
+    public PreparedStatementResult(ResultSet resultSet, PreparedStatement preparedStatement) {
         this.resultSet = resultSet;
         this.executeResult = null;
+        this.preparedStatement = preparedStatement;
     }
 
     public PreparedStatementResult(ExecuteResult executeResult) {
         this.resultSet = null;
         this.executeResult = executeResult;
+        this.preparedStatement = null;
     }
 
     public ResultSet getResultSet() {
@@ -46,7 +50,7 @@ public ExecuteResult getExecuteResult() {
 
     @Override
     public void close() throws ServerException, IOException, InterruptedException {
-        try (resultSet) {
+        try (preparedStatement; resultSet) {
             // close only
         }
     }
diff --git a/modules/tgsql/core/src/test/java/com/tsurugidb/tgsql/core/executor/engine/BasicEngineTest.java b/modules/tgsql/core/src/test/java/com/tsurugidb/tgsql/core/executor/engine/BasicEngineTest.java
index 522c100..43e5c66 100644
--- a/modules/tgsql/core/src/test/java/com/tsurugidb/tgsql/core/executor/engine/BasicEngineTest.java
+++ b/modules/tgsql/core/src/test/java/com/tsurugidb/tgsql/core/executor/engine/BasicEngineTest.java
@@ -169,7 +169,7 @@ void empty_statement() throws Exception {
     }
 
     @Test
-    void generic_staement_wo_result() throws Exception {
+    void generic_statement_wo_result() throws Exception {
         var reached = new AtomicBoolean();
         MockSqlProcessor sql = new MockSqlProcessor(true) {
             @Override
@@ -195,7 +195,7 @@ public Map<CounterType, Long> getCounters() {
     }
 
     @Test
-    void generic_staement_w_result() throws Exception {
+    void generic_statement_w_result() throws Exception {
         var reachedExec = new AtomicBoolean();
         MockSqlProcessor sql = new MockSqlProcessor(true) {
             @Override
@@ -205,7 +205,7 @@ public PreparedStatementResult execute(String statement, Region region) {
                 }
                 assertEquals("SELECT * FROM T", statement);
                 var rs = Relation.of(new Object[][] { { 1 } }).getResultSet(new ResultSetMetadataAdapter(SqlResponse.ResultSetMetadata.newBuilder().addColumns(Types.column(int.class)).build()));
-                return new PreparedStatementResult(rs);
+                return new PreparedStatementResult(rs, null);
             }
         };
         var reachedRs = new AtomicBoolean();
@@ -231,7 +231,7 @@ public long process(ResultSet target) throws ServerException, IOException, Inter
     }
 
     @Test
-    void call_staement_fall_through() throws Exception {
+    void call_statement_fall_through() throws Exception {
         var reached = new AtomicBoolean();
         MockSqlProcessor sql = new MockSqlProcessor(true) {
             @Override
@@ -257,7 +257,7 @@ public Map<CounterType, Long> getCounters() {
     }
 
     @Test
-    void generic_staement_inactive_tx() throws Exception {
+    void generic_statement_inactive_tx() throws Exception {
         MockSqlProcessor sql = new MockSqlProcessor(false);
         MockResultProcessor rs = new MockResultProcessor();
         var engine = newBasicEngine(sql, rs);