From 242c0ad2ecb4822e3feafce8f4a04666fefbb265 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Thu, 9 Jan 2025 16:45:18 +0100 Subject: [PATCH] Add grants support to Streamlit entity --- .../cli/_plugins/streamlit/manager.py | 8 +++ .../streamlit/streamlit_entity_model.py | 5 +- .../api/project/schemas/entities/common.py | 22 +++++++ tests/streamlit/test_streamlit_manager.py | 64 +++++++++++++++++++ 4 files changed, 98 insertions(+), 1 deletion(-) diff --git a/src/snowflake/cli/_plugins/streamlit/manager.py b/src/snowflake/cli/_plugins/streamlit/manager.py index 8d51c9f240..29bef6a0c7 100644 --- a/src/snowflake/cli/_plugins/streamlit/manager.py +++ b/src/snowflake/cli/_plugins/streamlit/manager.py @@ -214,8 +214,16 @@ def deploy(self, streamlit: StreamlitEntityModel, replace: bool = False): experimental=False, ) + self.grant_privileges(streamlit) + return self.get_url(streamlit_name=streamlit_id) + def grant_privileges(self, entity_model: StreamlitEntityModel): + if not entity_model.grants: + return + for grant in entity_model.grants: + self.execute_query(grant.get_grant_sql(entity_model)) + def get_url(self, streamlit_name: FQN) -> str: try: fqn = streamlit_name.using_connection(self._conn) diff --git a/src/snowflake/cli/_plugins/streamlit/streamlit_entity_model.py b/src/snowflake/cli/_plugins/streamlit/streamlit_entity_model.py index 55068adb5a..3eaff31a6f 100644 --- a/src/snowflake/cli/_plugins/streamlit/streamlit_entity_model.py +++ b/src/snowflake/cli/_plugins/streamlit/streamlit_entity_model.py @@ -20,6 +20,7 @@ from snowflake.cli.api.project.schemas.entities.common import ( EntityModelBase, ExternalAccessBaseModel, + GrantBaseModel, ImportsBaseModel, ) from snowflake.cli.api.project.schemas.updatable_model import ( @@ -27,7 +28,9 @@ ) -class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseModel): +class StreamlitEntityModel( + EntityModelBase, ExternalAccessBaseModel, ImportsBaseModel, GrantBaseModel +): type: Literal["streamlit"] = DiscriminatorField() # noqa: A003 title: Optional[str] = Field( title="Human-readable title for the Streamlit dashboard", default=None diff --git a/src/snowflake/cli/api/project/schemas/entities/common.py b/src/snowflake/cli/api/project/schemas/entities/common.py index d9036d9a4c..51008674e5 100644 --- a/src/snowflake/cli/api/project/schemas/entities/common.py +++ b/src/snowflake/cli/api/project/schemas/entities/common.py @@ -141,6 +141,28 @@ def get_imports_sql(self) -> str | None: return f"IMPORTS = ({imports})" +class Grant(UpdatableModel): + privileges: str | list[str] = Field(title="Required privileges") + role: str = Field(title="Role to which the privileges will be granted") + + def get_grant_sql(self, entity_model: EntityModelBase) -> str: + privileges = ( + ", ".join(self.privileges) + if isinstance(self.privileges, list) + else self.privileges + ) + return f"GRANT {privileges} ON {entity_model.get_type().upper()} {entity_model.fqn.sql_identifier} TO ROLE {self.role}" + + +class GrantBaseModel(UpdatableModel): + grants: Optional[List[Grant]] = Field(title="List of grants", default=None) + + def get_grant_sqls(self) -> list[str]: + return ( + [grant.get_grant_sql(self) for grant in self.grants] if self.grants else [] + ) + + class ExternalAccessBaseModel: external_access_integrations: Optional[List[str]] = Field( title="Names of external access integrations needed for this entity to access external networks", diff --git a/tests/streamlit/test_streamlit_manager.py b/tests/streamlit/test_streamlit_manager.py index 0fcfad4ec0..bf044ba314 100644 --- a/tests/streamlit/test_streamlit_manager.py +++ b/tests/streamlit/test_streamlit_manager.py @@ -172,6 +172,70 @@ def test_deploy_streamlit_with_default_warehouse( ) +@mock.patch("snowflake.cli._plugins.streamlit.manager.StageManager") +@mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager.get_url") +@mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager.execute_query") +@mock.patch( + "snowflake.cli._plugins.streamlit.manager.StreamlitManager.grant_privileges" +) +@mock_streamlit_exists +def test_deploy_streamlit_with_grants(mock_grants, _, __, mock_stage_manager, temp_dir): + mock_stage_manager().get_standard_stage_prefix.return_value = "stage_root" + + main_file = Path(temp_dir) / "main.py" + main_file.touch() + + st = StreamlitEntityModel( + type="streamlit", + identifier="my_streamlit_app", + title="MyStreamlit", + main_file=str(main_file), + artifacts=[main_file], + comment="This is a test comment", + grants=[{"privileges": ["USAGE"], "role": "FOO_BAR"}], + ) + + StreamlitManager(MagicMock(database="DB", schema="SH")).deploy( + streamlit=st, replace=False + ) + + mock_grants.assert_called_once_with(st) + + +@mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager.execute_query") +def test_grant_privileges_to_streamlit(mock_execute, temp_dir): + main_file = Path(temp_dir) / "main.py" + main_file.touch() + + st = StreamlitEntityModel( + type="streamlit", + identifier="my_streamlit_app", + title="MyStreamlit", + main_file="main.py", + artifacts=[main_file], + comment="This is a test comment", + grants=[ + {"privileges": ["AAAA"], "role": "FOO"}, + {"privileges": ["BBBB", "AAAA"], "role": "BAR"}, + ], + ) + + StreamlitManager(MagicMock(database="DB", schema="SH")).grant_privileges( + entity_model=st + ) + + mock_execute.assert_has_calls( + [ + mock.call( + "GRANT AAAA ON STREAMLIT IDENTIFIER('my_streamlit_app') TO ROLE FOO" + ), + mock.call( + "GRANT BBBB, AAAA ON STREAMLIT IDENTIFIER('my_streamlit_app') TO ROLE BAR" + ), + ] + ) + + @mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager.execute_query") @mock_streamlit_exists def test_execute_streamlit(mock_execute_query):