Tag-Based Masking in Snowflake Using Stored Procedures

Data protection crucial when you are handling your clients' personally identifiable information (PII).

Building a robust strategy to protect PII can sometimes prove challenging, as you need to strike a balance between very specific requirements and technical considerations. In this article, we will show you how to implement tag-based masking in Snowflake using stored procedures, without the need for any external applications.

PII challenges

Safeguarding clients' personal information, even from development teams, is a top priority. If your data resides in Snowflake, you have two options to implement tag-based masking to tackle this issue:

1. Assigning tags with dbt
2. Assigning tags with a Snowflake stored procedure

While a dbt solution would be simpler and quicker to implement, it carries the drawback of exposing PII data until it has been masked via dbt.

Let us consider an example to understand the delay in the dbt approach. Imagine using a third-party tool to ingest raw PII data into Snowflake. If we relied on dbt for masking, there would be a time lag while the dbt jobs run, particularly if they are scheduled for a specific time or day. Even if a third-party ingestion tool triggers the dbt jobs, running these jobs still takes a bit of time, creating a small window where PII data is exposed before it is masked.

With the tag-based in Snowflake, we can trigger the stored procedure as soon as the data is ingested by the tool. This process takes less time, significantly shortening the window of data exposure. Furthermore, this offers a sleeker solution, as everything remains within the data warehouse component. It is not always the case that dbt will be used within a project, and employing it solely for masking might not be viable.

Getting started with tag-based masking in Snowflake

Tag-based masking in Snowflake is an automated data protection mechanism. It dynamically applies masking policies based on tags assigned to database objects, such as columns. Rather than manually assigning a masking policy to each column, you simply apply a tag, and Snowflake automatically enforces the masking policy for all columns with that tag.

Implementing tag-based masking in Snowflake with stored procedures

The tag-based approach involves the creation of several components:

• Tags
• Masking policies
• Role to access the PII data
• Stored procedure
• Configuration table
• Logging table

We will now see how all of these elements work together.

To begin, we create the role, tags, and masking policies, and then apply the masking policies to the tag.

Creating this role is not strictly necessary; however, it is good practice to do so and then assign it to other roles or users.

If we instead use a set of roles, or a generic role such as ACCOUNTADMIN, we might still be able to access PII data when roles are chained. Your role may not be directly assigned, but if it has, for example, two or three chained roles, it could inadvertently provide the permissions to view the data. In situations where the infrastructure is complex and involves many objects, it is easy to make such mistakes.

Access roles

When implementing tag-based masking in Snowflake, you can specify the PII access role within the masking policy in two main ways:

• WHEN CURRENT_ROLE() IN ('PII_ACCESS') THEN val:
  • You can bypass the masking policy only by using the PII_ACCESS role directly;
  • You can see the PII data only by using the PII_ACCESS role;
• WHEN IS_ROLE_IN_SESSION('PII_ACCESS') THEN val:
  • If your user has been assigned the PII_ACCESS role and this role has been used within the current session or you are using a different role that has been granted the PII_ACCESS role to it, you will be able to bypass the masking policy;
  • You do not necessarily need to use the PII_ACCESS role directly to view PII data.

Next, we will create the tables:

The first table, METADATA_MASKING_CONFIG_TABLE, will contain the necessary configuration to mask a column. It will store the database, schema, table, and column name of the table, along with which tag to use for masking the column. It will also hold other metadata, which will be used later in the stored procedure for an efficient and clean implementation. We also need to populate this configuration table manually with the required PII data target column information.

Applying masking policies automatically with a stored procedure

We have covered the necessary steps to set up our process. Now, we will focus on the final component: the stored procedure that will bring everything together and apply the masking. The code for this stored procedure is quite lengthy and complex, so it is worth explaining in more detail what it does before sharing the SQL code.

At a high level, we are looking to accomplish the following:

1. Iterate through the configuration table, row by row;
2. Compare the actual data type of the column with the data type specified in the configuration table:
   1. If they match, create an ALTER TABLE statement using the values from each row of the configuration table.
   2. If they do not match, record the error as a data type mismatch;
3. Execute the ALTER TABLE statement to alter the column and apply the tag to it;
4. Check if the tag has been applied correctly by querying the tag location database and information schema, using the TAG_REFERENCES table function, and comparing the tag from the query with the tag from the configuration table:
   1. If they match, create an UPDATE table statement and update the configuration table;
   2. If they do not match, record the error as a tag mismatch;
5. Catch any other exceptions and record them;
6. Insert a log entry into the logs table.

If we were to translate the steps above into pseudocode, it would look something like this.

Lastly, here is the stored procedure's SQL code as a reference for those who need it.

Keep in mind

Here are a few important considerations:

EXECUTE AS CALLER vs. EXECUTE AS OWNER (default)

• EXECUTE AS CALLER executes the stored procedure with the caller’s permissions;
• EXECUTE AS OWNER executes the stored procedure with the permissions of the role that created it.

Executing through a Python connector

• Comments must be block comments (/* comment */); otherwise, a <EOF> error will occur;
• Use $$ at the beginning and end of the stored procedure definition; otherwise, it can throw a syntax error.

Triple encapsulating quotation marks

When concatenating variables in strings, use three quotation marks, for example: WHERE database_name = ''' || :loop_db || ''' ' ||.

Lastly, remember to call the stored procedure with the appropriate role, depending on your architecture, and then verify that the tags have been applied correctly:

CALL APPLY_MASKING_TAGS();

We hope you have found this article useful and can now confidently implement tag-based masking in Snowflake using stored procedures. 

You might also be interested in:

• A Practical Guide to Tag-Based Masking with Snowflake and dbt
• Automated Data Masking in Snowflake with Data Contracts

Make sure to visit the Infinite Lambda blog for more insights.

If you are struggling with data protection, drop us a line, we will be happy to have a look at your project.