Adding Support for Timestamp PKey. IT for fractional dateTime and fractional Timestamp (#2468)

Author: VardhanThigle

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/dialectadapter/mysql/MysqlDialectAdapter.java

@@ -357,7 +357,6 @@ private ImmutableMap<String, SourceColumnType> getTableCols(
       ImmutableMap.<String, SourceColumnIndexInfo.IndexType>builder()
           .put("BIGINT UNSIGNED", IndexType.BIG_INT_UNSIGNED)
           .put("BIGINT", IndexType.NUMERIC)
-          .put("DATETIME", IndexType.DATE_TIME)
           .put("INTEGER", IndexType.NUMERIC)
           .put("INTEGER UNSIGNED", IndexType.NUMERIC)
           .put("MEDIUMINT", IndexType.NUMERIC)
@@ -372,6 +371,8 @@ private ImmutableMap<String, SourceColumnType> getTableCols(
           .put("VARBINARY", IndexType.BINARY)
           .put("TINYBLOB", IndexType.BINARY)
           .put("TINYTEXT", IndexType.STRING)
+          .put("DATETIME", IndexType.TIME_STAMP)
+          .put("TIMESTAMP", IndexType.TIME_STAMP)
           .build();
 
   /**

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/dialectadapter/postgresql/PostgreSQLDialectAdapter.java

@@ -508,7 +508,7 @@ private SourceColumnIndexInfo.IndexType indexTypeFrom(String typeCategory) {
       case "N":
         return SourceColumnIndexInfo.IndexType.NUMERIC;
       case "D":
-        return SourceColumnIndexInfo.IndexType.DATE_TIME;
+        return SourceColumnIndexInfo.IndexType.TIME_STAMP;
       case "S":
         return SourceColumnIndexInfo.IndexType.STRING;
       default:

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/iowrapper/JdbcIoWrapper.java

@@ -311,7 +311,7 @@ private static TableConfig getTableConfig(
                 IndexType.STRING,
                 IndexType.BIG_INT_UNSIGNED,
                 IndexType.BINARY,
-                IndexType.DATE_TIME);
+                IndexType.TIME_STAMP);
         // As of now only Primary key index with Numeric type is supported.
         // TODO:
         //    1. support non-primary unique indexes.

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/uniformsplitter/range/BoundaryExtractorFactory.java

@@ -21,6 +21,8 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Timestamp;
+import java.util.Calendar;
+import java.util.TimeZone;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -158,6 +160,8 @@ private static Boundary<String> fromStrings(
         .build();
   }
 
+  private static final Calendar utcCalendar = Calendar.getInstance(TimeZone.getTimeZone("UTC"));
+
   private static Boundary<Timestamp> fromTimestamps(
       PartitionColumn partitionColumn,
       ResultSet resultSet,
@@ -167,8 +171,8 @@ private static Boundary<Timestamp> fromTimestamps(
     resultSet.next();
     return Boundary.<Timestamp>builder()
         .setPartitionColumn(partitionColumn)
-        .setStart(resultSet.getTimestamp(1))
-        .setEnd(resultSet.getTimestamp(2))
+        .setStart(resultSet.getTimestamp(1, utcCalendar))
+        .setEnd(resultSet.getTimestamp(2, utcCalendar))
         .setBoundarySplitter(BoundarySplitterFactory.create(Timestamp.class))
         .setBoundaryTypeMapper(boundaryTypeMapper)
         .build();

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/schema/SourceColumnIndexInfo.java

@@ -142,7 +142,7 @@ public enum IndexType {
     BIG_INT_UNSIGNED,
     BINARY,
     STRING,
-    DATE_TIME,
+    TIME_STAMP,
     OTHER
   };
 
@@ -153,5 +153,5 @@ public enum IndexType {
           IndexType.STRING, String.class,
           IndexType.BIG_INT_UNSIGNED, BigDecimal.class,
           IndexType.BINARY, BoundaryExtractorFactory.BYTE_ARRAY_CLASS,
-          IndexType.DATE_TIME, Timestamp.class);
+          IndexType.TIME_STAMP, Timestamp.class);
 }

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/dialectadapter/postgresql/PostgreSQLDialectAdapterTest.java

@@ -270,7 +270,7 @@ public void testDiscoverTableIndexes() throws SQLException, RetriableSchemaDisco
                     .setIsPrimary(false)
                     .setCardinality(2L)
                     .setOrdinalPosition(2L)
-                    .setIndexType(SourceColumnIndexInfo.IndexType.DATE_TIME)
+                    .setIndexType(SourceColumnIndexInfo.IndexType.TIME_STAMP)
                     .build()));
   }
 

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/uniformsplitter/range/BoundaryExtractorFactoryTest.java

@@ -18,6 +18,8 @@
 import static com.google.cloud.teleport.v2.source.reader.io.jdbc.uniformsplitter.range.BoundaryExtractorFactory.BYTE_ARRAY_CLASS;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.when;
 
@@ -246,9 +248,10 @@ public void testFromTimestamp() throws SQLException {
     BoundaryExtractor<Timestamp> extractor = BoundaryExtractorFactory.create(Timestamp.class);
     Timestamp start = Timestamp.valueOf("0000-01-01 00:00:00.000000000");
     Timestamp end = Timestamp.valueOf("2041-01-31 23:59:59.999999999");
+
     when(mockResultSet.next()).thenReturn(true);
-    when(mockResultSet.getTimestamp(1)).thenReturn(start);
-    when(mockResultSet.getTimestamp(2)).thenReturn(end);
+    when(mockResultSet.getTimestamp(eq(1), any())).thenReturn(start);
+    when(mockResultSet.getTimestamp(eq(2), any())).thenReturn(end);
     Boundary<Timestamp> boundary = extractor.getBoundary(partitionColumn, mockResultSet, null);
     assertThat(boundary.start()).isEqualTo(start);
     assertThat(boundary.end()).isEqualTo(end);

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/uniformsplitter/range/BoundarySplitterFactoryTest.java

@@ -302,6 +302,8 @@ public void testInstantMapping() {
   @Test
   public void testTimeStampSplitting() {
     assertThat(BoundarySplitterFactory.splitInstants(null, null)).isEqualTo(null);
+    assertThat(BoundarySplitterFactory.splitInstants(null, Instant.EPOCH))
+        .isEqualTo(BoundarySplitterFactory.splitInstants(Instant.EPOCH, null));
     assertThat(BoundarySplitterFactory.splitInstants(Instant.MIN, Instant.MAX))
         .isEqualTo(Instant.parse("0000-07-01T23:59:59.999999999Z"));
 

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/MySQLDataTypesIT.java

@@ -257,8 +257,22 @@ private Map<String, List<Map<String, Object>>> getExpectedData() {
             "date_time_pk",
             "1000-01-01T00:00:00Z",
             "1000-01-01T00:00:01Z",
+            "2001-01-01T00:01:54.123456000Z",
+            /* DateTime does not depend on time zone. */
+            "2005-01-01T05:31:54.123456000Z",
             "9999-12-30T23:59:59Z",
             "9999-12-31T23:59:59Z"));
+    expectedData.put(
+        "timestamp_pk",
+        createRows(
+            "timestamp_pk",
+            "1970-01-01T00:00:01Z",
+            "1970-01-01T00:00:02Z",
+            "2001-01-01T00:01:54.123456000Z",
+            /* Timestamp offsets by time zone. We always read in UTC. */
+            "2005-01-01T00:01:54.123456000Z",
+            "2037-12-30T23:59:59Z",
+            "2038-01-18T23:59:59Z"));
     return expectedData;
   }
 

v2/sourcedb-to-spanner/src/test/resources/DataTypesIT/mysql-data-types.sql

@@ -261,8 +261,14 @@ CREATE TABLE tiny_text_pk_table (
 );
 
 CREATE TABLE date_time_pk_table (
-                                    id DATETIME,
-                                    date_time_pk_col DATETIME NOT NULL,
+                                    id DATETIME(6),
+                                    date_time_pk_col DATETIME(6) NOT NULL,
+                                    CONSTRAINT PRIMARY KEY (id)
+);
+
+CREATE TABLE timestamp_pk_table (
+                                    id TIMESTAMP(6),
+                                    timestamp_pk_col TIMESTAMP(6) NOT NULL,
                                     CONSTRAINT PRIMARY KEY (id)
 );
 
@@ -403,7 +409,19 @@ INSERT INTO `char_pk_table` (`id`, `char_pk_col`) VALUES ('AA==', 'AA=='), ('gAA
 INSERT INTO `varchar_pk_table` (`id`, `varchar_pk_col`) VALUES ('AA==', 'AA=='), ('gAAAAAAAAAA=', 'gAAAAAAAAAA=');
 INSERT INTO `tiny_text_pk_table` (`id`, `tiny_text_pk_col`) VALUES ('AA==', 'AA=='), ('gAAAAAAAAAA=', 'gAAAAAAAAAA=');
 INSERT INTO `date_time_pk_table` (`id`, `date_time_pk_col`) VALUES ('1000-01-01 00:00:00', '1000-01-01 00:00:00'), ('1000-01-01 00:00:01', '1000-01-01 00:00:01'),
+                                                                   ('2001-01-01 00:01:54.123456', '2001-01-01 00:01:54.123456'),
                                                                    ('9999-12-30 23:59:59', '9999-12-30 23:59:59'),  ('9999-12-31 23:59:59', '9999-12-31 23:59:59');
+SET time_zone = 'Asia/Kolkata';
+INSERT INTO `date_time_pk_table` (`id`, `date_time_pk_col`) VALUES ('2005-01-01 05:31:54.123456', '2005-01-01 05:31:54.123456');
+SET time_zone = SYSTEM;
+
+SET time_zone = 'UTC';
+INSERT INTO `timestamp_pk_table` (`id`, `timestamp_pk_col`) VALUES ('1970-01-01 00:00:01', '1970-01-01 00:00:01'), ('1970-01-01 00:00:02', '1970-01-01 00:00:02'),
+                                                                   ('2001-01-01 00:01:54.123456', '2001-01-01 00:01:54.123456'),
+                                                                   ('2037-12-30 23:59:59', '2037-12-30 23:59:59'),  ('2038-01-18 23:59:59', '2038-01-18 23:59:59');
+SET time_zone = 'Asia/Kolkata';
+INSERT INTO `timestamp_pk_table` (`id`, `timestamp_pk_col`) VALUES ('2005-01-01 05:31:54.123456', '2005-01-01 05:31:54.123456');
+SET time_zone = SYSTEM;
 
 INSERT INTO `bigint_table` (`bigint_col`) VALUES (NULL);
 INSERT INTO `bigint_unsigned_table` (`bigint_unsigned_col`) VALUES (NULL);