Updated BigQueryReader so that BigQuery values of type TIMESTAMP are formatted similarly for local and service-based runs. This change updates the TIMESTAMP values produced in

chamikaramj · davorbonaci · commit 18c82ad5c749 · 2015-05-12T15:33:41.000-07:00
local runs to be formatted similar to the values produced by serviced-based runs. This closes apache#20 ----Release Notes---- This change modifies the formatting of the values of type TIMESTAMP produced by Dataflow BigQuery source. This change may only affect Dataflow jobs that are run in direct mode and Dataflow jobs that use BigQuery data as side inputs. This change will not affect Dataflow jobs that only use BigQuery data as a main input. [] ------------- Created by MOE: http://code.google.com/p/moe-java MOE_MIGRATED_REVID=93090065
diff --git a/sdk/src/main/java/com/google/cloud/dataflow/sdk/util/BigQueryTableRowIterator.java b/sdk/src/main/java/com/google/cloud/dataflow/sdk/util/BigQueryTableRowIterator.java
@@ -26,6 +26,9 @@
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.api.services.bigquery.model.TableSchema;
 
+import org.joda.time.format.DateTimeFormat;
+import org.joda.time.format.DateTimeFormatter;
+
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -82,6 +85,8 @@ public boolean hasNext() {
    *   <li> Record columns are {@link TableRow}s.
    *   <li> {@code BOOLEAN} columns are JSON booleans, hence Java {@link Boolean}s.
    *   <li> {@code FLOAT} columns are JSON floats, hence Java {@link Double}s.
+   *   <li> {@code TIMESTAMP} columns are {@link String}s that are of the format
+   *        {yyyy-MM-dd HH:mm:ss.SSS UTC}.
    *   <li> Every other atomic type is a {@link String}.
    * </ul></p>
    *
@@ -122,6 +127,14 @@ private Object getTypedCellValue(TableFieldSchema fieldSchema, Object v) {
       return Boolean.parseBoolean((String) v);
     }
 
+    if (fieldSchema.getType().equals("TIMESTAMP")) {
+      // Seconds to milliseconds
+      long milliSecs = (new Double(Double.parseDouble((String) v) * 1000)).longValue();
+      DateTimeFormatter formatter =
+          DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSS").withZoneUTC();
+      return formatter.print(milliSecs) + " UTC";
+    }
+
     return v;
   }
 
diff --git a/sdk/src/test/java/com/google/cloud/dataflow/sdk/util/BigQueryUtilTest.java b/sdk/src/test/java/com/google/cloud/dataflow/sdk/util/BigQueryUtilTest.java
@@ -166,6 +166,49 @@ private Table basicTableSchema() {
             )));
   }
 
+  private Table basicTableSchemaWithTime() {
+    return new Table()
+        .setSchema(new TableSchema()
+            .setFields(Arrays.asList(
+                new TableFieldSchema()
+                    .setName("name")
+                    .setType("STRING"),
+                new TableFieldSchema()
+                    .setName("time")
+                    .setType("TIMESTAMP"),
+                new TableFieldSchema()
+                    .setName("answer")
+                    .setType("INTEGER")
+            )));
+  }
+
+  @Test
+  public void testReadWithTime() throws IOException {
+    onTableGet(basicTableSchemaWithTime());
+
+    TableDataList dataList = rawDataList(rawRow("Arthur", "1.430397296789E9", 42));
+    onTableList(dataList);
+
+    BigQueryTableRowIterator iterator = new BigQueryTableRowIterator(
+        mockClient,
+        BigQueryIO.parseTableSpec("project:dataset.table"));
+
+    Assert.assertTrue(iterator.hasNext());
+    TableRow row = iterator.next();
+
+    Assert.assertTrue(row.containsKey("name"));
+    Assert.assertTrue(row.containsKey("time"));
+    Assert.assertTrue(row.containsKey("answer"));
+    Assert.assertEquals("Arthur", row.get("name"));
+    Assert.assertEquals("2015-04-30 12:34:56.789 UTC", row.get("time"));
+    Assert.assertEquals(42, row.get("answer"));
+
+    Assert.assertFalse(iterator.hasNext());
+
+    verifyTableGet();
+    verifyTabledataList();
+  }
+
   private TableRow rawRow(Object...args) {
     List<TableCell> cells = new LinkedList<>();
     for (Object a : args) {
