SNOW-2206349 Fix Streamlit Entity Grants

RELEASE-NOTES.md

@@ -23,6 +23,7 @@
 ## Fixes and improvements
 * Bumped `snowflake-connector-python==3.17.4`
 * Fixed `snow spcs logs` `IndexOutOfRange` error
+* 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(connection=self._conn).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(connection=self._conn).grant_privileges(self.model)

tests_integration/test_streamlit.py

@@ -132,3 +132,57 @@ 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 = snowflake_session.role
+    entity_id = "app_1"
+
+    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)
+
+
+@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 = snowflake_session.role
+    entity_id = "app_1"
+
+    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)

tests_integration/testing_utils/streamlit_utils.py

@@ -183,6 +183,15 @@ 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 verify_grants_applied(self, entity_id: str, test_role: str):
+        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"
+
 
 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()}"