Manually remove and clean up Immuta policy objects in Snowflake

Identify and gather the following:

• Immuta integration database name created in Snowflake

• Immuta system role created to manage the integration(e.g., <db name>.SYSTEM)

• Immuta system account user created to manage the integration (e.g., <db name>.SYSTEM_ACOUNT)

Below is the list of Snowflake SQLs can perform clean up of Immuta created policy objects in Snowflake:

• Review and understand the Snowflake script thoroughly; review the in-line comments about the overview of the script and the comments in each section.

• Update parameters in the SQL script/statements and make necessary changes to the database in referenced, i.e. <IMMUTA_DB> should be Immuta integration database in Snowflake.

• Call and execute both functions to iterate and disassociate policy-objects from tables.

• Thereafter, the Immuta database, system role, and system_account user can be dropped before re-integration.

--Manually Remove and clean up Immuta created Snowflake Policies--
-- OVERVIEW:
-- 1. grant the role needed to create clean up functions
-- 2. define and implement remove masking policy function
-- 3. define and implement remove row-access policy function
-- 4. run/execute functions to detach/disassociate policy objects
-- 5. drop database, drop system role, drop system_account user

--STEPS:
--1. Grant permissions to create Immuta functions
GRANT ROLE "<SYSTEM_ACCOUNT_ROLE>" TO ROLE "<CURRENT_ROLE>";

--2. Define function to remove masking policies
CREATE OR REPLACE PROCEDURE <IMMUTA_DB>.IMMUTA_PROCEDURES.DETACH_MASKING_POLICIES()
            RETURNS boolean
            LANGUAGE JAVASCRIPT
            EXECUTE AS CALLER
            AS
            $$
                    -- Get all masking policies in the Immuta database.
                    const maskingPolicies = snowflake.execute({
                        sqlText: `SHOW MASKING POLICIES IN <IMMUTA_DB>.IMMUTA_POLICIES`
                    });
                    const tableMap = {};
                    while (maskingPolicies.next())  {
                        const policyName = maskingPolicies.getColumnValue(2);
                        -- This query returns the database, schema, table, target type, and referenced column for current policy.
                        -- Policy names may contain single quotes. So need to escape them at the end.
                        const policyRecord = snowflake.execute({ sqlText: `SELECT REF_DATABASE_NAME, REF_SCHEMA_NAME, REF_ENTITY_NAME, REF_ENTITY_DOMAIN, REF_COLUMN_NAME FROM TABLE(<IMMUTA_DB>.INFORMATION_SCHEMA.POLICY_REFERENCES(policy_name => '<IMMUTA_DB>."IMMUTA_POLICIES"."${policyName.replace(/'/g, `\'`)}"'));` });
                        -- If previous query returns a result set, then call next() to get the row.
                        -- This should only run once in most cases, but we loop in order to handle detaching the policy from cloned tables.
                        while (policyRecord.getRowCount() > 0 && policyRecord.next()) {
                            const qualifiedTableName = `"${policyRecord.getColumnValue('REF_DATABASE_NAME')}"."${policyRecord.getColumnValue('REF_SCHEMA_NAME')}"."${policyRecord.getColumnValue('REF_ENTITY_NAME')}"`;
                            const targetType = policyRecord.getColumnValue('REF_ENTITY_DOMAIN');
                            const referencedColumn = policyRecord.getColumnValue('REF_COLUMN_NAME');
                            -- If the object contains the qualified table name as a key, that means a policy is applied to multiple columns in the table. Therefore, just add the column to its list.
                            -- Else, create a new key with the qualified table name and add the target type and referenced column to the object.
                            if (tableMap[qualifiedTableName]) {
                                tableMap[qualifiedTableName]['columns'].push(referencedColumn);
                            } else {
                                tableMap[qualifiedTableName] = { targetType, columns: [referencedColumn] };
                            }
                        }
                    }
                    -- Loop over the object, create a column clause for each column in the table, and execute an alter query to unset the masking policies from the columns.
                    for (const table in tableMap) {
                        const columnClauses = tableMap[table]['columns'].map((column) => {
                            return `COLUMN "${column}" UNSET MASKING POLICY`;
                        });
                        snowflake.execute({ sqlText: `ALTER ${tableMap[table]['targetType']} ${table} MODIFY ${columnClauses.join(',')}` });
                    }
                    return true;
            $$;

--3. Define function to remove row access policies
CREATE OR REPLACE PROCEDURE <IMMUTA_DB>.IMMUTA_PROCEDURES.DETACH_ROW_ACCESS_POLICIES()
            RETURNS boolean
            LANGUAGE JAVASCRIPT
            EXECUTE AS CALLER
            AS
            $$
                    -- Get all row access policies in the Immuta database.
                    const rowAccessPolicies = snowflake.execute({
                        sqlText: `SHOW ROW ACCESS POLICIES IN <IMMUTA_DB>."IMMUTA_POLICIES"`
                    });
                    while (rowAccessPolicies.next())  {
                        const policyName = rowAccessPolicies.getColumnValue(2);
                        -- This query returns the concatenated qualified table name and target type for current policy.
                        -- Policy names may contain single quotes. So need to escape them at the end.
                        const policyRecord = snowflake.execute({ sqlText: `SELECT REF_DATABASE_NAME, REF_SCHEMA_NAME, REF_ENTITY_NAME, REF_ENTITY_DOMAIN FROM TABLE(<IMMUTA_DB>.INFORMATION_SCHEMA.POLICY_REFERENCES(policy_name => '<IMMUTA_DB>."IMMUTA_POLICIES"."${policyName.replace(/'/g, `\'`)}"'))` });
                        -- If previous query returns a result set, then call next() to get the row.
                        -- This should only run once in most cases, but we loop in order to handle detaching the policy from cloned tables.
                        while(policyRecord.getRowCount() > 0 && policyRecord.next()) {
                            const qualifiedTableName = `"${policyRecord.getColumnValue('REF_DATABASE_NAME')}"."${policyRecord.getColumnValue('REF_SCHEMA_NAME')}"."${policyRecord.getColumnValue('REF_ENTITY_NAME')}"`;
                            const targetType = policyRecord.getColumnValue('REF_ENTITY_DOMAIN');
                            try {
                                -- Drop row access policy references on a table or view.
                                snowflake.execute({ sqlText: `ALTER ${targetType} ${qualifiedTableName} DROP ROW ACCESS POLICY <IMMUTA_DB>.IMMUTA_POLICIES."${policyName}"` });
                            } catch (err) {
                                -- If we failed to remove the RLS that probably means there was no policy.
                            }
                        }
                    }
                    return true;
                $$;

--4. Call and invoke functions and cleanup Immuta policy objects (same as cleanup script in Immuta UI)
-- Execute procedures to clean up masking and row access policies.
--
EXECUTE IMMEDIATE
$$
DECLARE 
  CURRENT_ROLE VARCHAR;
BEGIN
    -- Grant Immuta system role to the current executing role to access the policy stored procedures.
    CURRENT_ROLE := '"' || (SELECT CURRENT_ROLE()) || '"';   
    GRANT ROLE <SYSTEM_ACOUNT_ROLE> TO ROLE IDENTIFIER(:CURRENT_ROLE);
    CALL <IMMUTA_DB>.IMMUTA_PROCEDURES.DETACH_MASKING_POLICIES();
    CALL <IMMUTA_DB>.IMMUTA_PROCEDURES.DETACH_ROW_ACCESS_POLICIES();
    RETURN TRUE;
END;
$$
;


--5. After both masking and row-access policies have been detach/disassocated from table/view objects, drop the immuta database which contains the secure views for data sources and workspaces
--
DROP DATABASE IF EXISTS "<IMMUTA_DB>";
--
-- Drop the immuta_system role used by the native snowflake integration
--
-- NOTE: If you changed the default role as part of a manual bootstrap process this
--       role may need to be manually set.
--
DROP ROLE IF EXISTS "<SYSTEM_ACCOUNT_ROLE>";
--
-- Drop the Immuta system account if it exists (may not exist if a manual account was provided).
-- In the event a manual account was provided (and granted the immuta_system role) that account
-- can potentially be removed as well.
--
DROP USER IF EXISTS "<SYSTEM_ACCOUNT>";