Apply timezone to Timestamps for batch insert with bulkcopy

microsoft · Jan 10, 2024 · 0861bc5 · 0861bc5
1 parent c3343ce
commit 0861bc5
Show file tree

Hide file tree

Showing 2 changed files with 117 additions and 0 deletions.
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java
@@ -21,14 +21,17 @@
 import java.sql.SQLType;
 import java.sql.SQLXML;
 import java.sql.Statement;
+import java.text.DateFormat;
 import java.text.MessageFormat;
 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.Date;
 import java.util.logging.Level;
+import java.text.SimpleDateFormat;
 
 import com.microsoft.sqlserver.jdbc.SQLServerConnection.CityHash128Key;
 import com.microsoft.sqlserver.jdbc.SQLServerConnection.PreparedStatementHandle;
@@ -2938,6 +2941,44 @@ private boolean checkSQLLength(int length) {
         return true;
     }
 
+    private String getDateFormatPattern(int nanos) {
+        int sum = 0;
+        StringBuilder pattern = new StringBuilder("yyyy-MM-dd HH:mm:ss");
+        String[] figure = Integer.toString(nanos).split("(^0+|0+$)");
+
+        for(String digit: figure) {
+            sum += digit.length();
+        }
+
+        if (sum == 0) {
+            return pattern.toString();
+        }
+
+        if (sum > 7) {
+            sum = 7;
+        }
+
+        for (int i = 0; i < sum; i++) {
+            if (i == 0) {
+                pattern.append(".S");
+            }
+            pattern.append("S");
+        }
+
+        return pattern.toString();
+    }
+
+    private java.sql.Timestamp applyTimezoneTo(java.sql.Timestamp ts, Calendar cal) {
+        Date date = new Date(cal.getTimeInMillis());
+        int nanos = ts.getNanos();
+        DateFormat df = new SimpleDateFormat(getDateFormatPattern(nanos));
+        df.setTimeZone(cal.getTimeZone());
+        java.sql.Timestamp newTs = java.sql.Timestamp.valueOf(df.format(date));
+        newTs.setNanos(nanos);
+
+        return newTs;
+    }
+
     /**
      * Prepare statement batch execute command
      */
@@ -3416,6 +3457,11 @@ public final void setTimestamp(int n, java.sql.Timestamp x, java.util.Calendar c
         if (loggerExternal.isLoggable(java.util.logging.Level.FINER))
             loggerExternal.entering(getClassNameLogging(), "setTimestamp", new Object[] {n, x, cal});
         checkClosed();
+
+        if (this.useBulkCopyForBatchInsert && null != cal) {
+            x = applyTimezoneTo(x, cal);
+        }
+
         setValue(n, JDBCType.TIMESTAMP, x, JavaType.TIMESTAMP, cal, false);
         loggerExternal.exiting(getClassNameLogging(), "setTimestamp");
     }
@@ -3427,6 +3473,11 @@ public final void setTimestamp(int n, java.sql.Timestamp x, java.util.Calendar c
         if (loggerExternal.isLoggable(java.util.logging.Level.FINER))
             loggerExternal.entering(getClassNameLogging(), "setTimestamp", new Object[] {n, x, cal, forceEncrypt});
         checkClosed();
+
+        if (this.useBulkCopyForBatchInsert && null != cal) {
+            x = applyTimezoneTo(x, cal);
+        }
+
         setValue(n, JDBCType.TIMESTAMP, x, JavaType.TIMESTAMP, cal, forceEncrypt);
         loggerExternal.exiting(getClassNameLogging(), "setTimestamp");
     }

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java
@@ -13,11 +13,17 @@
 import java.sql.BatchUpdateException;
 import java.sql.CallableStatement;
 import java.sql.Connection;
+import java.sql.DriverManager;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.sql.Timestamp;
+import java.text.SimpleDateFormat;
 import java.util.Arrays;
+import java.util.Calendar;
+import java.util.Date;
+import java.util.TimeZone;
 
 import com.microsoft.sqlserver.jdbc.SQLServerPreparedStatement;
 import org.junit.jupiter.api.AfterAll;
@@ -54,6 +60,8 @@ public class BatchExecutionTest extends AbstractTest {
     private static String ctstable3;
     private static String ctstable4;
     private static String ctstable3Procedure1;
+    private static String timestampTable = AbstractSQLGenerator
+            .escapeIdentifier(RandomUtil.getIdentifier("timestampTableBatchInsert"));
 
     /**
      * This tests the updateCount when the error query does cause a SQL state HY008.
@@ -99,6 +107,64 @@ public void testBatchUpdateCountTrueOnFirstPstmtSpPrepare() throws Exception {
         testBatchUpdateCountWith(5, 4, true, "prepare", expectedUpdateCount);
     }
 
+    /**
+     * This tests the updateCount when the error query does cause a SQL state HY008.
+     *
+     * @throws Exception
+     */
+    @Test
+    public void testValidTimezoneForTimestampBatchInsertWithBulkCopy() throws Exception {
+        Calendar gmtCal = Calendar.getInstance(TimeZone.getTimeZone("GMT"));
+        SimpleDateFormat simpleDateFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
+        java.util.Date today = new Date();
+        String todayString = simpleDateFormat.format(today);
+        long ms = Timestamp.valueOf(todayString).getTime();
+
+        try (Connection con = DriverManager.getConnection(connectionString);
+                Statement stmt = con.createStatement();
+                PreparedStatement pstmt = con.prepareStatement(
+                        "INSERT INTO " + timestampTable + " VALUES(?)")) {
+
+            String dropSql = "IF EXISTS (SELECT * FROM dbo.sysobjects WHERE id = object_id(N'" + TestUtils.escapeSingleQuotes(timestampTable) + "') and OBJECTPROPERTY(id, N'IsUserTable') = 1) DROP TABLE " + timestampTable;
+            stmt.execute(dropSql);
+
+            String createSql = "CREATE TABLE" + timestampTable + " (c1 DATETIME2(3))";
+            stmt.execute(createSql);
+
+            Timestamp timestamp = new Timestamp(ms);
+
+            pstmt.setTimestamp(1, timestamp, gmtCal);
+            pstmt.addBatch();
+            pstmt.executeBatch();
+        }
+
+        try (Connection con = DriverManager.getConnection(connectionString + ";useBulkCopyForBatchInsert=true;sendTemporalDataTypesAsStringForBulkCopy=false;");
+                PreparedStatement pstmt = con.prepareStatement(
+                        "INSERT INTO " + timestampTable + " VALUES(?)")) {
+
+            Timestamp timestamp = new Timestamp(ms);
+
+            pstmt.setTimestamp(1, timestamp, gmtCal);
+            pstmt.addBatch();
+            pstmt.executeBatch();
+        }
+
+        try (Connection con = DriverManager.getConnection(connectionString);
+                Statement stmt = con.createStatement()) {
+            ResultSet rs = stmt.executeQuery("SELECT * FROM " + timestampTable);
+
+            Timestamp ts0; // Timestamp batch inserted without bulkcopy
+            Timestamp ts1; // Timestamp batch inserted with bulkcopy
+
+            rs.next();
+            ts0 = rs.getTimestamp(1);
+            rs.next();
+            ts1 = rs.getTimestamp(1);
+
+            assertEquals(ts0, ts1);
+        }
+    }
+
     /**
      * This tests the updateCount when the error query does not cause a SQL state HY008.
      *