Issue GoogleCloudDataproc#1043: fixed indirect write overwrite droppi…

…ng policy tags
vishalkarve15 · Jun 25, 2024 · 42b6e35 · 42b6e35
1 parent 27d36ae
commit 42b6e35
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 8 deletions.
diff --git a/...nnector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryUtil.java b/...nnector-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/...tor-common/src/test/java/com/google/cloud/bigquery/connector/common/BigQueryUtilTest.java b/...tor-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))