fix grants on multi word resource types

Titan-Systems · Nov 6, 2024 · 6403a76 · 6403a76
1 parent 60df92e
commit 6403a76
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 6 deletions.
diff --git a/tests/integration/test_lifecycle.py b/tests/integration/test_lifecycle.py
@@ -33,11 +33,12 @@ def test_create_drop_from_json(resource, cursor, suffix):
 
     # Not easily testable without flakiness
     if resource.__class__ in (
-        res.Service,
+        res.AccountParameter,
+        res.FutureGrant,
         res.Grant,
         res.RoleGrant,
-        res.FutureGrant,
         res.ScannerPackage,
+        res.Service,
     ):
         pytest.skip("Skipping")
 

diff --git a/tests/test_grant.py b/tests/test_grant.py
@@ -180,3 +180,60 @@ def test_grant_on_accepts_resource_name():
     grant = res.Grant(priv="usage", on_warehouse=wh.name, to="somerole")
     assert grant.on == "SOMEWH"
     assert grant.on_type == ResourceType.WAREHOUSE
+
+
+def test_grant_on_dynamic_tables():
+    grant = res.Grant(
+        priv="SELECT",
+        on_dynamic_table="somedb.someschema.sometable",
+        to="somerole",
+    )
+    assert grant._data.on == "SOMEDB.SOMESCHEMA.SOMETABLE"
+    assert grant._data.on_type == ResourceType.DYNAMIC_TABLE
+
+    dynamic_table = ResourcePointer(name="sometable", resource_type=ResourceType.DYNAMIC_TABLE)
+    grant = res.Grant(
+        priv="SELECT",
+        on=dynamic_table,
+        to="somerole",
+    )
+    assert grant._data.on == "SOMETABLE"
+    assert grant._data.on_type == ResourceType.DYNAMIC_TABLE
+
+    grant_on_all = res.GrantOnAll(
+        priv="SELECT",
+        on_all_dynamic_tables_in_schema="somedb.someschema",
+        to="somerole",
+    )
+    assert grant_on_all._data.in_name == "SOMEDB.SOMESCHEMA"
+    assert grant_on_all._data.in_type == ResourceType.SCHEMA
+    assert grant_on_all._data.on_type == ResourceType.DYNAMIC_TABLE
+
+    schema = res.Schema(name="somedb.someschema")
+    grant_on_all = res.GrantOnAll(
+        priv="SELECT",
+        on_all_dynamic_tables_in=schema,
+        to="somerole",
+    )
+    assert grant_on_all._data.in_name == "SOMEDB.SOMESCHEMA"
+    assert grant_on_all._data.in_type == ResourceType.SCHEMA
+    assert grant_on_all._data.on_type == ResourceType.DYNAMIC_TABLE
+
+    future_grant = res.FutureGrant(
+        priv="CREATE TABLE",
+        on_future_dynamic_tables_in_schema="somedb.someschema",
+        to="somerole",
+    )
+    assert future_grant._data.in_name == "SOMEDB.SOMESCHEMA"
+    assert future_grant._data.in_type == ResourceType.SCHEMA
+    assert future_grant._data.on_type == ResourceType.DYNAMIC_TABLE
+
+    schema = res.Schema(name="somedb.someschema")
+    future_grant = res.FutureGrant(
+        priv="CREATE TABLE",
+        on_future_dynamic_tables_in=schema,
+        to="somerole",
+    )
+    assert future_grant._data.in_name == "SOMEDB.SOMESCHEMA"
+    assert future_grant._data.in_type == ResourceType.SCHEMA
+    assert future_grant._data.on_type == ResourceType.DYNAMIC_TABLE
diff --git a/titan/resources/grant.py b/titan/resources/grant.py
@@ -355,7 +355,7 @@ def __init__(
                     # At some point we need to support _in_sometype=SomeType(blah)
 
                     if isinstance(arg, Resource):
-                        on_type = ResourceType(singularize(keyword[10:-3]))
+                        on_type = resource_type_for_label(singularize(keyword[10:-3]))
                         in_type = arg.resource_type
                         in_name = str(arg.fqn)
                         granted_in_ref = arg
@@ -499,15 +499,27 @@ def __init__(
             # Handle on_all_ kwargs
             if on_kwargs:
                 for keyword, arg in on_kwargs.items():
-                    on_keyword = keyword.split("_")[2]
-                    on_type = ResourceType(singularize(on_keyword))
                     if isinstance(arg, Resource):
+                        # In type inferred from Resource class
+                        # on_all_schemas_in=Database(name="somedb")
                         in_type = arg.resource_type
                         in_name = str(arg.fqn)
+                        on_type = resource_type_for_label(singularize(keyword[7:-3]))
                     else:
-                        in_stmt = keyword.split("_in_")[1]
+                        # on_all_schemas_in_database="somedb"
+                        on_stmt, in_stmt = keyword.split("_in_")
                         in_type = ResourceType(in_stmt)
                         in_name = arg
+                        on_type = resource_type_for_label(singularize(on_stmt[7:]))
+                    # on_keyword = keyword.split("_")[2]
+                    # on_type = ResourceType(singularize(on_keyword))
+                    # if isinstance(arg, Resource):
+                    #     in_type = arg.resource_type
+                    #     in_name = str(arg.fqn)
+                    # else:
+                    #     in_stmt = keyword.split("_in_")[1]
+                    #     in_type = ResourceType(in_stmt)
+                    #     in_name = arg
 
         super().__init__(**kwargs)
         self._data: _GrantOnAll = _GrantOnAll(