Add database role support

.changes/unreleased/Features-20241012-115830.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Support granting of Database Roles
+time: 2024-10-12T11:58:30.84775252+01:00
+custom:
+    Author: seediang
+    Issue: "1206"

dbt/adapters/snowflake/impl.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Mapping, Any, Optional, List, Union, Dict, FrozenSet, Tuple, TYPE_CHECKING
 
+
 from dbt.adapters.base.impl import AdapterConfig, ConstraintSupport
 from dbt.adapters.base.meta import available
 from dbt.adapters.capability import CapabilityDict, CapabilitySupport, Support, Capability
@@ -312,21 +313,120 @@ def quote_seed_column(self, column: str, quote_config: Optional[bool]) -> str:
         else:
             return column
 
+    GrantsDict = Dict[str, Dict[str, Any]]
+
     @available
-    def standardize_grants_dict(self, grants_table: "agate.Table") -> dict:
+    def standardize_grants_dict(self, grants_table: "agate.Table") -> GrantsDict:
         grants_dict: Dict[str, Any] = {}
+        self.AdapterSpecificConfigs
 
+        # granted_to maps to different object types like
+        # role, database_role and share
+        # create nest dictionaries [granted_to].[privilege][object_name]
         for row in grants_table:
             grantee = row["grantee_name"]
             granted_to = row["granted_to"]
             privilege = row["privilege"]
-            if privilege != "OWNERSHIP" and granted_to not in ["SHARE", "DATABASE_ROLE"]:
-                if privilege in grants_dict.keys():
-                    grants_dict[privilege].append(grantee)
-                else:
-                    grants_dict.update({privilege: [grantee]})
+            if privilege != "OWNERSHIP":
+                role_type_dict = grants_dict.setdefault(granted_to.lower(), {})
+                privilege_dict = role_type_dict.setdefault(privilege.lower(), [])
+                privilege_dict.append(grantee)
+
         return grants_dict
 
+    @available
+    def standardize_grant_config(self, grant_config: dict) -> GrantsDict:
+        """
+        Given a grants configuration object of either
+        Dict[str, Any] or Dict[str, [Dict[str, List[str] ] ]]]
+
+        Converts to Dict[str, Any] to [Dict["role", List[str]] ]]]
+
+        This enables use to handle new and older style grants
+        configurations.
+        """
+        grant_config_std: Dict[str, Any] = {}
+        self.AdapterSpecificConfigs
+
+        for grant_config_privilege, privilege_collection in grant_config.items():
+            # loop through the role entries and handle mapping, list & string entries
+            for privilege_item in privilege_collection:
+                # Assume old style list grants map to role
+                if not isinstance(privilege_item, dict):
+                    privilege_item = {"role": privilege_item}
+
+                for grantee_type, grantees in privilege_item.items():
+                    if grantees:
+                        # -- Make sure object_type is in grant_config_by_type
+                        grantee_type_privileges: Dict[str, Any] = grant_config_std.setdefault(
+                            grantee_type.lower(), {}
+                        )
+                        privilege_list = grantee_type_privileges.setdefault(
+                            grant_config_privilege.lower(), []
+                        )
+
+                        # -- convert string to array to make code simpler --#}
+                        if isinstance(grantees, str):
+                            grantees = [grantees]
+
+                        for grantee in grantees:
+                            # -- Only add the item if not already in the list --#}
+                            if grantee not in privilege_list:
+                                privilege_list.append(grantee)
+
+        return grant_config_std
+
+    @available
+    def diff_of_grants(self, grants_a: GrantsDict, grants_b: GrantsDict) -> GrantsDict:
+        """
+        Given two dictionaries of type Dict[str, Dict[str, List[str]]]:
+            grants_a = {'key_x': {'key_a': ['VALUE_1', 'VALUE_2']}, 'KEY_Y': {'key_b': ['value_2']}}
+            grants_b = {'KEY_x': {'key_a': ['VALUE_1']}, 'key_y': {'key_b': ['value_3']}}
+        Return the same dictionary representation of dict_a MINUS dict_b,
+        performing a case-insensitive comparison between the strings in each.
+        All keys returned will be in the original case of dict_a.
+            returns {"key_x": {'key_a': ['VALUE_2']},"KEY_Y":{"key_b": ["value_2"]}}
+        """
+
+        def diff_of_two_dicts(dict_a: dict, dict_b: dict) -> dict:
+            """
+            Given two dictionaries of type Dict[str, List[str]]:
+                dict_a = {'key_x': ['value_1', 'VALUE_2'], 'KEY_Y': ['value_3']}
+                dict_b = {'key_x': ['value_1'], 'key_z': ['value_4']}
+            Return the same dictionary representation of dict_a MINUS dict_b,
+            performing a case-insensitive comparison between the strings in each.
+            All keys returned will be in the original case of dict_a.
+                returns {'key_x': ['VALUE_2'], 'KEY_Y': ['value_3']}
+            """
+
+            dict_diff = {}
+            dict_b_lowered = {k.casefold(): [x.casefold() for x in v] for k, v in dict_b.items()}
+            for k in dict_a:
+                if k.casefold() in dict_b_lowered.keys():
+                    diff = []
+                    for v in dict_a[k]:
+                        if v.casefold() not in dict_b_lowered[k.casefold()]:
+                            diff.append(v)
+                    if diff:
+                        dict_diff.update({k: diff})
+                else:
+                    dict_diff.update({k: dict_a[k]})
+            return dict_diff
+
+        grants_diff: Dict[str, Any] = {}
+        grants_b_lowered = {
+            k.casefold(): {x.casefold(): y for x, y in v.items()} for k, v in grants_b.items()
+        }
+        for k in grants_a:
+            if k.casefold() in grants_b_lowered.keys():
+                diff = diff_of_two_dicts(grants_a[k], grants_b_lowered[k.casefold()])
+                if diff:
+                    grants_diff.update({k: diff})
+            else:
+                grants_diff.update({k: grants_a[k]})
+
+        return grants_diff
+
     def timestamp_add_sql(self, add_to: str, number: int = 1, interval: str = "hour") -> str:
         return f"DATEADD({interval}, {number}, {add_to})"
 

dbt/include/snowflake/macros/apply_grants.sql

@@ -6,3 +6,97 @@
 {%- macro snowflake__support_multiple_grantees_per_dcl_statement() -%}
     {{ return(False) }}
 {%- endmacro -%}
+
+{#
+  -- Create versions of get_grant_sql and get_revoke_sql that support an additional
+  -- object_type parameter
+#}
+
+{% macro get_grant_sql_by_type(relation, privilege, object_type, grantees) %}
+    {{ return(adapter.dispatch('get_grant_sql_by_type', 'dbt-snowflake')(relation, privilege, object_type, grantees)) }}
+{% endmacro %}
+
+{%- macro snowflake__get_grant_sql_by_type(relation, privilege, object_type, grantees) -%}
+    grant {{ privilege }} on {{ relation.render() }} to {{object_type | replace('_', ' ')}} {{ grantees | join(', ') }}
+{%- endmacro -%}
+
+{% macro get_revoke_sql_by_type(relation, privilege, object_type, grantees) %}
+    {{ return(adapter.dispatch('get_revoke_sql_by_type', 'dbt-snowflake')(relation, privilege, object_type, grantees)) }}
+{% endmacro %}
+
+{%- macro snowflake__get_revoke_sql_by_type(relation, privilege, object_type, grantees) -%}
+    revoke {{ privilege }} on {{ relation.render() }} from {{object_type | replace('_', ' ')}} {{ grantees | join(', ') }}
+{%- endmacro -%}
+
+
+{% macro get_dcl_statement_list_by_type(relation, grant_config_by_type, get_dcl_macro) %}
+    {{ return(adapter.dispatch('get_dcl_statement_list_by_type', 'dbt-snowflake')(relation, grant_config_by_type, get_dcl_macro)) }}
+{% endmacro %}
+
+
+{%- macro snowflake__get_dcl_statement_list_by_type(relation, grant_config_by_type, get_dcl_macro) -%}
+    {#
+      -- Unpack grant_config into specific privileges and the set of users who need them granted/revoked.
+      -- Depending on whether this database supports multiple grantees per statement, pass in the list of
+      -- all grantees per privilege, or (if not) template one statement per privilege-grantee pair.
+      -- `get_dcl_macro` will be either `get_grant_sql_by_type` or `get_revoke_sql_by_type`
+      --
+      -- grant_config_by_type should be in the following format  { grantee_type: { privilege: [grantee] } }
+    #}
+    {%- set dcl_statements = [] -%}
+    {%- for object_type, config in grant_config_by_type.items() %}
+        {%- for privilege, grantees in config.items() %}
+            {%- if support_multiple_grantees_per_dcl_statement() and grantees -%}
+            {%- set dcl = get_dcl_macro(relation, privilege, object_type, grantees) -%}
+            {%- do dcl_statements.append(dcl) -%}
+            {%- else -%}
+            {%- for grantee in grantees -%}
+                {% set dcl = get_dcl_macro(relation, privilege, object_type, [grantee]) %}
+                {%- do dcl_statements.append(dcl) -%}
+            {% endfor -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endfor -%}
+    {{ return(dcl_statements) }}
+{%- endmacro %}
+
+{% macro snowflake__apply_grants(relation, grant_config, should_revoke=True) %}
+    {#-- If grant_config is {} or None, this is a no-op --#}
+    {% if grant_config %}
+
+        {{ log('grant_config: ' ~ grant_config) }}
+        {#-- Check if we have defined new role type or are using default style --#}
+        {% set desired_grants_dict =  adapter.standardize_grant_config(grant_config) %}
+        {{ log('desired_grants_dict: ' ~ desired_grants_dict) }}
+
+
+        {% if should_revoke %}
+            {#-- We think previous grants may have carried over --#}
+            {#-- Show current grants and calculate diffs --#}
+            {% set current_grants_table = run_query(get_show_grant_sql(relation)) %}
+            {% set current_grants_dict = adapter.standardize_grants_dict(current_grants_table) %}
+
+            {% set needs_granting = adapter.diff_of_grants(desired_grants_dict, current_grants_dict) %}
+            {% set needs_revoking = adapter.diff_of_grants(current_grants_dict, desired_grants_dict) %}
+
+            {% if not (needs_granting or needs_revoking) %}
+                {{ log('On ' ~ relation.render() ~': All grants are in place, no revocation or granting needed.')}}
+            {% endif %}
+        {% else %}
+            {#-- We don't think there's any chance of previous grants having carried over. --#}
+            {#-- Jump straight to granting what the user has configured. --#}
+            {% set needs_revoking = {} %}
+            {% set needs_granting = desired_grants_dict %}
+        {% endif %}
+
+        {% if needs_granting or needs_revoking %}
+            {% set revoke_statement_list = get_dcl_statement_list_by_type(relation, needs_revoking, snowflake__get_revoke_sql_by_type) %}
+            {% set grant_statement_list = get_dcl_statement_list_by_type(relation, needs_granting, snowflake__get_grant_sql_by_type) %}
+            {% set dcl_statement_list = revoke_statement_list + grant_statement_list %}
+
+            {% if dcl_statement_list %}
+                {{ call_dcl_statements(dcl_statement_list) }}
+            {% endif %}
+        {% endif %}
+    {% endif %}
+{% endmacro %}

tests/functional/adapter/grants/base_grants.py

@@ -0,0 +1,67 @@
+import os
+import pytest
+import snowflake.connector
+from dbt.tests.adapter.grants.base_grants import BaseGrants as OrigBaseGrants
+
+
+class BaseGrantsSnowflakePatch:
+    """
+    Overides the adapter BaseGrants to use new adapter functions
+    for grants.
+    """
+
+    def assert_expected_grants_match_actual(self, project, relation_name, expected_grants):
+        adapter = project.adapter
+        actual_grants = self.get_grants_on_relation(project, relation_name)
+        expected_grants_std = adapter.standardize_grant_config(expected_grants)
+
+        # need a case-insensitive comparison -- this would not be true for all adapters
+        # so just a simple "assert expected == actual_grants" won't work
+        diff_a = adapter.diff_of_grants(actual_grants, expected_grants_std)
+        diff_b = adapter.diff_of_grants(expected_grants_std, actual_grants)
+        assert diff_a == diff_b == {}
+
+
+class BaseGrantsSnowflake(BaseGrantsSnowflakePatch, OrigBaseGrants):
+    @pytest.fixture(scope="session", autouse=True)
+    def ensure_database_roles(project):
+        """
+        We need to create database roles since test framework does not
+        have default database roles to work with. This has been patched
+        in with ta session scoped fixture and custom connection.
+        """
+        con = snowflake.connector.connect(
+            user=os.getenv("SNOWFLAKE_TEST_USER"),
+            password=os.getenv("SNOWFLAKE_TEST_PASSWORD"),
+            account=os.getenv("SNOWFLAKE_TEST_ACCOUNT"),
+            warehouse=os.getenv("SNOWFLAKE_TEST_WAREHOUSE"),
+            database=os.getenv("SNOWFLAKE_TEST_DATABASE"),
+        )
+
+        number_of_roles = 3
+
+        for index in range(1, number_of_roles + 1):
+            con.execute_string(f"CREATE DATABASE ROLE IF NOT EXISTS test_database_role_{index}")
+
+        yield
+
+        for index in range(1, number_of_roles + 1):
+            con.execute_string(f"DROP DATABASE ROLE test_database_role_{index}")
+
+
+class BaseCopyGrantsSnowflake:
+    # Try every test case without copy_grants enabled (default),
+    # and with copy_grants enabled (this base class)
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "models": {
+                "+copy_grants": True,
+            },
+            "seeds": {
+                "+copy_grants": True,
+            },
+            "snapshots": {
+                "+copy_grants": True,
+            },
+        }