From 8968b56a2dae339c7b0ad90ecafacd6cfb1d60a4 Mon Sep 17 00:00:00 2001 From: TJ Murphy Date: Mon, 25 Nov 2024 15:02:17 -0800 Subject: [PATCH] [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 | 1 + tests/fixtures/json/grant.json | 1 + .../data_provider/test_fetch_resource.py | 26 +++++++-- tests/integration/test_lifecycle.py | 20 +++++++ tests/integration/test_resources.py | 6 +-- tests/test_grant.py | 32 +++++++---- titan/data_provider.py | 53 +++++++++++++------ titan/lifecycle.py | 14 +++-- titan/resources/grant.py | 45 ++++++++++++++-- 9 files changed, 156 insertions(+), 42 deletions(-) diff --git a/tests/fixtures/json/future_grant.json b/tests/fixtures/json/future_grant.json index a46cc64e..1a4395a0 100644 --- a/tests/fixtures/json/future_grant.json +++ b/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 } \ No newline at end of file diff --git a/tests/fixtures/json/grant.json b/tests/fixtures/json/grant.json index 3c8daffc..72dd15a5 100644 --- a/tests/fixtures/json/grant.json +++ b/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": [ diff --git a/tests/integration/data_provider/test_fetch_resource.py b/tests/integration/data_provider/test_fetch_resource.py index f288a17e..71ca47d9 100644 --- a/tests/integration/data_provider/test_fetch_resource.py +++ b/tests/integration/data_provider/test_fetch_resource.py @@ -114,27 +114,31 @@ 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 @@ -142,6 +146,7 @@ def test_fetch_grant_all_on_resource(cursor): 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 diff --git a/tests/integration/test_lifecycle.py b/tests/integration/test_lifecycle.py index d59496c5..9731f148 100644 --- a/tests/integration/test_lifecycle.py +++ b/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) diff --git a/tests/integration/test_resources.py b/tests/integration/test_resources.py index 0934199b..23e9e901 100644 --- a/tests/integration/test_resources.py +++ b/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" diff --git a/tests/test_grant.py b/tests/test_grant.py index e11d3f86..69471f17 100644 --- a/tests/test_grant.py +++ b/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." + assert ( + str(URN.from_resource(grant)) + == "urn:::future_grant/FUTURE_GRANT?priv=CREATE VIEW&on=database/SOMEDB.&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." + assert ( + str(URN.from_resource(grant)) + == "urn:::future_grant/FUTURE_GRANT?priv=SELECT&on=schema/SOMESCHEMA.
&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.
" + str(URN.from_resource(grant)) + == "urn:::future_grant/FUTURE_GRANT?priv=SELECT&on=schema/SOMEDB.SOMESCHEMA.
&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.
" + str(URN.from_resource(grant)) + == "urn:::future_grant/FUTURE_GRANT?priv=SELECT&on=schema/SOMEDB.SOMESCHEMA.
&to=role/SOMEROLE" ) diff --git a/titan/data_provider.py b/titan/data_provider.py index be23d41b..12a49748 100644 --- a/titan/data_provider.py +++ b/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=), @@ -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, }, ) ) diff --git a/titan/lifecycle.py b/titan/lifecycle.py index d69dd23b..910d1c15 100644 --- a/titan/lifecycle.py +++ b/titan/lifecycle.py @@ -98,9 +98,9 @@ def create_future_grant(urn: URN, data: dict, props: Props, if_not_exists: bool) "IN", data["in_type"], data["in_name"], - "TO ROLE", - urn.fqn.name, - # props.render(data), #TODO grant option + "TO", + data["to_type"], + data["to"], ) @@ -116,7 +116,10 @@ def create_grant(urn: URN, data: dict, props: Props, if_not_exists: bool): "ON", on_type, data["on"], - props.render(data), + "TO", + data["to_type"], + data["to"], + "WITH GRANT OPTION" if data["grant_option"] else "", ) @@ -129,7 +132,8 @@ def create_grant_on_all(urn: URN, data: dict, props: Props, if_not_exists: bool) "IN", data["in_type"], data["in_name"], - "TO ROLE", + "TO", + data["to_type"], data["to"], ) diff --git a/titan/resources/grant.py b/titan/resources/grant.py index e8fd9d3f..5a711416 100644 --- a/titan/resources/grant.py +++ b/titan/resources/grant.py @@ -25,6 +25,7 @@ class _Grant(ResourceSpec): on: str on_type: ResourceType to: RoleRef + to_type: ResourceType = None grant_option: bool = False owner: Role = field(default=None, metadata={"fetchable": False}) _privs: list[str] = field(default_factory=list, metadata={"triggers_create": True}) @@ -44,6 +45,8 @@ def __post_init__(self): else: self._privs = [self.priv] + self.to_type = self.to.resource_type + class Grant(Resource): """ @@ -108,6 +111,11 @@ def __init__( ): kwargs.pop("_privs", None) + to_type = kwargs.pop("to_type", None) + + if all([to_type, to]): + to_type = ResourceType(to_type) + to = ResourcePointer(name=to, resource_type=to_type) priv = priv.value if isinstance(priv, ParseableEnum) else priv @@ -210,6 +218,10 @@ def on_type(self) -> ResourceType: def to(self): return self._data.to + @property + def to_type(self) -> ResourceType: + return self._data.to_type + @property def priv(self): return self._data.priv @@ -217,11 +229,13 @@ def priv(self): def grant_fqn(grant: _Grant): on = f"{resource_label_for_type(grant.on_type)}/{grant.on}" + to = f"{resource_label_for_type(grant.to_type)}/{grant.to.fqn}" return FQN( - name=grant.to.name, + name=ResourceName("GRANT"), params={ "priv": grant.priv, "on": on, + "to": to, }, ) @@ -232,7 +246,8 @@ def grant_yaml(data: dict): return { "priv": grant.priv, f"on_{resource_label}": grant.on, - "to": grant.to.name, + "to": str(grant.to.fqn), + "to_type": str(grant.to_type), "grant_option": grant.grant_option, } @@ -244,12 +259,14 @@ class _FutureGrant(ResourceSpec): in_type: ResourceType in_name: ResourceName to: RoleRef + to_type: ResourceType = None grant_option: bool = False def __post_init__(self): super().__post_init__() if isinstance(self.priv, str): self.priv = self.priv.upper() + self.to_type = self.to.resource_type class FutureGrant(Resource): @@ -327,6 +344,7 @@ def __init__( on_type = kwargs.pop("on_type", None) in_type = kwargs.pop("in_type", None) in_name = kwargs.pop("in_name", None) + to_type = kwargs.pop("to_type", None) granted_in_ref = None if all([on_type, in_type, in_name]): @@ -334,6 +352,10 @@ def __init__( on_type = ResourceType(on_type) granted_in_ref = ResourcePointer(name=in_name, resource_type=in_type) + if all([to_type, to]): + to_type = ResourceType(to_type) + to = ResourcePointer(name=to, resource_type=to_type) + else: # Collect on_ kwargs @@ -405,17 +427,23 @@ def in_name(self) -> str: def to(self): return self._data.to + @property + def to_type(self) -> ResourceType: + return self._data.to_type + def future_grant_fqn(data: _FutureGrant): in_type = resource_label_for_type(data.in_type) in_name = data.in_name on_type = resource_label_for_type(data.on_type).upper() collection = format_collection_string({"in_name": in_name, "in_type": in_type, "on_type": on_type}) + to = f"{resource_label_for_type(data.to_type)}/{data.to.fqn}" return FQN( - name=data.to.name, + name=ResourceName("FUTURE_GRANT"), params={ "priv": data.priv, "on": f"{in_type}/{collection}", + "to": to, }, ) @@ -427,12 +455,14 @@ class _GrantOnAll(ResourceSpec): in_type: ResourceType in_name: ResourceName to: RoleRef + to_type: ResourceType = None grant_option: bool = False def __post_init__(self): super().__post_init__() if self.in_type not in [ResourceType.DATABASE, ResourceType.SCHEMA]: raise ValueError(f"in_type must be either DATABASE or SCHEMA, not {self.in_type}") + self.to_type = self.to.resource_type class GrantOnAll(Resource): @@ -514,6 +544,11 @@ def __init__( on_type = kwargs.pop("on_type", None) in_type = kwargs.pop("in_type", None) in_name = kwargs.pop("in_name", None) + to_type = kwargs.pop("to_type", None) + + if all([to_type, to]): + to_type = ResourceType(to_type) + to = ResourcePointer(name=to, resource_type=to_type) _owner = kwargs.pop("owner", None) if _owner is not None: @@ -577,11 +612,13 @@ def grant_on_all_fqn(data: _GrantOnAll): in_name = data.in_name on_type = resource_label_for_type(data.on_type).upper() collection = format_collection_string({"in_name": in_name, "in_type": in_type, "on_type": on_type}) + to = f"{resource_label_for_type(data.to_type)}/{data.to.fqn}" return FQN( - name=data.to.name, + name=ResourceName("GRANT_ON_ALL"), params={ "priv": data.priv, "on": f"{in_type}/{collection}", + "to": to, }, )