diff --git a/v2/googlecloud-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/schemautils/TypesUtils.java b/v2/googlecloud-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/schemautils/TypesUtils.java index b853eb21dc..d1480eb69f 100644 --- a/v2/googlecloud-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/schemautils/TypesUtils.java +++ b/v2/googlecloud-to-googlecloud/src/main/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/schemautils/TypesUtils.java @@ -16,6 +16,8 @@ package com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.schemautils; import com.google.cloud.spanner.Type; +import java.util.regex.Matcher; +import java.util.regex.Pattern; import org.json.JSONObject; /** @@ -24,6 +26,9 @@ */ public class TypesUtils { + private static final Pattern POSTGRES_VECTOR_TYPE_PATTERN = + Pattern.compile("^(.+)\\[\\]\\s+VECTOR\\s+LENGTH\\s+[1-9][0-9]*$", Pattern.CASE_INSENSITIVE); + public static Type informationSchemaGoogleSQLTypeToSpannerType(String type) { type = cleanInformationSchemaType(type); switch (type) { @@ -62,6 +67,12 @@ public static Type informationSchemaGoogleSQLTypeToSpannerType(String type) { } public static Type informationSchemaPostgreSQLTypeToSpannerType(String type) { + Matcher vectorTypeMatcher = POSTGRES_VECTOR_TYPE_PATTERN.matcher(type); + if (vectorTypeMatcher.matches()) { + Type itemType = informationSchemaPostgreSQLTypeToSpannerType(vectorTypeMatcher.group(1)); + return Type.array(itemType); + } + boolean isPostgresArray = isPostgresArray(type); String cleanedType = ""; if (isPostgresArray) { diff --git a/v2/googlecloud-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/TypesUtilsTest.java b/v2/googlecloud-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/TypesUtilsTest.java index 818dc2d20a..6c96f595c8 100644 --- a/v2/googlecloud-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/TypesUtilsTest.java +++ b/v2/googlecloud-to-googlecloud/src/test/java/com/google/cloud/teleport/v2/templates/spannerchangestreamstobigquery/TypesUtilsTest.java @@ -82,10 +82,37 @@ public void testInformationSchemaPostgreSQLTypeToSpannerType() { .isEqualTo(Type.array(Type.string())); assertThat(TypesUtils.informationSchemaPostgreSQLTypeToSpannerType("JSONB[]")) .isEqualTo(Type.array(Type.pgJsonb())); + assertThat(TypesUtils.informationSchemaPostgreSQLTypeToSpannerType("real[]")) + .isEqualTo(Type.array(Type.float32())); assertThat(TypesUtils.informationSchemaPostgreSQLTypeToSpannerType("real")) .isEqualTo(Type.float32()); } + @Test + public void testInformationSchemaPostgreSQLVectorTypeToSpannerType() { + assertThat(TypesUtils.informationSchemaPostgreSQLTypeToSpannerType("real[] vector length 512")) + .isEqualTo(Type.array(Type.float32())); + assertThat( + TypesUtils.informationSchemaPostgreSQLTypeToSpannerType( + "DOUBLE PRECISION[] VECTOR LENGTH 4")) + .isEqualTo(Type.array(Type.float64())); + } + + @Test + public void testInformationSchemaPostgreSQLVectorTypeRejectsInvalidTypes() { + assertThrows( + IllegalArgumentException.class, + () -> + TypesUtils.informationSchemaPostgreSQLTypeToSpannerType( + "real[] vector length invalid")); + assertThrows( + IllegalArgumentException.class, + () -> TypesUtils.informationSchemaPostgreSQLTypeToSpannerType("real[] vector length 0")); + assertThrows( + IllegalArgumentException.class, + () -> TypesUtils.informationSchemaPostgreSQLTypeToSpannerType("integer[] vector length 4")); + } + @Test public void testExtractTypeFromInvalidTypeCode() { final JSONObject invalidJsonObject = new JSONObject("{\"type_code\":\"STRING\"}");