feat: Adds support for policy tags

googleapis · Aug 31, 2023 · 730a30e · 730a30e
1 parent 8037c35
commit 730a30e
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 6 deletions.
diff --git a/apis/Google.Cloud.BigQuery.V2/Google.Cloud.BigQuery.V2.IntegrationTests/TablesTest.cs b/apis/Google.Cloud.BigQuery.V2/Google.Cloud.BigQuery.V2.IntegrationTests/TablesTest.cs
@@ -21,6 +21,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
+using static Google.Apis.Bigquery.v2.Data.TableFieldSchema;
 
 namespace Google.Cloud.BigQuery.V2.IntegrationTests
 {
@@ -821,5 +822,26 @@ public void TestTableIamPermissions()
             var response = client.TestTableIamPermissions(_fixture.DatasetId, _fixture.HighScoreTableId, new List<string> { "bigquery.tables.get" });
             Assert.Collection(response.Permissions, role => Assert.Equal("bigquery.tables.get", role));
         }
+
+        [Fact]
+        public void CreateTableWithPolicyTagsFields()
+        {
+            var client = BigQueryClient.Create(_fixture.ProjectId);
+
+            string sampleTag = $"projects/{_fixture.ProjectId}/locations/us/taxonomies/1/policyTags/2";
+            var dataset = client.GetDataset(_fixture.DatasetId);
+            var tableId = _fixture.CreateTableId();
+
+            var schema = new TableSchemaBuilder
+            {
+                { "full_name", BigQueryDbType.String, BigQueryFieldMode.Nullable, "field with policy tag" }
+            }.ModifyField("full_name", field => field.PolicyTags = new PolicyTagsData { Names = new[] { sampleTag } })
+            .Build();
+
+            dataset.CreateTable(tableId, schema);
+            var table = dataset.GetTable(tableId);
+
+            Assert.Contains(sampleTag, table.Schema.Fields.First().PolicyTags.Names);
+        }
     }
 }
diff --git a/apis/Google.Cloud.BigQuery.V2/Google.Cloud.BigQuery.V2.Tests/TableSchemaBuilderTest.cs b/apis/Google.Cloud.BigQuery.V2/Google.Cloud.BigQuery.V2.Tests/TableSchemaBuilderTest.cs
@@ -1,4 +1,4 @@
-// Copyright 2016 Google Inc. All Rights Reserved.
+// Copyright 2016 Google Inc. All Rights Reserved.
 // 
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -13,6 +13,7 @@
 // limitations under the License.
 using System;
 using Xunit;
+using static Google.Apis.Bigquery.v2.Data.TableFieldSchema;
 
 namespace Google.Cloud.BigQuery.V2.Tests
 {
@@ -79,6 +80,23 @@ public void ValidateFieldName_Invalid(string name)
             Assert.Throws<ArgumentException>(() => TableSchemaBuilder.ValidateFieldName(name, "field"));
         }
 
+        [Fact]
+        public void AddWithPolicyTags()
+        {
+            var builder = new TableSchemaBuilder
+            {
+                { "field", BigQueryDbType.Int64, BigQueryFieldMode.Repeated, "My field" }
+            }.ModifyField("field", field => field.PolicyTags = new PolicyTagsData { Names = new[] { "policyName" } });
+
+            var schema = builder.Build();
+            var field = schema.Fields[0];
+            Assert.Equal("field", field.Name);
+            Assert.Equal("INTEGER", field.Type);
+            Assert.Equal("REPEATED", field.Mode);
+            Assert.Equal("My field", field.Description);
+            Assert.Contains("policyName", field.PolicyTags.Names);
+        }
+
         [Fact]
         public void ValidateFieldName_InvalidTooLong()
         {

diff --git a/apis/Google.Cloud.BigQuery.V2/Google.Cloud.BigQuery.V2/TableSchemaBuilder.cs b/apis/Google.Cloud.BigQuery.V2/Google.Cloud.BigQuery.V2/TableSchemaBuilder.cs
@@ -1,4 +1,4 @@
-// Copyright 2016 Google Inc. All Rights Reserved.
+// Copyright 2016 Google Inc. All Rights Reserved.
 // 
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,13 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+using Google.Api.Gax;
+using Google.Api.Gax.Rest;
 using Google.Apis.Bigquery.v2.Data;
+using System;
 using System.Collections;
 using System.Collections.Generic;
-using System;
+using System.Linq;
 using System.Text.RegularExpressions;
-using Google.Api.Gax.Rest;
-using Google.Api.Gax;
 
 namespace Google.Cloud.BigQuery.V2
 {
@@ -40,10 +41,27 @@ public sealed class TableSchemaBuilder : IEnumerable
         public void Add(TableFieldSchema field)
         {
             GaxPreconditions.CheckNotNull(field, nameof(field));
-
             _fields.Add(field);
         }
 
+        /// <summary>
+        /// Modifies the field with <paramref name="fieldName"/> by running the given action.
+        /// </summary>
+        /// <param name="fieldName">The name of the field to be modified. Must not be null.</param>
+        /// <param name="fieldModifier">An action to be performed on the given field.</param>
+        /// <returns>
+        /// This <see cref="TableSchemaBuilder"/> after having applied the modification. For the purposes of method chaining.
+        /// </returns>
+        /// <exception cref="ArgumentException"><paramref name="fieldName"/> is not present within the schema being built.</exception>
+        public TableSchemaBuilder ModifyField(string fieldName, Action<TableFieldSchema> fieldModifier)
+        {
+            GaxPreconditions.CheckNotNull(fieldName, nameof(fieldName));
+            GaxPreconditions.CheckNotNull(fieldModifier, nameof(fieldModifier));
+            var field = _fields.FirstOrDefault(field => field.Name == fieldName) ?? throw new ArgumentException($"Field '{fieldName}' not found in the schema.", nameof(fieldName));
+            fieldModifier(field);
+            return this;
+        }
+
         /// <summary>
         /// Creates a field with the specified details, and adds it to the schema being built.
         /// </summary>