[BUGFIX] Some grants don't work with database roles (#163)

* fix grants to work with database roles

---------

Co-authored-by: TJ Murphy <1796+teej@users.noreply.github.com>

tests/fixtures/json/future_grant.json

@@ -4,5 +4,6 @@
     "in_type": "DATABASE",
     "in_name": "STATIC_DATABASE",
     "to": "STATIC_ROLE",
+    "to_type": "ROLE",
     "grant_option": false
 }

tests/fixtures/json/grant.json

@@ -3,6 +3,7 @@
     "on_type": "DATABASE",
     "on": "STATIC_DATABASE",
     "to": "STATIC_ROLE",
+    "to_type": "ROLE",
     "owner": "SYSADMIN",
     "grant_option": false,
     "_privs": [

tests/integration/data_provider/test_fetch_resource.py

@@ -114,34 +114,39 @@ def test_fetch_grant_on_account(cursor, suffix):
     cursor.execute(f"GRANT BIND SERVICE ENDPOINT ON ACCOUNT TO ROLE {role.name}")
 
     try:
-        bind_service_urn = parse_URN(f"urn:::grant/{role.name}?priv=BIND SERVICE ENDPOINT&on=account/ACCOUNT")
+        bind_service_urn = parse_URN(
+            f"urn:::grant/GRANT?priv=BIND SERVICE ENDPOINT&on=account/ACCOUNT&to=role/{role.name}"
+        )
         bind_service_grant = safe_fetch(cursor, bind_service_urn)
         assert bind_service_grant is not None
         assert bind_service_grant["priv"] == "BIND SERVICE ENDPOINT"
         assert bind_service_grant["on"] == "ACCOUNT"
         assert bind_service_grant["on_type"] == "ACCOUNT"
         assert bind_service_grant["to"] == role.name
-        audit_urn = parse_URN(f"urn:::grant/{role.name}?priv=AUDIT&on=account/ACCOUNT")
+        assert bind_service_grant["to_type"] == "ROLE"
+        audit_urn = parse_URN(f"urn:::grant/GRANT?priv=AUDIT&on=account/ACCOUNT&to=role/{role.name}")
         audit_grant = safe_fetch(cursor, audit_urn)
         assert audit_grant is not None
         assert audit_grant["priv"] == "AUDIT"
         assert audit_grant["on"] == "ACCOUNT"
         assert audit_grant["on_type"] == "ACCOUNT"
         assert audit_grant["to"] == role.name
+        assert audit_grant["to_type"] == "ROLE"
     finally:
         cursor.execute(role.drop_sql(if_exists=True))
 
 
 def test_fetch_grant_all_on_resource(cursor):
     cursor.execute("GRANT ALL ON WAREHOUSE STATIC_WAREHOUSE TO ROLE STATIC_ROLE")
-    grant_all_urn = parse_URN("urn:::grant/STATIC_ROLE?priv=ALL&on=warehouse/STATIC_WAREHOUSE")
+    grant_all_urn = parse_URN("urn:::grant/GRANT_ON_ALL?priv=ALL&on=warehouse/STATIC_WAREHOUSE&to=role/STATIC_ROLE")
     try:
         grant = safe_fetch(cursor, grant_all_urn)
         assert grant is not None
         assert grant["priv"] == "ALL"
         assert grant["on_type"] == "WAREHOUSE"
         assert grant["on"] == "STATIC_WAREHOUSE"
         assert grant["to"] == "STATIC_ROLE"
+        assert grant["to_type"] == "ROLE"
         assert grant["owner"] == "SYSADMIN"
         assert grant["grant_option"] is False
         assert grant["_privs"] == ["APPLYBUDGET", "MODIFY", "MONITOR", "OPERATE", "USAGE"]
@@ -724,3 +729,18 @@ def test_fetch_task_predecessor(cursor, suffix, marked_for_cleanup):
     result = clean_resource_data(res.Task.spec, result)
     data = clean_resource_data(res.Task.spec, child_task.to_dict())
     assert result == data
+
+
+def test_fetch_database_role_grant(cursor, suffix, marked_for_cleanup):
+    role = res.DatabaseRole(name=f"TEST_FETCH_DATABASE_ROLE_GRANT_{suffix}", database="STATIC_DATABASE")
+    create(cursor, role)
+    marked_for_cleanup.append(role)
+
+    grant = res.Grant(priv="USAGE", on_schema="STATIC_DATABASE.PUBLIC", to=role)
+    create(cursor, grant)
+
+    result = safe_fetch(cursor, grant.urn)
+    assert result is not None
+    result = clean_resource_data(res.Grant.spec, result)
+    data = clean_resource_data(res.Grant.spec, grant.to_dict())
+    assert result == data

tests/integration/test_lifecycle.py

@@ -210,3 +210,23 @@ def test_task_lifecycle_remove_predecessor(cursor, suffix, marked_for_cleanup):
     assert len(plan) == 1
     assert isinstance(plan[0], UpdateResource)
     blueprint.apply(cursor.connection, plan)
+
+
+def test_database_role_grants(cursor, suffix, marked_for_cleanup):
+    db = res.Database(name="whatever")
+    role = res.DatabaseRole(name="whatever_role", database=db)
+    grant = res.Grant(priv="USAGE", on_schema=db.public_schema.fqn, to=role)
+    future_grant = res.FutureGrant(priv="SELECT", on_future_tables_in=db, to=role)
+
+    marked_for_cleanup.append(db)
+    marked_for_cleanup.append(role)
+    marked_for_cleanup.append(grant)
+    marked_for_cleanup.append(future_grant)
+
+    bp = Blueprint(
+        resources=[db, role, grant, future_grant],
+    )
+    plan = bp.plan(cursor.connection)
+    assert len(plan) == 4
+    assert all(isinstance(r, CreateResource) for r in plan)
+    bp.apply(cursor.connection, plan)

tests/integration/test_resources.py

@@ -87,7 +87,7 @@ def test_grant_on_all(cursor, suffix, marked_for_cleanup):
     cursor.execute(grant.create_sql())
 
     schema_1_usage_grant = safe_fetch(
-        cursor, parse_URN(f"urn:::grant/STATIC_ROLE?priv=USAGE&on=schema/{test_db}.SCHEMA_1")
+        cursor, parse_URN(f"urn:::grant/GRANT?priv=USAGE&on=schema/{test_db}.SCHEMA_1&to=role/STATIC_ROLE")
     )
     assert schema_1_usage_grant is not None
     assert schema_1_usage_grant["priv"] == "USAGE"
@@ -96,7 +96,7 @@ def test_grant_on_all(cursor, suffix, marked_for_cleanup):
     assert schema_1_usage_grant["on_type"] == "SCHEMA"
 
     schema_2_usage_grant = safe_fetch(
-        cursor, parse_URN(f"urn:::grant/STATIC_ROLE?priv=USAGE&on=schema/{test_db}.SCHEMA_2")
+        cursor, parse_URN(f"urn:::grant/GRANT?priv=USAGE&on=schema/{test_db}.SCHEMA_2&to=role/STATIC_ROLE")
     )
     assert schema_2_usage_grant is not None
     assert schema_2_usage_grant["priv"] == "USAGE"
@@ -105,7 +105,7 @@ def test_grant_on_all(cursor, suffix, marked_for_cleanup):
     assert schema_2_usage_grant["on_type"] == "SCHEMA"
 
     schema_3_usage_grant = safe_fetch(
-        cursor, parse_URN(f"urn:::grant/STATIC_ROLE?priv=USAGE&on=schema/{test_db}.SCHEMA_3")
+        cursor, parse_URN(f"urn:::grant/GRANT?priv=USAGE&on=schema/{test_db}.SCHEMA_3&to=role/STATIC_ROLE")
     )
     assert schema_3_usage_grant is not None
     assert schema_3_usage_grant["priv"] == "USAGE"

tests/test_grant.py

@@ -13,8 +13,10 @@ def test_grant_global_priv():
     assert grant.priv == "CREATE WAREHOUSE"
     assert grant.on == "ACCOUNT"
     assert grant.to.name == "somerole"
-    assert str(URN.from_resource(grant)) == "urn:::grant/SOMEROLE?priv=CREATE WAREHOUSE&on=account/ACCOUNT"
-    assert grant.create_sql() == "GRANT CREATE WAREHOUSE ON ACCOUNT TO SOMEROLE"
+    assert (
+        str(URN.from_resource(grant)) == "urn:::grant/GRANT?priv=CREATE WAREHOUSE&on=account/ACCOUNT&to=role/SOMEROLE"
+    )
+    assert grant.create_sql() == "GRANT CREATE WAREHOUSE ON ACCOUNT TO ROLE SOMEROLE"
 
 
 def test_grant_account_obj_priv_with_resource():
@@ -24,7 +26,7 @@ def test_grant_account_obj_priv_with_resource():
     assert grant.on == "SOMEWH"
     assert grant.on_type == ResourceType.WAREHOUSE
     assert grant.to.name == "SOMEROLE"
-    assert str(URN.from_resource(grant)) == "urn:::grant/SOMEROLE?priv=MODIFY&on=warehouse/SOMEWH"
+    assert str(URN.from_resource(grant)) == "urn:::grant/GRANT?priv=MODIFY&on=warehouse/SOMEWH&to=role/SOMEROLE"
 
 
 def test_grant_account_obj_priv_with_kwarg():
@@ -33,7 +35,7 @@ def test_grant_account_obj_priv_with_kwarg():
     assert grant.on == "SOMEWH"
     assert grant.on_type == ResourceType.WAREHOUSE
     assert grant.to.name == "SOMEROLE"
-    assert str(URN.from_resource(grant)) == "urn:::grant/SOMEROLE?priv=MODIFY&on=warehouse/SOMEWH"
+    assert str(URN.from_resource(grant)) == "urn:::grant/GRANT?priv=MODIFY&on=warehouse/SOMEWH&to=role/SOMEROLE"
 
 
 def test_grant_schema_priv_with_resource():
@@ -43,7 +45,7 @@ def test_grant_schema_priv_with_resource():
     assert grant.on == "SOMESCHEMA"
     assert grant.on_type == ResourceType.SCHEMA
     assert grant.to.name == "SOMEROLE"
-    assert str(URN.from_resource(grant)) == "urn:::grant/SOMEROLE?priv=CREATE VIEW&on=schema/SOMESCHEMA"
+    assert str(URN.from_resource(grant)) == "urn:::grant/GRANT?priv=CREATE VIEW&on=schema/SOMESCHEMA&to=role/SOMEROLE"
 
 
 def test_grant_schema_priv_with_kwarg():
@@ -52,7 +54,7 @@ def test_grant_schema_priv_with_kwarg():
     assert grant.on == "SOMESCHEMA"
     assert grant.on_type == ResourceType.SCHEMA
     assert grant.to.name == "SOMEROLE"
-    assert str(URN.from_resource(grant)) == "urn:::grant/SOMEROLE?priv=CREATE VIEW&on=schema/SOMESCHEMA"
+    assert str(URN.from_resource(grant)) == "urn:::grant/GRANT?priv=CREATE VIEW&on=schema/SOMESCHEMA&to=role/SOMEROLE"
 
 
 def test_grant_all():
@@ -62,7 +64,7 @@ def test_grant_all():
     assert grant.on_type == ResourceType.WAREHOUSE
     assert grant.to.name == "SOMEROLE"
     assert grant._data._privs == all_privs_for_resource_type(ResourceType.WAREHOUSE)
-    assert str(URN.from_resource(grant)) == "urn:::grant/SOMEROLE?priv=ALL&on=warehouse/SOMEWH"
+    assert str(URN.from_resource(grant)) == "urn:::grant/GRANT?priv=ALL&on=warehouse/SOMEWH&to=role/SOMEROLE"
 
 
 def test_future_grant_schemas_priv():
@@ -72,7 +74,10 @@ def test_future_grant_schemas_priv():
     assert grant.in_type == ResourceType.DATABASE
     assert grant.in_name == "SOMEDB"
     assert grant.to.name == "SOMEROLE"
-    assert str(URN.from_resource(grant)) == "urn:::future_grant/SOMEROLE?priv=CREATE VIEW&on=database/SOMEDB.<SCHEMA>"
+    assert (
+        str(URN.from_resource(grant))
+        == "urn:::future_grant/FUTURE_GRANT?priv=CREATE VIEW&on=database/SOMEDB.<SCHEMA>&to=role/SOMEROLE"
+    )
 
 
 def test_future_grant_anonymous_target():
@@ -82,7 +87,10 @@ def test_future_grant_anonymous_target():
     assert grant.in_type == ResourceType.SCHEMA
     assert grant.in_name == "SOMESCHEMA"
     assert grant.to.name == "SOMEROLE"
-    assert str(URN.from_resource(grant)) == "urn:::future_grant/SOMEROLE?priv=SELECT&on=schema/SOMESCHEMA.<TABLE>"
+    assert (
+        str(URN.from_resource(grant))
+        == "urn:::future_grant/FUTURE_GRANT?priv=SELECT&on=schema/SOMESCHEMA.<TABLE>&to=role/SOMEROLE"
+    )
 
 
 def test_future_grant_anonymous_nested_target():
@@ -93,7 +101,8 @@ def test_future_grant_anonymous_nested_target():
     assert grant.in_name == "somedb.SOMESCHEMA"
     assert grant.to.name == "SOMEROLE"
     assert (
-        str(URN.from_resource(grant)) == "urn:::future_grant/SOMEROLE?priv=SELECT&on=schema/SOMEDB.SOMESCHEMA.<TABLE>"
+        str(URN.from_resource(grant))
+        == "urn:::future_grant/FUTURE_GRANT?priv=SELECT&on=schema/SOMEDB.SOMESCHEMA.<TABLE>&to=role/SOMEROLE"
     )
 
 
@@ -106,7 +115,8 @@ def test_future_grant_referenced_inferred_target():
     assert grant.in_name == "somedb.SOMESCHEMA"
     assert grant.to.name == "SOMEROLE"
     assert (
-        str(URN.from_resource(grant)) == "urn:::future_grant/SOMEROLE?priv=SELECT&on=schema/SOMEDB.SOMESCHEMA.<TABLE>"
+        str(URN.from_resource(grant))
+        == "urn:::future_grant/FUTURE_GRANT?priv=SELECT&on=schema/SOMEDB.SOMESCHEMA.<TABLE>&to=role/SOMEROLE"
     )
 
 

titan/data_provider.py

@@ -163,9 +163,14 @@ def _fail_if_not_granted(result, *args):
 
 
 def _fetch_grant_to_role(
-    session: SnowflakeConnection, role: ResourceName, granted_on: str, on_name: str, privilege: str
+    session: SnowflakeConnection,
+    role: ResourceName,
+    granted_on: str,
+    on_name: str,
+    privilege: str,
+    role_type: ResourceType = ResourceType.ROLE,
 ):
-    grants = _show_grants_to_role(session, role, cacheable=True)
+    grants = _show_grants_to_role(session, role, role_type=role_type, cacheable=True)
     if id(grants) not in _INDEX:
         local_index: dict[tuple[str, str, str], dict[str, Any]] = {}
         _INDEX[id(grants)] = local_index
@@ -536,7 +541,10 @@ def _get_account_privilege_roles(session: SnowflakeConnection) -> dict[str, list
 
 
 def _show_grants_to_role(
-    session: SnowflakeConnection, role: ResourceName, cacheable: bool = False
+    session: SnowflakeConnection,
+    role: ResourceName,
+    role_type: ResourceType = ResourceType.ROLE,
+    cacheable: bool = False,
 ) -> list[dict[str, Any]]:
     """
     {
@@ -552,7 +560,7 @@ def _show_grants_to_role(
     """
     grants = execute(
         session,
-        f"SHOW GRANTS TO ROLE {role}",
+        f"SHOW GRANTS TO {role_type} {role}",
         cacheable=cacheable,
         empty_response_codes=[DOES_NOT_EXIST_ERR],
     )
@@ -1154,8 +1162,12 @@ def fetch_function(session: SnowflakeConnection, fqn: FQN):
 
 
 def fetch_future_grant(session: SnowflakeConnection, fqn: FQN):
+
+    to_type, to = fqn.params["to"].split("/", 1)
+    to_type = resource_type_for_label(to_type)
+
     try:
-        show_result = execute(session, f"SHOW FUTURE GRANTS TO ROLE {fqn.name}", cacheable=True)
+        show_result = execute(session, f"SHOW FUTURE GRANTS TO {to_type} {to}", cacheable=True)
         """
         {
             'created_on': datetime.datetime(2024, 2, 5, 19, 39, 50, 146000, tzinfo=<DstTzInfo 'America/Los_Angeles' PST-1 day, 16:00:00 STD>),
@@ -1192,8 +1204,8 @@ def fetch_future_grant(session: SnowflakeConnection, fqn: FQN):
         show_result,
         privilege=fqn.params["priv"],
         name=collection_str,
-        grant_to="ROLE",
-        grantee_name=fqn.name,
+        grant_to=str(to_type),
+        grantee_name=to,
     )
 
     if len(grants) == 0:
@@ -1208,16 +1220,20 @@ def fetch_future_grant(session: SnowflakeConnection, fqn: FQN):
         "on_type": str(resource_type_for_label(data["grant_on"])),
         "in_type": collection["in_type"].upper(),
         "in_name": collection["in_name"],
-        "to": data["grantee_name"],
+        "to": to,
+        "to_type": resource_type_for_label(data["grant_to"]),
         "grant_option": data["grant_option"] == "true",
     }
 
 
 def fetch_grant(session: SnowflakeConnection, fqn: FQN):
     priv = fqn.params["priv"]
-    on_type, on = fqn.params["on"].split("/")
+    on_type, on = fqn.params["on"].split("/", 1)
     on_type = on_type.upper()
 
+    to_type, to = fqn.params["to"].split("/", 1)
+    to_type = resource_type_for_label(to_type)
+
     if priv == "ALL":
 
         filters = {
@@ -1227,7 +1243,7 @@ def fetch_grant(session: SnowflakeConnection, fqn: FQN):
         if on_type != "ACCOUNT":
             filters["name"] = on
 
-        grants = _show_grants_to_role(session, fqn.name, cacheable=True)
+        grants = _show_grants_to_role(session, to, role_type=to_type, cacheable=True)
         grants = _filter_result(grants, **filters)
 
         if len(grants) == 0:
@@ -1239,10 +1255,11 @@ def fetch_grant(session: SnowflakeConnection, fqn: FQN):
     else:
         data = _fetch_grant_to_role(
             session,
-            role=fqn.name,
+            role=to,
             granted_on=on_type,
             on_name=on,
             privilege=priv,
+            role_type=to_type,
         )
         if data is None:
             return None
@@ -1258,7 +1275,8 @@ def fetch_grant(session: SnowflakeConnection, fqn: FQN):
         "priv": priv,
         "on": "ACCOUNT" if on_type == "ACCOUNT" else data["name"],
         "on_type": data["granted_on"].replace("_", " "),
-        "to": data["grantee_name"],
+        "to": to,
+        "to_type": resource_type_for_label(data["granted_to"]),
         "grant_option": data["grant_option"] == "true",
         "owner": data["granted_by"],
         "_privs": privs,
@@ -2463,12 +2481,14 @@ def list_future_grants(session: SnowflakeConnection) -> list[FQN]:
         for data in grant_data:
             in_type = "database" if data["grant_on"] == "SCHEMA" else "schema"
             collection = data["name"]
+            to = f"role/{role_name}"
             grants.append(
                 FQN(
-                    name=role_name,
+                    name=ResourceName("FUTURE_GRANT"),
                     params={
                         "priv": data["privilege"],
                         "on": f"{in_type}/{collection}",
+                        "to": to,
                     },
                 )
             )
@@ -2495,10 +2515,9 @@ def list_grants(session: SnowflakeConnection) -> list[FQN]:
         role_name = resource_name_from_snowflake_metadata(role["name"])
         if role_name in SYSTEM_ROLES:
             continue
-        grant_data = _show_grants_to_role(session, role_name, cacheable=False)
+        grant_data = _show_grants_to_role(session, role_name, role_type=ResourceType.ROLE, cacheable=False)
         for data in grant_data:
             if data["granted_on"] == "ROLE":
-                # raise Exception(f"Role grants are not supported yet: {data}")
                 continue
 
             # Titan Grants don't support OWNERSHIP privilege
@@ -2513,12 +2532,14 @@ def list_grants(session: SnowflakeConnection) -> list[FQN]:
             if data["granted_on"] == "ACCOUNT":
                 name = "ACCOUNT"
             on = f"{data['granted_on'].lower()}/{name}"
+            to = f"role/{role_name}"
             grants.append(
                 FQN(
-                    name=role_name,
+                    name=ResourceName("GRANT"),
                     params={
                         "priv": data["privilege"],
                         "on": on,
+                        "to": to,
                     },
                 )
             )