changes to add conditional masking policy feature (#62)

entechlog · Apr 1, 2023 · 41d079f · 41d079f
1 parent e0206e7
commit 41d079f
Show file tree

Hide file tree

Showing 11 changed files with 138 additions and 50 deletions.
diff --git a/integration_tests/macros/snow-mask-ddl/create_masking_policy_mp_conditional_contacts_pii.sql b/integration_tests/macros/snow-mask-ddl/create_masking_policy_mp_conditional_contacts_pii.sql
@@ -0,0 +1,16 @@
+{% macro create_masking_policy_mp_conditional_contacts_pii(node_database, node_schema, masked_column) %}
+
+    CREATE MASKING POLICY IF NOT EXISTS {{node_database}}.{{node_schema}}.mp_conditional_contacts_pii AS (
+        {{masked_column}} string,
+        last_name string
+    ) RETURNS string ->
+        CASE
+            WHEN CURRENT_ROLE() IN ('ANALYST') THEN {{masked_column}}
+            WHEN CURRENT_ROLE() IN ('DEVELOPER') AND last_name like 'A%' THEN {{masked_column}}
+            WHEN CURRENT_ROLE() IN ('DEVELOPER') AND last_name like 'B%' THEN SHA2({{masked_column}})
+            WHEN CURRENT_ROLE() IN ('DEVELOPER') AND last_name='Skeffington' THEN '*TARGETED_MASKING*'
+             WHEN CURRENT_ROLE() IN ('SYSADMIN') THEN SHA2({{masked_column}})
+        ELSE '**********'
+        END
+
+{% endmacro %}
diff --git a/integration_tests/macros/snow-mask-ddl/create_masking_policy_mp_conditional_customer_pii.sql b/integration_tests/macros/snow-mask-ddl/create_masking_policy_mp_conditional_customer_pii.sql
@@ -0,0 +1,16 @@
+{% macro create_masking_policy_mp_conditional_customer_pii(node_database, node_schema, masked_column) %}
+
+    CREATE MASKING POLICY IF NOT EXISTS {{node_database}}.{{node_schema}}.mp_conditional_customer_pii AS (
+        {{masked_column}} string,
+        store_id int,
+        active int
+    ) RETURNS string ->
+        CASE
+            WHEN CURRENT_ROLE() IN ('ANALYST') AND active=1 AND store_id=1 THEN {{masked_column}}
+            WHEN CURRENT_ROLE() IN ('ANALYST') AND active=1 AND store_id=2 THEN SHA2({{masked_column}}) 
+            WHEN CURRENT_ROLE() IN ('ANALYST') AND active=0 THEN '**********'
+            WHEN CURRENT_ROLE() IN ('SYSADMIN') THEN SHA2({{masked_column}})
+        ELSE '**********'
+        END
+
+{% endmacro %}
diff --git a/integration_tests/models/sources/sakila/src_sakila.yml b/integration_tests/models/sources/sakila/src_sakila.yml
@@ -7,4 +7,7 @@ sources:
         columns:
           - name: first_name
             meta:
-                masking_policy: mp_encrypt_pii
+              masking_policy: mp_conditional_customer_pii
+              mp_conditional_columns: 
+                - store_id 
+                - active
diff --git a/integration_tests/models/staging/pii/stg_contacts.yml b/integration_tests/models/staging/pii/stg_contacts.yml
@@ -1,27 +1,31 @@
 version: 2
 
 models:
-- name: stg_contacts
-  description: ""
-  columns:
-  - name: contact_id
-    description: ""
-  - name: first_name
-    description: ""
-  - name: last_name
-    description: ""
-  - name: email
-    description: "Email Address"
-    meta:
-      masking_policy: mp_email
-  - name: gender
+  - name: stg_contacts
     description: ""
+    columns:
+      - name: contact_id
+        description: ""
+      - name: first_name
+        description: ""
+        meta:
+          masking_policy: mp_conditional_contacts_pii
+          mp_conditional_columns: 
+            - last_name 
+      - name: last_name
+        description: ""
+      - name: email
+        description: "Email Address"
+        meta:
+          masking_policy: mp_email
+      - name: gender
+        description: ""
 
-  - name: ip_address
-    description: ""
+      - name: ip_address
+        description: ""
 
-  - name: ssn
-    description: ""
+      - name: ssn
+        description: ""
 
-  - name: phone
-    description: ""
+      - name: phone
+        description: ""
diff --git a/macros/snow-mask-ddl/create_masking_policy_mp_conditional_pii.sql b/macros/snow-mask-ddl/create_masking_policy_mp_conditional_pii.sql
@@ -0,0 +1,15 @@
+{% macro create_masking_policy_mp_conditional_pii(node_database, node_schema, masked_column) %}
+
+    CREATE MASKING POLICY IF NOT EXISTS {{node_database}}.{{node_schema}}.mp_conditional_pii AS (
+        {{masked_column}} string,
+        my_conditional_col_1 string,
+        my_conditional_col_2 string
+    ) RETURNS string ->
+        CASE 
+            WHEN CURRENT_ROLE() IN ('ANALYST') AND my_conditional_col_1='foo' THEN {{masked_column}}
+            WHEN CURRENT_ROLE() IN ('ANALYST') AND my_conditional_col_2='bar' THEN SHA2({{masked_column}})
+             WHEN CURRENT_ROLE() IN ('SYSADMIN') THEN SHA2({{masked_column}})
+        ELSE '**********'
+        END
+
+{% endmacro %}
diff --git a/macros/snow-mask/apply-policy/apply_masking_policy_list_for_models.sql b/macros/snow-mask/apply-policy/apply_masking_policy_list_for_models.sql
@@ -50,15 +50,19 @@
             {% endif %}
 
             {%- for meta_tuple in meta_columns if meta_columns | length > 0 %}
-                {% set column   = meta_tuple[0] %}
+                {% set column               = meta_tuple[0] %}
                 {% set masking_policy_name  = meta_tuple[1] %}
-                    {% if masking_policy_name is not none %}
+                {% set conditional_columns  = meta_tuple[2] %}
+
+                {% if masking_policy_name is not none %}
 
                     {% for masking_policy_in_db in masking_policy_list['MASKING_POLICY'] %}
                         {% if masking_policy_db|upper ~ '.' ~ masking_policy_schema|upper ~ '.' ~ masking_policy_name|upper == masking_policy_in_db %}
                             {{ log(modules.datetime.datetime.now().strftime("%H:%M:%S") ~ " | " ~ operation_type ~ "ing masking policy to model  : " ~ masking_policy_db|upper ~ '.' ~ masking_policy_schema|upper ~ '.' ~ masking_policy_name|upper ~ " on " ~ database ~ '.' ~ schema ~ '.' ~ alias ~ '.' ~ column ~ ' [force = ' ~ var('use_force_applying_masking_policy','False') ~ ']', info=True) }}
                             {% set query %}
-                            alter {{materialization}}  {{database}}.{{schema}}.{{alias}} modify column  {{column}} set masking policy {{masking_policy_db}}.{{masking_policy_schema}}.{{masking_policy_name}} {% if var('use_force_applying_masking_policy','False')|upper in ['TRUE','YES'] %} force {% endif %};
+                            alter {{materialization}} {{database}}.{{schema}}.{{alias}}
+                            modify column {{column}}
+                            set masking policy {{masking_policy_db}}.{{masking_policy_schema}}.{{masking_policy_name}} {% if conditional_columns | length > 0 %}using ({{column}}, {{conditional_columns|join(', ')}}){% endif %} {% if var('use_force_applying_masking_policy','False')|upper in ['TRUE','YES'] %} force {% endif %};
                             {% endset %}
                             {% do run_query(query) %}
                         {% endif %}

diff --git a/macros/snow-mask/apply-policy/apply_masking_policy_list_for_sources.sql b/macros/snow-mask/apply-policy/apply_masking_policy_list_for_sources.sql
@@ -51,8 +51,9 @@
         {% endif %}
 
         {%- for meta_tuple in meta_columns if meta_columns | length > 0 %}
-            {% set column   = meta_tuple[0] %}
+            {% set column               = meta_tuple[0] %}
             {% set masking_policy_name  = meta_tuple[1] %}
+            {% set conditional_columns  = meta_tuple[2] %}
 
             {% if masking_policy_name is not none %}
 
@@ -61,7 +62,9 @@
                         {{ log(modules.datetime.datetime.now().strftime("%H:%M:%S") ~ " | " ~ operation_type ~ "ing masking policy to source : " ~ masking_policy_db|upper ~ '.' ~ masking_policy_schema|upper ~ '.' ~ masking_policy_name|upper ~ " on " ~ database ~ '.' ~ schema ~ '.' ~ identifier ~ '.' ~ column ~ ' [force = ' ~ var('use_force_applying_masking_policy','False') ~ ']', info=True) }}
                         {% set query %}
                             {% if operation_type == "apply" %}
-                                alter {{materialization}}  {{database}}.{{schema}}.{{identifier}} modify column  {{column}} set masking policy  {{masking_policy_db}}.{{masking_policy_schema}}.{{masking_policy_name}} {% if var('use_force_applying_masking_policy','False')|upper in ['TRUE','YES'] %} force {% endif %}
+                                alter {{materialization}} {{database}}.{{schema}}.{{identifier}}
+                                modify column {{column}}
+                                set masking policy  {{masking_policy_db}}.{{masking_policy_schema}}.{{masking_policy_name}} {% if conditional_columns | length > 0 %}using ({{column}}, {{conditional_columns|join(', ')}}){% endif %} {% if var('use_force_applying_masking_policy','False')|upper in ['TRUE','YES'] %} force {% endif %}
                             {% elif operation_type == "unapply" %}
                                 alter {{materialization}}  {{database}}.{{schema}}.{{identifier}} modify column  {{column}} unset masking policy
                             {% endif %}

diff --git a/macros/snow-mask/create-policy/create_masking_policy.sql b/macros/snow-mask/create-policy/create_masking_policy.sql
@@ -30,14 +30,19 @@
             {% endif %}
         {% endif %}
 
-        {% set current_policy_name = masking_policy[2] | string  %}
+        {% set current_policy_name          = masking_policy[2] | string  %}
+        {% set conditionally_masked_column  = masking_policy[3] %}
 
         {%- if (var('create_masking_policy_schema', 'True')|upper in ['TRUE','YES']) -%}
             {% do adapter.create_schema(api.Relation.create(database=masking_policy_db, schema=masking_policy_schema)) %}
         {% endif %}
 
         {% set call_masking_policy_macro = context["create_masking_policy_" | string ~ current_policy_name | string]  %}
-        {% set result = run_query(call_masking_policy_macro(masking_policy_db, masking_policy_schema)) %}
+        {% if conditionally_masked_column is not none %}
+            {% set result = run_query(call_masking_policy_macro(masking_policy_db, masking_policy_schema, conditionally_masked_column)) %}
+        {% else %}
+            {% set result = run_query(call_masking_policy_macro(masking_policy_db, masking_policy_schema)) %}
+        {% endif %}
     {% endfor %}
 
 {% endif %}

diff --git a/macros/snow-mask/create-policy/get_masking_policy_list_for_models.sql b/macros/snow-mask/create-policy/get_masking_policy_list_for_models.sql
@@ -6,21 +6,26 @@
 
         {{ log(modules.datetime.datetime.now().strftime("%H:%M:%S") ~ " | macro - now processing            : " ~ node.unique_id | string , info=False) }}
 
-        {% set node_database = node.database | string %}
-        {% set node_schema   = node.schema | string %}
-        {% set node_unique_id = node.unique_id | string %}
-        {% set node_resource_type = node.resource_type | string %}
+        {% set node_database        = node.database | string %}
+        {% set node_schema          = node.schema | string %}
+        {% set node_unique_id       = node.unique_id | string %}
+        {% set node_resource_type   = node.resource_type | string %}
 
         {% set meta_columns = dbt_snow_mask.get_meta_objects(node_unique_id,meta_key,node_resource_type) %}
 
         {%- for meta_tuple in meta_columns if meta_columns | length > 0 %}
             {{ log(modules.datetime.datetime.now().strftime("%H:%M:%S") ~ " | macro - meta_columns               : " ~ node_unique_id ~ " has " ~ meta_columns | string ~ " masking tags set", info=False) }}
 
-            {% set column   = meta_tuple[0] %}
-            {% set masking_policy_name  = meta_tuple[1] %}
-
+            {% set column                           = meta_tuple[0] %}
+            {% set masking_policy_name              = meta_tuple[1] %}
+            {% set conditional_columns              = meta_tuple[2] %}
+            {% if conditional_columns | length > 0 %}
+                {% set conditionally_masked_column = column %}
+            {% else %}
+                {% set conditionally_masked_column = none %}
+            {% endif %}
             {% if masking_policy_name is not none %}
-                {% set masking_policy_tuple = (node_database, node_schema, meta_tuple[1]) %}
+                {% set masking_policy_tuple = (node_database, node_schema, masking_policy_name, conditionally_masked_column) %}
                 {% do masking_policies.append(masking_policy_tuple) %}
             {% endif %}
 

diff --git a/macros/snow-mask/create-policy/get_masking_policy_list_for_sources.sql b/macros/snow-mask/create-policy/get_masking_policy_list_for_sources.sql
@@ -6,21 +6,26 @@
 
         {{ log(modules.datetime.datetime.now().strftime("%H:%M:%S") ~ " | macro - now processing            : " ~ node.unique_id | string , info=False) }}
 
-        {% set node_database = node.database | string %}
-        {% set node_schema   = node.schema | string %}
-        {% set node_unique_id = node.unique_id | string %}
-        {% set node_resource_type = node.resource_type | string %}
+        {% set node_database        = node.database | string %}
+        {% set node_schema          = node.schema | string %}
+        {% set node_unique_id       = node.unique_id | string %}
+        {% set node_resource_type   = node.resource_type | string %}
 
         {% set meta_columns = dbt_snow_mask.get_meta_objects(node_unique_id,meta_key,node_resource_type) %}
 
         {%- for meta_tuple in meta_columns if meta_columns | length > 0 %}
             {{ log(modules.datetime.datetime.now().strftime("%H:%M:%S") ~ " | macro - meta_columns               : " ~ node_unique_id ~ " has " ~ meta_columns | string ~ " masking tags set", info=False) }}
 
-            {% set column   = meta_tuple[0] %}
-            {% set masking_policy_name  = meta_tuple[1] %}
-
+            {% set column                       = meta_tuple[0] %}
+            {% set masking_policy_name          = meta_tuple[1] %}
+            {% set conditional_columns          = meta_tuple[2] %}
+            {% if conditional_columns | length > 0 %}
+                {% set conditionally_masked_column = column %}
+            {% else %}
+                {% set conditionally_masked_column = none %}
+            {% endif %}
             {% if masking_policy_name is not none %}
-                {% set masking_policy_tuple = (node_database, node_schema, meta_tuple[1]) %}
+                {% set masking_policy_tuple = (node_database, node_schema, masking_policy_name, conditionally_masked_column) %}
                 {% do masking_policies.append(masking_policy_tuple) %}
             {% endif %}
 

diff --git a/macros/snow-mask/get_meta_objects.sql b/macros/snow-mask/get_meta_objects.sql
@@ -12,18 +12,30 @@
             {% if node_resource_type == "source" %} 
                 {% for column in columns if graph.sources[node_unique_id]['columns'][column]['meta'][meta_key] | length > 0 %}
                     {% set meta_dict = graph.sources[node_unique_id]['columns'][column]['meta'] %}
-                    {% for key, value in meta_dict.items() if key == meta_key %}
-                        {% set meta_tuple = (column ,value ) %}
+                    {% if meta_key in meta_dict %}
+                        {% set policy_name = meta_dict[meta_key] %}
+                        {% if "mp_conditional_columns" in meta_dict %}
+                            {% set conditional_columns = meta_dict['mp_conditional_columns'] %}
+                        {% else %}
+                            {% set conditional_columns = [] %}
+                        {% endif %}
+                        {% set meta_tuple = (column, policy_name, conditional_columns) %}
                         {% do meta_columns.append(meta_tuple) %}
-                    {% endfor %}
+                    {% endif %}
                 {% endfor %}
             {% else %}
                 {% for column in columns if graph.nodes[node_unique_id]['columns'][column]['meta'][meta_key] | length > 0 %}
                     {% set meta_dict = graph.nodes[node_unique_id]['columns'][column]['meta'] %}
-                    {% for key, value in meta_dict.items() if key == meta_key %}
-                        {% set meta_tuple = (column ,value ) %}
+                    {% if meta_key in meta_dict %}
+                        {% set policy_name = meta_dict[meta_key] %}
+                        {% if "mp_conditional_columns" in meta_dict %}
+                            {% set conditional_columns = meta_dict['mp_conditional_columns'] %}
+                        {% else %}
+                            {% set conditional_columns = [] %}
+                        {% endif %}
+                        {% set meta_tuple = (column, policy_name, conditional_columns) %}
                         {% do meta_columns.append(meta_tuple) %}
-                    {% endfor %}
+                    {% endif %}
                 {% endfor %}
             {% endif %}
         {% else %}