Skip to content

Commit 6b26c9e

Browse files
pratickchokhanipratickchokhani
committed
feat: add support for UUID and PG_UUID types across Spanner migration utilities
1 parent 466301e commit 6b26c9e

10 files changed

Lines changed: 156 additions & 2 deletions

File tree

‎v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/ddl/Column.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,10 @@ private static String typeString(Type type, Integer size) {
164164
return Type.Code.JSON.getName();
165165
case PG_JSONB:
166166
return Type.Code.PG_JSONB.getName();
167+
case UUID:
168+
return Type.Code.UUID.getName();
169+
case PG_UUID:
170+
return Type.Code.PG_UUID.getName();
167171
case TOKENLIST:
168172
return Type.Code.TOKENLIST.getName();
169173
case ARRAY:
@@ -250,6 +254,10 @@ public Builder bytes() {
250254
return type(Type.bytes());
251255
}
252256

257+
public Builder uuid() {
258+
return type(Type.uuid());
259+
}
260+
253261
public Builder timestamp() {
254262
return type(Type.timestamp());
255263
}
@@ -306,6 +314,10 @@ public Builder pgJsonb() {
306314
return type(Type.pgJsonb());
307315
}
308316

317+
public Builder pgUuid() {
318+
return type(Type.pgUuid());
319+
}
320+
309321
public Builder max() {
310322
return size(-1);
311323
}
@@ -377,6 +389,9 @@ private static SizedType parseSpannerType(String spannerType, Dialect dialect) {
377389
if (spannerType.equals(Type.Code.NUMERIC.getName())) {
378390
return t(Type.numeric(), null);
379391
}
392+
if (spannerType.equals(Type.Code.UUID.getName())) {
393+
return t(Type.uuid(), null);
394+
}
380395
if (spannerType.equals(Type.Code.JSON.getName())) {
381396
return t(Type.json(), null);
382397
}
@@ -432,6 +447,9 @@ private static SizedType parseSpannerType(String spannerType, Dialect dialect) {
432447
if (spannerType.equals(Type.Code.PG_NUMERIC.getName())) {
433448
return t(Type.pgNumeric(), null);
434449
}
450+
if (spannerType.equals(Type.Code.PG_UUID.getName())) {
451+
return t(Type.pgUuid(), null);
452+
}
435453
if (spannerType.equals(Type.Code.PG_JSONB.getName())) {
436454
return t(Type.pgJsonb(), null);
437455
}

‎v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/migrations/convertors/ChangeEventSpannerConvertor.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,8 @@ public static com.google.cloud.spanner.Key changeEventToPrimaryKey(
111111
case STRING:
112112
case PG_VARCHAR:
113113
case PG_TEXT:
114+
case UUID:
115+
case PG_UUID:
114116
pk.append(
115117
ChangeEventTypeConvertor.toString(
116118
changeEvent, keyColName, /* requiredField= */ true));

‎v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/migrations/convertors/ChangeEventTypeConvertor.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,8 @@ public static Value toValue(
7171
case STRING:
7272
case PG_VARCHAR:
7373
case PG_TEXT:
74+
case UUID:
75+
case PG_UUID:
7476
return Value.string(toString(changeEvent, key, requiredField));
7577
case NUMERIC:
7678
case PG_NUMERIC:

‎v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/migrations/spanner/SpannerReadUtils.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ private static void bindGoogleSqlValue(
111111
break;
112112
case STRING:
113113
case JSON:
114+
case UUID:
114115
stmtBuilder.bind(bindName).to((String) value);
115116
break;
116117
case NUMERIC:
@@ -153,6 +154,7 @@ private static void bindPgValue(
153154
case PG_VARCHAR:
154155
case PG_TEXT:
155156
case PG_JSONB:
157+
case PG_UUID:
156158
stmtBuilder.bind(bindName).to((String) value);
157159
break;
158160
case PG_NUMERIC:

‎v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/type/Type.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public final class Type implements Serializable {
4444
private static final Type TYPE_BYTES = new Type(Type.Code.BYTES, null, null);
4545
private static final Type TYPE_TIMESTAMP = new Type(Type.Code.TIMESTAMP, null, null);
4646
private static final Type TYPE_DATE = new Type(Type.Code.DATE, null, null);
47+
private static final Type TYPE_UUID = new Type(Type.Code.UUID, null, null);
4748
private static final Type TYPE_ARRAY_BOOL = new Type(Type.Code.ARRAY, TYPE_BOOL, null);
4849
private static final Type TYPE_ARRAY_INT64 = new Type(Type.Code.ARRAY, TYPE_INT64, null);
4950
private static final Type TYPE_ARRAY_FLOAT32 = new Type(Type.Code.ARRAY, TYPE_FLOAT32, null);
@@ -66,6 +67,7 @@ public final class Type implements Serializable {
6667
private static final Type TYPE_PG_BYTEA = new Type(Type.Code.PG_BYTEA, null, null);
6768
private static final Type TYPE_PG_TIMESTAMPTZ = new Type(Type.Code.PG_TIMESTAMPTZ, null, null);
6869
private static final Type TYPE_PG_DATE = new Type(Type.Code.PG_DATE, null, null);
70+
private static final Type TYPE_PG_UUID = new Type(Type.Code.PG_UUID, null, null);
6971
private static final Type TYPE_PG_ARRAY_BOOL = new Type(Type.Code.PG_ARRAY, TYPE_PG_BOOL, null);
7072
private static final Type TYPE_PG_ARRAY_INT8 = new Type(Type.Code.PG_ARRAY, TYPE_PG_INT8, null);
7173
private static final Type TYPE_PG_ARRAY_FLOAT4 =
@@ -160,6 +162,11 @@ public static Type date() {
160162
return TYPE_DATE;
161163
}
162164

165+
/** Returns the descriptor for the {@code UUID} type. */
166+
public static Type uuid() {
167+
return TYPE_UUID;
168+
}
169+
163170
public static Type pgBool() {
164171
return TYPE_PG_BOOL;
165172
}
@@ -204,6 +211,10 @@ public static Type pgDate() {
204211
return TYPE_PG_DATE;
205212
}
206213

214+
public static Type pgUuid() {
215+
return TYPE_PG_UUID;
216+
}
217+
207218
public static Type pgCommitTimestamp() {
208219
return TYPE_PG_COMMIT_TIMESTAMP;
209220
}
@@ -317,6 +328,7 @@ public enum Code {
317328
BYTES("BYTES", Dialect.GOOGLE_STANDARD_SQL),
318329
TIMESTAMP("TIMESTAMP", Dialect.GOOGLE_STANDARD_SQL),
319330
DATE("DATE", Dialect.GOOGLE_STANDARD_SQL),
331+
UUID("UUID", Dialect.GOOGLE_STANDARD_SQL),
320332
ARRAY("ARRAY", Dialect.GOOGLE_STANDARD_SQL),
321333
STRUCT("STRUCT", Dialect.GOOGLE_STANDARD_SQL),
322334
PG_BOOL("boolean", Dialect.POSTGRESQL),
@@ -330,6 +342,7 @@ public enum Code {
330342
PG_BYTEA("bytea", Dialect.POSTGRESQL),
331343
PG_TIMESTAMPTZ("timestamp with time zone", Dialect.POSTGRESQL),
332344
PG_DATE("date", Dialect.POSTGRESQL),
345+
PG_UUID("uuid", Dialect.POSTGRESQL),
333346
PG_ARRAY("array", Dialect.POSTGRESQL),
334347
PG_COMMIT_TIMESTAMP("spanner.commit_timestamp", Dialect.POSTGRESQL);
335348

‎v2/spanner-common/src/test/java/com/google/cloud/teleport/v2/spanner/ddl/ColumnTest.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ public void testTypeString_GoogleStandardSQL() {
5151
assertEquals("JSON", Column.builder().name("col").type(Type.json()).autoBuild().typeString());
5252
assertEquals(
5353
"TOKENLIST", Column.builder().name("col").type(Type.tokenlist()).autoBuild().typeString());
54+
assertEquals("UUID", Column.builder().name("col").type(Type.uuid()).autoBuild().typeString());
5455
assertEquals(
5556
"ARRAY<INT64>",
5657
Column.builder().name("col").type(Type.array(Type.int64())).autoBuild().typeString());
@@ -151,6 +152,13 @@ public void testTypeString_PostgreSQL() {
151152
.type(Type.pgCommitTimestamp())
152153
.autoBuild()
153154
.typeString());
155+
assertEquals(
156+
"uuid",
157+
Column.builder(Dialect.POSTGRESQL)
158+
.name("col")
159+
.type(Type.pgUuid())
160+
.autoBuild()
161+
.typeString());
154162
assertEquals(
155163
"bigint[]",
156164
Column.builder(Dialect.POSTGRESQL)
@@ -180,6 +188,8 @@ public void testParseSpannerType_GoogleStandardSQL() {
180188
assertEquals(Type.json(), Column.builder().name("col").parseType("JSON").autoBuild().type());
181189
assertEquals(
182190
Type.tokenlist(), Column.builder().name("col").parseType("TOKENLIST").autoBuild().type());
191+
assertEquals(
192+
Type.uuid(), Column.builder().name("col").parseType("UUID").autoBuild().type());
183193
assertEquals(
184194
Type.array(Type.int64()),
185195
Column.builder().name("col").parseType("ARRAY<INT64>").autoBuild().type());
@@ -239,6 +249,13 @@ public void testParseSpannerType_PostgreSQL() {
239249
.parseType("spanner.commit_timestamp")
240250
.autoBuild()
241251
.type());
252+
assertEquals(
253+
Type.pgUuid(),
254+
Column.builder(Dialect.POSTGRESQL)
255+
.name("col")
256+
.parseType("uuid")
257+
.autoBuild()
258+
.type());
242259
assertEquals(
243260
Type.pgArray(Type.pgInt8()),
244261
Column.builder(Dialect.POSTGRESQL).name("col").parseType("bigint[]").autoBuild().type());

‎v2/spanner-common/src/test/java/com/google/cloud/teleport/v2/spanner/migrations/convertors/ChangeEventSpannerConvertorTest.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,4 +432,75 @@ public void mutationFromEventWithGeneratedColumn() throws ChangeEventConvertorEx
432432
Truth.assertThat(actual.get("last_name").getAsString()).isEqualTo("B");
433433
Truth.assertThat(actual.containsKey("full_name")).isFalse();
434434
}
435+
436+
@Test
437+
public void canConvertChangeEventWithUuidToPrimaryKey() throws Exception {
438+
Ddl ddl =
439+
Ddl.builder()
440+
.createTable("UsersUuid")
441+
.column("uuid_field")
442+
.type(Type.uuid())
443+
.endColumn()
444+
.column("pg_uuid_field")
445+
.type(Type.pgUuid())
446+
.endColumn()
447+
.primaryKey()
448+
.asc("uuid_field")
449+
.asc("pg_uuid_field")
450+
.end()
451+
.endTable()
452+
.build();
453+
454+
JSONObject changeEvent = new JSONObject();
455+
changeEvent.put("uuid_field", "550e8400-e29b-41d4-a716-446655440000");
456+
changeEvent.put("pg_uuid_field", "123e4567-e89b-12d3-a456-426614174000");
457+
changeEvent.put(Constants.EVENT_TABLE_NAME_KEY, "UsersUuid");
458+
JsonNode ce = parseChangeEvent(changeEvent.toString());
459+
460+
Key key =
461+
ChangeEventSpannerConvertor.changeEventToPrimaryKey(
462+
"UsersUuid", ddl, ce, /* convertNameToLowerCase= */ true);
463+
464+
Iterable<Object> keyParts = key.getParts();
465+
ArrayList<Object> expectedKeyParts = new ArrayList<>();
466+
expectedKeyParts.add("550e8400-e29b-41d4-a716-446655440000");
467+
expectedKeyParts.add("123e4567-e89b-12d3-a456-426614174000");
468+
469+
assertThat(keyParts, is(expectedKeyParts));
470+
}
471+
472+
@Test
473+
public void mutationFromEventWithUuid() throws Exception {
474+
Ddl ddl =
475+
Ddl.builder()
476+
.createTable("UsersUuid")
477+
.column("uuid_field")
478+
.type(Type.uuid())
479+
.endColumn()
480+
.column("pg_uuid_field")
481+
.type(Type.pgUuid())
482+
.endColumn()
483+
.primaryKey()
484+
.asc("uuid_field")
485+
.end()
486+
.endTable()
487+
.build();
488+
489+
JSONObject changeEvent = new JSONObject();
490+
changeEvent.put("uuid_field", "550e8400-e29b-41d4-a716-446655440000");
491+
changeEvent.put("pg_uuid_field", "123e4567-e89b-12d3-a456-426614174000");
492+
changeEvent.put(Constants.EVENT_TABLE_NAME_KEY, "UsersUuid");
493+
JsonNode ce = parseChangeEvent(changeEvent.toString());
494+
495+
Mutation mutation =
496+
ChangeEventSpannerConvertor.mutationFromEvent(
497+
ddl.table("UsersUuid"),
498+
ce,
499+
List.of("uuid_field", "pg_uuid_field"),
500+
Set.of("uuid_field"));
501+
502+
Map<String, Value> actual = mutation.asMap();
503+
Truth.assertThat(actual.get("uuid_field").getAsString()).isEqualTo("550e8400-e29b-41d4-a716-446655440000");
504+
Truth.assertThat(actual.get("pg_uuid_field").getAsString()).isEqualTo("123e4567-e89b-12d3-a456-426614174000");
505+
}
435506
}

‎v2/spanner-common/src/test/java/com/google/cloud/teleport/v2/spanner/migrations/convertors/ChangeEventTypeConvertorTest.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -772,4 +772,20 @@ public void canConvertNullStrings() throws Exception {
772772
assertNull(ChangeEventTypeConvertor.toDate(ce, "date_field", true));
773773
assertEquals("NULL", ChangeEventTypeConvertor.toString(ce, "string_field", true));
774774
}
775+
776+
@Test
777+
public void canConvertToValue() throws Exception {
778+
JSONObject changeEvent = new JSONObject();
779+
changeEvent.put("uuid_field", "550e8400-e29b-41d4-a716-446655440000");
780+
changeEvent.put("pg_uuid_field", "123e4567-e89b-12d3-a456-426614174000");
781+
JsonNode ce = getJsonNode(changeEvent.toString());
782+
783+
assertEquals(
784+
com.google.cloud.spanner.Value.string("550e8400-e29b-41d4-a716-446655440000"),
785+
ChangeEventTypeConvertor.toValue(ce, com.google.cloud.teleport.v2.spanner.type.Type.uuid(), "uuid_field", true));
786+
787+
assertEquals(
788+
com.google.cloud.spanner.Value.string("123e4567-e89b-12d3-a456-426614174000"),
789+
ChangeEventTypeConvertor.toValue(ce, com.google.cloud.teleport.v2.spanner.type.Type.pgUuid(), "pg_uuid_field", true));
790+
}
775791
}

‎v2/spanner-common/src/test/java/com/google/cloud/teleport/v2/spanner/migrations/spanner/SpannerReadUtilsTest.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,9 @@ public void testGenerateReadSQLAllTypesPostgres() {
121121
.column("date_field")
122122
.pgDate()
123123
.endColumn()
124+
.column("pg_uuid_field")
125+
.pgUuid()
126+
.endColumn()
124127
.column("version")
125128
.pgInt8()
126129
.endColumn()
@@ -133,6 +136,7 @@ public void testGenerateReadSQLAllTypesPostgres() {
133136
.asc("bytea_field")
134137
.asc("timestamptz_field")
135138
.asc("date_field")
139+
.asc("pg_uuid_field")
136140
.end()
137141
.endTable()
138142
.build();
@@ -146,6 +150,7 @@ public void testGenerateReadSQLAllTypesPostgres() {
146150
ByteArray bytesValue = ByteArray.copyFrom("test_bytes");
147151
Timestamp timestampValue = Timestamp.ofTimeMicroseconds(1234567);
148152
Date dateValue = Date.fromYearMonthDay(2024, 1, 1);
153+
String uuidValue = "123e4567-e89b-12d3-a456-426614174000";
149154

150155
Key primaryKey =
151156
Key.of(
@@ -156,7 +161,8 @@ public void testGenerateReadSQLAllTypesPostgres() {
156161
textValue,
157162
bytesValue,
158163
timestampValue,
159-
dateValue);
164+
dateValue,
165+
uuidValue);
160166

161167
// 3. Generate the Statement
162168
Statement stmt =
@@ -171,7 +177,7 @@ public void testGenerateReadSQLAllTypesPostgres() {
171177
+ "\" WHERE "
172178
+ "\"bool_field\"=$1 AND \"int_field\"=$2 AND \"float_field\"=$3 AND "
173179
+ "\"numeric_field\"=$4 AND \"text_field\"=$5 AND \"bytea_field\"=$6 AND "
174-
+ "\"timestamptz_field\"=$7 AND \"date_field\"=$8";
180+
+ "\"timestamptz_field\"=$7 AND \"date_field\"=$8 AND \"pg_uuid_field\"=$9";
175181

176182
assertEquals(expectedSql, stmt.getSql());
177183

@@ -186,6 +192,7 @@ public void testGenerateReadSQLAllTypesPostgres() {
186192
assertEquals(Value.bytes(bytesValue), params.get("p6"));
187193
assertEquals(Value.timestamp(timestampValue), params.get("p7"));
188194
assertEquals(Value.date(dateValue), params.get("p8"));
195+
assertEquals(Value.string(uuidValue), params.get("p9"));
189196
}
190197

191198
@Test

‎v2/spanner-to-sourcedb/src/main/java/com/google/cloud/teleport/v2/templates/transforms/AssignShardIdFn.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -586,6 +586,8 @@ private com.google.cloud.spanner.Key generateKey(String tableName, JsonNode keys
586586
case STRING:
587587
case PG_VARCHAR:
588588
case PG_TEXT:
589+
case UUID:
590+
case PG_UUID:
589591
pk.append(
590592
DataChangeRecordTypeConvertor.toString(
591593
keysJson, keyColName, /* requiredField= */ true));
@@ -645,6 +647,8 @@ private Object getColumnValueFromJson(Column column, JsonNode valuesJson) throws
645647
case STRING:
646648
case PG_VARCHAR:
647649
case PG_TEXT:
650+
case UUID:
651+
case PG_UUID:
648652
return DataChangeRecordTypeConvertor.toString(valuesJson, colName, false);
649653
case NUMERIC:
650654
case PG_NUMERIC:
@@ -690,6 +694,8 @@ private Object getColumnValueFromRow(Column column, Value value) throws Exceptio
690694
case STRING:
691695
case PG_VARCHAR:
692696
case PG_TEXT:
697+
case UUID:
698+
case PG_UUID:
693699
return value.getString();
694700
case NUMERIC:
695701
case PG_NUMERIC:

0 commit comments

Comments
 (0)