Issue#2550 - Fixed getGeneratedKeys functionality for execute API (#2554

) * Issue#2550 - Fixed getGeneratedKeys functionality for execute API * Adapted the test for working with jdk8 * Fixed indenetation * Enable thre new TCGenKeys tests for AzureDW * Incorporated review comments. * Incorporated review comments * Incorporated review comments * Add a test for PreparedStatement * Adding a fix and test case for issue # 2587 * Added a new test for execute API with set no count
microsoft · Jan 30, 2025 · 25ea8da · 25ea8da
1 parent 03cfcfd
commit 25ea8da
Show file tree

Hide file tree

Showing 4 changed files with 419 additions and 12 deletions.
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
@@ -705,7 +705,7 @@ final void doExecutePreparedStatement(PrepStmtExecCmd command) throws SQLServerE
         if (EXECUTE_QUERY == executeMethod && null == resultSet) {
             SQLServerException.makeFromDriverError(connection, this, SQLServerException.getErrString("R_noResultset"),
                     null, true);
-        } else if (EXECUTE_UPDATE == executeMethod && null != resultSet) {
+        } else if ((EXECUTE_UPDATE == executeMethod) && (null != resultSet) && !bRequestedGeneratedKeys) {
             SQLServerException.makeFromDriverError(connection, this,
                     SQLServerException.getErrString("R_resultsetGeneratedForUpdate"), null, false);
         }

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java
@@ -1601,9 +1601,15 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                         if (null != procedureName)
                             return false;
 
+                        //For Insert, we must fetch additional TDS_DONE token that comes with the actual update count
+                        if ((StreamDone.CMD_INSERT == doneToken.getCurCmd()) &&  (-1 != doneToken.getUpdateCount()) && EXECUTE == executeMethod) {
+                            return true;
+                        }
+
                         // Always return all update counts from statements executed through Statement.execute()
-                        if (EXECUTE == executeMethod)
-                            return false;
+                        if (EXECUTE == executeMethod) {
+                           return false;
+                        }
 
                         // Statement.executeUpdate() may or may not return this update count depending on the
                         // setting of the lastUpdateCount connection property:
@@ -2357,17 +2363,27 @@ public final ResultSet getGeneratedKeys() throws SQLServerException {
 
         if (null == autoGeneratedKeys) {
             long orgUpd = updateCount;
+            //
+            //A case of SET NOCOUNT ON and GENERATED KEYS requested
+            //where we may not have received update count but would have already read the resultset
+            //so directly consume it.
+            //
+            if ((executeMethod != EXECUTE_QUERY) && bRequestedGeneratedKeys && (resultSet != null)) {
+                autoGeneratedKeys = resultSet;
+                updateCount = orgUpd;
+            } else {
 
-            // Generated keys are returned in a ResultSet result right after the update count.
-            // Try to get that ResultSet. If there are no more results after the update count,
-            // or if the next result isn't a ResultSet, then something is wrong.
-            if (!getNextResult(true) || null == resultSet) {
-                SQLServerException.makeFromDriverError(connection, this,
-                        SQLServerException.getErrString("R_statementMustBeExecuted"), null, false);
+                // Generated keys are returned in a ResultSet result right after the update count.
+                // Try to get that ResultSet. If there are no more results after the update count,
+                // or if the next result isn't a ResultSet, then something is wrong.
+                if (!getNextResult(true) || null == resultSet) {
+                    SQLServerException.makeFromDriverError(connection, this,
+                            SQLServerException.getErrString("R_statementMustBeExecuted"), null, false);
+                }
+                autoGeneratedKeys = resultSet;
+                updateCount = orgUpd;
             }
 
-            autoGeneratedKeys = resultSet;
-            updateCount = orgUpd;
         }
         loggerExternal.exiting(getClassNameLogging(), "getGeneratedKeys", autoGeneratedKeys);
         return autoGeneratedKeys;
@@ -2616,6 +2632,11 @@ SQLServerColumnEncryptionKeyStoreProvider getColumnEncryptionKeyStoreProvider(
             lock.unlock();
         }
     }
+
+    protected void setAutoGeneratedKey(SQLServerResultSet rs) {
+        autoGeneratedKeys = rs;
+    }
+
 }
 
 

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/StreamDone.java b/src/main/java/com/microsoft/sqlserver/jdbc/StreamDone.java
@@ -232,7 +232,7 @@ final long getUpdateCount() {
     }
 
     final boolean cmdIsDMLOrDDL() {
-        switch (curCmd) {
+    	switch (curCmd) {
             case CMD_INSERT:
             case CMD_BULKINSERT:
             case CMD_DELETE: