Support multiple policy tags

config/locales/en.yml

@@ -32,9 +32,11 @@ en:
           description: |
             The name of the BigQuery airbyte dataset we're writing to.
           default: ENV['BIGQUERY_AIRBYTE_DATASET']
-        bigquery_hidden_policy_tag:
-          description: The name of the BigQuery hidden policy tag applied to PII in raw airbyte tables.
-          default: ENV['BIGQUERY_HIDDEN_POLICY_TAG']
+        bigquery_policy_tags:
+          description: |
+            The BigQuery policy tags JSON with multiple policy tags to be
+            applied to hidden and sensitive data in raw airbyte tables.
+          default: ENV['BIGQUERY_POLICY_TAGS']
         bigquery_retries:
           description: |
             Passed directly to the retries: option on the BigQuery client

lib/dfe/analytics.rb

@@ -165,7 +165,10 @@ def self.extract_model_attributes(model, attributes = nil)
       table_name = model.class.table_name.to_sym
 
       exportable_attrs = (allowlist[table_name].presence || []).map(&:to_sym)
-      hidden_pii_attrs = (hidden_pii[table_name].presence || []).map(&:to_sym)
+      hidden_pii_attrs = (hidden_pii[table_name].presence || []).map do |field|
+        field = field.keys.first.to_sym if field.is_a?(Hash)
+        field.to_sym
+      end
       exportable_hidden_pii_attrs = exportable_attrs & hidden_pii_attrs
 
       # Exclude hidden pii attributes from allowed_attributes

lib/dfe/analytics/big_query_api.rb

@@ -9,20 +9,9 @@ class BigQueryApi
       RETRY_INITIAL_BASE_INTERVAL = 15
       RETRY_MAX_INTERVAL = 60
       RETRY_INTERVAL_MULTIPLIER = 2
-      BIGQUERY_MANDATORY_CONFIG = %i[
-        bigquery_project_id
-        bigquery_table_name
-        bigquery_dataset
-        azure_client_id
-        azure_token_path
-        azure_scope
-        gcp_scope
-        google_cloud_credentials
-      ].freeze
-      BIGQUERY_AIRBYTE_MANDATORY_CONFIG = %i[
-        bigquery_airbyte_dataset
-        bigquery_hidden_policy_tag
-      ].freeze
+      BIGQUERY_MANDATORY_CONFIG =
+        %i[bigquery_project_id bigquery_table_name bigquery_dataset azure_client_id azure_token_path azure_scope gcp_scope google_cloud_credentials].freeze
+      BIGQUERY_AIRBYTE_MANDATORY_CONFIG = %i[bigquery_airbyte_dataset].freeze
 
       def self.client
         @client ||= begin
@@ -81,8 +70,8 @@ def self.error_message_for(response)
         "DfE::Analytics BigQuery API insert error for #{response.insert_errors.length} event(s):\n#{message}"
       end
 
-      def self.apply_policy_tags(dataset, tables, policy_tag)
-        tables.each do |table_name, column_names|
+      def self.apply_policy_tags(dataset, tables, policy_tags)
+        tables.each do |table_name, columns|
           begin
             table = client.get_table(
               DfE::Analytics.config.bigquery_project_id,
@@ -96,7 +85,17 @@ def self.apply_policy_tags(dataset, tables, policy_tag)
           end
 
           updated_fields = table.schema.fields.map do |field|
-            field.policy_tags = Google::Apis::BigqueryV2::TableFieldSchema::PolicyTags.new(names: [policy_tag]) if column_names.include?(field.name)
+            column = columns.find do |column_config|
+              column_name = column_config.is_a?(Hash) ? column_config.keys.first : column_config
+              column_name.to_s == field.name
+            end
+
+            if column
+              policy_tag_key = column.is_a?(Hash) ? column.values.first.to_sym : :hidden
+
+              field.policy_tags = Google::Apis::BigqueryV2::TableFieldSchema::PolicyTags.new(names: [policy_tags.symbolize_keys.fetch(policy_tag_key)])
+            end
+
             field
           end
 

lib/dfe/analytics/config.rb

@@ -16,6 +16,7 @@ module Config
         bigquery_airbyte_dataset
         bigquery_api_json_key
         bigquery_hidden_policy_tag
+        bigquery_policy_tags
         bigquery_retries
         bigquery_timeout
         enable_analytics
@@ -56,6 +57,9 @@ def self.configure(config)
         config.bigquery_airbyte_dataset         ||= ENV.fetch('BIGQUERY_AIRBYTE_DATASET', nil)
         config.bigquery_api_json_key            ||= ENV.fetch('BIGQUERY_API_JSON_KEY', nil)
         config.bigquery_hidden_policy_tag       ||= ENV.fetch('BIGQUERY_HIDDEN_POLICY_TAG', nil)
+        config.bigquery_policy_tags             ||=
+          JSON.parse(ENV.fetch('BIGQUERY_POLICY_TAGS', '{}')).transform_keys(&:underscore).symbolize_keys
+        config.bigquery_policy_tags             ||= { hidden: config.bigquery_hidden_policy_tag } if config.bigquery_hidden_policy_tag.present?
         config.bigquery_retries                 ||= 3
         config.bigquery_timeout                 ||= 120
         config.environment                      ||= ENV.fetch('RAILS_ENV', 'development')

lib/dfe/analytics/jobs/big_query_apply_policy_tags_job.rb

@@ -3,25 +3,25 @@
 module DfE
   module Analytics
     module Jobs
-      # Applies BigQuery hidden policy tags to PII fields in the airbyte tables
+      # Applies BigQuery policy tags to hidden and sensitive fields in the airbyte tables
       class BigQueryApplyPolicyTagsJob < AnalyticsJob
-        def self.do(dataset:, tables:, policy_tag:, delay_in_minutes: 0)
+        def self.do(dataset:, tables:, policy_tags:, delay_in_minutes: 0)
           if delay_in_minutes.zero?
-            perform_later(dataset, tables, policy_tag)
+            perform_later(dataset, tables, policy_tags)
           else
             time_to_run = Time.zone.now + delay_in_minutes.minutes
 
-            set(wait_until: time_to_run).perform_later(dataset, tables, policy_tag)
+            set(wait_until: time_to_run).perform_later(dataset, tables, policy_tags)
           end
         end
 
-        def perform(dataset, tables, policy_tag)
+        def perform(dataset, tables, policy_tags)
           unless DfE::Analytics.airbyte_enabled?
             Rails.logger.warn('DfE::Analytics::BigQueryApplyPolicyTags.perform called but airbyte is disabled. Please check DfE::Analytics.airbyte_enabled? before applying policy tags in BigQuery')
             return
           end
 
-          DfE::Analytics::BigQueryApi.apply_policy_tags(dataset, tables, policy_tag)
+          DfE::Analytics::BigQueryApi.apply_policy_tags(dataset, tables, policy_tags)
         end
       end
     end

lib/dfe/analytics/services.rb

@@ -3,7 +3,6 @@
 require 'dfe/analytics/shared/service_pattern'
 
 require 'dfe/analytics/services/apply_airbyte_final_tables_policy_tags'
-require 'dfe/analytics/services/apply_airbyte_internal_tables_policy_tags'
 require 'dfe/analytics/services/checksum_calculator'
 require 'dfe/analytics/services/entity_table_checks'
 require 'dfe/analytics/services/generic_checksum_calculator'

lib/dfe/analytics/services/apply_airbyte_final_tables_policy_tags.rb

@@ -3,7 +3,7 @@
 module DfE
   module Analytics
     module Services
-      # Apply hidden policy tags to the final airbyte table columns in the hidden pii config
+      # Apply policy tags to the final airbyte table columns configured as hidden or sensitive
       class ApplyAirbyteFinalTablesPolicyTags
         include ServicePattern
 
@@ -16,7 +16,7 @@ def call
             delay_in_minutes: @delay_in_minutes,
             dataset: DfE::Analytics.config.bigquery_airbyte_dataset,
             tables: DfE::Analytics.hidden_pii,
-            policy_tag: DfE::Analytics.config.bigquery_hidden_policy_tag
+            policy_tags: DfE::Analytics.config.bigquery_policy_tags
           )
         end
       end

lib/dfe/analytics/tasks/big_query_apply_policy_tags.rake

@@ -6,10 +6,6 @@ namespace :dfe do
 
       puts "Calling DfE::Analytics::Services::ApplyAirbyteFinalTablesPolicyTags.call(delay_in_minutes: #{delay})"
       DfE::Analytics::Services::ApplyAirbyteFinalTablesPolicyTags.call(delay_in_minutes: delay)
-
-      # Internal tables no longer created by airbyte - remove following code following testing
-      # puts "Calling DfE::Analytics::Services::ApplyAirbyteInternalTablesPolicyTags.call(delay_in_minutes: #{delay})"
-      # DfE::Analytics::Services::ApplyAirbyteInternalTablesPolicyTags.call(delay_in_minutes: delay)
     end
   end
 end

spec/dfe/analytics/big_query_api_spec.rb

@@ -87,55 +87,132 @@
 
   describe '.apply_policy_tags' do
     let(:dataset) { 'airbyte_dataset' }
-    let(:tables) { { users: %w[email name] } }
-    let(:policy_tag) { 'projects/my-project/locations/eu/taxonomies/123/policyTags/abc' }
+
+    let(:hidden_policy_tag) do
+      'projects/my-project/locations/eu/taxonomies/123/policyTags/hidden'
+    end
+
+    let(:sensitive_hidden_policy_tag) do
+      'projects/my-project/locations/eu/taxonomies/123/policyTags/sensitive-hidden'
+    end
+
+    let(:policy_tags) do
+      {
+        hidden: hidden_policy_tag,
+        sensitive_hidden: sensitive_hidden_policy_tag
+      }
+    end
+
     let(:fields) do
       [
         instance_double(Google::Apis::BigqueryV2::TableFieldSchema, name: 'email'),
         instance_double(Google::Apis::BigqueryV2::TableFieldSchema, name: 'name'),
         instance_double(Google::Apis::BigqueryV2::TableFieldSchema, name: 'created_at')
       ]
     end
-    let(:schema) { instance_double(Google::Apis::BigqueryV2::TableSchema, fields: fields) }
-    let(:table) { instance_double(Google::Apis::BigqueryV2::Table, schema: schema) }
+
+    let(:schema) do
+      instance_double(
+        Google::Apis::BigqueryV2::TableSchema,
+        fields: fields
+      )
+    end
+
+    let(:table) do
+      instance_double(
+        Google::Apis::BigqueryV2::Table,
+        schema: schema
+      )
+    end
 
     before do
-      allow(fields[0]).to receive(:policy_tags=)
-      allow(fields[1]).to receive(:policy_tags=)
-      allow(fields[2]).to receive(:policy_tags=)
+      fields.each do |field|
+        allow(field).to receive(:policy_tags=)
+      end
 
       allow(client).to receive(:get_table).and_return(table)
       allow(client).to receive(:patch_table)
     end
 
-    it 'sets policy tags on specified columns only' do
-      with_analytics_config(test_dummy_config) do
-        described_class.apply_policy_tags(dataset, tables, policy_tag)
+    context 'when no policy tag is specified for the column' do
+      let(:tables) do
+        {
+          users: %w[email name]
+        }
+      end
+
+      it 'applies the hidden policy tag by default' do
+        with_analytics_config(test_dummy_config) do
+          described_class.apply_policy_tags(dataset, tables, policy_tags)
 
-        expect(fields[0]).to have_received(:policy_tags=)
-        expect(fields[1]).to have_received(:policy_tags=)
-        expect(fields[2]).not_to have_received(:policy_tags=)
+          expect(fields[0]).to have_received(:policy_tags=).with(
+            an_object_having_attributes(names: [hidden_policy_tag])
+          )
+
+          expect(fields[1]).to have_received(:policy_tags=).with(
+            an_object_having_attributes(names: [hidden_policy_tag])
+          )
+
+          expect(fields[2]).not_to have_received(:policy_tags=)
+        end
+      end
+    end
+
+    context 'when a policy tag is specified for a column' do
+      let(:tables) do
+        {
+          users: [
+            { 'email' => 'sensitive_hidden' },
+            'name'
+          ]
+        }
+      end
+
+      it 'applies the specified policy tag to that column' do
+        with_analytics_config(test_dummy_config) do
+          described_class.apply_policy_tags(dataset, tables, policy_tags)
+
+          expect(fields[0]).to have_received(:policy_tags=).with(
+            an_object_having_attributes(names: [sensitive_hidden_policy_tag])
+          )
+
+          expect(fields[1]).to have_received(:policy_tags=).with(
+            an_object_having_attributes(names: [hidden_policy_tag])
+          )
+
+          expect(fields[2]).not_to have_received(:policy_tags=)
+        end
       end
     end
 
     it 'raises and logs error if get_table fails' do
+      tables = { users: %w[email name] }
+
       with_analytics_config(test_dummy_config) do
-        allow(client).to receive(:get_table).and_raise(Google::Apis::ClientError.new('not found'))
+        allow(client)
+          .to receive(:get_table)
+          .and_raise(Google::Apis::ClientError.new('not found'))
 
         expect(Rails.logger).to receive(:error).with(/Failed to retrieve table/)
+
         expect do
-          described_class.apply_policy_tags(dataset, tables, policy_tag)
+          described_class.apply_policy_tags(dataset, tables, policy_tags)
         end.to raise_error(DfE::Analytics::BigQueryApi::PolicyTagError)
       end
     end
 
     it 'raises and logs error if patch_table fails' do
+      tables = { users: %w[email name] }
+
       with_analytics_config(test_dummy_config) do
-        allow(client).to receive(:patch_table).and_raise(Google::Apis::ClientError.new('invalid'))
+        allow(client)
+          .to receive(:patch_table)
+          .and_raise(Google::Apis::ClientError.new('invalid'))
 
         expect(Rails.logger).to receive(:error).with(/Failed to update table/)
+
         expect do
-          described_class.apply_policy_tags(dataset, tables, policy_tag)
+          described_class.apply_policy_tags(dataset, tables, policy_tags)
         end.to raise_error(DfE::Analytics::BigQueryApi::PolicyTagError)
       end
     end

spec/dfe/analytics/config_spec.rb

@@ -41,6 +41,7 @@
         allow(ENV).to receive(:fetch).with('BIGQUERY_AIRBYTE_DATASET', nil).and_return('my_airbyte_dataset')
         allow(ENV).to receive(:fetch).with('BIGQUERY_API_JSON_KEY', nil).and_return(nil)
         allow(ENV).to receive(:fetch).with('BIGQUERY_HIDDEN_POLICY_TAG', nil).and_return('my_policy_tag')
+        allow(ENV).to receive(:fetch).with('BIGQUERY_POLICY_TAGS', '{}').and_return('{"key":"val"}')
         allow(ENV).to receive(:fetch).with('BIGQUERY_MAINTENANCE_WINDOW', nil).and_return(nil)
         allow(ENV).to receive(:fetch).with('RAILS_ENV', 'development').and_return('test')
         allow(ENV).to receive(:fetch).with('AIRBYTE_INTERNAL_DATASET', 'airbyte_internal').and_return('my_airbyte_internal_dataset')