Addressed PR review comments.

Author: Ananya2

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java

@@ -795,7 +795,7 @@ boolean onRetValue(TDSReader tdsReader) throws SQLServerException {
      * PreparedStatement does not require this additional token processing.
      */
     @Override
-    protected boolean hasUpdateCountTDSToken(StreamDone doneToken, int executeMethod) {
+    protected boolean hasUpdateCountTDSTokenForInsertCmd() {
         return false;
     }
 

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java

@@ -1602,7 +1602,8 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                             return false;
 
                         // For Insert operations, check if additional TDS_DONE token processing is required.
-                        if (hasUpdateCountTDSToken(doneToken, executeMethod)) {
+                        if (hasUpdateCountTDSTokenForInsertCmd() && (StreamDone.CMD_INSERT == doneToken.getCurCmd()) && (-1 != doneToken.getUpdateCount())
+                                && EXECUTE == executeMethod) {
                             return true;
                         }
 
@@ -1846,20 +1847,15 @@ boolean consumeExecOutParam(TDSReader tdsReader) throws SQLServerException {
 
     /**
      * Determines whether to continue processing additional TDS_DONE tokens for INSERT statements.
-     * For INSERT operations, the driver must fetch an additional TDS_DONE token that contains
+     * For INSERT operations, regular Statement requires reading an additional TDS_DONE token that contains
      * the actual update count. This method can be overridden by subclasses to customize 
      * TDS token processing behavior.
      * 
-     * @param doneToken The current DONE token being processed
-     * @param executeMethod The execution method used
-     * @return true to continue processing more tokens to get the actual update count, 
-     *         false to stop and return this token
+     * @return true to continue processing more tokens to get the actual update count for INSERT operations
      */
-    protected boolean hasUpdateCountTDSToken(StreamDone doneToken, int executeMethod) {
+    protected boolean hasUpdateCountTDSTokenForInsertCmd() {
         // For Insert, we must fetch additional TDS_DONE token that comes with the actual update count
-        return (StreamDone.CMD_INSERT == doneToken.getCurCmd()) &&
-               (-1 != doneToken.getUpdateCount()) &&
-               (EXECUTE == executeMethod);
+        return true;
     }
 
     // --------------------------JDBC 2.0-----------------------------

src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/StatementTest.java

@@ -3150,7 +3150,9 @@ public void testMultiStatementPreparedStatementLoopResults() {
         @Test
         public void testPreparedStatementExecuteInsertAndSelect() {
             try (Connection con = getConnection()) {
-                try (PreparedStatement ps = con.prepareStatement("INSERT INTO " + tableName + " (NAME) VALUES(?) SELECT NAME FROM " + tableName + " WHERE ID = 1")) {
+                String sql = "INSERT INTO " + tableName + " (NAME) VALUES(?) " +
+                            "SELECT NAME FROM " + tableName + " WHERE ID = 1";
+                try (PreparedStatement ps = con.prepareStatement(sql)) {
                     ps.setString(1, "test");
                     boolean retval = ps.execute();
                     do {
@@ -3187,7 +3189,12 @@ public void testPreparedStatementExecuteInsertAndSelect() {
         @Test
         public void testPreparedStatementExecuteMergeAndSelect() {
             try (Connection con = getConnection()) {
-                try (PreparedStatement ps = con.prepareStatement("MERGE INTO " + tableName + " AS target USING (VALUES (?)) AS source (name) ON target.name = source.name WHEN NOT MATCHED THEN INSERT (name) VALUES (?); SELECT NAME FROM " + tableName + " WHERE ID = 1")) {
+                String sql = "MERGE INTO " + tableName + " AS target " +
+                            "USING (VALUES (?)) AS source (name) " +
+                            "ON target.name = source.name " +
+                            "WHEN NOT MATCHED THEN INSERT (name) VALUES (?); " +
+                            "SELECT NAME FROM " + tableName + " WHERE ID = 1";
+                try (PreparedStatement ps = con.prepareStatement(sql)) {
                     ps.setString(1, "test1");
                     ps.setString(2, "test1");
                     boolean retval = ps.execute();