From 7522a5fe91d1cf57201d3cf49592c424c9081ce0 Mon Sep 17 00:00:00 2001 From: TJ Murphy <1796+teej@users.noreply.github.com> Date: Tue, 17 Dec 2024 21:11:26 -0800 Subject: [PATCH 1/2] fix vars for role names --- tests/test_blueprint.py | 9 +++++++ tests/test_vars.py | 11 +++++++++ titan/blueprint.py | 7 ++++++ titan/resources/grant.py | 5 +--- titan/resources/resource.py | 49 ++++++++++++++++++++++++++----------- titan/role_ref.py | 8 +++--- 6 files changed, 68 insertions(+), 21 deletions(-) diff --git a/tests/test_blueprint.py b/tests/test_blueprint.py index c9e28297..94499713 100644 --- a/tests/test_blueprint.py +++ b/tests/test_blueprint.py @@ -638,6 +638,15 @@ def test_blueprint_vars_spec(session_ctx): blueprint.generate_manifest(session_ctx) +def test_blueprint_vars_in_owner(session_ctx): + blueprint = Blueprint( + resources=[res.Schema(name="schema", owner="role_{{ var.role_name }}", database="STATIC_DATABASE")], + vars={"role_name": "role123"}, + ) + manifest = blueprint.generate_manifest(session_ctx) + # assert manifest.resources[1].data["owner"] == "role_{{ var.role_name }}" + + def test_blueprint_allowlist(session_ctx, remote_state): blueprint = Blueprint( resources=[res.Role(name="role1")], diff --git a/tests/test_vars.py b/tests/test_vars.py index de17f75d..36daef75 100644 --- a/tests/test_vars.py +++ b/tests/test_vars.py @@ -9,3 +9,14 @@ def test_blueprint_vars_comparison_with_system_names(): schema = res.Schema(name=var.schema_name) assert isinstance(schema.name, VarString) + + +def test_vars_in_owner(): + schema = res.Schema(name="schema", owner="role_{{ var.role_name }}") + assert isinstance(schema._data.owner, VarString) + + +def test_vars_database_role(): + role = res.DatabaseRole(name="role_{{ var.role_name }}", database="db_{{ var.db_name }}") + assert isinstance(role._data.name, VarString) + assert isinstance(role._data.database, VarString) diff --git a/titan/blueprint.py b/titan/blueprint.py index 8fda91df..d77429c9 100644 --- a/titan/blueprint.py +++ b/titan/blueprint.py @@ -675,6 +675,12 @@ def _resolve_vars(self): for resource in self._staged: resource._resolve_vars(self._config.vars) + def _resolve_role_refs(self): + for resource in _walk(self._root): + if isinstance(resource, ResourcePointer): + continue + resource._resolve_role_refs() + def _build_resource_graph(self, session_ctx: SessionContext) -> None: """ Convert the staged resources into a directed graph of resources @@ -921,6 +927,7 @@ def _finalize(self, session_ctx: SessionContext) -> None: self._finalized = True self._resolve_vars() self._build_resource_graph(session_ctx) + self._resolve_role_refs() self._create_tag_references() self._create_ownership_refs(session_ctx) self._create_grandparent_refs() diff --git a/titan/resources/grant.py b/titan/resources/grant.py index 582db8c9..0406299e 100644 --- a/titan/resources/grant.py +++ b/titan/resources/grant.py @@ -185,7 +185,6 @@ def __init__( owner=owner, ) - self.requires(self._data.to) granted_on = None if on_type: granted_on = ResourcePointer(name=on, resource_type=on_type) @@ -336,7 +335,7 @@ class FutureGrant(Resource): def __init__( self, priv: str, - to: Role, + to: Union[Role, DatabaseRole], grant_option: bool = False, **kwargs, ): @@ -391,7 +390,6 @@ def __init__( to=to, grant_option=grant_option, ) - self.requires(self._data.to) if granted_in_ref: self.requires(granted_in_ref) @@ -592,7 +590,6 @@ def __init__( to=to, grant_option=grant_option, ) - self.requires(self._data.to) @classmethod def from_sql(cls, sql): diff --git a/titan/resources/resource.py b/titan/resources/resource.py index 2143a9bc..a9c73dab 100644 --- a/titan/resources/resource.py +++ b/titan/resources/resource.py @@ -89,7 +89,13 @@ def _coerce_resource_field(field_value, field_type): return {k: _coerce_resource_field(v, field_type=dict_types[1]) for k, v in field_value.items()} elif field_type is RoleRef: - return convert_role_ref(field_value) + if isinstance(field_value, str) and string_contains_var(field_value): + return VarString(field_value) + elif isinstance(field_value, (Resource, VarString, str)): + return convert_role_ref(field_value) + # return field_value + else: + raise TypeError # Check for field_value's type in a Union elif get_origin(field_type) == Union: @@ -139,14 +145,7 @@ def _coerce_resource_field(field_value, field_type): elif field_type is ResourceTags: return ResourceTags(field_value) elif field_type is str: - if isinstance(field_value, str) and string_contains_var(field_value): - return VarString(field_value) - elif isinstance(field_value, VarString): - return field_value - elif not isinstance(field_value, str): - raise TypeError - else: - return field_value + return convert_to_varstring(field_value) elif field_type is float: if isinstance(field_value, float): return field_value @@ -181,6 +180,9 @@ def to_dict(self, account_edition: AccountEdition): def _serialize_field(field, value): if field.name == "owner": + # if isinstance(value, str): + # return value + # else: return str(value.fqn) elif isinstance(value, ResourcePointer): return str(value.fqn) @@ -477,6 +479,15 @@ def _render_vars(field_value): if isinstance(self, NamedResource) and isinstance(self._name, VarString): self._name = ResourceName(self._name.to_string(vars)) + def _resolve_role_refs(self): + for f in fields(self._data): + if f.type == RoleRef: + field_value = getattr(self._data, f.name) + new_value = convert_role_ref(field_value) + setattr(self._data, f.name, new_value) + if new_value.name != "": + self.requires(new_value) + def to_pointer(self): return ResourcePointer( name=str(self.fqn), @@ -639,7 +650,7 @@ def container(self): @property def database(self): if isinstance(self.scope, DatabaseScope): - return self.container.name + return self.container.name # type: ignore else: raise ValueError("ResourcePointer does not have a database") @@ -703,16 +714,15 @@ def convert_to_resource( def convert_role_ref(role_ref: RoleRef) -> Resource: if role_ref.__class__.__name__ == "Role": - return role_ref + return role_ref # type: ignore elif role_ref.__class__.__name__ == "DatabaseRole": - return role_ref + return role_ref # type: ignore elif isinstance(role_ref, ResourcePointer) and role_ref.resource_type in ( ResourceType.DATABASE_ROLE, ResourceType.ROLE, ): return role_ref - - elif isinstance(role_ref, str) or isinstance(role_ref, ResourceName): + elif isinstance(role_ref, (str, ResourceName)): return ResourcePointer(name=role_ref, resource_type=infer_role_type_from_name(role_ref)) else: raise TypeError @@ -728,3 +738,14 @@ def infer_role_type_from_name(name: Union[str, ResourceName]) -> ResourceType: return ResourceType.DATABASE_ROLE else: return ResourceType.ROLE + + +def convert_to_varstring(value: Union[str, ResourceName]) -> Union[VarString, str]: + if isinstance(value, str) and string_contains_var(value): + return VarString(value) + elif isinstance(value, VarString): + return value + elif not isinstance(value, str): + raise TypeError + else: + return value diff --git a/titan/role_ref.py b/titan/role_ref.py index afafda0b..87083f84 100644 --- a/titan/role_ref.py +++ b/titan/role_ref.py @@ -1,6 +1,8 @@ -from typing import Union, TYPE_CHECKING +from typing import TYPE_CHECKING, Union + +from .var import VarString if TYPE_CHECKING: - from titan.resources.role import Role, DatabaseRole + from titan.resources.role import DatabaseRole, Role -RoleRef = Union["Role", "DatabaseRole", str] +RoleRef = Union["Role", "DatabaseRole", VarString, str] From 74d75187a3f0d610e000bea291adff75ae8eab1f Mon Sep 17 00:00:00 2001 From: TJ Murphy <1796+teej@users.noreply.github.com> Date: Tue, 17 Dec 2024 21:12:52 -0800 Subject: [PATCH 2/2] final touches --- tests/test_blueprint.py | 3 +-- titan/resources/resource.py | 4 ---- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/tests/test_blueprint.py b/tests/test_blueprint.py index 94499713..ad3fe82e 100644 --- a/tests/test_blueprint.py +++ b/tests/test_blueprint.py @@ -643,8 +643,7 @@ def test_blueprint_vars_in_owner(session_ctx): resources=[res.Schema(name="schema", owner="role_{{ var.role_name }}", database="STATIC_DATABASE")], vars={"role_name": "role123"}, ) - manifest = blueprint.generate_manifest(session_ctx) - # assert manifest.resources[1].data["owner"] == "role_{{ var.role_name }}" + assert blueprint.generate_manifest(session_ctx) def test_blueprint_allowlist(session_ctx, remote_state): diff --git a/titan/resources/resource.py b/titan/resources/resource.py index a9c73dab..282bf1a2 100644 --- a/titan/resources/resource.py +++ b/titan/resources/resource.py @@ -93,7 +93,6 @@ def _coerce_resource_field(field_value, field_type): return VarString(field_value) elif isinstance(field_value, (Resource, VarString, str)): return convert_role_ref(field_value) - # return field_value else: raise TypeError @@ -180,9 +179,6 @@ def to_dict(self, account_edition: AccountEdition): def _serialize_field(field, value): if field.name == "owner": - # if isinstance(value, str): - # return value - # else: return str(value.fqn) elif isinstance(value, ResourcePointer): return str(value.fqn)