Use TableReference instead of String for table identification in BigQueryStorageApiInsertError

Following reviewer feedback, this change updates BigQueryStorageApiInsertError
to use TableReference instead of String tableUrn for better consistency with
BigQueryInsertError.

Changes:
- Changed BigQueryStorageApiInsertError to use TableReference field instead of String tableUrn
- Updated BigQueryStorageApiInsertErrorCoder to use BigQueryHelpers.toTableSpec() and parseTableSpec()
- Added null safety checks in coder to prevent NullPointerException when table is unknown
- Updated all calling sites to pass TableReference from TableDestination.getTableReference():
- StorageApiWriteUnshardedRecords.java (3 locations)
- StorageApiWritesShardedRecords.java (3 locations + added tableReference variable)
- StorageApiConvertMessages.java (1 location)
- Added TableReference imports to modified files

The null checks in the coder ensure job stability even when table information
is unavailable during error handling, preventing pipeline failures in error
reporting scenarios.

Author: hekk-kaori-maeda

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.io.gcp.bigquery;
 
+import com.google.api.services.bigquery.model.TableReference;
 import com.google.api.services.bigquery.model.TableRow;
 import javax.annotation.Nullable;
 
@@ -25,9 +26,7 @@ public class BigQueryStorageApiInsertError {
 
   private @Nullable String errorMessage;
 
-  private @Nullable String tableUrn;
-
-  private @Nullable String[] parsedParts;
+  private @Nullable TableReference table;
 
   public BigQueryStorageApiInsertError(TableRow row) {
     this(row, null, null);
@@ -38,10 +37,10 @@ public BigQueryStorageApiInsertError(TableRow row, @Nullable String errorMessage
   }
 
   public BigQueryStorageApiInsertError(
-      TableRow row, @Nullable String errorMessage, @Nullable String tableUrn) {
+      TableRow row, @Nullable String errorMessage, @Nullable TableReference table) {
     this.row = row;
     this.errorMessage = errorMessage;
-    this.tableUrn = tableUrn;
+    this.table = table;
   }
 
   public TableRow getRow() {
@@ -54,37 +53,8 @@ public String getErrorMessage() {
   }
 
   @Nullable
-  public String getTableUrn() {
-    return tableUrn;
-  }
-
-  @Nullable
-  public String getProjectId() {
-    return getParsedPart(1);
-  }
-
-  @Nullable
-  public String getDatasetId() {
-    return getParsedPart(3);
-  }
-
-  @Nullable
-  public String getTableId() {
-    return getParsedPart(5);
-  }
-
-  @Nullable
-  private String getParsedPart(int index) {
-    if (parsedParts == null && tableUrn != null && !tableUrn.isEmpty()) {
-      String[] parts = tableUrn.split("/");
-      if (parts.length == 6
-          && "projects".equals(parts[0])
-          && "datasets".equals(parts[2])
-          && "tables".equals(parts[4])) {
-        parsedParts = parts;
-      }
-    }
-    return parsedParts != null ? parsedParts[index] : null;
+  public TableReference getTable() {
+    return table;
   }
 
   @Override

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.io.gcp.bigquery;
 
+import com.google.api.services.bigquery.model.TableReference;
 import com.google.api.services.bigquery.model.TableRow;
 import java.io.IOException;
 import java.io.InputStream;
@@ -42,15 +43,18 @@ public void encode(BigQueryStorageApiInsertError value, OutputStream outStream)
       throws IOException {
     TABLE_ROW_CODER.encode(value.getRow(), outStream);
     STRING_CODER.encode(value.getErrorMessage(), outStream);
-    STRING_CODER.encode(value.getTableUrn(), outStream);
+    TableReference table = value.getTable();
+    String tableSpec = table != null ? BigQueryHelpers.toTableSpec(table) : null;
+    STRING_CODER.encode(tableSpec, outStream);
   }
 
   @Override
   public BigQueryStorageApiInsertError decode(InputStream inStream)
       throws CoderException, IOException {
-    return new BigQueryStorageApiInsertError(
-        TABLE_ROW_CODER.decode(inStream),
-        STRING_CODER.decode(inStream),
-        STRING_CODER.decode(inStream));
+    TableRow row = TABLE_ROW_CODER.decode(inStream);
+    String errorMessage = STRING_CODER.decode(inStream);
+    String tableSpec = STRING_CODER.decode(inStream);
+    TableReference table = tableSpec != null ? BigQueryHelpers.parseTableSpec(tableSpec) : null;
+    return new BigQueryStorageApiInsertError(row, errorMessage, table);
   }
 }

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

@@ -19,6 +19,7 @@
 
 import static org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.BAD_RECORD_TAG;
 
+import com.google.api.services.bigquery.model.TableReference;
 import com.google.api.services.bigquery.model.TableRow;
 import java.io.IOException;
 import org.apache.beam.sdk.coders.Coder;
@@ -186,15 +187,15 @@ public void processElement(
           badRecordRouter.route(o, element, elementCoder, e, "Unable to convert value to TableRow");
           return;
         }
-        String tableUrn = null;
+        TableReference tableReference = null;
         TableDestination tableDestination = dynamicDestinations.getTable(element.getKey());
         if (tableDestination != null) {
-          tableUrn = tableDestination.getTableUrn(pipelineOptions.as(BigQueryOptions.class));
+          tableReference = tableDestination.getTableReference();
         }
         o.get(failedWritesTag)
             .output(
                 new BigQueryStorageApiInsertError(
-                    failsafeTableRow, conversionException.toString(), tableUrn));
+                    failsafeTableRow, conversionException.toString(), tableReference));
       } catch (Exception e) {
         badRecordRouter.route(
             o, element, elementCoder, e, "Unable to convert value to StorageWriteApiPayload");

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

@@ -608,7 +608,8 @@ void addMessage(
               org.joda.time.Instant timestamp = payload.getTimestamp();
               rowsSentToFailedRowsCollection.inc();
               failedRowsReceiver.outputWithTimestamp(
-                  new BigQueryStorageApiInsertError(tableRow, e.toString(), tableUrn),
+                  new BigQueryStorageApiInsertError(
+                      tableRow, e.toString(), tableDestination.getTableReference()),
                   timestamp != null ? timestamp : elementTs);
               return;
             }
@@ -668,7 +669,9 @@ long flush(
             org.joda.time.Instant timestamp = insertTimestamps.get(i);
             failedRowsReceiver.outputWithTimestamp(
                 new BigQueryStorageApiInsertError(
-                    failedRow, "Row payload too large. Maximum size " + maxRequestSize, tableUrn),
+                    failedRow,
+                    "Row payload too large. Maximum size " + maxRequestSize,
+                    tableDestination.getTableReference()),
                 timestamp);
           }
           int numRowsFailed = inserts.getSerializedRowsCount();
@@ -755,7 +758,7 @@ long flush(
                         new BigQueryStorageApiInsertError(
                             failedRow,
                             error.getRowIndexToErrorMessage().get(failedIndex),
-                            tableUrn);
+                            tableDestination.getTableReference());
                   } catch (Exception e) {
                     LOG.error("Failed to insert row and could not parse the result!", e);
                   }

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

@@ -21,6 +21,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
+import com.google.api.services.bigquery.model.TableReference;
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.cloud.bigquery.storage.v1.AppendRowsRequest;
 import com.google.cloud.bigquery.storage.v1.AppendRowsResponse;
@@ -481,6 +482,7 @@ public void process(
               });
       final String tableId = tableDestination.getTableUrn(bigQueryOptions);
       final String shortTableId = tableDestination.getShortTableUrn();
+      final TableReference tableReference = tableDestination.getTableReference();
       final DatasetService datasetService = getDatasetService(pipelineOptions);
       final WriteStreamService writeStreamService = getWriteStreamService(pipelineOptions);
 
@@ -620,7 +622,7 @@ public void process(
                 o.get(failedRowsTag)
                     .outputWithTimestamp(
                         new BigQueryStorageApiInsertError(
-                            failedRow.getValue(), errorMessage, tableId),
+                            failedRow.getValue(), errorMessage, tableReference),
                         failedRow.getTimestamp());
                 rowsSentToFailedRowsCollection.inc();
                 BigQuerySinkMetrics.appendRowsRowStatusCounter(
@@ -740,7 +742,9 @@ public void process(
                 o.get(failedRowsTag)
                     .outputWithTimestamp(
                         new BigQueryStorageApiInsertError(
-                            failedRow, error.getRowIndexToErrorMessage().get(failedIndex), tableId),
+                            failedRow,
+                            error.getRowIndexToErrorMessage().get(failedIndex),
+                            tableReference),
                         timestamp);
               }
               int failedRows = failedRowIndices.size();
@@ -913,7 +917,7 @@ public void process(
                     new BigQueryStorageApiInsertError(
                         failedRow,
                         "Row payload too large. Maximum size " + maxRequestSize,
-                        tableId),
+                        tableReference),
                     timestamp);
           }
           int numRowsFailed = splitValue.getProtoRows().getSerializedRowsCount();