Add support for multiple BQ policy tags

src/Dfe.Analytics.EFCore/AnalyticsDeployer.cs

@@ -21,6 +21,7 @@ public async Task DeployAsync(
         string airbyteConnectionId,
         string? airbyteSyncMode,
         string hiddenPolicyTagName,
+        IReadOnlyDictionary<string, string> additionalPolicyTags,
         IProgressReporter? progressReporter = null,
         CancellationToken cancellationToken = default)
     {
@@ -46,7 +47,7 @@ await WithProgressReportingAsync(
             "Waiting for sync to complete");
 
         await WithProgressReportingAsync(
-            () => UpdateBigQueryPolicyTagsAsync(configuration, hiddenPolicyTagName, progressReporter, cancellationToken),
+            () => UpdateBigQueryPolicyTagsAsync(configuration, hiddenPolicyTagName, additionalPolicyTags, progressReporter, cancellationToken),
             progressReporter,
             "Updating BigQuery policy tags");
     }
@@ -207,6 +208,7 @@ async Task WaitForCompletionAsync()
     internal async Task UpdateBigQueryPolicyTagsAsync(
         DatabaseSyncConfiguration configuration,
         string hiddenPolicyTagName,
+        IReadOnlyDictionary<string, string> additionalPolicyTags,
         IProgressReporter progressReporter,
         CancellationToken cancellationToken = default)
     {
@@ -269,7 +271,13 @@ await WithProgressReportingAsync(
 
                 if (column.Hidden)
                 {
-                    bqField.PolicyTags.Names.Add(hiddenPolicyTagName);
+                    var policyTagName = column.PolicyTag is string policyTag && policyTag != "hidden"
+                        ? additionalPolicyTags.TryGetValue(column.PolicyTag, out var tag)
+                            ? tag
+                            : throw new InvalidOperationException($"Missing policy tag mapping for '{policyTag}'.")
+                        : hiddenPolicyTagName;
+
+                    bqField.PolicyTags.Names.Add(policyTagName);
                 }
 
                 var policyTagNamesChanged = !existingPolicyTagNames.SetEquals(bqField.PolicyTags.Names);

src/Dfe.Analytics.EFCore/Cli/Commands.Config.Apply.cs

@@ -1,3 +1,4 @@
+using System.Collections.Frozen;
 using System.CommandLine;
 using Dfe.Analytics.EFCore.Configuration;
 using Microsoft.EntityFrameworkCore;
@@ -19,6 +20,35 @@ public static Command GetConfigApplyCommand()
         var projectIdOption = new Option<string>("--project-id") { Required = true };
         var datasetIdOption = new Option<string>("--dataset-id") { Required = true };
         var hiddenPolicyTagNameOption = new Option<string>("--hidden-policy-tag-name") { Required = true };
+        var policyTagOption = new Option<IReadOnlyDictionary<string, string>>("--policy-tag")
+        {
+            Required = false,
+            Arity = new ArgumentArity(2, ArgumentArity.OneOrMore.MaximumNumberOfValues),
+            AllowMultipleArgumentsPerToken = true,
+            CustomParser = result =>
+            {
+                var policyTagMapping = new Dictionary<string, string>();
+
+                if (result.Tokens.Count % 2 != 0)
+                {
+                    result.AddError("--policy-tag must be specified as an alias followed by a name");
+                    return policyTagMapping;
+                }
+
+                for (var i = 0; i < result.Tokens.Count; i += 2)
+                {
+                    var alias = result.Tokens[i].Value;
+                    var name = result.Tokens[i + 1].Value;
+
+                    if (!policyTagMapping.TryAdd(alias, name))
+                    {
+                        result.AddError($"--policy-tag with alias '{alias}' has already been specified");
+                    }
+                }
+
+                return policyTagMapping;
+            }
+        };
 
         // Airbyte options
         var airbyteApiBaseAddressOption = new Option<string>("--airbyte-api-base-address") { Required = true };
@@ -40,6 +70,7 @@ public static Command GetConfigApplyCommand()
             projectIdOption,
             datasetIdOption,
             hiddenPolicyTagNameOption,
+            policyTagOption,
             airbyteApiBaseAddressOption,
             airbyteClientIdOption,
             airbyteClientSecretOption,
@@ -59,6 +90,7 @@ public static Command GetConfigApplyCommand()
             var projectId = parseResult.GetRequiredValue(projectIdOption);
             var datasetId = parseResult.GetRequiredValue(datasetIdOption);
             var hiddenPolicyTagName = parseResult.GetRequiredValue(hiddenPolicyTagNameOption);
+            var additionalPolicyTags = parseResult.GetValue(policyTagOption) ?? FrozenDictionary<string, string>.Empty;
             var airbyteApiBaseAddress = parseResult.GetRequiredValue(airbyteApiBaseAddressOption);
             var airbyteClientId = parseResult.GetRequiredValue(airbyteClientIdOption);
             var airbyteClientSecret = parseResult.GetRequiredValue(airbyteClientSecretOption);
@@ -100,6 +132,7 @@ await analyticsDeployer.DeployAsync(
                 airbyteConnectionId,
                 airbyteSyncMode,
                 hiddenPolicyTagName,
+                additionalPolicyTags,
                 cancellationToken: cts.Token);
         });
 

src/Dfe.Analytics.EFCore/Configuration/AnalyticsConfigurationProvider.cs

@@ -118,7 +118,10 @@ private static IReadOnlyCollection<ColumnSyncInfo> GetColumns(IEntityType entity
                 columnSyncInfos.Add(new ColumnSyncInfo
                 {
                     Name = columnName,
-                    Hidden = hidden
+                    Hidden = hidden,
+                    PolicyTag = hidden ?
+                        columnSyncMetadata?.PolicyTag ?? tableSyncMetadata.DefaultColumnSettings.PolicyTag :
+                        null
                 });
             }
         }

src/Dfe.Analytics.EFCore/Configuration/ColumnSyncInfo.cs

@@ -4,6 +4,7 @@ public sealed class ColumnSyncInfo : IEquatable<ColumnSyncInfo>
 {
     public required string Name { get; init; }
     public required bool Hidden { get; init; }
+    public required string? PolicyTag { get; init; }
 
     public override bool Equals(object? obj)
     {
@@ -22,7 +23,7 @@ public override bool Equals(object? obj)
 
     public override int GetHashCode()
     {
-        return HashCode.Combine(Name, Hidden);
+        return HashCode.Combine(Name, Hidden, PolicyTag);
     }
 
     public bool Equals(ColumnSyncInfo? other)
@@ -37,6 +38,6 @@ public bool Equals(ColumnSyncInfo? other)
             return true;
         }
 
-        return Name == other.Name && Hidden == other.Hidden;
+        return Name == other.Name && Hidden == other.Hidden && PolicyTag == other.PolicyTag;
     }
 }

src/Dfe.Analytics.EFCore/DbContextConfigurationExtensions.cs

@@ -5,6 +5,8 @@ namespace Dfe.Analytics.EFCore;
 
 public static class DbContextConfigurationExtensions
 {
+    private const string DefaultPolicyTagName = "hidden";
+
     // ReSharper disable once UnusedMethodReturnValue.Global
     public static T IncludeInAnalyticsSync<T>(this T builder, bool includeAllColumns = true, bool? hidden = null)
         where T : EntityTypeBuilder
@@ -13,19 +15,24 @@ public static T IncludeInAnalyticsSync<T>(this T builder, bool includeAllColumns
 
         builder.HasAnnotation(
             AnnotationKeys.TableAnalyticsSyncMetadata,
-            new TableSyncMetadata(new ColumnSyncMetadata(includeAllColumns, hidden)));
+            new TableSyncMetadata(new ColumnSyncMetadata(includeAllColumns, hidden, PolicyTag: DefaultPolicyTagName)));
 
         return builder;
     }
 
-    public static T ConfigureAnalyticsSync<T>(this T builder, bool included = true, bool? hidden = null)
+    public static T ConfigureAnalyticsSync<T>(this T builder, bool included = true, bool? hidden = null, string? policyTag = null)
         where T : PropertyBuilder
     {
         ArgumentNullException.ThrowIfNull(builder);
 
+        if (hidden is null && policyTag is not null)
+        {
+            hidden = true;
+        }
+
         builder.HasAnnotation(
             AnnotationKeys.ColumnAnalyticsSyncMetadata,
-            new ColumnSyncMetadata(included, hidden));
+            new ColumnSyncMetadata(included, hidden, policyTag ?? DefaultPolicyTagName));
 
         return builder;
     }

src/Dfe.Analytics.EFCore/Description/ColumnSyncMetadata.cs

@@ -1,3 +1,3 @@
 namespace Dfe.Analytics.EFCore.Description;
 
-internal record ColumnSyncMetadata(bool Included, bool? Hidden);
+internal record ColumnSyncMetadata(bool Included, bool? Hidden, string? PolicyTag);

tests/Dfe.Analytics.EFCore.Tests/AnalyticsDeployerTests.cs

@@ -17,6 +17,9 @@ public class AnalyticsDeployerTests
     private const string DatasetId = "dummy-dataset";
 
     private const string HiddenPolicyTagName = "projects/dummy-project/locations/us/taxonomies/dummy-taxonomy/policyTags/dummy-policy-tag";
+    private const string SensitivePolicyTagName = "projects/dummy-project/locations/us/taxonomies/dummy-taxonomy/policyTags/dummy-sensitive-policy-tag";
+
+    private static readonly IReadOnlyDictionary<string, string> NoAdditionalPolicyTags = new Dictionary<string, string>();
 
     [Fact]
     public async Task ApplyAirbyteConfigurationAsync_UpdatesAirbyteConnectionWithExpectedPayload()
@@ -192,7 +195,7 @@ public async Task UpdateBigQueryPolicyTagsAsync_ThrowsIfColumnIsMissingFromBqTab
 
         // Act
         var ex = await Record.ExceptionAsync(
-            () => deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, progressReporter, TestContext.Current.CancellationToken));
+            () => deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, NoAdditionalPolicyTags, progressReporter, TestContext.Current.CancellationToken));
 
         // Assert
         Assert.IsType<InvalidOperationException>(ex);
@@ -255,12 +258,128 @@ public async Task UpdateBigQueryPolicyTagsAsync_ColumnIsMissingHiddenPolicyTag_A
         var deployer = CreateDeployer(httpClient, bigQueryClientMock.Object);
 
         // Act
-        await deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, progressReporter, TestContext.Current.CancellationToken);
+        await deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, NoAdditionalPolicyTags, progressReporter, TestContext.Current.CancellationToken);
 
         // Assert
         bigQueryClientMock.Verify();
     }
 
+    [Fact]
+    public async Task UpdateBigQueryPolicyTagsAsync_ColumnHasPolicyTag_AddsMappedPolicyTagToSchema()
+    {
+        // Arrange
+        var configuration = GetConfiguration(namePolicyTag: "sensitive");
+        var additionalPolicyTags = new Dictionary<string, string> { ["sensitive"] = SensitivePolicyTagName };
+        var progressReporter = new RecordingProgressReporter();
+
+        var tableName = configuration.Tables.Select(t => t.Name).Single();
+
+        using var httpClient = new HttpClient();
+
+        var bigQueryClientMock = new Mock<BigQueryClient>();
+
+        bigQueryClientMock
+            .Setup(mock => mock.ListTablesAsync(ProjectId, DatasetId, It.IsAny<ListTablesOptions>()))
+            .Returns(new TestablePagedAsyncEnumerable<TableList, BigQueryTable>([
+                new BigQueryTable(bigQueryClientMock.Object, new Table { TableReference = new TableReference { TableId = tableName } })
+            ]));
+
+        bigQueryClientMock
+            .Setup(mock => mock.GetTableAsync(ProjectId, DatasetId, tableName, It.IsAny<GetTableOptions>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(() =>
+            {
+                var table = new BigQueryTable(
+                    bigQueryClientMock.Object,
+                    new Table
+                    {
+                        Schema = new TableSchema
+                        {
+                            Fields =
+                            [
+                                new TableFieldSchema { Name = "TestEntityId", Type = "INTEGER" },
+                                new TableFieldSchema { Name = "Name", Type = "STRING" },
+                                new TableFieldSchema { Name = "DateOfBirth", Type = "DATE" }
+                            ]
+                        }
+                    });
+
+                return table;
+            });
+
+        bigQueryClientMock
+            .Setup(mock => mock.PatchTableAsync(
+                ProjectId,
+                DatasetId,
+                tableName,
+                It.Is<Table>(t =>
+                    t.Schema.Fields.Single(f => f.Name == "Name").PolicyTags.Names.Contains(SensitivePolicyTagName) &&
+                    !t.Schema.Fields.Single(f => f.Name == "Name").PolicyTags.Names.Contains(HiddenPolicyTagName)),
+                It.IsAny<PatchTableOptions>(),
+                It.IsAny<CancellationToken>()))
+            .ReturnsAsync(() => null)
+            .Verifiable();
+
+        var deployer = CreateDeployer(httpClient, bigQueryClientMock.Object);
+
+        // Act
+        await deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, additionalPolicyTags, progressReporter, TestContext.Current.CancellationToken);
+
+        // Assert
+        bigQueryClientMock.Verify();
+    }
+
+    [Fact]
+    public async Task UpdateBigQueryPolicyTagsAsync_ColumnHasPolicyTagWithNoMapping_ThrowsInvalidOperationException()
+    {
+        // Arrange
+        var configuration = GetConfiguration(namePolicyTag: "sensitive");
+        var progressReporter = new RecordingProgressReporter();
+
+        var tableName = configuration.Tables.Select(t => t.Name).Single();
+
+        using var httpClient = new HttpClient();
+
+        var bigQueryClientMock = new Mock<BigQueryClient>();
+
+        bigQueryClientMock
+            .Setup(mock => mock.ListTablesAsync(ProjectId, DatasetId, It.IsAny<ListTablesOptions>()))
+            .Returns(new TestablePagedAsyncEnumerable<TableList, BigQueryTable>([
+                new BigQueryTable(bigQueryClientMock.Object, new Table { TableReference = new TableReference { TableId = tableName } })
+            ]));
+
+        bigQueryClientMock
+            .Setup(mock => mock.GetTableAsync(ProjectId, DatasetId, tableName, It.IsAny<GetTableOptions>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(() =>
+            {
+                var table = new BigQueryTable(
+                    bigQueryClientMock.Object,
+                    new Table
+                    {
+                        Schema = new TableSchema
+                        {
+                            Fields =
+                            [
+                                new TableFieldSchema { Name = "TestEntityId", Type = "INTEGER" },
+                                new TableFieldSchema { Name = "Name", Type = "STRING" },
+                                new TableFieldSchema { Name = "DateOfBirth", Type = "DATE" }
+                            ]
+                        }
+                    });
+
+                return table;
+            });
+
+        var deployer = CreateDeployer(httpClient, bigQueryClientMock.Object);
+
+        // Act
+        var ex = await Record.ExceptionAsync(
+            () => deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, NoAdditionalPolicyTags, progressReporter, TestContext.Current.CancellationToken));
+
+        // Assert
+        Assert.IsType<InvalidOperationException>(ex);
+        Assert.Equal("Missing policy tag mapping for 'sensitive'.", ex.Message);
+    }
+
     [Fact]
     public async Task UpdateBigQueryPolicyTagsAsync_ColumnHasHiddenPolicyTagButSchemaDoesNotHaveHiddenFlag_RemovesHiddenPolicyTagFromSchema()
     {
@@ -334,7 +453,7 @@ public async Task UpdateBigQueryPolicyTagsAsync_ColumnHasHiddenPolicyTagButSchem
         var deployer = CreateDeployer(httpClient, bigQueryClientMock.Object);
 
         // Act
-        await deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, progressReporter, TestContext.Current.CancellationToken);
+        await deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, NoAdditionalPolicyTags, progressReporter, TestContext.Current.CancellationToken);
 
         // Assert
         bigQueryClientMock.Verify();
@@ -392,7 +511,7 @@ public async Task UpdateBigQueryPolicyTagsAsync_BigQuerySchemaMatchesConfigurati
         var deployer = CreateDeployer(httpClient, bigQueryClientMock.Object);
 
         // Act
-        await deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, progressReporter, TestContext.Current.CancellationToken);
+        await deployer.UpdateBigQueryPolicyTagsAsync(configuration, HiddenPolicyTagName, NoAdditionalPolicyTags, progressReporter, TestContext.Current.CancellationToken);
 
         // Assert
         bigQueryClientMock.Verify(
@@ -478,7 +597,7 @@ private AnalyticsDeployer CreateDeployer(
         return new AnalyticsDeployer(airbyteApiClient, options);
     }
 
-    private DatabaseSyncConfiguration GetConfiguration() =>
+    private DatabaseSyncConfiguration GetConfiguration(string? namePolicyTag = null) =>
         new()
         {
             DbContextName = typeof(TestDbContext).AssemblyQualifiedName!,
@@ -490,9 +609,9 @@ private DatabaseSyncConfiguration GetConfiguration() =>
                     PrimaryKey = new TablePrimaryKeySyncInfo { ColumnNames = ["TestEntityId"] },
                     Columns =
                     [
-                        new ColumnSyncInfo { Name = "TestEntityId", Hidden = false },
-                        new ColumnSyncInfo { Name = "Name", Hidden = true },
-                        new ColumnSyncInfo { Name = "DateOfBirth", Hidden = false }
+                        new ColumnSyncInfo { Name = "TestEntityId", Hidden = false, PolicyTag = null },
+                        new ColumnSyncInfo { Name = "Name", Hidden = true, PolicyTag = namePolicyTag },
+                        new ColumnSyncInfo { Name = "DateOfBirth", Hidden = false, PolicyTag = null }
                     ]
                 }
             ]