Repository: microsoft/mssql-jdbc
Pull request #2085 (2023-02-28 15:56:23)

Fix shared timer race #2084


Patch:
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
index a15a9e68..c8bc87ab 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -7953,16 +7953,19 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
             return supportsTransactions;
         }
 
+        supportsTransactions = true;
+
         try {
             this.connectionCommand("SELECT @@TRANCOUNT", "SQLServerConnection.supportsTransactions");
         } catch (SQLServerException e) {
             if (e.getMessage().equals(SQLServerException.getErrString("R_transactionsNotSupported"))) {
-                return Boolean.FALSE.equals(supportsTransactions);
+                supportsTransactions = false;
+                return false;
             }
             throw e;
         }
 
-        return Boolean.TRUE.equals(supportsTransactions);
+        return supportsTransactions;
     }
 
     /**
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
index 4743407b..b6e789b5 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
@@ -26,6 +26,7 @@ import java.util.ArrayList;
 import java.util.Calendar;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Vector;
 import java.util.logging.Level;
 
@@ -2087,6 +2088,7 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                             SQLServerResultSet rs = stmt
                                     .executeQueryInternal("sp_executesql N'SET FMTONLY ON SELECT * FROM "
                                             + Util.escapeSingleQuotes(tableName) + " '")) {
+                        Map<Integer, Integer> columnMappings = null;
                         if (null != columnList && !columnList.isEmpty()) {
                             if (columnList.size() != valueList.size()) {
 
@@ -2095,6 +2097,7 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                                 Object[] msgArgs = {columnList.size(), valueList.size()};
                                 throw new IllegalArgumentException(form.format(msgArgs));
                             }
+                            columnMappings = new HashMap<>(columnList.size());
                         } else {
                             if (rs.getColumnCount() != valueList.size()) {
                                 MessageFormat form = new MessageFormat(
@@ -2118,8 +2121,17 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                             } else {
                                 jdbctype = ti.getSSType().getJDBCType().getIntValue();
                             }
-                            batchRecord.addColumnMetadata(i, c.getColumnName(), jdbctype, ti.getPrecision(),
-                                    ti.getScale());
+                            if (null != columnList && columnList.size() > 0) {
+                                int columnIndex = columnList.indexOf(c.getColumnName());
+                                if (columnIndex > -1) {
+                                    columnMappings.put(columnIndex + 1, i);
+                                    batchRecord.addColumnMetadata(columnIndex + 1, c.getColumnName(), jdbctype,
+                                            ti.getPrecision(), ti.getScale());
+                                }
+                            } else {
+                                batchRecord.addColumnMetadata(i, c.getColumnName(), jdbctype, ti.getPrecision(),
+                                        ti.getScale());
+                            }
                         }
 
                         SQLServerBulkCopy bcOperation = new SQLServerBulkCopy(connection);
@@ -2127,6 +2139,11 @@ public class SQLServerPreparedStatement extends SQLServerStatement implements IS
                         option.setBulkCopyTimeout(queryTimeout);
                         bcOperation.setBulkCopyOptions(option);
                         bcOperation.setDestinationTableName(tableName);
+                        if (columnMappings != null) {
+                            for (Entry<Integer, Integer> pair : columnMappings.entrySet()) {
+                                bcOperation.addColumnMapping(pair.getKey(), pair.getValue());
+                            }
+                        }
                         bcOperation.setStmtColumnEncriptionSetting(this.getStmtColumnEncriptionSetting());
                         bcOperation.setDestinationTableMetadata(rs);
                         bcOperation.writeToServer(batchRecord);
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
index c936a572..d4634ae8 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResource.java
@@ -532,6 +532,7 @@ public final class SQLServerResource extends ListResourceBundle {
         {"R_ALPNFailed", "Failed to negotiate Application-Layer Protocol {0}. Server returned: {1}."},
         {"R_serverError", "An error occurred during the current command (Done status {0}). {1}"},
         {"R_ManagedIdentityTokenAcquisitionFail", "Failed to acquire managed identity token. Request for the token succeeded, but no token was returned. The token is null."},
+        {"R_AmbiguousRowUpdate", "Failed to execute updateRow(). The update is attempting an ambiguous update on tables \"{0}\" and \"{1}\". Ensure all columns being updated prior to the updateRow() call belong to the same table."},
         {"R_InvalidSqlQuery", "Invalid SQL Query: {0}"}
     };
 }
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
index cc82f40d..1cb5167a 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
@@ -4893,7 +4893,7 @@ public class SQLServerResultSet implements ISQLServerResultSet, java.io.Serializ
         tdsWriter.writeRPCInt(null, serverCursorId, false);
         tdsWriter.writeRPCInt(null, TDS.SP_CURSOR_OP_UPDATE | TDS.SP_CURSOR_OP_SETPOSITION, false);
         tdsWriter.writeRPCInt(null, fetchBufferGetRow(), false);
-        tdsWriter.writeRPCStringUnicode("");
+        tdsWriter.writeRPCStringUnicode(getUpdatedColumnTableName());
 
         assert hasUpdatedColumns();
 
@@ -4912,6 +4912,23 @@ public class SQLServerResultSet implements ISQLServerResultSet, java.io.Serializ
         return false;
     }
 
+    private String getUpdatedColumnTableName() throws SQLServerException {
+        String columnTableName = "";
+
+        for (Column column : columns) {
+            if (column.hasUpdates() && columnTableName.isEmpty()) {
+                columnTableName = column.getTableName().asEscapedString();
+            } else if (column.hasUpdates()
+                    && !columnTableName.equals(column.getTableName().asEscapedString())) {
+                MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_AmbiguousRowUpdate"));
+                Object[] msgArgs = {columnTableName, column.getTableName().asEscapedString()};
+                SQLServerException.makeFromDriverError(stmt.connection, stmt, form.format(msgArgs), null, false);
+            }
+        }
+
+        return columnTableName;
+    }
+
     @Override
     public void deleteRow() throws SQLException {
         loggerExternal.entering(getClassNameLogging(), "deleteRow");
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerXAResource.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerXAResource.java
index a4ca6bbd..8829072d 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerXAResource.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerXAResource.java
@@ -24,6 +24,7 @@ import javax.transaction.xa.XAException;
 import javax.transaction.xa.XAResource;
 import javax.transaction.xa.Xid;
 
+import com.microsoft.sqlserver.jdbc.SQLServerError.TransientError;
 
 /**
  * Implements Transaction id used to recover transactions.
@@ -751,9 +752,24 @@ public final class SQLServerXAResource implements javax.transaction.xa.XAResourc
                     }
                 }
             }
-        } catch (SQLServerException | SQLTimeoutException ex) {
+        } catch (SQLTimeoutException ex) {
             if (xaLogger.isLoggable(Level.FINER))
                 xaLogger.finer(toString() + " exception:" + ex);
+            XAException e = new XAException(ex.toString());
+            e.errorCode = XAException.XAER_RMFAIL;
+            throw e;
+
+        } catch (SQLServerException ex) {
+            if (xaLogger.isLoggable(Level.FINER))
+                xaLogger.finer(toString() + " exception:" + ex);
+
+            if (ex.getMessage().equals(SQLServerException.getErrString("R_noServerResponse"))
+                    || TransientError.isTransientError(ex.getSQLServerError())) {
+                XAException e = new XAException(ex.toString());
+                e.errorCode = XAException.XAER_RMFAIL;
+                throw e;
+            }
+
             XAException e = new XAException(ex.toString());
             e.errorCode = XAException.XAER_RMERR;
             throw e;
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SharedTimer.java b/src/main/java/com/microsoft/sqlserver/jdbc/SharedTimer.java
index 986a0bbb..e46e0c83 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/SharedTimer.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/SharedTimer.java
@@ -111,20 +111,16 @@ class SharedTimer implements Serializable {
      * When the caller is finished with the SharedTimer it must be released via {@link#removeRef}
      */
     public static SharedTimer getTimer() {
-        SharedTimer result = instance;
-        if (result == null) {
-            LOCK.lock();
-            try {
-                result = instance;
-                if (result == null) {
-                    instance = result = new SharedTimer();
-                }
-            } finally {
-                LOCK.unlock();
+        LOCK.lock();
+        try {
+            if (instance == null) {
+                instance = new SharedTimer();
             }
+            instance.refCount.getAndIncrement();
+        } finally {
+            LOCK.unlock();
         }
-        result.refCount.getAndIncrement();
-        return result;
+        return instance;
     }
 
     /**


Test patch:
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/SharedTimerTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/SharedTimerTest.java
new file mode 100644
index 00000000..d7a8c5ce
--- /dev/null
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/SharedTimerTest.java
@@ -0,0 +1,26 @@
+package com.microsoft.sqlserver.jdbc;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.concurrent.*;
+
+class SharedTimerTest {
+
+    @Test
+    void getTimer() throws InterruptedException, ExecutionException, TimeoutException {
+        final int iterations = 500;
+
+        ExecutorService executor = Executors.newFixedThreadPool(2);
+        try {
+            ArrayList<CompletableFuture<?>> futures = new ArrayList<>(iterations);
+            for (int i = 0; i < iterations; i++) {
+                futures.add(CompletableFuture.runAsync(() -> SharedTimer.getTimer().removeRef(), executor));
+            }
+
+            CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(2, TimeUnit.MINUTES);
+        } finally {
+            executor.shutdown();
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/preparedStatement/BatchExecutionWithBulkCopyTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/preparedStatement/BatchExecutionWithBulkCopyTest.java
index 4178c15d..dbccf4e7 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/preparedStatement/BatchExecutionWithBulkCopyTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/preparedStatement/BatchExecutionWithBulkCopyTest.java
@@ -721,6 +721,45 @@ public class BatchExecutionWithBulkCopyTest extends AbstractTest {
         }
     }
 
+    @Test
+    public void testReverseColumnOrder() throws Exception {
+        String valid = "insert into " + AbstractSQLGenerator.escapeIdentifier(tableName) + " (c2, c1) values "
+                + "(" + "?, " + "? " + ")";
+
+        try (Connection connection = PrepUtil.getConnection(connectionString + ";useBulkCopyForBatchInsert=true;");
+                SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) connection.prepareStatement(valid);
+                Statement stmt = connection.createStatement()) {
+            Field f1 = SQLServerConnection.class.getDeclaredField("isAzureDW");
+            f1.setAccessible(true);
+            f1.set(connection, true);
+
+            TestUtils.dropTableIfExists(AbstractSQLGenerator.escapeIdentifier(tableName), stmt);
+            String createTable = "create table " + AbstractSQLGenerator.escapeIdentifier(tableName)
+                    + " (c1 varchar(1), c2 varchar(3))";
+            stmt.execute(createTable);
+
+            pstmt.setString(1, "One");
+            pstmt.setString(2, "1");
+            pstmt.addBatch();
+
+            pstmt.executeBatch();
+
+            try (ResultSet rs = stmt
+                    .executeQuery("select c1, c2 from " + AbstractSQLGenerator.escapeIdentifier(tableName))) {
+
+                Object[] expected = new Object[2];
+
+                expected[0] = "1";
+                expected[1] = "One";
+                rs.next();
+
+                for (int i = 0; i < expected.length; i++) {
+                    assertEquals(expected[i], rs.getObject(i + 1));
+                }
+            }
+        }
+    }
+
     @BeforeAll
     public static void setupTests() throws Exception {
         setConnection();
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resultset/ResultSetTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/resultset/ResultSetTest.java
index ac285bb9..9ad05409 100644
--- a/src/test/java/com/microsoft/sqlserver/jdbc/resultset/ResultSetTest.java
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/resultset/ResultSetTest.java
@@ -31,6 +31,7 @@ import java.util.TimeZone;
 import java.util.UUID;
 
 import com.microsoft.sqlserver.jdbc.SQLServerConnection;
+import com.microsoft.sqlserver.jdbc.SQLServerException;
 import com.microsoft.sqlserver.jdbc.TestResource;
 import com.microsoft.sqlserver.testframework.PrepUtil;
 import org.junit.jupiter.api.AfterEach;
@@ -52,6 +53,10 @@ import com.microsoft.sqlserver.testframework.Constants;
 @RunWith(JUnitPlatform.class)
 public class ResultSetTest extends AbstractTest {
     private static final String tableName = RandomUtil.getIdentifier("StatementParam");
+    private static final String tableName1 = AbstractSQLGenerator
+            .escapeIdentifier(RandomUtil.getIdentifier("AmbiguousRs1"));
+    private static final String tableName2 = AbstractSQLGenerator
+            .escapeIdentifier(RandomUtil.getIdentifier("AmbiguousRs2"));
 
     private static final String expectedSqlState = "S0001";
 
@@ -68,6 +73,8 @@ public class ResultSetTest extends AbstractTest {
     public void init() throws Exception {
         try (Connection conn = getConnection(); Statement stmt = conn.createStatement()) {
             TestUtils.dropTableIfExists(AbstractSQLGenerator.escapeIdentifier(tableName), stmt);
+            TestUtils.dropTableIfExists(tableName1, stmt);
+            TestUtils.dropTableIfExists(tableName2, stmt);
         }
     }
 
@@ -75,6 +82,8 @@ public class ResultSetTest extends AbstractTest {
     public void cleanUp() throws Exception {
         try (Connection conn = getConnection(); Statement stmt = conn.createStatement()) {
             TestUtils.dropTableIfExists(AbstractSQLGenerator.escapeIdentifier(tableName), stmt);
+            TestUtils.dropTableIfExists(tableName1, stmt);
+            TestUtils.dropTableIfExists(tableName2, stmt);
         }
     }
 
@@ -283,6 +292,55 @@ public class ResultSetTest extends AbstractTest {
         }
     }
 
+    @Test
+    @Tag(Constants.xAzureSQLDW)
+    public void testUpdateOnColumnForAssociatedTable() throws SQLException {
+        String query1 = "SELECT t.* FROM " + tableName1 + " k INNER JOIN " + tableName2 + " t ON  t.i = k.i";
+        String query2 = "SELECT * FROM " + tableName2;
+        String data = "NEW EDIT";
+
+        try (SQLServerConnection conn = PrepUtil.getConnection(connectionString)) {
+
+            ambiguousUpdateRowTestSetup(conn);
+
+            try (Statement stmt = conn.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE);
+                    ResultSet rs = stmt.executeQuery(query1);) {
+                rs.first();
+                rs.updateString(3, data);
+                rs.updateRow();
+            }
+
+            try (Statement createStatement = conn.createStatement();
+                    ResultSet rs = createStatement.executeQuery(query2);) {
+                rs.next();
+                assertEquals(data, rs.getString(3));
+            }
+        }
+    }
+
+    @Test
+    @Tag(Constants.xAzureSQLDW)
+    public void testErrorOnAmbiguousUpdate() throws SQLException {
+        String query1 = "SELECT t.*, k.* FROM " + tableName1 + " k INNER JOIN " + tableName2 + " t ON  t.i = k.i";
+        String data = "NEW EDIT";
+
+        try (SQLServerConnection conn = PrepUtil.getConnection(connectionString)) {
+
+            ambiguousUpdateRowTestSetup(conn);
+
+            try (Statement stmt = conn.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE);
+                    ResultSet rs = stmt.executeQuery(query1);) {
+                rs.first();
+                rs.updateString(3, data);
+                rs.updateString(5, data);
+                rs.updateRow();
+            }
+            fail(TestResource.getResource("R_expectedExceptionNotThrown"));
+        } catch (SQLServerException se) {
+            assertTrue(se.getMessage().matches(TestUtils.formatErrorMsg("R_AmbiguousRowUpdate")));
+        }
+    }
+
     /**
      * Tests getObject(n, java.time.LocalDateTime.class).
      * 
@@ -634,13 +692,12 @@ public class ResultSetTest extends AbstractTest {
     }
 
     @Test
-    public void testResultSetClientCursorInitializerSqlErrorState() throws Exception {
+    public void testResultSetClientCursorInitializerSqlErrorState() {
         try (Connection con = PrepUtil.getConnection(connectionString); Statement stmt = con.createStatement()) {
-            stmt.executeUpdate("create table " + AbstractSQLGenerator.escapeIdentifier(tableName)
-                    + " (col1 int)");
-            boolean hasResults = stmt.execute("select * from " + AbstractSQLGenerator.escapeIdentifier(tableName)
-                    + "; select 1/0");
-            while(hasResults) {
+            stmt.executeUpdate("create table " + AbstractSQLGenerator.escapeIdentifier(tableName) + " (col1 int)");
+            boolean hasResults = stmt
+                    .execute("select * from " + AbstractSQLGenerator.escapeIdentifier(tableName) + "; select 1/0");
+            while (hasResults) {
                 ResultSet rs = stmt.getResultSet();
                 while (rs.next()) {}
                 hasResults = stmt.getMoreResults();
@@ -651,4 +708,13 @@ public class ResultSetTest extends AbstractTest {
             assertEquals(expectedErrorCode, e.getErrorCode());
         }
     }
+
+    private void ambiguousUpdateRowTestSetup(Connection conn) throws SQLException {
+        try (Statement stmt = conn.createStatement()) {
+            stmt.execute("CREATE TABLE " + tableName1 + " (i INT, data VARCHAR(30))");
+            stmt.execute("CREATE TABLE " + tableName2 + " (i INT, row INT, data VARCHAR(30))");
+            stmt.execute("INSERT INTO " + tableName1 + " SELECT 1, 'EDIT'");
+            stmt.execute("INSERT INTO " + tableName2 + " SELECT 1, 1, 'TEST1' UNION SELECT 1, 2, 'TEST2'");
+        }
+    }
 }