diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Field.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Field.java index 8259cc7e7..1213bc4e8 100644 --- a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Field.java +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Field.java @@ -58,6 +58,7 @@ public TableFieldSchema apply(Field field) { private final FieldList subFields; private final String mode; private final String description; + private final PolicyTags policyTags; /** * Mode for a BigQuery Table field. {@link Mode#NULLABLE} fields can be set to {@code null}, @@ -77,6 +78,7 @@ public static final class Builder { private FieldList subFields; private String mode; private String description; + private PolicyTags policyTags; private Builder() {} @@ -86,6 +88,7 @@ private Builder(Field field) { this.subFields = field.subFields; this.mode = field.mode; this.description = field.description; + this.policyTags = field.policyTags; } /** @@ -190,6 +193,12 @@ public Builder setDescription(String description) { return this; } + /** Sets the policy tags for the field. */ + public Builder setPolicyTags(PolicyTags policyTags) { + this.policyTags = policyTags; + return this; + } + /** Creates a {@code Field} object. */ public Field build() { return new Field(this); @@ -202,6 +211,7 @@ private Field(Builder builder) { this.subFields = builder.subFields; this.mode = builder.mode; this.description = builder.description; + this.policyTags = builder.policyTags; } /** Returns the field name. */ @@ -229,6 +239,11 @@ public String getDescription() { return Data.isNull(description) ? null : description; } + /** Returns the policy tags for the field. */ + public PolicyTags getPolicyTags() { + return policyTags; + } + /** * Returns the list of sub-fields if {@link #getType()} is a {@link LegacySQLTypeName#RECORD}. * Returns {@code null} otherwise. @@ -249,12 +264,13 @@ public String toString() { .add("type", type) .add("mode", mode) .add("description", description) + .add("policyTags", policyTags) .toString(); } @Override public int hashCode() { - return Objects.hash(name, type, mode, description); + return Objects.hash(name, type, mode, description, policyTags); } @Override @@ -316,6 +332,9 @@ TableFieldSchema toPb() { if (description != null) { fieldSchemaPb.setDescription(description); } + if (policyTags != null) { + fieldSchemaPb.setPolicyTags(policyTags.toPb()); + } if (getSubFields() != null) { List fieldsPb = Lists.transform(getSubFields(), TO_PB_FUNCTION); fieldSchemaPb.setFields(fieldsPb); @@ -332,6 +351,9 @@ static Field fromPb(TableFieldSchema fieldSchemaPb) { if (fieldSchemaPb.getDescription() != null) { fieldBuilder.setDescription(fieldSchemaPb.getDescription()); } + if (fieldSchemaPb.getPolicyTags() != null) { + fieldBuilder.setPolicyTags(PolicyTags.fromPb(fieldSchemaPb.getPolicyTags())); + } FieldList subFields = fieldSchemaPb.getFields() != null ? FieldList.of(Lists.transform(fieldSchemaPb.getFields(), FROM_PB_FUNCTION)) diff --git a/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/PolicyTags.java b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/PolicyTags.java new file mode 100644 index 000000000..a78fa92fb --- /dev/null +++ b/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/PolicyTags.java @@ -0,0 +1,65 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigquery; + +import com.google.auto.value.AutoValue; +import com.google.common.collect.ImmutableList; +import java.io.Serializable; +import java.util.List; +import javax.annotation.Nullable; + +@AutoValue +public abstract class PolicyTags implements Serializable { + + private static final long serialVersionUID = 1L; + + @Nullable + abstract ImmutableList getNamesImmut(); + + public List getNames() { + return getNamesImmut(); + } + + public abstract Builder toBuilder(); + + @AutoValue.Builder + public abstract static class Builder { + + abstract Builder setNamesImmut(ImmutableList namesImmut); + + public Builder setNames(List names) { + return setNamesImmut(ImmutableList.copyOf(names)); + } + + public abstract PolicyTags build(); + } + + public static Builder newBuilder() { + return new AutoValue_PolicyTags.Builder(); + } + + com.google.api.services.bigquery.model.TableFieldSchema.PolicyTags toPb() { + com.google.api.services.bigquery.model.TableFieldSchema.PolicyTags tagPb = + new com.google.api.services.bigquery.model.TableFieldSchema.PolicyTags(); + tagPb.setNames(getNames()); + return tagPb; + } + + static PolicyTags fromPb( + com.google.api.services.bigquery.model.TableFieldSchema.PolicyTags tagPb) { + return newBuilder().setNames(tagPb.getNames()).build(); + } +} diff --git a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SchemaTest.java b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SchemaTest.java index 090bb8a27..63c6752d7 100644 --- a/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SchemaTest.java +++ b/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/SchemaTest.java @@ -25,6 +25,9 @@ public class SchemaTest { + private static PolicyTags POLICY_TAGS = + PolicyTags.newBuilder().setNames(ImmutableList.of("someTag")).build(); + private static final Field FIELD_SCHEMA1 = Field.newBuilder("StringField", LegacySQLTypeName.STRING) .setMode(Field.Mode.NULLABLE) @@ -34,12 +37,14 @@ public class SchemaTest { Field.newBuilder("IntegerField", LegacySQLTypeName.INTEGER) .setMode(Field.Mode.REPEATED) .setDescription("FieldDescription2") + .setPolicyTags(POLICY_TAGS) .build(); private static final Field FIELD_SCHEMA3 = Field.newBuilder("RecordField", LegacySQLTypeName.RECORD, FIELD_SCHEMA1, FIELD_SCHEMA2) .setMode(Field.Mode.REQUIRED) .setDescription("FieldDescription3") .build(); + private static final List FIELDS = ImmutableList.of(FIELD_SCHEMA1, FIELD_SCHEMA2, FIELD_SCHEMA3); private static final Schema TABLE_SCHEMA = Schema.of(FIELDS); 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 index 03b55dd4a..472d8ecf8 100644 --- 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 @@ -70,6 +70,7 @@ import com.google.cloud.bigquery.Model; import com.google.cloud.bigquery.ModelId; import com.google.cloud.bigquery.ModelInfo; +import com.google.cloud.bigquery.PolicyTags; import com.google.cloud.bigquery.QueryJobConfiguration; import com.google.cloud.bigquery.QueryParameterValue; import com.google.cloud.bigquery.RangePartitioning; @@ -141,6 +142,10 @@ public class ITBigQueryTest { ImmutableMap.of( "example-label1", "example-value1", "example-label2", "example-value2"); + private static final String sampleTag = + String.format("projects/%s/locations/us/taxonomies/1/policyTags/2", PROJECT_ID); + private static final PolicyTags POLICY_TAGS = + PolicyTags.newBuilder().setNames(ImmutableList.of(sampleTag)).build(); private static final Field TIMESTAMP_FIELD_SCHEMA = Field.newBuilder("TimestampField", LegacySQLTypeName.TIMESTAMP) .setMode(Field.Mode.NULLABLE) @@ -198,6 +203,12 @@ public class ITBigQueryTest { .setMode(Field.Mode.NULLABLE) .setDescription("NumericDescription") .build(); + private static final Field STRING_FIELD_SCHEMA_WITH_POLICY = + Field.newBuilder("StringFieldWithPolicy", LegacySQLTypeName.STRING) + .setMode(Field.Mode.NULLABLE) + .setDescription("field has a policy") + .setPolicyTags(POLICY_TAGS) + .build(); private static final Schema TABLE_SCHEMA = Schema.of( TIMESTAMP_FIELD_SCHEMA, @@ -211,6 +222,8 @@ public class ITBigQueryTest { GEOGRAPHY_FIELD_SCHEMA, NUMERIC_FIELD_SCHEMA); private static final Schema SIMPLE_SCHEMA = Schema.of(STRING_FIELD_SCHEMA); + private static final Schema POLICY_SCHEMA = + Schema.of(STRING_FIELD_SCHEMA, STRING_FIELD_SCHEMA_WITH_POLICY, INTEGER_FIELD_SCHEMA); private static final Schema QUERY_RESULT_SCHEMA = Schema.of( Field.newBuilder("TimestampField", LegacySQLTypeName.TIMESTAMP) @@ -514,6 +527,21 @@ public void testCreateTableWithRangePartitioning() { } } + public void testCreateTableWithPolicyTags() { + String tableName = "test_create_table_policytags"; + TableId tableId = TableId.of(DATASET, tableName); + try { + StandardTableDefinition tableDefinition = + StandardTableDefinition.newBuilder().setSchema(POLICY_SCHEMA).build(); + Table createdTable = bigquery.create(TableInfo.of(tableId, tableDefinition)); + assertNotNull(createdTable); + Table remoteTable = bigquery.getTable(DATASET, tableName); + assertEquals(POLICY_SCHEMA, remoteTable.getDefinition().getSchema()); + } finally { + bigquery.delete(tableId); + } + } + @Test public void testCreateAndGetTable() { String tableName = "test_create_and_get_table";