feat: add max staleness to ExternalTableDefinition (#3499)

PhongChuong · web-flow · commit f1ebd5be5877 · 2024-10-01T13:58:34.000-04:00
* feat: add max staleness

TODO: Remove the manual testing anchors

* test: add max staleness to ExternalTableDefinition

* Format fix

* Remove unintentional changes
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExternalTableDefinition.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExternalTableDefinition.java
@@ -200,6 +200,20 @@ public Builder setMetadataCacheMode(String metadataCacheMode) {
 
     abstract Builder setMetadataCacheModeInner(String metadataCacheMode);
 
+    /**
+     * [Optional] Metadata Cache Mode for the table. Set this to enable caching of metadata from
+     * external data source.
+     *
+     * @see <a
+     *     href="https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource:-table">
+     *     MaxStaleness</a>
+     */
+    public Builder setMaxStaleness(String maxStaleness) {
+      return setMaxStalenessInner(maxStaleness);
+    }
+
+    abstract Builder setMaxStalenessInner(String maxStaleness);
+
     /** Creates an {@code ExternalTableDefinition} object. */
     @Override
     public abstract ExternalTableDefinition build();
@@ -305,6 +319,22 @@ public String getMetadataCacheMode() {
   @Nullable
   abstract String getMetadataCacheModeInner();
 
+  /**
+   * Returns the maximum staleness of data that could be returned when the table is queried.
+   * Staleness encoded as a string encoding of sql IntervalValue type.
+   *
+   * @see <a
+   *     href="hhttps://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource:-table">
+   *     MaxStaleness</a>
+   */
+  @Nullable
+  public String getMaxStaleness() {
+    return getMaxStalenessInner();
+  }
+
+  @Nullable
+  abstract String getMaxStalenessInner();
+
   /**
    * Returns the source format, and possibly some parsing options, of the external data. Supported
    * formats are {@code CSV} and {@code NEWLINE_DELIMITED_JSON}.
@@ -351,6 +381,9 @@ public HivePartitioningOptions getHivePartitioningOptions() {
   com.google.api.services.bigquery.model.Table toPb() {
     Table tablePb = super.toPb();
     tablePb.setExternalDataConfiguration(toExternalDataConfigurationPb());
+    if (getMaxStaleness() != null) {
+      tablePb.setMaxStaleness(getMaxStaleness());
+    }
     return tablePb;
   }
 
@@ -616,6 +649,9 @@ static ExternalTableDefinition fromPb(Table tablePb) {
       if (externalDataConfiguration.getMetadataCacheMode() != null) {
         builder.setMetadataCacheMode(externalDataConfiguration.getMetadataCacheMode());
       }
+      if (tablePb.getMaxStaleness() != null) {
+        builder.setMaxStaleness(tablePb.getMaxStaleness());
+      }
     }
     return builder.build();
   }
diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExternalTableDefinitionTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExternalTableDefinitionTest.java
@@ -59,8 +59,8 @@ public class ExternalTableDefinitionTest {
           .setSourceUriPrefix(SOURCE_URIS.get(0))
           .build();
   private static final String OBJECT_METADATA = "SIMPLE";
-
   private static final String METADATA_CACHE_MODE = "AUTOMATIC";
+  private static final String MAX_STALENESS = "INTERVAL 15 MINUTE";
   private static final ExternalTableDefinition EXTERNAL_TABLE_DEFINITION =
       ExternalTableDefinition.newBuilder(SOURCE_URIS, TABLE_SCHEMA, CSV_OPTIONS)
           .setFileSetSpecType("FILE_SET_SPEC_TYPE_FILE_SYSTEM_MATCH")
@@ -73,6 +73,7 @@ public class ExternalTableDefinitionTest {
           .setHivePartitioningOptions(HIVE_PARTITIONING_OPTIONS)
           .setObjectMetadata(OBJECT_METADATA)
           .setMetadataCacheMode(METADATA_CACHE_MODE)
+          .setMaxStaleness(MAX_STALENESS)
           .build();
 
   private static final ExternalTableDefinition EXTERNAL_TABLE_DEFINITION_AVRO =
@@ -174,5 +175,6 @@ private void compareExternalTableDefinition(
     assertEquals(expected.getHivePartitioningOptions(), value.getHivePartitioningOptions());
     assertEquals(expected.getObjectMetadata(), value.getObjectMetadata());
     assertEquals(expected.getMetadataCacheMode(), value.getMetadataCacheMode());
+    assertEquals(expected.getMaxStaleness(), value.getMaxStaleness());
   }
 }
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
@@ -2111,9 +2111,13 @@ public void testCreateAndGetTableWithSelectedField() {
   public void testCreateExternalTable() throws InterruptedException {
     String tableName = "test_create_external_table";
     TableId tableId = TableId.of(DATASET, tableName);
+
     ExternalTableDefinition externalTableDefinition =
         ExternalTableDefinition.of(
-            "gs://" + BUCKET + "/" + JSON_LOAD_FILE, TABLE_SCHEMA, FormatOptions.json());
+                "gs://" + BUCKET + "/" + JSON_LOAD_FILE, TABLE_SCHEMA, FormatOptions.json())
+            .toBuilder()
+            .setMaxStaleness("INTERVAL 15 MINUTE")
+            .build();
     TableInfo tableInfo = TableInfo.of(tableId, externalTableDefinition);
     Table createdTable = bigquery.create(tableInfo);
     assertNotNull(createdTable);
