Skip to content

Commit 5445cd4

Browse files
Ananya2Ananya2
committed
Fix update count handling for multi-statement queries executed via PreparedStatement.
1 parent 1ed161f commit 5445cd4

File tree

3 files changed

+271
-5
lines changed

3 files changed

+271
-5
lines changed

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -788,6 +788,17 @@ boolean onRetValue(TDSReader tdsReader) throws SQLServerException {
788788
return false;
789789
}
790790

791+
/**
792+
* Override TDS token processing behavior for PreparedStatement.
793+
* For regular Statement, the execute API for INSERT requires reading an additional explicit
794+
* TDS_DONE token that contains the actual update count returned by the server.
795+
* PreparedStatement does not require this additional token processing.
796+
*/
797+
@Override
798+
protected boolean hasUpdateCountTDSToken(StreamDone doneToken, int executeMethod) {
799+
return false;
800+
}
801+
791802
/**
792803
* Sends the statement parameters by RPC.
793804
*/

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

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1601,9 +1601,8 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
16011601
if (null != procedureName)
16021602
return false;
16031603

1604-
// For Insert, we must fetch additional TDS_DONE token that comes with the actual update count
1605-
if ((StreamDone.CMD_INSERT == doneToken.getCurCmd()) && (-1 != doneToken.getUpdateCount())
1606-
&& EXECUTE == executeMethod) {
1604+
// For Insert operations, check if additional TDS_DONE token processing is required.
1605+
if (hasUpdateCountTDSToken(doneToken, executeMethod)) {
16071606
return true;
16081607
}
16091608

@@ -1845,6 +1844,24 @@ boolean consumeExecOutParam(TDSReader tdsReader) throws SQLServerException {
18451844
return false;
18461845
}
18471846

1847+
/**
1848+
* Determines whether to continue processing additional TDS_DONE tokens for INSERT statements.
1849+
* For INSERT operations, the driver must fetch an additional TDS_DONE token that contains
1850+
* the actual update count. This method can be overridden by subclasses to customize
1851+
* TDS token processing behavior.
1852+
*
1853+
* @param doneToken The current DONE token being processed
1854+
* @param executeMethod The execution method used
1855+
* @return true to continue processing more tokens to get the actual update count,
1856+
* false to stop and return this token
1857+
*/
1858+
protected boolean hasUpdateCountTDSToken(StreamDone doneToken, int executeMethod) {
1859+
// For Insert, we must fetch additional TDS_DONE token that comes with the actual update count
1860+
return (StreamDone.CMD_INSERT == doneToken.getCurCmd()) &&
1861+
(-1 != doneToken.getUpdateCount()) &&
1862+
(EXECUTE == executeMethod);
1863+
}
1864+
18481865
// --------------------------JDBC 2.0-----------------------------
18491866

18501867
@Override

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

Lines changed: 240 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2961,7 +2961,7 @@ public void testExecuteInsertManyRowsAndSelect() {
29612961
// no more results
29622962
break;
29632963
} else {
2964-
assertEquals(count, 3, "update count should have been 6");
2964+
assertEquals(count, 3, "update count should have been 3");
29652965
}
29662966
} else {
29672967
// process ResultSet
@@ -2998,7 +2998,7 @@ public void testExecuteTwoInsertsRowsAndSelect() {
29982998
// no more results
29992999
break;
30003000
} else {
3001-
assertEquals(count, 1, "update count should have been 2");
3001+
assertEquals(count, 1, "update count should have been 1");
30023002
}
30033003
} else {
30043004
// process ResultSet
@@ -3091,6 +3091,244 @@ public void testExecuteDelAndSelect() {
30913091
}
30923092
}
30933093

3094+
/**
3095+
* Tests multi-statement PreparedStatement with loop to process all results
3096+
*
3097+
* @throws Exception
3098+
*/
3099+
@Test
3100+
public void testMultiStatementPreparedStatementLoopResults() {
3101+
try (Connection con = getConnection()) {
3102+
try (PreparedStatement ps = con.prepareStatement("DELETE FROM " + tableName + " " +
3103+
"INSERT INTO " + tableName + " (NAME) VALUES (?) " +
3104+
"INSERT INTO " + tableName + " (NAME) VALUES (?) " +
3105+
"UPDATE " + tableName + " SET NAME = 'updated' " +
3106+
"INSERT INTO " + tableName + " (NAME) VALUES (?) " +
3107+
"INSERT INTO " + tableName + " (NAME) VALUES (?) " +
3108+
"SELECT * FROM " + tableName)) {
3109+
3110+
ps.setString(1, "test1");
3111+
ps.setString(2, "test2");
3112+
ps.setString(3, "test3");
3113+
ps.setString(4, "test4");
3114+
3115+
boolean retval = ps.execute();
3116+
do {
3117+
if (!retval) {
3118+
int count = ps.getUpdateCount();
3119+
if (count == -1) {
3120+
// no more results
3121+
break;
3122+
} else {
3123+
assertTrue(count >= 0, "update count should be non-negative: " + count);
3124+
}
3125+
} else {
3126+
// process ResultSet
3127+
try (ResultSet rs = ps.getResultSet()) {
3128+
int rowCount = 0;
3129+
while (rs.next()) {
3130+
String name = rs.getString("NAME");
3131+
assertTrue(name != null, "name should not be null");
3132+
rowCount++;
3133+
}
3134+
assertEquals(4, rowCount, "Expected 4 rows in result set");
3135+
}
3136+
}
3137+
retval = ps.getMoreResults();
3138+
} while (true);
3139+
}
3140+
} catch (SQLException e) {
3141+
fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
3142+
}
3143+
}
3144+
3145+
/**
3146+
* Tests PreparedStatement execute for Insert followed by select
3147+
*
3148+
* @throws Exception
3149+
*/
3150+
@Test
3151+
public void testPreparedStatementExecuteInsertAndSelect() {
3152+
try (Connection con = getConnection()) {
3153+
try (PreparedStatement ps = con.prepareStatement("INSERT INTO " + tableName + " (NAME) VALUES(?) SELECT NAME FROM " + tableName + " WHERE ID = 1")) {
3154+
ps.setString(1, "test");
3155+
boolean retval = ps.execute();
3156+
do {
3157+
if (!retval) {
3158+
int count = ps.getUpdateCount();
3159+
if (count == -1) {
3160+
// no more results
3161+
break;
3162+
} else {
3163+
assertEquals(count, 1, "update count should have been 1");
3164+
}
3165+
} else {
3166+
// process ResultSet
3167+
try (ResultSet rs = ps.getResultSet()) {
3168+
if (rs.next()) {
3169+
String val = rs.getString(1);
3170+
assertEquals(val, "test", "read value should have been 'test'");
3171+
}
3172+
}
3173+
}
3174+
retval = ps.getMoreResults();
3175+
} while (true);
3176+
}
3177+
} catch (SQLException e) {
3178+
fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
3179+
}
3180+
}
3181+
3182+
/**
3183+
* Tests PreparedStatement execute for Merge followed by select
3184+
*
3185+
* @throws Exception
3186+
*/
3187+
@Test
3188+
public void testPreparedStatementExecuteMergeAndSelect() {
3189+
try (Connection con = getConnection()) {
3190+
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")) {
3191+
ps.setString(1, "test1");
3192+
ps.setString(2, "test1");
3193+
boolean retval = ps.execute();
3194+
do {
3195+
if (!retval) {
3196+
int count = ps.getUpdateCount();
3197+
if (count == -1) {
3198+
// no more results
3199+
break;
3200+
} else {
3201+
assertEquals(count, 1, "update count should have been 1");
3202+
}
3203+
} else {
3204+
// process ResultSet
3205+
try (ResultSet rs = ps.getResultSet()) {
3206+
if (rs.next()) {
3207+
String val = rs.getString(1);
3208+
assertEquals(val, "test", "read value should have been 'test'");
3209+
}
3210+
}
3211+
}
3212+
retval = ps.getMoreResults();
3213+
} while (true);
3214+
}
3215+
} catch (SQLException e) {
3216+
fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
3217+
}
3218+
}
3219+
3220+
/**
3221+
* Tests PreparedStatement execute two Inserts followed by select
3222+
*
3223+
* @throws Exception
3224+
*/
3225+
@Test
3226+
public void testPreparedStatementExecuteTwoInsertsRowsAndSelect() {
3227+
try (Connection con = getConnection()) {
3228+
try (PreparedStatement ps = con.prepareStatement("INSERT INTO " + tableName + " (NAME) VALUES(?) INSERT INTO " + tableName + " (NAME) VALUES(?) SELECT NAME from " + tableName + " WHERE ID = 1")) {
3229+
ps.setString(1, "test");
3230+
ps.setString(2, "test");
3231+
boolean retval = ps.execute();
3232+
do {
3233+
if (!retval) {
3234+
int count = ps.getUpdateCount();
3235+
if (count == -1) {
3236+
// no more results
3237+
break;
3238+
} else {
3239+
assertEquals(count, 1, "update count should have been 1");
3240+
}
3241+
} else {
3242+
// process ResultSet
3243+
try (ResultSet rs = ps.getResultSet()) {
3244+
if (rs.next()) {
3245+
String val = rs.getString(1);
3246+
assertEquals(val, "test", "read value should have been 'test'");
3247+
}
3248+
}
3249+
}
3250+
retval = ps.getMoreResults();
3251+
} while (true);
3252+
}
3253+
} catch (SQLException e) {
3254+
fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
3255+
}
3256+
}
3257+
3258+
/**
3259+
* Tests PreparedStatement execute for Update followed by select
3260+
*
3261+
* @throws Exception
3262+
*/
3263+
@Test
3264+
public void testPreparedStatementExecuteUpdAndSelect() {
3265+
try (Connection con = getConnection()) {
3266+
try (PreparedStatement ps = con.prepareStatement("UPDATE " + tableName + " SET NAME = ? SELECT NAME FROM " + tableName + " WHERE ID = 1")) {
3267+
ps.setString(1, "test");
3268+
boolean retval = ps.execute();
3269+
do {
3270+
if (!retval) {
3271+
int count = ps.getUpdateCount();
3272+
if (count == -1) {
3273+
// no more results
3274+
break;
3275+
} else {
3276+
assertEquals(count, 3, "update count should have been 3");
3277+
}
3278+
} else {
3279+
// process ResultSet
3280+
try (ResultSet rs = ps.getResultSet()) {
3281+
if (rs.next()) {
3282+
String val = rs.getString(1);
3283+
assertEquals(val, "test", "read value should have been 'test'");
3284+
}
3285+
}
3286+
}
3287+
retval = ps.getMoreResults();
3288+
} while (true);
3289+
}
3290+
} catch (SQLException e) {
3291+
fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
3292+
}
3293+
}
3294+
3295+
/**
3296+
* Tests PreparedStatement execute for Delete followed by select
3297+
*
3298+
* @throws Exception
3299+
*/
3300+
@Test
3301+
public void testPreparedStatementExecuteDelAndSelect() {
3302+
try (Connection con = getConnection()) {
3303+
try (PreparedStatement ps = con.prepareStatement("DELETE FROM " + tableName + " WHERE ID = ? SELECT NAME FROM " + tableName + " WHERE ID = 2")) {
3304+
ps.setInt(1, 1);
3305+
boolean retval = ps.execute();
3306+
do {
3307+
if (!retval) {
3308+
int count = ps.getUpdateCount();
3309+
if (count == -1) {
3310+
// no more results
3311+
break;
3312+
} else {
3313+
assertEquals(count, 1, "update count should have been 1");
3314+
}
3315+
} else {
3316+
// process ResultSet
3317+
try (ResultSet rs = ps.getResultSet()) {
3318+
if (rs.next()) {
3319+
String val = rs.getString(1);
3320+
assertEquals(val, "test", "read value should have been 'test'");
3321+
}
3322+
}
3323+
}
3324+
retval = ps.getMoreResults();
3325+
} while (true);
3326+
}
3327+
} catch (SQLException e) {
3328+
fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
3329+
}
3330+
}
3331+
30943332
@AfterEach
30953333
public void terminate() {
30963334
try (Connection con = getConnection(); Statement stmt = con.createStatement()) {

0 commit comments

Comments
 (0)