From 6412fb2c03fc6b2c6d220d72bbd89e47922bf970 Mon Sep 17 00:00:00 2001
From: Stephanie Wang <stephaniewang526@users.noreply.github.com>
Date: Mon, 14 Mar 2022 18:30:16 -0400
Subject: [PATCH] feat: add JsonToProtoMessage support for TIMESTAMP (#1574)

Fixes #1515
---
 .../storage/v1/JsonToProtoMessage.java         | 18 ++++++++++++++++++
 .../v1/it/ITBigQueryWriteManualClientTest.java | 15 +++++++++++++++
 2 files changed, 33 insertions(+)

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
index 50e57b551c..668a447dbf 100644
--- 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
@@ -25,6 +25,7 @@
 import com.google.protobuf.Message;
 import com.google.protobuf.UninitializedMessageException;
 import java.math.BigDecimal;
+import java.sql.Timestamp;
 import java.time.LocalDate;
 import java.util.List;
 import java.util.logging.Logger;
@@ -294,6 +295,14 @@ private static void fillField(
               protoMsg.setField(fieldDescriptor, (Long) val);
               return;
             }
+          } else if (fieldSchema.getType() == TableFieldSchema.Type.TIMESTAMP) {
+            if (val instanceof String) {
+              protoMsg.setField(fieldDescriptor, Timestamp.valueOf((String) val).getTime());
+              return;
+            } else if (val instanceof Long) {
+              protoMsg.setField(fieldDescriptor, (Long) val);
+              return;
+            }
           }
         }
         if (val instanceof Integer) {
@@ -469,6 +478,15 @@ private static void fillRepeatedField(
             } else {
               fail = true;
             }
+          } else if (fieldSchema != null
+              && fieldSchema.getType() == TableFieldSchema.Type.TIMESTAMP) {
+            if (val instanceof String) {
+              protoMsg.addRepeatedField(fieldDescriptor, Timestamp.valueOf((String) val).getTime());
+            } else if (val instanceof Long) {
+              protoMsg.addRepeatedField(fieldDescriptor, (Long) val);
+            } else {
+              fail = true;
+            }
           } else if (val instanceof Integer) {
             protoMsg.addRepeatedField(fieldDescriptor, new Long((Integer) val));
           } else if (val instanceof Long) {
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
index 89e2c80e27..13d0f1a697 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
@@ -39,6 +39,7 @@
 import io.grpc.Status.Code;
 import java.io.IOException;
 import java.math.BigDecimal;
+import java.sql.Timestamp;
 import java.util.*;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -339,12 +340,19 @@ public void testJsonStreamWriterWithDefaultStream()
             .setMode(TableFieldSchema.Mode.REPEATED)
             .setName("test_bytestring_repeated")
             .build();
+    TableFieldSchema TEST_TIMESTAMP =
+        TableFieldSchema.newBuilder()
+            .setName("test_timeStamp")
+            .setType(TableFieldSchema.Type.TIMESTAMP)
+            .setMode(TableFieldSchema.Mode.NULLABLE)
+            .build();
     TableSchema tableSchema =
         TableSchema.newBuilder()
             .addFields(0, TEST_STRING)
             .addFields(1, TEST_DATE)
             .addFields(2, TEST_NUMERIC)
             .addFields(3, TEST_REPEATED_BYTESTRING)
+            .addFields(4, TEST_TIMESTAMP)
             .build();
     TableInfo tableInfo =
         TableInfo.newBuilder(
@@ -364,6 +372,9 @@ public void testJsonStreamWriterWithDefaultStream()
                         com.google.cloud.bigquery.Field.newBuilder(
                                 "test_bytestring_repeated", StandardSQLTypeName.BYTES)
                             .setMode(Field.Mode.REPEATED)
+                            .build(),
+                        com.google.cloud.bigquery.Field.newBuilder(
+                                "test_timestamp", StandardSQLTypeName.TIMESTAMP)
                             .build())))
             .build();
 
@@ -396,6 +407,7 @@ public void testJsonStreamWriterWithDefaultStream()
                 ByteString.copyFromUtf8("a").toByteArray(),
                 ByteString.copyFromUtf8("b").toByteArray()
               }));
+      row1.put("test_timestamp", "2022-02-06 07:24:47.84");
       JSONArray jsonArr1 = new JSONArray(new JSONObject[] {row1});
 
       ApiFuture<AppendRowsResponse> response1 = jsonStreamWriter.append(jsonArr1, -1);
@@ -444,6 +456,9 @@ public void testJsonStreamWriterWithDefaultStream()
       assertEquals("2020-10-01T12:00:00", currentRow.get(2).getStringValue());
       assertEquals(2, currentRow.get(3).getRepeatedValue().size());
       assertEquals("Yg==", currentRow.get(3).getRepeatedValue().get(1).getStringValue());
+      assertEquals(
+          Timestamp.valueOf("2022-02-06 07:24:47.84").getTime(),
+          currentRow.get(4).getTimestampValue()); // timestamp long of "2022-02-06 07:24:47.84"
       assertEquals("bbb", iter.next().get(0).getStringValue());
       assertEquals("ccc", iter.next().get(0).getStringValue());
       assertEquals("ddd", iter.next().get(0).getStringValue());