Add grants support to Streamlit entity

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

@@ -247,8 +247,16 @@ def deploy(
                 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)

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

@@ -20,6 +20,7 @@
     Artifacts,
     EntityModelBaseWithArtifacts,
     ExternalAccessBaseModel,
+    GrantBaseModel,
     ImportsBaseModel,
 )
 from snowflake.cli.api.project.schemas.updatable_model import (
@@ -28,7 +29,10 @@
 
 
 class StreamlitEntityModel(
-    EntityModelBaseWithArtifacts, ExternalAccessBaseModel, ImportsBaseModel
+    EntityModelBaseWithArtifacts,
+    ExternalAccessBaseModel,
+    ImportsBaseModel,
+    GrantBaseModel,
 ):
     type: Literal["streamlit"] = DiscriminatorField()  # noqa: A003
     title: Optional[str] = Field(

src/snowflake/cli/api/project/schemas/entities/common.py

@@ -141,6 +141,23 @@ def get_imports_sql(self) -> str | None:
         return f"IMPORTS = ({imports})"
 
 
+class Grant(UpdatableModel):
+    privilege: 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:
+        return f"GRANT {self.privilege} 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",

tests/streamlit/test_streamlit_manager.py

@@ -183,6 +183,69 @@ 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.py"],
+        comment="This is a test comment",
+        grants=[{"privilege": "USAGE", "role": "FOO_BAR"}],
+    )
+
+    StreamlitManager(MagicMock(database="DB", schema="SH")).deploy(
+        streamlit=st,
+        replace=False,
+        streamlit_project_paths=StreamlitProjectPaths(Path(temp_dir)),
+    )
+
+    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):
+    st = StreamlitEntityModel(
+        type="streamlit",
+        identifier="my_streamlit_app",
+        title="MyStreamlit",
+        main_file="main.py",
+        artifacts=["main.py"],
+        comment="This is a test comment",
+        grants=[
+            {"privilege": "AAAA", "role": "FOO"},
+            {"privilege": "BBBB", "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 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):

tests_integration/test_data/projects/streamlit_v2/snowflake.yml

@@ -10,3 +10,6 @@ entities:
     main_file: streamlit_app.py
     artifacts:
       - streamlit_app.py
+    grants:
+      - privilege: "USAGE"
+        role: "ACCOUNTADMIN"