feat: add several types to PGAdapter (GoogleCloudPlatform#1887)

* feat: add multiple type mapping for PG adapter

Adds 60+ different type mappings for the PostgreSQL adapter.

* trigger build

* trigger build

* feat: remove unsupported type mappings

Simplifies code and add testing for unsupported types.

* tests: remove time type expectation

* trigger build

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/source/reader/io/jdbc/rowmapper/provider/PostgreSQLJdbcValueMappings.java

@@ -19,33 +19,138 @@
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.rowmapper.JdbcValueMappingsProvider;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.rowmapper.ResultSetValueExtractor;
 import com.google.cloud.teleport.v2.source.reader.io.jdbc.rowmapper.ResultSetValueMapper;
+import com.google.cloud.teleport.v2.source.reader.io.schema.typemapping.provider.unified.CustomSchema.TimeStampTz;
 import com.google.common.collect.ImmutableMap;
+import java.nio.ByteBuffer;
 import java.sql.ResultSet;
+import java.time.Instant;
+import java.time.OffsetDateTime;
+import java.time.OffsetTime;
+import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
+import java.time.format.DateTimeFormatterBuilder;
+import java.time.temporal.ChronoField;
 import java.util.Calendar;
 import java.util.Map;
 import java.util.TimeZone;
+import java.util.concurrent.TimeUnit;
+import org.apache.avro.generic.GenericRecordBuilder;
 import org.apache.commons.lang3.tuple.Pair;
 
 /** PostgreSQL data type mapping to AVRO types. */
 public class PostgreSQLJdbcValueMappings implements JdbcValueMappingsProvider {
+
+  private static final Calendar UTC_CALENDAR =
+      Calendar.getInstance(TimeZone.getTimeZone(ZoneOffset.UTC));
+
+  private static final DateTimeFormatter TIMESTAMPTZ_FORMAT =
+      new DateTimeFormatterBuilder()
+          .appendPattern("yyyy-MM-dd HH:mm:ss")
+          .optionalStart()
+          .appendFraction(ChronoField.NANO_OF_SECOND, 1, 6, true)
+          .optionalEnd()
+          .appendOffset("+HH:mm", "+00")
+          .toFormatter();
+
+  private static long toMicros(Instant instant) {
+    return TimeUnit.SECONDS.toMicros(instant.getEpochSecond())
+        + TimeUnit.NANOSECONDS.toMicros(instant.getNano());
+  }
+
+  private static long toMicros(OffsetTime offsetTime) {
+    return TimeUnit.HOURS.toMicros(offsetTime.getHour())
+        + TimeUnit.MINUTES.toMicros(offsetTime.getMinute())
+        + TimeUnit.SECONDS.toMicros(offsetTime.getSecond())
+        + TimeUnit.NANOSECONDS.toMicros(offsetTime.getNano());
+  }
+
   private static final ResultSetValueMapper<?> valuePassThrough = (value, schema) -> value;
 
-  private static final Calendar utcCalendar = Calendar.getInstance(TimeZone.getTimeZone("UTC"));
+  private static final ResultSetValueExtractor<ByteBuffer> bytesExtractor =
+      (rs, fieldName) -> {
+        byte[] bytes = rs.getBytes(fieldName);
+        if (bytes == null) {
+          return null;
+        }
+        return ByteBuffer.wrap(bytes);
+      };
+
+  private static final ResultSetValueExtractor<java.sql.Date> dateExtractor =
+      (rs, fieldName) -> rs.getDate(fieldName, UTC_CALENDAR);
+
+  private static final ResultSetValueExtractor<java.sql.Timestamp> timestampExtractor =
+      (rs, fieldName) -> rs.getTimestamp(fieldName, UTC_CALENDAR);
 
-  private static final ResultSetValueExtractor<java.sql.Date> utcDateExtractor =
-      (rs, fieldName) -> rs.getDate(fieldName, utcCalendar);
+  private static final ResultSetValueExtractor<OffsetDateTime> timestamptzExtractor =
+      (rs, fieldName) -> {
+        String timestampTz = rs.getString(fieldName);
+        if (timestampTz == null) {
+          return null;
+        }
+        return OffsetDateTime.parse(timestampTz, TIMESTAMPTZ_FORMAT);
+      };
 
-  private static final ResultSetValueMapper<java.sql.Date> sqlDateToAvro =
+  // Value might be a Double.NaN or a valid BigDecimal
+  private static final ResultSetValueMapper<Number> numericToAvro =
+      (value, schema) -> value.toString();
+
+  private static final ResultSetValueMapper<java.sql.Date> dateToAvro =
       (value, schema) -> (int) value.toLocalDate().toEpochDay();
 
-  // TODO(thiagotnunes): Add missing type mappings
+  private static final ResultSetValueMapper<java.sql.Timestamp> timestampToAvro =
+      (value, schema) -> toMicros(value.toInstant());
+
+  private static final ResultSetValueMapper<OffsetDateTime> timestamptzToAvro =
+      (value, schema) ->
+          new GenericRecordBuilder(TimeStampTz.SCHEMA)
+              .set(TimeStampTz.TIMESTAMP_FIELD_NAME, toMicros(value.toInstant()))
+              .set(
+                  TimeStampTz.OFFSET_FIELD_NAME,
+                  TimeUnit.SECONDS.toMillis(value.getOffset().getTotalSeconds()))
+              .build();
+
   private static final ImmutableMap<String, JdbcValueMapper<?>> SCHEMA_MAPPINGS =
       ImmutableMap.<String, Pair<ResultSetValueExtractor<?>, ResultSetValueMapper<?>>>builder()
           .put("BIGINT", Pair.of(ResultSet::getLong, valuePassThrough))
+          .put("BIGSERIAL", Pair.of(ResultSet::getLong, valuePassThrough))
+          .put("BIT", Pair.of(bytesExtractor, valuePassThrough))
+          .put("BIT VARYING", Pair.of(bytesExtractor, valuePassThrough))
+          .put("BOOL", Pair.of(ResultSet::getBoolean, valuePassThrough))
+          .put("BOOLEAN", Pair.of(ResultSet::getBoolean, valuePassThrough))
+          .put("BYTEA", Pair.of(bytesExtractor, valuePassThrough))
+          .put("CHAR", Pair.of(ResultSet::getString, valuePassThrough))
+          .put("CHARACTER", Pair.of(ResultSet::getString, valuePassThrough))
           .put("CHARACTER VARYING", Pair.of(ResultSet::getString, valuePassThrough))
-          .put("DATE", Pair.of(utcDateExtractor, sqlDateToAvro))
+          .put("CITEXT", Pair.of(ResultSet::getString, valuePassThrough))
+          .put("DATE", Pair.of(dateExtractor, dateToAvro))
+          .put("DECIMAL", Pair.of(ResultSet::getObject, numericToAvro))
+          .put("DOUBLE PRECISION", Pair.of(ResultSet::getDouble, valuePassThrough))
+          .put("FLOAT4", Pair.of(ResultSet::getFloat, valuePassThrough))
+          .put("FLOAT8", Pair.of(ResultSet::getDouble, valuePassThrough))
+          .put("INT", Pair.of(ResultSet::getInt, valuePassThrough))
+          .put("INTEGER", Pair.of(ResultSet::getInt, valuePassThrough))
+          .put("INT2", Pair.of(ResultSet::getInt, valuePassThrough))
+          .put("INT4", Pair.of(ResultSet::getInt, valuePassThrough))
           .put("INT8", Pair.of(ResultSet::getLong, valuePassThrough))
+          .put("JSON", Pair.of(ResultSet::getString, valuePassThrough))
+          .put("JSONB", Pair.of(ResultSet::getString, valuePassThrough))
+          .put("MONEY", Pair.of(ResultSet::getDouble, valuePassThrough))
+          .put("NUMERIC", Pair.of(ResultSet::getObject, numericToAvro))
+          .put("OID", Pair.of(ResultSet::getLong, valuePassThrough))
+          .put("REAL", Pair.of(ResultSet::getFloat, valuePassThrough))
+          .put("SERIAL", Pair.of(ResultSet::getInt, valuePassThrough))
+          .put("SERIAL2", Pair.of(ResultSet::getInt, valuePassThrough))
+          .put("SERIAL4", Pair.of(ResultSet::getInt, valuePassThrough))
+          .put("SERIAL8", Pair.of(ResultSet::getLong, valuePassThrough))
+          .put("SMALLINT", Pair.of(ResultSet::getInt, valuePassThrough))
+          .put("SMALLSERIAL", Pair.of(ResultSet::getInt, valuePassThrough))
           .put("TEXT", Pair.of(ResultSet::getString, valuePassThrough))
+          .put("TIMESTAMP", Pair.of(timestampExtractor, timestampToAvro))
+          .put("TIMESTAMPTZ", Pair.of(timestamptzExtractor, timestamptzToAvro))
+          .put("TIMESTAMP WITH TIME ZONE", Pair.of(timestamptzExtractor, timestamptzToAvro))
+          .put("TIMESTAMP WITHOUT TIME ZONE", Pair.of(timestampExtractor, timestampToAvro))
+          .put("UUID", Pair.of(ResultSet::getString, valuePassThrough))
+          .put("VARBIT", Pair.of(bytesExtractor, valuePassThrough))
           .put("VARCHAR", Pair.of(ResultSet::getString, valuePassThrough))
           .build()
           .entrySet()

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

@@ -33,20 +33,67 @@ public final class PostgreSQLMappingProvider {
   private static final ImmutableMap<String, UnifiedTypeMapping> MAPPING =
       ImmutableMap.<String, UnifiedMappingProvider.Type>builder()
           .put("BIGINT", UnifiedMappingProvider.Type.LONG)
+          .put("BIGSERIAL", UnifiedMappingProvider.Type.LONG)
+          .put("BIT", UnifiedMappingProvider.Type.BYTES)
+          .put("BIT VARYING", UnifiedMappingProvider.Type.BYTES)
+          .put("BOOL", UnifiedMappingProvider.Type.BOOLEAN)
+          .put("BOOLEAN", UnifiedMappingProvider.Type.BOOLEAN)
+          .put("BYTEA", UnifiedMappingProvider.Type.BYTES)
+          // TODO(thiagotnunes): Refine mapping type according to
+          // https://cloud.google.com/datastream/docs/unified-types#map-psql (if there is a limit
+          // for length we should use varchar instead)
+          .put("CHAR", UnifiedMappingProvider.Type.STRING)
+          // TODO(thiagotnunes): Refine mapping type according to
+          // https://cloud.google.com/datastream/docs/unified-types#map-psql (if there is a limit
+          // for length we should use varchar instead)
+          .put("CHARACTER", UnifiedMappingProvider.Type.STRING)
           // TODO(thiagotnunes): Refine mapping type according to
           // https://cloud.google.com/datastream/docs/unified-types#map-psql (if there is a limit
           // for length we should use varchar instead)
           .put("CHARACTER VARYING", UnifiedMappingProvider.Type.STRING)
+          .put("CITEXT", UnifiedMappingProvider.Type.STRING)
           .put("DATE", UnifiedMappingProvider.Type.DATE)
+          // TODO(thiagotnunes): Refine mapping type according to
+          // https://cloud.google.com/datastream/docs/unified-types#map-psql (if there is a
+          // precision and scale are >= 0, map to DECIMAL)
+          .put("DECIMAL", UnifiedMappingProvider.Type.NUMBER)
+          .put("DOUBLE PRECISION", UnifiedMappingProvider.Type.DOUBLE)
+          .put("FLOAT4", UnifiedMappingProvider.Type.FLOAT)
+          .put("FLOAT8", UnifiedMappingProvider.Type.DOUBLE)
+          .put("INT", UnifiedMappingProvider.Type.INTEGER)
+          .put("INTEGER", UnifiedMappingProvider.Type.INTEGER)
+          .put("INT2", UnifiedMappingProvider.Type.INTEGER)
+          .put("INT4", UnifiedMappingProvider.Type.INTEGER)
           .put("INT8", UnifiedMappingProvider.Type.LONG)
+          .put("JSON", UnifiedMappingProvider.Type.JSON)
+          .put("JSONB", UnifiedMappingProvider.Type.JSON)
+          .put("MONEY", UnifiedMappingProvider.Type.DOUBLE)
+          // https://cloud.google.com/datastream/docs/unified-types#map-psql (if there is a
+          // precision and scale are >= 0, map to DECIMAL)
+          .put("NUMERIC", UnifiedMappingProvider.Type.NUMBER)
+          .put("OID", UnifiedMappingProvider.Type.LONG)
+          .put("REAL", UnifiedMappingProvider.Type.FLOAT)
+          .put("SERIAL", UnifiedMappingProvider.Type.INTEGER)
+          .put("SERIAL2", UnifiedMappingProvider.Type.INTEGER)
+          .put("SERIAL4", UnifiedMappingProvider.Type.INTEGER)
+          .put("SERIAL8", UnifiedMappingProvider.Type.LONG)
+          .put("SMALLINT", UnifiedMappingProvider.Type.INTEGER)
+          .put("SMALLSERIAL", UnifiedMappingProvider.Type.INTEGER)
           // TODO(thiagotnunes): Refine mapping type according to
           // https://cloud.google.com/datastream/docs/unified-types#map-psql (if there is a limit
           // for length we should use varchar instead)
           .put("TEXT", UnifiedMappingProvider.Type.STRING)
+          .put("TIMESTAMP", UnifiedMappingProvider.Type.TIMESTAMP)
+          .put("TIMESTAMPTZ", UnifiedMappingProvider.Type.TIMESTAMP_WITH_TIME_ZONE)
+          .put("TIMESTAMP WITH TIME ZONE", UnifiedMappingProvider.Type.TIMESTAMP_WITH_TIME_ZONE)
+          .put("TIMESTAMP WITHOUT TIME ZONE", UnifiedMappingProvider.Type.TIMESTAMP)
+          .put("UUID", UnifiedMappingProvider.Type.STRING)
+          .put("VARBIT", UnifiedMappingProvider.Type.BYTES)
           // TODO(thiagotnunes): Refine mapping type according to
           // https://cloud.google.com/datastream/docs/unified-types#map-psql (if there is a limit
           // for length we should use varchar instead)
           .put("VARCHAR", UnifiedMappingProvider.Type.STRING)
+          .put("UNSUPPORTED", UnifiedMappingProvider.Type.UNSUPPORTED)
           .build()
           .entrySet()
           .stream()

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/source/reader/io/CustomAsserts.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.google.cloud.teleport.v2.source.reader.io;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+
+import java.nio.ByteBuffer;
+
+public class CustomAsserts {
+
+  private static final double DELTA = 0.0001f;
+
+  private CustomAsserts() {}
+
+  public static void assertColumnEquals(String message, Object expected, Object actual) {
+    if (expected instanceof byte[] && actual instanceof byte[]) {
+      assertArrayEquals(message, (byte[]) expected, (byte[]) actual);
+    } else if (expected instanceof byte[] && actual instanceof ByteBuffer) {
+      assertArrayEquals(message, (byte[]) expected, ((ByteBuffer) actual).array());
+    } else if (expected instanceof Float && actual instanceof Float) {
+      assertEquals(message, (float) expected, (float) actual, DELTA);
+    } else if (expected instanceof Double && actual instanceof Double) {
+      assertEquals(message, (double) expected, (double) actual, DELTA);
+    } else {
+      assertEquals(message, expected, actual);
+    }
+  }
+}