feat: add support for byte and short (#2789)

agrawal-siddharth · web-flow · commit 98a714fbdd40 · 2024-12-05T10:38:45.000-08:00
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessage.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessage.java
@@ -486,8 +486,9 @@ private FieldDescriptorAndFieldTableSchema computeDescriptorAndSchema(
    * Fills a non-repetaed protoField with the json data.
    *
    * @param protoMsg The protocol buffer message being constructed
-   * @param fieldDescriptor
-   * @param fieldSchema
+   * @param fieldDescriptor Proto format to be transmitted over the wire (derived from table schema
+   *     via BQTableSchemaToProtoDescriptor.BQTableSchemaModeMap)
+   * @param fieldSchema Actual table column schema type if available
    * @param json
    * @param exactJsonKeyName Exact key name in JSONObject instead of lowercased version
    * @param currentScope Debugging purposes
@@ -647,6 +648,12 @@ private void fillField(
         } else if (val instanceof Long) {
           protoMsg.setField(fieldDescriptor, val);
           return;
+        } else if (val instanceof Byte) {
+          protoMsg.setField(fieldDescriptor, Long.valueOf((Byte) val));
+          return;
+        } else if (val instanceof Short) {
+          protoMsg.setField(fieldDescriptor, Long.valueOf((Short) val));
+          return;
         }
         if (val instanceof String) {
           Long parsed = Longs.tryParse((String) val);
@@ -727,8 +734,9 @@ private void fillField(
    * Fills a repeated protoField with the json data.
    *
    * @param protoMsg The protocol buffer message being constructed
-   * @param fieldDescriptor
-   * @param fieldSchema
+   * @param fieldDescriptor Proto format to be transmitted over the wire (derived from table schema
+   *     via BQTableSchemaToProtoDescriptor.BQTableSchemaModeMap)
+   * @param fieldSchema Actual table column schema type if available
    * @param json If root level has no matching fields, throws exception.
    * @param exactJsonKeyName Exact key name in JSONObject instead of lowercased version
    * @param currentScope Debugging purposes
@@ -912,6 +920,10 @@ private void fillRepeatedField(
             protoMsg.addRepeatedField(fieldDescriptor, Long.valueOf((Integer) val));
           } else if (val instanceof Long) {
             protoMsg.addRepeatedField(fieldDescriptor, val);
+          } else if (val instanceof Byte) {
+            protoMsg.addRepeatedField(fieldDescriptor, Long.valueOf((Byte) val));
+          } else if (val instanceof Short) {
+            protoMsg.addRepeatedField(fieldDescriptor, Long.valueOf((Short) val));
           } else if (val instanceof String) {
             Long parsed = Longs.tryParse((String) val);
             if (parsed != null) {
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessageTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonToProtoMessageTest.java
@@ -30,7 +30,9 @@
 import java.math.BigDecimal;
 import java.time.LocalTime;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.logging.Logger;
@@ -579,8 +581,8 @@ public void testInt64() throws Exception {
     TestInt64 expectedProto =
         TestInt64.newBuilder().setByte(1).setShort(1).setInt(1).setLong(1).setString(1).build();
     JSONObject json = new JSONObject();
-    json.put("byte", (byte) 1);
-    json.put("short", (short) 1);
+    json.put("byte", (byte) 1); // This does NOT actually verify byte as it is converted to int
+    json.put("short", (short) 1); // This does NOT actually verify short as it is converted to int
     json.put("int", 1);
     json.put("long", 1L);
     json.put("string", "1");
@@ -589,6 +591,46 @@ public void testInt64() throws Exception {
     assertEquals(expectedProto, protoMsg);
   }
 
+  @Test
+  public void testInt64Extended() throws Exception {
+    TestInt64 expectedProto =
+        TestInt64.newBuilder().setByte(1).setShort(1).setInt(1).setLong(1).setString(1).build();
+    Map map = new HashMap();
+    map.put("byte", (byte) 1);
+    map.put("short", (short) 1);
+    map.put("int", (int) 1);
+    map.put("long", (long) 1);
+    map.put("string", "1");
+    JSONObject json = new JSONObject(map);
+    DynamicMessage protoMsg =
+        JsonToProtoMessage.INSTANCE.convertToProtoMessage(TestInt64.getDescriptor(), json);
+    assertEquals(expectedProto, protoMsg);
+  }
+
+  @Test
+  public void testInt64Repeated() throws Exception {
+    RepeatedInt64 expectedProto =
+        RepeatedInt64.newBuilder()
+            .addTestRepeated(1)
+            .addTestRepeated(1)
+            .addTestRepeated(1)
+            .addTestRepeated(1)
+            .addTestRepeated(1)
+            .build();
+    Collection collection = new ArrayList();
+    collection.add((byte) 1);
+    collection.add((short) 1);
+    collection.add((int) 1);
+    collection.add((long) 1);
+    collection.add("1");
+    JSONArray array = new JSONArray(collection);
+    JSONObject json = new JSONObject();
+    json.put("test_repeated", array);
+    DynamicMessage protoMsg =
+        JsonToProtoMessage.INSTANCE.convertToProtoMessage(RepeatedInt64.getDescriptor(), json);
+    assertEquals(expectedProto, protoMsg);
+  }
+
   @Test
   public void testInt32() throws Exception {
     TestInt32 expectedProto =
