Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add grants support to Streamlit entity #1981

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/snowflake/cli/_plugins/streamlit/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,17 @@
from snowflake.cli.api.project.schemas.entities.common import (
EntityModelBase,
ExternalAccessBaseModel,
GrantBaseModel,
ImportsBaseModel,
)
from snowflake.cli.api.project.schemas.updatable_model import (
DiscriminatorField,
)


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
Expand Down
22 changes: 22 additions & 0 deletions src/snowflake/cli/api/project/schemas/entities/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

title="Privileges that will be granted by a role"

Copy link
Contributor

@sfc-gh-astus sfc-gh-astus Jan 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should stay with only list type? WDYT?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wondering if we even need to support list here. Without list you can still do

grants:
  - role: A
     grant: B
  - role: A
     grant: C

One can argue that this is even better

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",
Expand Down
64 changes: 64 additions & 0 deletions tests/streamlit/test_streamlit_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading