Enable functionality to use common schema in model database

README.md

@@ -46,21 +46,51 @@ This dbt package contains macros that can be (re)used across dbt projects with s
 
 # How to configure database and schema for the masking policy ?
 
-By default this process creates the masking policies in same directory as the database objects. You can change this default behavior by using following parameters in your `dbt_project.yml` 
+By default this process creates the masking policies in same directory as the database objects. There are 2 methods for changing the default behavior by using the following parameters in your `dbt_project.yml` 
 
+**Method 1 : Use a common database**
+
+To change the database that your masking polices are created in set the following parameters:
 * `use_common_masking_policy_db` (optional): Flag to enable the usage of a common db/schema for all masking policies. Valid values are “True” OR "False"
 * `common_masking_policy_db` (optional): The database name for creating masking policies
 * `common_masking_policy_schema` (optional): The schema name for creating masking policies
 
-**Example** : var block in dbt_project.yml
-
+**Example** : var block in dbt_project.yml to enable using a common masking policy database
 ```yaml
 vars:
   use_common_masking_policy_db: "True"
   common_masking_policy_db: "DEMO_DB"
   common_masking_policy_schema: "COMPLIANCE"
 ```
 
+**Method 2 : Use a common schema (in the current database)**
+
+To change only the schema (so that a common masking policy schema is used in the same database as your model) set the following parameters:
+* `use_common_masking_policy_schema_only` (optional): Flag to enable the usage of a common schema in the current database for all masking policies. Valid values are “True” OR "False"
+* `common_masking_policy_schema` (optional): The schema name for creating masking policies
+
+**Example** : var block in dbt_project.yml to enable using a common masking policy schema (in the current database)
+
+```yaml
+vars:
+  use_common_masking_policy_schema_only: "True"
+  common_masking_policy_schema: "COMPLIANCE"
+```
+
+> ✅ If both `use_common_masking_policy_db` and `use_common_masking_policy_schema_only` are set to True, then `use_common_masking_policy_db` will supercede `use_common_masking_policy_schema_only`.
+
+**Allow Custom Materializations**
+
+To enable dbt_snow_mask to apply masking policies to models generated from custom materializations in dbt, configure the following parameter:
+* `custom_materializations_map` (optional): A dictionary containing key-value pairs mapping custom materializations in dbt to the objects they generate in Snowflake. For each pair, the key must be the name of the custom_materialization and the value must be either `table` or `view`. 
+
+**Example** : var block in dbt_project.yml to enable application of masking policies to a model generated using a custom materialiazition that ends up as a table in Snowflake.
+
+```yaml
+vars:
+  custom_materializations_map: '{ "custom_incremental": "table" }'
+```
+
 # 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. 

integration_tests/dbt_project.yml

@@ -21,6 +21,9 @@ clean-targets:
 #  common_masking_policy_db: "DEMO_DB"
 #  common_masking_policy_schema: "COMPLIANCE"
 
+#   use_common_masking_policy_schema_only: "True"
+#   common_masking_policy_schema: "COMPLIANCE"
+
 seeds:
   dbt_snow_mask_integration_tests:
     sakila:

macros/snow-mask/apply-policy/apply_masking_policy_list_for_models.sql

@@ -6,23 +6,34 @@
         {% set alias    = model.alias %}    
         {% set database = model.database %}
         {% set schema   = model.schema %}
-        {% set materialization = model.config.get("materialized") %}
-        {% if materialization == "incremental" %}
-            {% set materialization = "table" %}
-        {% endif %}
+
+        {# This dictionary stores a mapping between materializations in dbt and the objects they will generate in Snowflake  #}
+        {% set materialization_map = {"table": "table", "view": "view", "incremental": "table"} %}
+
+        {# Append custom materializations to the list of standard materializations  #}
+        {% do materialization_map.update(fromjson(var('custom_materializations_map', '{}'))) %}
+
+        {% set materialization = materialization_map[model.config.get("materialized")] %}
         {% set meta_columns = dbt_snow_mask.get_meta_objects(model_id,meta_key) %}
 
         {% set masking_policy_db = model.database %}
         {% set masking_policy_schema = model.schema %}
 
-        {# Override the database and schema name when use common_masking_policy_db flag is set #}
-        {%- if (var('use_common_masking_policy_db', 'False')|upper == 'TRUE') or (var('use_common_masking_policy_db', 'False')|upper == 'YES') -%}
-            {% if var('common_masking_policy_db') and var('common_masking_policy_schema') %}
+        {# Override the database and schema name when use_common_masking_policy_db flag is set #}
+        {%- if (var('use_common_masking_policy_db', 'False')|upper in ['TRUE','YES']) -%}
+            {% if (var('common_masking_policy_db') and var('common_masking_policy_schema')) %}
                 {% set masking_policy_db = var('common_masking_policy_db') | string  %}
                 {% set masking_policy_schema = var('common_masking_policy_schema') | string  %}
             {% endif %}
         {% endif %}
 
+        {# Override the schema name (in the masking_policy_db) when use_common_masking_policy_schema_only flag is set #}
+        {%- if (var('use_common_masking_policy_schema_only', 'False')|upper in ['TRUE','YES']) and (var('use_common_masking_policy_db', 'False')|upper in ['FALSE','NO']) -%}
+            {% if var('common_masking_policy_schema') %}
+                {% set masking_policy_schema = var('common_masking_policy_schema') | string  %}
+            {% endif %}
+        {% endif %}
+
         {% set masking_policy_list_sql %}     
             show masking policies in {{masking_policy_db}}.{{masking_policy_schema}};
             select $3||'.'||$4||'.'||$2 as masking_policy from table(result_scan(last_query_id()));
@@ -39,7 +50,7 @@
                     {% 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, info=True) }}
                         {% set query %}
-                        alter {{materialization}}  {{database}}.{{schema}}.{{alias}} modify column  {{column}} set masking policy  {{masking_policy_db}}.{{masking_policy_schema}}.{{masking_policy_name}};
+                        alter {{materialization}}  {{database}}.{{schema}}.{{alias}} modify column  {{column}} set masking policy {{masking_policy_db}}.{{masking_policy_schema}}.{{masking_policy_name}};
                         {% endset %}
                         {% do run_query(query) %}
                     {% endif %}
@@ -56,10 +67,12 @@
             {% 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 %}
-            {% if materialization == "incremental" %}
-                {% set materialization = "table" %}
-            {% endif %}
+            {% set materialization_map = {"table": "table", "view": "view", "incremental": "table"} %}
+
+            {# Append custom materializations to the list of standard materializations  #}
+            {% do materialization_map.update(fromjson(var('custom_materializations_map', '{}'))) %}
+
+            {% set materialization = materialization_map[model.config.get("materialized")] %}
             {% set alias    = node.alias %}
 
             {% set meta_columns = dbt_snow_mask.get_meta_objects(node_unique_id,meta_key,node_resource_type) %}

macros/snow-mask/apply-policy/apply_masking_policy_list_for_sources.sql

@@ -23,14 +23,21 @@
         {% set masking_policy_db = node.database %}
         {% set masking_policy_schema = node.schema %}
 
-        {# Override the database and schema name when use common_masking_policy_db flag is set #}
-        {%- if (var('use_common_masking_policy_db', 'False')|upper == 'TRUE') or (var('use_common_masking_policy_db', 'False')|upper == 'YES') -%}
-            {% if var('common_masking_policy_db') and var('common_masking_policy_schema') %}
+        {# Override the database and schema name when use_common_masking_policy_db flag is set #}
+        {%- if (var('use_common_masking_policy_db', 'False')|upper in ['TRUE','YES']) -%}
+            {% if (var('common_masking_policy_db') and var('common_masking_policy_schema')) %}
                 {% set masking_policy_db = var('common_masking_policy_db') | string  %}
                 {% set masking_policy_schema = var('common_masking_policy_schema') | string  %}
             {% endif %}
         {% endif %}
 
+        {# Override the schema name (in the masking_policy_db) when use_common_masking_policy_schema_only flag is set #}
+        {%- if (var('use_common_masking_policy_schema_only', 'False')|upper in ['TRUE','YES']) and (var('use_common_masking_policy_db', 'False')|upper in ['FALSE','NO']) -%}
+            {% if var('common_masking_policy_schema') %}
+                {% set masking_policy_schema = var('common_masking_policy_schema') | string  %}
+            {% endif %}
+        {% endif %}
+
         {% set masking_policy_list_sql %}
             show masking policies in {{masking_policy_db}}.{{masking_policy_schema}};
             select $3||'.'||$4||'.'||$2 as masking_policy from table(result_scan(last_query_id()));

macros/snow-mask/create-policy/create_masking_policy.sql

@@ -12,24 +12,30 @@
 
     {% for masking_policy in masking_policies | unique -%}
 
-        {% set current_database = masking_policy[0] | string  %}
-        {% set current_schema = masking_policy[1] | string  %}
-
-        {# Override the database and schema name when use common_masking_policy_db flag is set #}
-        {%- if (var('use_common_masking_policy_db', 'False')|upper == 'TRUE') or (var('use_common_masking_policy_db', 'False')|upper == 'YES') -%}
-            {% if var('common_masking_policy_db') and var('common_masking_policy_schema') %}
-                {% set current_database = var('common_masking_policy_db') | string  %}
-                {% set current_schema = var('common_masking_policy_schema') | string  %}
+        {% set masking_policy_db = masking_policy[0] | string  %}
+        {% set masking_policy_schema = masking_policy[1] | string  %}
+
+        {# Override the database and schema name when use_common_masking_policy_db flag is set #}
+        {%- if (var('use_common_masking_policy_db', 'False')|upper in ['TRUE','YES']) -%}
+            {% if (var('common_masking_policy_db') and var('common_masking_policy_schema')) %}
+                {% set masking_policy_db = var('common_masking_policy_db') | string  %}
+                {% set masking_policy_schema = var('common_masking_policy_schema') | string  %}
             {% endif %}
         {% endif %}
 
-        {% set current_policy_name = masking_policy[2] | string  %}
-        {{ log(modules.datetime.datetime.now().strftime("%H:%M:%S") ~ " | creating masking policy           : " ~ current_database|upper ~ '.' ~ current_schema|upper ~ '.' ~ current_policy_name|upper , info=True) }}
+        {# Override the schema name (in the masking_policy_db) when use_common_masking_policy_schema_only flag is set #}
+        {%- if (var('use_common_masking_policy_schema_only', 'False')|upper in ['TRUE','YES']) and (var('use_common_masking_policy_db', 'False')|upper in ['FALSE','NO']) -%}
+            {% if var('common_masking_policy_schema') %}
+                {% set masking_policy_schema = var('common_masking_policy_schema') | string  %}
+            {% endif %}
+        {% endif %}
 
-        {% do adapter.create_schema(api.Relation.create(database=current_database, schema=current_schema)) %}
+        {% set current_policy_name = masking_policy[2] | string  %}
+
+        {% do adapter.create_schema(api.Relation.create(database=masking_policy_db, schema=masking_policy_schema)) %}
 
         {% set call_masking_policy_macro = context["create_masking_policy_" | string ~ current_policy_name | string]  %}
-        {% set result = run_query(call_masking_policy_macro(current_database, current_schema)) %}
+        {% set result = run_query(call_masking_policy_macro(masking_policy_db, masking_policy_schema)) %}
     {% endfor %}
 
 {% endif %}