From 42b6e356d94b5450b980b7b3d457fb7755f6e3bc Mon Sep 17 00:00:00 2001
From: Vishal Karve <vishalkarve15@gmail.com>
Date: Tue, 25 Jun 2024 20:08:39 +0530
Subject: [PATCH] Issue #1043: fixed indirect write overwrite dropping policy
 tags

---
 .../connector/common/BigQueryUtil.java        | 20 +++++++++++--------
 .../connector/common/BigQueryUtilTest.java    | 17 ++++++++++++++++
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryUtil.java b/bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryUtil.java
index f7e3cfed1..847d6662e 100644
--- a/bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryUtil.java
+++ b/bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryUtil.java
@@ -571,14 +571,18 @@ public static Schema adjustSchemaIfNeeded(Schema wantedSchema, Schema existingTa
    */
   @VisibleForTesting
   static Field adjustField(Field field, @Nullable Field existingField) {
+    if (existingField == null) {
+      return field;
+    }
+    Field.Builder fieldBuilder = field.toBuilder();
+    if (existingField.getPolicyTags() != null) {
+      fieldBuilder.setPolicyTags(existingField.getPolicyTags());
+    }
     if (field.getType().equals(LegacySQLTypeName.NUMERIC)
-        && existingField != null
         && existingField.getType().equals(LegacySQLTypeName.BIGNUMERIC)) {
       // convert type
-      return field.toBuilder().setType(LegacySQLTypeName.BIGNUMERIC).build();
-    }
-    if (field.getType().equals(LegacySQLTypeName.RECORD)
-        && existingField != null
+      return fieldBuilder.setType(LegacySQLTypeName.BIGNUMERIC).build();
+    } else if (field.getType().equals(LegacySQLTypeName.RECORD)
         && existingField.getType().equals(LegacySQLTypeName.RECORD)) {
       // need to go recursively
       FieldList subFields = field.getSubFields();
@@ -592,10 +596,10 @@ static Field adjustField(Field field, @Nullable Field existingField) {
                       subField ->
                           adjustField(subField, existingSubFieldsMap.get(subField.getName())))
                   .collect(Collectors.toList()));
-      return field.toBuilder().setType(LegacySQLTypeName.RECORD, adjustedSubFields).build();
+      return fieldBuilder.setType(LegacySQLTypeName.RECORD, adjustedSubFields).build();
+    } else {
+      return fieldBuilder.build();
     }
-    // no adjustment
-    return field;
   }
 
   public static String prepareQueryForLog(String query, int maxLength) {
diff --git a/bigquery-connector-common/src/test/java/com/google/cloud/bigquery/connector/common/BigQueryUtilTest.java b/bigquery-connector-common/src/test/java/com/google/cloud/bigquery/connector/common/BigQueryUtilTest.java
index ea48907a5..38aaa4c57 100644
--- a/bigquery-connector-common/src/test/java/com/google/cloud/bigquery/connector/common/BigQueryUtilTest.java
+++ b/bigquery-connector-common/src/test/java/com/google/cloud/bigquery/connector/common/BigQueryUtilTest.java
@@ -31,6 +31,7 @@
 import com.google.cloud.bigquery.FormatOptions;
 import com.google.cloud.bigquery.HivePartitioningOptions;
 import com.google.cloud.bigquery.LegacySQLTypeName;
+import com.google.cloud.bigquery.PolicyTags;
 import com.google.cloud.bigquery.RangePartitioning;
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.StandardSQLTypeName;
@@ -642,6 +643,22 @@ public void testAdjustField_nullExistingFieldWithRecordType() {
     assertThat(adjustedField.getType()).isEqualTo(LegacySQLTypeName.RECORD);
   }
 
+  @Test
+  public void testAdjustField_policyTagsExistingField() {
+    Field field = Field.of("f", LegacySQLTypeName.BOOLEAN).toBuilder().build();
+    PolicyTags existingPolicyTags =
+        PolicyTags.newBuilder().setNames(Arrays.asList("test-tag-1")).build();
+    Field existingField =
+        Field.of("f", LegacySQLTypeName.BOOLEAN)
+            .toBuilder()
+            .setPolicyTags(existingPolicyTags)
+            .build();
+    Field adjustedField = BigQueryUtil.adjustField(field, existingField);
+    assertThat(adjustedField.getType()).isEqualTo(LegacySQLTypeName.BOOLEAN);
+    assertThat(adjustedField.getPolicyTags()).isNotNull();
+    assertThat(adjustedField.getPolicyTags()).isEqualTo(existingPolicyTags);
+  }
+
   @Test
   public void testPrepareQueryForLog_withNewLine() {
     assertThat(BigQueryUtil.prepareQueryForLog("SELECT a\nFROM table", 40))