entechlog · entechlog · Jun 14, 2021 · Jun 14, 2021
diff --git a/README.md b/README.md
@@ -1,12 +1,16 @@
 - [Overview](#overview)
-- [How to use this package ?](#how-to-use-this-package-)
-- [Credits](#credits)
+- [Installation Instructions](#installation-instructions)
+- [How to apply masking policy ?](#how-to-apply-masking-policy-)
+- [How to remove masking policy ?](#how-to-remove-masking-policy-)
+- [How to validate masking policy ?](#how-to-validate-masking-policy-)
 - [Future Enhancements](#future-enhancements)
+- [Credits](#credits)
+- [References](#references)
 
 # Overview
 This dbt package contains macros that can be (re)used across dbt projects with snowflake. `dbt_snow_mask` will help to apply [Dynamic Data Masking](https://docs.snowflake.com/en/user-guide/security-column-ddm-use.html) using [dbt meta](https://docs.getdbt.com/reference/resource-properties/meta).
 
-# How to use this package ?
+# Installation Instructions
 
 - Add the package into your project.
 
@@ -16,7 +20,10 @@ This dbt package contains macros that can be (re)used across dbt projects with s
     - git: "https://github.com/entechlog/dbt-snow-mask.git"
     revision: 0.1.1
   ```
-> Please refer to the release version for the latest revision
+
+> Please refer to the release version of this repo/dbt hub for the latest revision
+
+# How to apply masking policy ?
 
 - Masking is controlled by [meta](https://docs.getdbt.com/reference/resource-properties/meta) in [dbt resource properties](https://docs.getdbt.com/reference/declaring-properties) for sources and models. 
 
@@ -62,9 +69,10 @@ This dbt package contains macros that can be (re)used across dbt projects with s
   {% endmacro %}
   ```
 
-> Its also good to keep the masking policy ddl organized in a directory say `\macros\snow-mask-ddl`
+> Its good to keep the masking policy ddl organized in a directory say `\macros\snow-mask-ddl`
+
+- Create the masking policies by running below command  
 
-- Create the masking policies by running below command
 
   | Resource Type | Command                                                                         |
   | ------------- | ------------------------------------------------------------------------------- |
@@ -73,25 +81,58 @@ This dbt package contains macros that can be (re)used across dbt projects with s
 
 - Add post-hook to `dbt_project.yml`
 
-  ***Example** : dbt_project.yml
+  **Example** : dbt_project.yml
 
   ```bash
   models:
   post-hook: 
     - "{{ dbt_snow_mask.apply_masking_policy() }}"
   ```
 
-- Apply the masking policy by running below commands
+- Apply the masking policy by running below commands  
+
 
   | Resource Type | Command                                                                        |
   | ------------- | ------------------------------------------------------------------------------ |
   | sources       | `dbt run-operation apply_masking_policy --args '{"resource_type": "sources"}'` |
   | models        | `dbt run -- model <model-name>`                                                |
-
-# Credits
-This package was created using example macros from [Serge](https://getdbt.slack.com/archives/CJN7XRF1B/p1609177817234800)
+
+# How to remove masking policy ?
+
+- Remove the masking policy applied by this package by running below commands  
+
+
+  | Resource Type | Command                                                                          |
+  | ------------- | -------------------------------------------------------------------------------- |
+  | sources       | `dbt run-operation unapply_masking_policy --args '{"resource_type": "sources"}'` |
+  | models        | `dbt run-operation unapply_masking_policy --args '{"resource_type": "models"}'`  |
+
+# How to validate masking policy ?
+
+```sql
+-- Show masking policy
+SHOW MASKING POLICIES;
+
+-- Describe masking policy
+DESCRIBE MASKING POLICY <masking-policy-name>;
+
+-- Show masking policy references
+USE DATABASE <database-name>;
+
+USE SCHEMA INFORMATION_SCHEMA;
+
+SELECT *
+  FROM TABLE(INFORMATION_SCHEMA.POLICY_REFERENCES(POLICY_NAME => '<database-name>.<schema-name>.<masking-policy-name>'));
+```
 
 # Future Enhancements
 - Optimize macros & reduce number of lines in macros
 - `apply_masking_policy_list_for_sources` needs changes to find the `materialization` 
 - Add support for `CREATE OR REPLACE MASKING POLICY`. This needs unset to happen before replacing the existing policy to avoid `SQL compilation error: Policy TEMP cannot be dropped/replaced as it is associated with one or more entities.` error
+
+# Credits
+This package was created using examples from [Serge](https://www.linkedin.com/in/serge-gekker-912b9928/) and [Matt](https://www.linkedin.com/in/matt-winkler-4024263a/)
+
+# References
+- https://docs.snowflake.com/en/user-guide/security-column-ddm-intro.html
+- https://getdbt.slack.com/archives/CJN7XRF1B/p1609177817234800
diff --git a/macros/snow-mask/apply_masking_policy_list_for_models.sql b/macros/snow-mask/apply_masking_policy_list_for_models.sql
@@ -1,5 +1,7 @@
-{% macro apply_masking_policy_list_for_models(meta_key) %}
+{% macro apply_masking_policy_list_for_models(meta_key,operation_type="apply") %}
 
+    {% if operation_type == "apply" %}
+
         {% set model_id = model.unique_id | string %}
         {% set alias    = model.alias %}    
         {% set database = model.database %}
@@ -24,7 +26,7 @@
 
                 {% for masking_policy_in_db in masking_policy_list['MASKING_POLICY'] %}
                     {% if database|upper ~ '.' ~ schema|upper ~ '.' ~ masking_policy_name|upper == masking_policy_in_db %}
-                        {{ log(modules.datetime.time() ~ " | applying masking policy (model)   : " ~ database|upper ~ '.' ~ schema|upper ~ '.' ~ masking_policy_name|upper ~ " on " ~ database ~ '.' ~ schema ~ '.' ~ alias ~ '.' ~ column, info=True) }}
+                        {{ log(modules.datetime.time() ~ " | " ~ operation_type ~ "ing masking policy to model  : " ~ database|upper ~ '.' ~ schema|upper ~ '.' ~ masking_policy_name|upper ~ " on " ~ database ~ '.' ~ schema ~ '.' ~ alias ~ '.' ~ column, info=True) }}
                         {% set query %}
                         alter {{materialization}}  {{database}}.{{schema}}.{{alias}} modify column  {{column}} set masking policy  {{database}}.{{schema}}.{{masking_policy_name}};
                         {% endset %}
@@ -35,4 +37,35 @@
                {% endif %}
         {% endfor %}
 
-    {% endmacro %}
+    {% elif operation_type == "unapply" %}
+
+        {% for node in graph.nodes.values() -%}
+
+            {% set database = node.database | string %}
+            {% set schema   = node.schema | string %}
+            {% set node_unique_id = node.unique_id | string %}
+            {% set node_resource_type = node.resource_type | string %}
+            {% set materialization = node.config.materialized | string %}
+            {% set alias    = node.alias %}
+
+            {% 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 %}
+                {% set column   = meta_tuple[0] %}
+                {% set masking_policy_name  = meta_tuple[1] %}
+
+                {% if masking_policy_name is not none %}
+                    {{ log(modules.datetime.time() ~ " | " ~ operation_type ~ "ing masking policy to model  : " ~ database|upper ~ '.' ~ schema|upper ~ '.' ~ masking_policy_name|upper ~ " on " ~ database ~ '.' ~ schema ~ '.' ~ alias ~ '.' ~ column, info=True) }}
+                    {% set query %}
+                        alter {{materialization}}  {{database}}.{{schema}}.{{alias}} modify column  {{column}} unset masking policy
+                    {% endset %}
+                    {% do run_query(query) %}
+                {% endif %}
+
+            {% endfor %}
+
+        {% endfor %}
+
+    {% endif %}
+
+{% endmacro %}
diff --git a/macros/snow-mask/apply_masking_policy_list_for_sources.sql b/macros/snow-mask/apply_masking_policy_list_for_sources.sql
@@ -1,4 +1,4 @@
-{% macro apply_masking_policy_list_for_sources(meta_key) %}
+{% macro apply_masking_policy_list_for_sources(meta_key,operation_type="apply") %}
 
     {% for node in graph.sources.values() -%}
 
@@ -25,9 +25,13 @@
 
                 {% for masking_policy_in_db in masking_policy_list['MASKING_POLICY'] %}
                     {% if database|upper ~ '.' ~ schema|upper ~ '.' ~ masking_policy_name|upper == masking_policy_in_db %}
-                        {{ log(modules.datetime.time() ~ " | applying masking policy (source)  : " ~ database|upper ~ '.' ~ schema|upper ~ '.' ~ masking_policy_name|upper ~ " on " ~ database ~ '.' ~ schema ~ '.' ~ name ~ '.' ~ column, info=True) }}
+                        {{ log(modules.datetime.time() ~ " | " ~ operation_type ~ "ing masking policy to source : " ~ database|upper ~ '.' ~ schema|upper ~ '.' ~ masking_policy_name|upper ~ " on " ~ database ~ '.' ~ schema ~ '.' ~ name ~ '.' ~ column, info=True) }}
                         {% set query %}
-                        alter {{materialization}}  {{database}}.{{schema}}.{{name}} modify column  {{column}} set masking policy  {{database}}.{{schema}}.{{masking_policy_name}}
+                            {% if operation_type == "apply" %}
+                                alter {{materialization}}  {{database}}.{{schema}}.{{name}} modify column  {{column}} set masking policy  {{database}}.{{schema}}.{{masking_policy_name}}
+                            {% elif operation_type == "unapply" %}
+                                alter {{materialization}}  {{database}}.{{schema}}.{{name}} modify column  {{column}} unset masking policy
+                            {% endif %}
                         {% endset %}
                         {% do run_query(query) %}
                     {% endif %}

diff --git a/macros/snow-mask/create_masking_policy.sql b/macros/snow-mask/create_masking_policy.sql
@@ -14,7 +14,7 @@
         {% set current_database = masking_policy[0] | string  %}
         {% set current_schema = masking_policy[1] | string  %}
         {% set current_policy_name = masking_policy[2] | string  %}
-        {{ log(modules.datetime.time() ~ " | creating masking policies for     : " ~ current_database|upper ~ '.' ~ current_schema|upper ~ '.' ~ current_policy_name|upper , info=True) }}
+        {{ log(modules.datetime.time() ~ " | creating masking policy           : " ~ current_database|upper ~ '.' ~ current_schema|upper ~ '.' ~ current_policy_name|upper , info=True) }}
         {% set call_masking_policy_macro = context["create_masking_policy_" | string ~ current_policy_name | string]  %}
         {{ run_query(call_masking_policy_macro(current_database, current_schema)) }}
     {% endfor %}

diff --git a/macros/snow-mask/unapply_masking_policy.sql b/macros/snow-mask/unapply_masking_policy.sql
@@ -0,0 +1,13 @@
+{% macro unapply_masking_policy(resource_type="models",meta_key="masking_policy",operation_type="unapply") %}
+
+    {% if execute %}
+
+        {% if resource_type == "sources" %}
+            {{ dbt_snow_mask.apply_masking_policy_list_for_sources(meta_key,operation_type) }}
+        {% else %}
+            {{ dbt_snow_mask.apply_masking_policy_list_for_models(meta_key,operation_type) }}
+        {% endif %}
+
+    {% endif %}
+
+{% endmacro %}