fix: properly handle external table schema on table update (#2236)

the-other-tim-brown · web-flow · commit 460ef318297f · 2022-10-17T23:28:13.000Z
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [x] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigquery/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) Fixes #2235 ☕️ If you write sample code, please follow the [samples format]( https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
@@ -288,12 +288,7 @@ public Table create(TableInfo tableInfo, TableOption... options) {
                     ? getOptions().getProjectId()
                     : tableInfo.getTableId().getProject())
             .toPb();
-    // Set schema on the Table for permanent external table
-    if (tablePb.getExternalDataConfiguration() != null) {
-      tablePb.setSchema(tablePb.getExternalDataConfiguration().getSchema());
-      // clear table schema on ExternalDataConfiguration
-      tablePb.getExternalDataConfiguration().setSchema(null);
-    }
+    handleExternalTableSchema(tablePb);
     final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
     try {
       return Table.fromPb(
@@ -313,6 +308,16 @@ public com.google.api.services.bigquery.model.Table call() {
     }
   }
 
+  private void handleExternalTableSchema(
+      final com.google.api.services.bigquery.model.Table tablePb) {
+    // Set schema on the Table for permanent external table
+    if (tablePb.getExternalDataConfiguration() != null) {
+      tablePb.setSchema(tablePb.getExternalDataConfiguration().getSchema());
+      // clear table schema on ExternalDataConfiguration
+      tablePb.getExternalDataConfiguration().setSchema(null);
+    }
+  }
+
   @Override
   public Routine create(RoutineInfo routineInfo, RoutineOption... options) {
     final com.google.api.services.bigquery.model.Routine routinePb =
@@ -679,6 +684,7 @@ public Table update(TableInfo tableInfo, TableOption... options) {
                     ? getOptions().getProjectId()
                     : tableInfo.getTableId().getProject())
             .toPb();
+    handleExternalTableSchema(tablePb);
     final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
     try {
       return Table.fromPb(
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
@@ -812,6 +812,25 @@ public void testCreateTable() {
     verify(bigqueryRpcMock).create(tableInfo.toPb(), EMPTY_RPC_OPTIONS);
   }
 
+  @Test
+  public void tesCreateExternalTable() {
+    TableInfo createTableInfo =
+        TableInfo.of(TABLE_ID, ExternalTableDefinition.newBuilder().setSchema(TABLE_SCHEMA).build())
+            .setProjectId(OTHER_PROJECT);
+
+    com.google.api.services.bigquery.model.Table expectedCreateInput =
+        createTableInfo.toPb().setSchema(TABLE_SCHEMA.toPb());
+    expectedCreateInput.getExternalDataConfiguration().setSchema(null);
+    when(bigqueryRpcMock.create(expectedCreateInput, EMPTY_RPC_OPTIONS))
+        .thenReturn(createTableInfo.toPb());
+    BigQueryOptions bigQueryOptions =
+        createBigQueryOptionsForProject(OTHER_PROJECT, rpcFactoryMock);
+    bigquery = bigQueryOptions.getService();
+    Table table = bigquery.create(createTableInfo);
+    assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(createTableInfo)), table);
+    verify(bigqueryRpcMock).create(expectedCreateInput, EMPTY_RPC_OPTIONS);
+  }
+
   @Test
   public void testCreateTableWithoutProject() {
     TableInfo tableInfo = TABLE_INFO.setProjectId(PROJECT);
@@ -1189,6 +1208,25 @@ public void testUpdateTable() {
     verify(bigqueryRpcMock).patch(updatedTableInfo.toPb(), EMPTY_RPC_OPTIONS);
   }
 
+  @Test
+  public void testUpdateExternalTableWithNewSchema() {
+    TableInfo updatedTableInfo =
+        TableInfo.of(TABLE_ID, ExternalTableDefinition.newBuilder().setSchema(TABLE_SCHEMA).build())
+            .setProjectId(OTHER_PROJECT);
+
+    com.google.api.services.bigquery.model.Table expectedPatchInput =
+        updatedTableInfo.toPb().setSchema(TABLE_SCHEMA.toPb());
+    expectedPatchInput.getExternalDataConfiguration().setSchema(null);
+    when(bigqueryRpcMock.patch(expectedPatchInput, EMPTY_RPC_OPTIONS))
+        .thenReturn(updatedTableInfo.toPb());
+    BigQueryOptions bigQueryOptions =
+        createBigQueryOptionsForProject(OTHER_PROJECT, rpcFactoryMock);
+    bigquery = bigQueryOptions.getService();
+    Table table = bigquery.update(updatedTableInfo);
+    assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(updatedTableInfo)), table);
+    verify(bigqueryRpcMock).patch(expectedPatchInput, EMPTY_RPC_OPTIONS);
+  }
+
   @Test
   public void testUpdateTableWithoutProject() {
     TableInfo tableInfo = TABLE_INFO.setProjectId(PROJECT);
