Add UUID support in SpannerSchema (#34034)

* Add UUID support in Spanner Schema * Add test
apache · Feb 20, 2025 · 8031f09 · 8031f09
1 parent 6430e07
commit 8031f09
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java b/...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java
@@ -181,6 +181,9 @@ private static Type parseSpannerType(String spannerType, Dialect dialect) {
           if (spannerType.startsWith("STRING")) {
             return Type.string();
           }
+          if ("UUID".equals(spannerType)) {
+            return Type.string();
+          }
           if (spannerType.startsWith("BYTES")) {
             return Type.bytes();
           }
@@ -260,6 +263,9 @@ private static Type parseSpannerType(String spannerType, Dialect dialect) {
           if (spannerType.startsWith("JSONB")) {
             return Type.pgJsonb();
           }
+          if ("UUID".equals(spannerType)) {
+            return Type.string();
+          }
           throw new IllegalArgumentException("Unknown spanner type " + spannerType);
         default:
           throw new IllegalArgumentException("Unrecognized dialect: " + dialect.name());

diff --git a/...le-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java b/...le-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java
@@ -41,17 +41,19 @@ public void testSingleTable() throws Exception {
             .addColumn("test", "protoVal", "PROTO<customer.app.TestMessage>")
             .addColumn("test", "enumVal", "ENUM<customer.app.TestEnum>")
             .addColumn("test", "tokens", "TOKENLIST")
+            .addColumn("test", "uuidCol", "UUID")
             .build();
 
     assertEquals(1, schema.getTables().size());
-    assertEquals(7, schema.getColumns("test").size());
+    assertEquals(8, schema.getColumns("test").size());
     assertEquals(1, schema.getKeyParts("test").size());
     assertEquals(Type.json(), schema.getColumns("test").get(3).getType());
     assertEquals(
         Type.proto("customer.app.TestMessage"), schema.getColumns("test").get(4).getType());
     assertEquals(
         Type.protoEnum("customer.app.TestEnum"), schema.getColumns("test").get(5).getType());
     assertEquals(Type.bytes(), schema.getColumns("test").get(6).getType());
+    assertEquals(Type.string(), schema.getColumns("test").get(7).getType());
   }
 
   @Test
@@ -84,12 +86,14 @@ public void testSinglePgTable() throws Exception {
             .addColumn("test", "numericVal", "numeric")
             .addColumn("test", "commitTime", "spanner.commit_timestamp")
             .addColumn("test", "jsonbCol", "jsonb")
+            .addColumn("test", "uuidCol", "uuid")
             .build();
 
     assertEquals(1, schema.getTables().size());
-    assertEquals(5, schema.getColumns("test").size());
+    assertEquals(6, schema.getColumns("test").size());
     assertEquals(1, schema.getKeyParts("test").size());
     assertEquals(Type.timestamp(), schema.getColumns("test").get(3).getType());
+    assertEquals(Type.string(), schema.getColumns("test").get(5).getType());
   }
 
   @Test