SNOW-2206349 Fix Streamlit Entity Grants

Author: sfc-gh-daniszewski

RELEASE-NOTES.md

@@ -21,6 +21,7 @@
 ## New additions
 
 ## Fixes and improvements
+* Grant privileges defined in `snowflake.yml` after deploying Streamlit
 
 
 # v3.12.0

src/snowflake/cli/_plugins/streamlit/streamlit_entity.py

@@ -6,6 +6,7 @@
 from snowflake.cli._plugins.connection.util import make_snowsight_url
 from snowflake.cli._plugins.nativeapp.artifacts import build_bundle
 from snowflake.cli._plugins.stage.manager import StageManager
+from snowflake.cli._plugins.streamlit.manager import StreamlitManager
 from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
     StreamlitEntityModel,
 )
@@ -132,6 +133,8 @@ def deploy(
                 self.get_deploy_sql(replace=replace, from_stage_name=stage_root)
             )
 
+            StreamlitManager().grant_privileges(self.model)
+
         return self.perform(EntityActions.GET_URL, action_context, *args, **kwargs)
 
     def describe(self) -> SnowflakeCursor:
@@ -256,3 +259,5 @@ def _deploy_experimental(
             print_diff=True,
             force_overwrite=True,  # files copied to streamlit vstage need to be overwritten
         )
+
+        StreamlitManager().grant_privileges(self.model)

tests_integration/test_streamlit.py

@@ -1,3 +1,5 @@
+import uuid
+
 import pytest
 
 from tests_integration.testing_utils import FlowTestSetup
@@ -132,3 +134,65 @@ def _test_setup(
 @pytest.fixture
 def _streamlit_test_steps(_test_setup):
     return StreamlitTestSteps(_test_setup)
+
+
+@pytest.mark.integration
+def test_streamlit_grants_flow(
+    _streamlit_test_steps,
+    project_directory,
+    snowflake_session,
+    alter_snowflake_yml,
+):
+    """Test that streamlit grants are properly applied during deployment."""
+    test_role = f"snowcli_streamlit_grants_test_{uuid.uuid4().hex[:8]}"
+    entity_id = "app_1"
+
+    _streamlit_test_steps.create_test_role(test_role)
+
+    with project_directory("streamlit_v2"):
+        alter_snowflake_yml(
+            "snowflake.yml",
+            "entities.app_1.grants",
+            [{"privilege": "USAGE", "role": test_role}],
+        )
+
+        _streamlit_test_steps.deploy_with_entity_id_specified_should_succeed(
+            entity_id, snowflake_session, experimental=False
+        )
+
+        _streamlit_test_steps.verify_grants_applied(entity_id, test_role)
+
+        _streamlit_test_steps.drop_should_succeed(entity_id, snowflake_session)
+
+    _streamlit_test_steps.cleanup_test_role(test_role)
+
+
+@pytest.mark.integration
+def test_streamlit_grants_experimental_flow(
+    _streamlit_test_steps,
+    project_directory,
+    snowflake_session,
+    alter_snowflake_yml,
+):
+    """Test that streamlit grants are properly applied during experimental deployment."""
+    test_role = f"snowcli_streamlit_grants_exp_test_{uuid.uuid4().hex[:8]}"
+    entity_id = "app_1"
+
+    _streamlit_test_steps.create_test_role(test_role)
+
+    with project_directory("streamlit_v2"):
+        alter_snowflake_yml(
+            "snowflake.yml",
+            "entities.app_1.grants",
+            [{"privilege": "USAGE", "role": test_role}],
+        )
+
+        _streamlit_test_steps.deploy_with_entity_id_specified_should_succeed(
+            entity_id, snowflake_session, experimental=True
+        )
+
+        _streamlit_test_steps.verify_grants_applied(entity_id, test_role)
+
+        _streamlit_test_steps.drop_should_succeed(entity_id, snowflake_session)
+
+    _streamlit_test_steps.cleanup_test_role(test_role)

tests_integration/testing_utils/streamlit_utils.py

@@ -183,6 +183,30 @@ def assert_proper_url_is_returned(
         assert message.startswith("Streamlit successfully deployed and available under")
         assert message.endswith(create_expected_url_suffix(entity_id, session))
 
+    def create_test_role(self, role_name: str):
+        self.setup.sql_test_helper.execute_single_sql(
+            f"set user = (select current_user()); "
+            f"create role {role_name}; "
+            f"grant role {role_name} to user IDENTIFIER($USER)"
+        )
+
+    def cleanup_test_role(self, role_name: str):
+        self.setup.sql_test_helper.execute_single_sql(
+            f"DROP ROLE IF EXISTS {role_name}"
+        )
+
+    def verify_grants_applied(self, entity_id: str, test_role: str):
+        try:
+            self.setup.sql_test_helper.execute_single_sql(f"USE ROLE {test_role}")
+            streamlits_with_role = self.setup.sql_test_helper.execute_single_sql(
+                f"SHOW STREAMLITS LIKE '{entity_id}'"
+            )
+            assert (
+                len(streamlits_with_role) == 1
+            ), f"Role {test_role} should have USAGE access to the streamlit"
+        finally:
+            self.setup.sql_test_helper.execute_single_sql("USE ROLE ACCOUNTADMIN")
+
 
 def create_expected_url_suffix(entity_id: str, session: SnowflakeConnection):
     return f".snowflake.com/SFENGINEERING/{get_account(session)}/#/streamlit-apps/{session.database.upper()}.{session.schema.upper()}.{entity_id.upper()}"