Snowflake masking policy: can the input be a constant string variable? - snowflake-cloud-data-platform

Question

I am trying to create a masking policy with tagging:
CREATE OR REPLACE MASKING POLICY TAGS_MASKING
AS (val VARCHAR, col_name STRING) RETURNS VARCHAR ->
CASE
WHEN CURRENT_ROLE() IN ('ADMIN_ROLE') THEN val
WHEN CURRENT_ROLE() IN ('ANALYST_ROLE') AND (SELECT SYSTEM$GET_TAG('TAG_NAME', col_name , 'COLUMN') = 'PUBLIC') THEN val
WHEN CURRENT_ROLE() IN ('ANALYST_ROLE') AND (SELECT SYSTEM$GET_TAG('TAG_NAME', col_name , 'COLUMN') IN ('PROTECTED')) THEN '****MASKED****'
END;
Here, col_name is a string (e.g. 'mytable.col1'), so that I can assign this masking policy to any columns I want to. But when I used the following query to assign it to one column of one table, it failed:
ALTER TABLE IF EXISTS db.masking.mytable MODIFY COLUMN col1
SET MASKING POLICY TAGS_MASKING using (col1, 'mytable.col1');
The error message is:
Syntax error: unexpected "mytable.col1"
How should I figure this out? Thanks!

Answer 1

I have not found a way to parametrize the column name(passing it as optional second parameter) so I have used a differnt approach.
It automatically creates a masking policy per colum using Snowflake Scripting.
Setup:
CREATE OR REPLACE TAG TAG_NAME;
CREATE OR REPLACE TABLE mytable(col1 STRING);
ALTER TABLE mytable SET TAG TAG_NAME='PUBLIC';
INSERT INTO mytable(col1) VALUES ('Test');
SELECT * FROM mytable;
-- Test
Procedure:
CREATE OR REPLACE PROCEDURE test(schema_name STRING, tab_name STRING, col_name STRING)
RETURNS STRING
LANGUAGE SQL
AS
$$
DECLARE
sql_masking_policy STRING;
sql_alter_table STRING;
masking_policy_name STRING := CONCAT_WS('_', 'TAGS_MASKING_', SCHEMA_NAME, TAB_NAME, COL_NAME);
BEGIN
sql_masking_policy := '
CREATE OR REPLACE MASKING POLICY <masking_policy_name>
AS (val VARCHAR) RETURNS VARCHAR ->
CASE
WHEN CURRENT_ROLE() IN (''ADMIN_ROLE'') THEN val
WHEN CURRENT_ROLE() IN (''ANALYST_ROLE'') AND (SYSTEM$GET_TAG(''TAG_NAME'', ''<col_name>'', ''COLUMN'') = ''PUBLIC'') THEN val
WHEN CURRENT_ROLE() IN (''ANALYST_ROLE'') AND (SYSTEM$GET_TAG(''TAG_NAME'', ''<col_name>'', ''COLUMN'') IN (''PROTECTED'')) THEN ''****MASKED****''
END;';
sql_alter_table := 'ALTER TABLE IF EXISTS <tab_name> MODIFY COLUMN <col_name>
SET MASKING POLICY <masking_policy_name>;';
sql_masking_policy := REPLACE(sql_masking_policy, '<masking_policy_name>', :masking_policy_name);
sql_masking_policy := REPLACE(sql_masking_policy, '<col_name>', CONCAT_WS('.', schema_name, tab_name, col_name));
sql_alter_table := REPLACE(sql_alter_table, '<masking_policy_name>', :masking_policy_name);
sql_alter_table := REPLACE(sql_alter_table, '<tab_name>', CONCAT_WS('.', schema_name, tab_name));
sql_alter_table := REPLACE(sql_alter_table, '<col_name>', col_name);
EXECUTE IMMEDIATE :sql_masking_policy;
EXECUTE IMMEDIATE :sql_alter_table;
RETURN sql_masking_policy || CHR(10) || sql_alter_table;
END;
$$;
Call:
CALL test('public', 'mytable', 'col1');
Output:
CREATE OR REPLACE MASKING POLICY TAGS_MASKING__public_mytable_col1
AS (val VARCHAR) RETURNS VARCHAR ->
CASE
WHEN CURRENT_ROLE() IN ('ADMIN_ROLE') THEN val
WHEN CURRENT_ROLE() IN ('ANALYST_ROLE') AND (SYSTEM$GET_TAG('TAG_NAME', 'public.mytable.col1', 'COLUMN') = 'PUBLIC') THEN val
WHEN CURRENT_ROLE() IN ('ANALYST_ROLE') AND (SYSTEM$GET_TAG('TAG_NAME', 'public.mytable.col1', 'COLUMN') IN ('PROTECTED')) THEN '****MASKED****'
END;
ALTER TABLE IF EXISTS public.mytable MODIFY COLUMN col1 SET MASKING POLICY TAGS_MASKING__public_mytable_col1;
Check:
SHOW MASKING POLICIES;
Output:
Test of select using POLICY_CONTEXT:
execute using policy_context(current_role => 'PUBLIC')
AS
SELECT * FROM public.mytable;
-- NULL
execute using policy_context(current_role => 'ADMIN_ROLE')
AS
SELECT * FROM public.mytable;
-- Test
execute using policy_context(current_role => 'ANALYST_ROLE')
AS
SELECT * FROM public.mytable;
-- Test
ALTER TABLE mytable SET TAG TAG_NAME='PROTECTED';
execute using policy_context(current_role => 'ANALYST_ROLE')
AS
SELECT * FROM public.mytable;
-- ****MASKED****

Answer 2

The previous answer could be improved by using new feature: Tag-based Masking Policies:
A tag-based masking policy combines the object tagging and masking policy features to allow a masking policy to be set on a tag using an ALTER TAG command. When the data type in the masking policy signature and the data type of the column match, the tagged column is automatically protected by the conditions in the masking policy. This simplifies the data protection efforts because column data that should be protected no longer needs a masking policy manually applied to the column to protect the data. A column can be protected by a masking policy directly assigned to a column and a tag-based masking policy
SYSTEM$GET_TAG_ON_CURRENT_COLUMN:
SYSTEM$GET_TAG_ON_CURRENT_COLUMN( '<tag_name>' )
Returns the tag string value assigned to the column based upon the specified tag or NULL if a tag is not assigned to the specified column.
For this scenario:
CREATE OR REPLACE TAG TAG_NAME;
CREATE OR REPLACE MASKING POLICY TAGS_MASKING
AS (val VARCHAR) RETURNS VARCHAR ->
CASE
WHEN CURRENT_ROLE() IN ('ADMIN_ROLE') THEN val
WHEN CURRENT_ROLE() IN ('ANALYST_ROLE')
AND SYSTEM$GET_TAG_ON_CURRENT_COLUMN('TAG_NAME') = 'PUBLIC' THEN val
WHEN CURRENT_ROLE() IN ('ANALYST_ROLE')
AND SYSTEM$GET_TAG_ON_CURRENT_COLUMN('TAG_NAME') IN ('PROTECTED') THEN '****MASKED****'
END;
ALTER TAG TAG_NAME SET MASKING POLICY TAGS_MASKING;
Table:
CREATE OR REPLACE TABLE mytable(col1 STRING);
INSERT INTO mytable(col1) VALUES ('Test');
SELECT * FROM mytable;
-- COL1
-- Test
Assigning tag:
ALTER TABLE mytable ALTER COLUMN col1 SET TAG TAG_NAME='PROTECTED';
Role outside ADMIN_ROLE/ANALYST_ROLE:
SELECT * FROM mytable;
-- COL1
-- null
Switching to ANALYST_ROLE(column tagged with PROTECTED value):
USE ROLE ANALYST_ROLE;
SELECT * FROM mytable;
-- COL1
-- ****MASKED****