From 069240fc4ebfc54d39d84fc3a9cc31d093cfa70a Mon Sep 17 00:00:00 2001
From: Praful Makani <praful@qlogic.io>
Date: Fri, 17 Jan 2020 22:20:38 +0530
Subject: [PATCH] feat: add range partitioning field to tableslist (#114)

---
 .../bigquery/spi/v2/HttpBigQueryRpc.java      |  3 +-
 .../cloud/bigquery/BigQueryImplTest.java      | 27 +++++++++++++++++
 .../cloud/bigquery/it/ITBigQueryTest.java     | 30 +++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
index efebd53a4..391209fcb 100644
--- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
+++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
@@ -299,7 +299,8 @@ public Table apply(TableList.Tables tablePb) {
                       .setTableReference(tablePb.getTableReference())
                       .setType(tablePb.getType())
                       .setCreationTime(tablePb.getCreationTime())
-                      .setTimePartitioning(tablePb.getTimePartitioning());
+                      .setTimePartitioning(tablePb.getTimePartitioning())
+                      .setRangePartitioning(tablePb.getRangePartitioning());
                 }
               }));
     } catch (IOException ex) {
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java
index b1e1a5d36..7ab3e8ef2 100644
--- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java
@@ -121,6 +121,17 @@ public class BigQueryImplTest {
           .setSchema(TABLE_SCHEMA)
           .setTimePartitioning(TIME_PARTITIONING)
           .build();
+  private static final RangePartitioning.Range RANGE =
+      RangePartitioning.Range.newBuilder().setStart(1L).setInterval(2L).setEnd(10L).build();
+  private static final RangePartitioning RANGE_PARTITIONING =
+      RangePartitioning.newBuilder().setField("IntegerField").setRange(RANGE).build();
+  private static final StandardTableDefinition TABLE_DEFINITION_WITH_RANGE_PARTITIONING =
+      StandardTableDefinition.newBuilder()
+          .setSchema(TABLE_SCHEMA)
+          .setRangePartitioning(RANGE_PARTITIONING)
+          .build();
+  private static final TableInfo TABLE_INFO_RANGE_PARTITIONING =
+      TableInfo.of(TABLE_ID, TABLE_DEFINITION_WITH_RANGE_PARTITIONING);
   private static final TableInfo TABLE_INFO = TableInfo.of(TABLE_ID, TABLE_DEFINITION);
   private static final TableInfo OTHER_TABLE_INFO = TableInfo.of(OTHER_TABLE_ID, TABLE_DEFINITION);
   private static final TableInfo TABLE_INFO_WITH_PROJECT =
@@ -901,6 +912,22 @@ public void testListTablesReturnedParameters() {
     assertArrayEquals(tableList.toArray(), Iterables.toArray(page.getValues(), Table.class));
   }
 
+  @Test
+  public void testListTablesWithRangePartitioning() {
+    bigquery = options.getService();
+    ImmutableList<Table> tableList =
+        ImmutableList.of(
+            new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_RANGE_PARTITIONING)));
+    Tuple<String, Iterable<com.google.api.services.bigquery.model.Table>> result =
+        Tuple.of(CURSOR, Iterables.transform(tableList, TableInfo.TO_PB_FUNCTION));
+    EasyMock.expect(bigqueryRpcMock.listTables(PROJECT, DATASET, TABLE_LIST_OPTIONS))
+        .andReturn(result);
+    EasyMock.replay(bigqueryRpcMock);
+    Page<Table> page = bigquery.listTables(DATASET, TABLE_LIST_PAGE_SIZE, TABLE_LIST_PAGE_TOKEN);
+    assertEquals(CURSOR, page.getNextPageToken());
+    assertArrayEquals(tableList.toArray(), Iterables.toArray(page.getValues(), Table.class));
+  }
+
   @Test
   public void testListTablesFromDatasetId() {
     bigquery = options.getService();
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
index ade2e65cf..547c0892e 100644
--- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
+++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -712,6 +712,36 @@ public void testListTablesWithPartitioning() {
     }
   }
 
+  @Test
+  public void testListTablesWithRangePartitioning() {
+    String tableName = "test_list_tables_range_partitioning";
+    StandardTableDefinition tableDefinition =
+        StandardTableDefinition.newBuilder()
+            .setSchema(TABLE_SCHEMA)
+            .setRangePartitioning(RANGE_PARTITIONING)
+            .build();
+    TableInfo tableInfo = TableInfo.of(TableId.of(DATASET, tableName), tableDefinition);
+    Table createdRangePartitioningTable = bigquery.create(tableInfo);
+    assertNotNull(createdRangePartitioningTable);
+    try {
+      Page<Table> tables = bigquery.listTables(DATASET);
+      boolean found = false;
+      Iterator<Table> tableIterator = tables.getValues().iterator();
+      while (tableIterator.hasNext() && !found) {
+        StandardTableDefinition standardTableDefinition = tableIterator.next().getDefinition();
+        if (standardTableDefinition.getRangePartitioning() != null) {
+          assertEquals(RANGE_PARTITIONING, standardTableDefinition.getRangePartitioning());
+          assertEquals(RANGE, standardTableDefinition.getRangePartitioning().getRange());
+          assertEquals("IntegerField", standardTableDefinition.getRangePartitioning().getField());
+          found = true;
+        }
+      }
+      assertTrue(found);
+    } finally {
+      createdRangePartitioningTable.delete();
+    }
+  }
+
   @Test
   public void testListPartitions() throws InterruptedException {
     String tableName = "test_table_partitions";