From 346525e5f0c8877dd5d42438b23a7936f18de013 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EA=B9=80=EC=84=A0=EC=9A=B0?= <seonwoo960000@toss.im>
Date: Sun, 24 Aug 2025 09:25:59 +0900
Subject: [PATCH] Update test

---
 .../java/tech/turso/jdbc4/JDBC4Statement.java | 44 ++++++++++++++-----
 .../tech/turso/jdbc4/JDBC4StatementTest.java  |  4 +-
 2 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/bindings/java/src/main/java/tech/turso/jdbc4/JDBC4Statement.java b/bindings/java/src/main/java/tech/turso/jdbc4/JDBC4Statement.java
index 67225aac5..2230b1f96 100644
--- a/bindings/java/src/main/java/tech/turso/jdbc4/JDBC4Statement.java
+++ b/bindings/java/src/main/java/tech/turso/jdbc4/JDBC4Statement.java
@@ -4,15 +4,14 @@ import static java.util.Objects.requireNonNull;
 
 import java.sql.BatchUpdateException;
 import java.sql.Connection;
-import java.sql.DatabaseMetaData;
 import java.sql.ResultSet;
 import java.sql.SQLException;
-import java.sql.SQLTimeoutException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.locks.ReentrantLock;
+
 import tech.turso.annotations.Nullable;
 import tech.turso.annotations.SkipNullableCheck;
 import tech.turso.core.TursoResultSet;
@@ -268,13 +267,9 @@ public class JDBC4Statement implements Statement {
     for (int i = 0; i < batchCommands.size(); i++) {
       String sql = batchCommands.get(i);
       try {
-        // Check if the statement returns a ResultSet (SELECT statements)
-        // In batch processing, SELECT statements should throw an exception
-        if (execute(sql)) {
-          // This means the statement returned a ResultSet, which is not allowed in batch
+        if (!isBatchCompatibleStatement(sql)) {
           failedCommands.add(sql);
           updateCounts[i] = EXECUTE_FAILED;
-          // Create a BatchUpdateException for the failed command
           BatchUpdateException bue =
               new BatchUpdateException(
                   "Batch entry "
@@ -289,12 +284,12 @@ public class JDBC4Statement implements Statement {
           // Clear the batch after failure
           clearBatch();
           throw bue;
-        } else {
-          // For DML statements, get the update count
-          updateCounts[i] = getUpdateCount();
         }
+
+        execute(sql);
+        // For DML statements, get the update count
+        updateCounts[i] = getUpdateCount();
       } catch (SQLException e) {
-        // Handle SQL exceptions during batch execution
         failedCommands.add(sql);
         updateCounts[i] = EXECUTE_FAILED;
 
@@ -317,6 +312,33 @@ public class JDBC4Statement implements Statement {
     return updateCounts;
   }
 
+  /**
+   * Checks if a SQL statement is compatible with batch execution. Only INSERT, UPDATE, DELETE, and
+   * DDL statements are allowed in batch. SELECT and other query statements are not allowed.
+   *
+   * @param sql The SQL statement to check
+   * @return true if the statement is batch-compatible, false otherwise
+   */
+  private boolean isBatchCompatibleStatement(String sql) {
+    if (sql == null || sql.trim().isEmpty()) {
+      return false;
+    }
+
+    // Trim and convert to uppercase for case-insensitive comparison
+    String trimmedSql = sql.trim().toUpperCase();
+
+    // Check if it starts with batch-compatible keywords
+    return trimmedSql.startsWith("INSERT")
+        || trimmedSql.startsWith("UPDATE")
+        || trimmedSql.startsWith("DELETE")
+        || trimmedSql.startsWith("CREATE")
+        || trimmedSql.startsWith("DROP")
+        || trimmedSql.startsWith("ALTER")
+        || trimmedSql.startsWith("TRUNCATE")
+        || trimmedSql.startsWith("REPLACE")
+        || trimmedSql.startsWith("MERGE");
+  }
+
   @Override
   public Connection getConnection() {
     return connection;
diff --git a/bindings/java/src/test/java/tech/turso/jdbc4/JDBC4StatementTest.java b/bindings/java/src/test/java/tech/turso/jdbc4/JDBC4StatementTest.java
index 53572d3be..4fe1d32ed 100644
--- a/bindings/java/src/test/java/tech/turso/jdbc4/JDBC4StatementTest.java
+++ b/bindings/java/src/test/java/tech/turso/jdbc4/JDBC4StatementTest.java
@@ -176,7 +176,7 @@ class JDBC4StatementTest {
         "INSERT INTO batch_test VALUES (1, 'initial1'), (2, 'initial2'), (3, 'initial3');");
 
     // Add batch commands with different operations
-    stmt.addBatch("UPDATE batch_test SET value = 'updated' WHERE id = 1;");
+    stmt.addBatch("UPDATE batch_test SET value = 'updated';");
     stmt.addBatch("DELETE FROM batch_test WHERE id = 2;");
     stmt.addBatch("INSERT INTO batch_test VALUES (4, 'new');");
 
@@ -185,7 +185,7 @@ class JDBC4StatementTest {
 
     // Verify update counts
     assertThat(updateCounts).hasSize(3);
-    assertThat(updateCounts[0]).isEqualTo(1); // UPDATE affected 1 row
+    assertThat(updateCounts[0]).isEqualTo(3); // UPDATE affected 3 row
     assertThat(updateCounts[1]).isEqualTo(1); // DELETE affected 1 row
     assertThat(updateCounts[2]).isEqualTo(1); // INSERT affected 1 row