Adding support for pg Jsonb datatype in SpannerIO

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java

@@ -126,6 +126,8 @@ private static long estimatePrimitiveValue(Value v) {
         return v.isNull() ? 0 : v.getNumeric().toString().length();
       case JSON:
         return v.isNull() ? 0 : v.getJson().length();
+      case PG_JSONB:
+        return v.isNull() ? 0 : v.getPgJsonb().length();
       default:
         throw new IllegalArgumentException("Unsupported type " + v.getType());
     }
@@ -188,6 +190,15 @@ private static long estimateArrayValue(Value v) {
           totalLength += s.length();
         }
         return totalLength;
+      case PG_JSONB:
+        totalLength = 0;
+        for (String s : v.getPgJsonbArray()) {
+          if (s == null) {
+            continue;
+          }
+          totalLength += s.length();
+        }
+        return totalLength;
       default:
         throw new IllegalArgumentException("Unsupported type " + v.getType());
     }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchema.java

@@ -233,6 +233,9 @@ private static Type parseSpannerType(String spannerType, Dialect dialect) {
           if ("SPANNER.COMMIT_TIMESTAMP".equals(spannerType)) {
             return Type.timestamp();
           }
+          if (spannerType.startsWith("JSONB")) {
+            return Type.pgJsonb();
+          }
           throw new IllegalArgumentException("Unknown spanner type " + spannerType);
         default:
           throw new IllegalArgumentException("Unrecognized dialect: " + dialect.name());

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimatorTest.java

@@ -65,6 +65,11 @@ public void primitives() throws Exception {
             .to(Value.json("{\"key1\":\"value1\", \"key2\":\"value2\"}"))
             .build();
     Mutation deleteDouble = Mutation.delete("test", Key.of(1223.));
+    Mutation jsonb =
+        Mutation.newInsertOrUpdateBuilder("test")
+            .set("one")
+            .to(Value.pgJsonb("{\"key123\":\"value123\", \"key321\":\"value321\"}"))
+            .build();
 
     assertThat(MutationSizeEstimator.sizeOf(int64), is(8L));
     assertThat(MutationSizeEstimator.sizeOf(float64), is(8L));
@@ -74,6 +79,7 @@ public void primitives() throws Exception {
     assertThat(MutationSizeEstimator.sizeOf(pgNumericNaN), is(3L));
     assertThat(MutationSizeEstimator.sizeOf(json), is(34L));
     assertThat(MutationSizeEstimator.sizeOf(deleteDouble), is(8L));
+    assertThat(MutationSizeEstimator.sizeOf(jsonb), is(42L));
   }
 
   @Test
@@ -131,13 +137,22 @@ public void primitiveArrays() throws Exception {
                     ByteArray.copyFrom("some_bytes".getBytes(UTF_8)),
                     ByteArray.copyFrom("some_bytes".getBytes(UTF_8))))
             .build();
+    Mutation jsonb =
+        Mutation.newInsertOrUpdateBuilder("test")
+            .set("one")
+            .toPgJsonbArray(
+                ImmutableList.of(
+                    "{\"key123\":\"value123\", \"key321\":\"value321\"}",
+                    "{\"key456\":\"value456\", \"key789\":600}"))
+            .build();
     assertThat(MutationSizeEstimator.sizeOf(int64), is(24L));
     assertThat(MutationSizeEstimator.sizeOf(float64), is(16L));
     assertThat(MutationSizeEstimator.sizeOf(bool), is(4L));
     assertThat(MutationSizeEstimator.sizeOf(numeric), is(153L));
     assertThat(MutationSizeEstimator.sizeOf(pgNumeric), is(156L));
     assertThat(MutationSizeEstimator.sizeOf(json), is(62L));
     assertThat(MutationSizeEstimator.sizeOf(bytes), is(20L));
+    assertThat(MutationSizeEstimator.sizeOf(jsonb), is(77L));
   }
 
   @Test
@@ -162,13 +177,16 @@ public void nullPrimitiveArrays() throws Exception {
             .toPgNumericArray((Iterable<String>) null)
             .build();
     Mutation json = Mutation.newInsertOrUpdateBuilder("test").set("one").toJsonArray(null).build();
+    Mutation jsonb =
+        Mutation.newInsertOrUpdateBuilder("test").set("one").toPgJsonbArray(null).build();
 
     assertThat(MutationSizeEstimator.sizeOf(int64), is(0L));
     assertThat(MutationSizeEstimator.sizeOf(float64), is(0L));
     assertThat(MutationSizeEstimator.sizeOf(bool), is(0L));
     assertThat(MutationSizeEstimator.sizeOf(numeric), is(0L));
     assertThat(MutationSizeEstimator.sizeOf(pgNumeric), is(0L));
     assertThat(MutationSizeEstimator.sizeOf(json), is(0L));
+    assertThat(MutationSizeEstimator.sizeOf(jsonb), is(0L));
   }
 
   @Test
@@ -237,6 +255,29 @@ public void jsons() throws Exception {
     assertThat(MutationSizeEstimator.sizeOf(nullArray), is(0L));
   }
 
+  @Test
+  public void pgJsonb() throws Exception {
+    Mutation empty =
+        Mutation.newInsertOrUpdateBuilder("test").set("one").to(Value.pgJsonb("{}")).build();
+    Mutation nullValue =
+        Mutation.newInsertOrUpdateBuilder("test")
+            .set("one")
+            .to(Value.pgJsonb((String) null))
+            .build();
+    Mutation sample =
+        Mutation.newInsertOrUpdateBuilder("test")
+            .set("one")
+            .to(Value.pgJsonb("{\"type_name\":\"number\",\"value\":12345.123}"))
+            .build();
+    Mutation nullArray =
+        Mutation.newInsertOrUpdateBuilder("test").set("one").toPgJsonbArray(null).build();
+
+    assertThat(MutationSizeEstimator.sizeOf(empty), is(2L));
+    assertThat(MutationSizeEstimator.sizeOf(nullValue), is(0L));
+    assertThat(MutationSizeEstimator.sizeOf(sample), is(40L));
+    assertThat(MutationSizeEstimator.sizeOf(nullArray), is(0L));
+  }
+
   @Test
   public void dates() throws Exception {
     Mutation timestamp =

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerSchemaTest.java

@@ -75,10 +75,11 @@ public void testSinglePgTable() throws Exception {
             .addColumn("test", "maxKey", "character varying")
             .addColumn("test", "numericVal", "numeric")
             .addColumn("test", "commitTime", "spanner.commit_timestamp")
+            .addColumn("test", "jsonbCol", "jsonb")
             .build();
 
     assertEquals(1, schema.getTables().size());
-    assertEquals(4, schema.getColumns("test").size());
+    assertEquals(5, schema.getColumns("test").size());
     assertEquals(1, schema.getKeyParts("test").size());
     assertEquals(Type.timestamp(), schema.getColumns("test").get(3).getType());
   }
@@ -90,14 +91,15 @@ public void testTwoPgTables() throws Exception {
             .addColumn("test", "pk", "character varying(48)")
             .addKeyPart("test", "pk", false)
             .addColumn("test", "maxKey", "character varying")
+            .addColumn("test", "jsonbCol", "jsonb")
             .addColumn("other", "pk", "bigint")
             .addKeyPart("other", "pk", true)
             .addColumn("other", "maxKey", "character varying")
             .addColumn("other", "commitTime", "spanner.commit_timestamp")
             .build();
 
     assertEquals(2, schema.getTables().size());
-    assertEquals(2, schema.getColumns("test").size());
+    assertEquals(3, schema.getColumns("test").size());
     assertEquals(1, schema.getKeyParts("test").size());
 
     assertEquals(3, schema.getColumns("other").size());