Add model table type for the new BigQuery ML models (#3529)

Fixes listTables on datasets containing models.

Author: gsabhnani

google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ModelTableDefinition.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.bigquery;
+
+import com.google.api.core.BetaApi;
+import com.google.api.services.bigquery.model.Table;
+import com.google.auto.value.AutoValue;
+import javax.annotation.Nullable;
+
+/**
+ * A Google BigQuery Model table definition. This definition is used to represent a BigQuery
+ * ML model.
+ *
+ * @see <a href="https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create#models_in_bqml_name">BigQuery ML Model</a>
+ */
+@AutoValue
+@BetaApi
+public abstract class ModelTableDefinition extends TableDefinition {
+
+  private static final long serialVersionUID = 2113445776046717900L;
+
+  @AutoValue.Builder
+  public abstract static class Builder
+      extends TableDefinition.Builder<ModelTableDefinition, Builder> {
+
+    public abstract Builder setNumBytes(Long numBytes);
+
+    public abstract Builder setLocation(String location);
+
+    public abstract Builder setType(Type type);
+
+    /** Creates a {@code ModelTableDefinition} object. */
+    public abstract ModelTableDefinition build();
+  }
+
+  /** Returns the size of this table in bytes, excluding any data in the streaming buffer. */
+  @Nullable
+  public abstract Long getNumBytes();
+
+  /**
+   * Returns the geographic location where the table should reside. This value is inherited from the
+   * dataset.
+   *
+   * @see <a
+   *     href="https://cloud.google.com/bigquery/docs/managing_jobs_datasets_projects#dataset-location">
+   *     Dataset Location</a>
+   */
+  @Nullable
+  public abstract String getLocation();
+
+  /**
+   * Returns a builder for a BigQuery ML model table definition.
+   */
+  public static Builder newBuilder() {
+    return new AutoValue_ModelTableDefinition.Builder().setType(Type.MODEL);
+  }
+
+  /** Returns a builder for the {@code <ModelTableDefinition} object. */
+  public abstract Builder toBuilder();
+
+  @Override
+  Table toPb() {
+    Table tablePb = super.toPb();
+    tablePb.setNumBytes(getNumBytes());
+    tablePb.setLocation(getLocation());
+    return tablePb;
+  }
+
+  @SuppressWarnings("unchecked")
+  static ModelTableDefinition fromPb(Table tablePb) {
+    Builder builder = newBuilder().table(tablePb);
+    return builder.setNumBytes(tablePb.getNumBytes()).setLocation(tablePb.getLocation()).build();
+  }
+}

google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableDefinition.java

@@ -71,6 +71,14 @@ public Type apply(String constant) {
      */
     public static final Type EXTERNAL = type.createAndRegister("EXTERNAL");
 
+    /**
+     * A BigQuery table representing BigQuery ML Model.
+     *
+     * @see <a href="https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create#models_in_bqml_name">
+     *     BigQuery ML Model</a>
+     */
+    public static final Type MODEL = type.createAndRegister("MODEL");
+
     private Type(String constant) {
       super(constant);
     }
@@ -157,6 +165,8 @@ static <T extends TableDefinition> T fromPb(Table tablePb) {
         return (T) ViewDefinition.fromPb(tablePb);
       case "EXTERNAL":
         return (T) ExternalTableDefinition.fromPb(tablePb);
+      case "MODEL":
+        return (T) ModelTableDefinition.fromPb(tablePb);
       default:
         // never reached
         throw new IllegalArgumentException("Format " + tablePb.getType() + " is not supported");

google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/TableInfo.java

@@ -32,7 +32,8 @@
 /**
  * Google BigQuery table information. Use {@link StandardTableDefinition} to create simple BigQuery
  * table. Use {@link ViewDefinition} to create a BigQuery view. Use {@link ExternalTableDefinition}
- * to create a BigQuery a table backed by external data.
+ * to create a BigQuery a table backed by external data. Use {@link ModelDefinition} to create a 
+ * BigQuery ML model.
  *
  * @see <a href="https://cloud.google.com/bigquery/docs/tables">Managing Tables</a>
  */

google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -108,10 +108,14 @@ public class BigQueryImplTest {
   private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA1, FIELD_SCHEMA2);
   private static final StandardTableDefinition TABLE_DEFINITION =
       StandardTableDefinition.of(TABLE_SCHEMA);
+  private static final ModelTableDefinition MODEL_TABLE_DEFINITION =
+      ModelTableDefinition.newBuilder().build();
   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 =
       TableInfo.of(TABLE_ID_WITH_PROJECT, TABLE_DEFINITION);
+  private static final TableInfo MODEL_TABLE_INFO_WITH_PROJECT =
+      TableInfo.of(TABLE_ID_WITH_PROJECT, MODEL_TABLE_DEFINITION);
   private static final LoadJobConfiguration LOAD_JOB_CONFIGURATION =
       LoadJobConfiguration.of(TABLE_ID, "URI");
   private static final LoadJobConfiguration LOAD_JOB_CONFIGURATION_WITH_PROJECT =
@@ -613,7 +617,8 @@ public void testListTables() {
     ImmutableList<Table> tableList =
         ImmutableList.of(
             new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)),
-            new Table(bigquery, new TableInfo.BuilderImpl(OTHER_TABLE_INFO)));
+            new Table(bigquery, new TableInfo.BuilderImpl(OTHER_TABLE_INFO)),
+            new Table(bigquery, new TableInfo.BuilderImpl(MODEL_TABLE_INFO_WITH_PROJECT)));
     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, EMPTY_RPC_OPTIONS))