Merge pull request #21938 from [cherry-pick][release-2.40.0] BigQueryIO: Adding the BASIC view setting to getTable request (#21879)

[cherry-pick][release-2.40.0] BigQueryIO: Adding the BASIC view setting to getTable request  (#21879)

Author: pabloem

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java

@@ -33,6 +33,7 @@
 import java.io.IOException;
 import java.math.BigInteger;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -516,7 +517,9 @@ static String randomUUIDString() {
 
   static void verifyTableNotExistOrEmpty(DatasetService datasetService, TableReference tableRef) {
     try {
-      if (datasetService.getTable(tableRef) != null) {
+      if (datasetService.getTable(
+              tableRef, Collections.emptyList(), DatasetService.TableMetadataView.BASIC)
+          != null) {
         checkState(
             datasetService.isTableEmpty(tableRef),
             "BigQuery table is not empty: %s.",

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java

@@ -111,6 +111,16 @@ JobStatistics dryRunQuery(
 
   /** An interface to get, create and delete Cloud BigQuery datasets and tables. */
   public interface DatasetService extends AutoCloseable {
+
+    // maps the values at
+    // https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/get#TableMetadataView
+    enum TableMetadataView {
+      TABLE_METADATA_VIEW_UNSPECIFIED,
+      BASIC,
+      STORAGE_STATS,
+      FULL;
+    };
+
     /**
      * Gets the specified {@link Table} resource by table ID.
      *
@@ -123,6 +133,10 @@ public interface DatasetService extends AutoCloseable {
     Table getTable(TableReference tableRef, List<String> selectedFields)
         throws InterruptedException, IOException;
 
+    @Nullable
+    Table getTable(TableReference tableRef, List<String> selectedFields, TableMetadataView view)
+        throws InterruptedException, IOException;
+
     /** Creates the specified table if it does not exist. */
     void createTable(Table table) throws InterruptedException, IOException;
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java

@@ -616,13 +616,24 @@ private DatasetServiceImpl(BigQueryOptions bqOptions) {
     @Override
     public @Nullable Table getTable(TableReference tableRef, List<String> selectedFields)
         throws IOException, InterruptedException {
-      return getTable(tableRef, selectedFields, createDefaultBackoff(), Sleeper.DEFAULT);
+      return getTable(tableRef, selectedFields, TableMetadataView.STORAGE_STATS);
+    }
+
+    @Override
+    public @Nullable Table getTable(
+        TableReference tableRef, List<String> selectedFields, TableMetadataView view)
+        throws IOException, InterruptedException {
+      return getTable(tableRef, selectedFields, view, createDefaultBackoff(), Sleeper.DEFAULT);
     }
 
     @VisibleForTesting
     @Nullable
     Table getTable(
-        TableReference ref, List<String> selectedFields, BackOff backoff, Sleeper sleeper)
+        TableReference ref,
+        List<String> selectedFields,
+        TableMetadataView view,
+        BackOff backoff,
+        Sleeper sleeper)
         throws IOException, InterruptedException {
       Tables.Get get =
           client
@@ -632,6 +643,9 @@ Table getTable(
       if (!selectedFields.isEmpty()) {
         get.setSelectedFields(String.join(",", selectedFields));
       }
+      if (view != null) {
+        get.set("view", view.name());
+      }
       try {
         return executeWithRetries(
             get,

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java

@@ -112,7 +112,9 @@ private static void tryCreateTable(
     tableReference.setTableId(BigQueryHelpers.stripPartitionDecorator(tableReference.getTableId()));
     try (DatasetService datasetService =
         bqServices.getDatasetService(context.getPipelineOptions().as(BigQueryOptions.class))) {
-      if (datasetService.getTable(tableReference) == null) {
+      if (datasetService.getTable(
+              tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC)
+          == null) {
         TableSchema tableSchema = schemaSupplier.get();
         Preconditions.checkArgumentNotNull(
             tableSchema,

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableSchemaCache.java

@@ -22,6 +22,7 @@
 import com.google.api.services.bigquery.model.TableSchema;
 import com.google.auto.value.AutoValue;
 import java.io.IOException;
+import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
@@ -166,7 +167,12 @@ public TableSchema getSchema(TableReference tableReference, DatasetService datas
     if (!schemaHolder.isPresent()) {
       // Not initialized. Query the new schema with the monitor released and then update the cache.
       try {
-        @Nullable Table table = datasetService.getTable(tableReference);
+        // requesting the BASIC view will prevent BQ backend to run calculations
+        // related with storage stats that are not needed here
+        @Nullable
+        Table table =
+            datasetService.getTable(
+                tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC);
         schemaHolder =
             Optional.ofNullable((table == null) ? null : SchemaHolder.of(table.getSchema(), 0));
       } catch (Exception e) {
@@ -298,7 +304,12 @@ private Map<String, TableSchema> refreshAll(Map<String, Refresh> tables)
     Map<String, TableSchema> schemas = Maps.newHashMapWithExpectedSize(tables.size());
     for (Map.Entry<String, Refresh> entry : tables.entrySet()) {
       TableReference tableReference = BigQueryHelpers.parseTableSpec(entry.getKey());
-      Table table = entry.getValue().getDatasetService().getTable(tableReference);
+      Table table =
+          entry
+              .getValue()
+              .getDatasetService()
+              .getTable(
+                  tableReference, Collections.emptyList(), DatasetService.TableMetadataView.BASIC);
       if (table == null) {
         throw new RuntimeException("Did not get value for table " + tableReference);
       }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java

@@ -173,6 +173,23 @@ public Table getTable(TableReference tableRef) throws InterruptedException, IOEx
   @Override
   public Table getTable(TableReference tableRef, @Nullable List<String> selectedFields)
       throws InterruptedException, IOException {
+    return getTable(tableRef, selectedFields, null);
+  }
+
+  @Override
+  public Table getTable(
+      TableReference tableRef,
+      @Nullable List<String> selectedFields,
+      @Nullable TableMetadataView view)
+      throws InterruptedException, IOException {
+    return getTableImpl(tableRef, selectedFields, view);
+  }
+
+  public Table getTableImpl(
+      TableReference tableRef,
+      @Nullable List<String> selectedFields,
+      @Nullable TableMetadataView view)
+      throws InterruptedException, IOException {
     synchronized (FakeDatasetService.class) {
       Map<String, TableContainer> dataset =
           tables.get(tableRef.getProjectId(), tableRef.getDatasetId());

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java

@@ -474,7 +474,7 @@ public void testGetTableSucceeds() throws Exception {
 
     Table table =
         datasetService.getTable(
-            tableRef, Collections.emptyList(), BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
+            tableRef, Collections.emptyList(), null, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
 
     assertEquals(testTable, table);
     verifyAllResponsesAreRead();
@@ -499,7 +499,7 @@ public void testGetTableNotFound() throws IOException, InterruptedException {
             .setTableId("tableId");
     Table table =
         datasetService.getTable(
-            tableRef, Collections.emptyList(), BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
+            tableRef, Collections.emptyList(), null, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
 
     assertNull(table);
     verifyAllResponsesAreRead();
@@ -526,7 +526,7 @@ public void testGetTableThrows() throws Exception {
         new BigQueryServicesImpl.DatasetServiceImpl(
             bigquery, null, PipelineOptionsFactory.create());
     datasetService.getTable(
-        tableRef, Collections.emptyList(), BackOff.STOP_BACKOFF, Sleeper.DEFAULT);
+        tableRef, Collections.emptyList(), null, BackOff.STOP_BACKOFF, Sleeper.DEFAULT);
   }
 
   @Test

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryUtilTest.java

@@ -125,13 +125,15 @@ private void onTableGet(Table table) throws IOException {
     when(mockClient.tables()).thenReturn(mockTables);
     when(mockTables.get(anyString(), anyString(), anyString())).thenReturn(mockTablesGet);
     when(mockTablesGet.setPrettyPrint(false)).thenReturn(mockTablesGet);
+    when(mockTablesGet.set(anyString(), anyString())).thenReturn(mockTablesGet);
     when(mockTablesGet.execute()).thenReturn(table);
   }
 
   private void verifyTableGet() throws IOException {
     verify(mockClient).tables();
     verify(mockTables).get("project", "dataset", "table");
     verify(mockTablesGet, atLeastOnce()).setPrettyPrint(false);
+    verify(mockTablesGet, atLeastOnce()).set(anyString(), anyString());
     verify(mockTablesGet, atLeastOnce()).execute();
   }