From 6bdb1af67f2011e425e0be79f6aaa4214fe50564 Mon Sep 17 00:00:00 2001 From: Jan Sikorski Date: Fri, 29 Nov 2024 15:28:42 +0100 Subject: [PATCH] Actions --- .../_plugins/streamlit/streamlit_entity.py | 89 +++++++++++++- .../__snapshots__/test_commands.ambr | 11 +- tests/streamlit/test_actions.py | 114 ++++++++++++++++++ tests/streamlit/test_commands.py | 2 +- .../example_streamlit_v2/snowflake.yml | 2 +- .../example_streamlit_v2/utils/utils.py | 0 6 files changed, 207 insertions(+), 11 deletions(-) create mode 100644 tests/streamlit/test_actions.py create mode 100644 tests/test_data/projects/example_streamlit_v2/utils/utils.py diff --git a/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py b/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py index 6def772525..8f4e32fcc1 100644 --- a/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py +++ b/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py @@ -1,7 +1,13 @@ +import functools +import shutil + +from snowflake.cli._plugins.connection.util import make_snowsight_url from snowflake.cli._plugins.streamlit.streamlit_entity_model import ( StreamlitEntityModel, ) -from snowflake.cli.api.entities.common import EntityBase +from snowflake.cli._plugins.workspace.context import ActionContext +from snowflake.cli.api.entities.common import EntityBase, get_sql_executor +from snowflake.connector.cursor import SnowflakeCursor class StreamlitEntity(EntityBase[StreamlitEntityModel]): @@ -9,4 +15,83 @@ class StreamlitEntity(EntityBase[StreamlitEntityModel]): A Streamlit app. """ - pass + @property + def root(self): + return self._workspace_ctx.project_root + + @property + def artifacts(self): + return self._entity_model.artifacts + + @functools.cached_property + def _sql_executor(self): + return get_sql_executor() + + @functools.cached_property + def _conn(self): + return self._sql_executor._conn # noqa + + def action_bundle(self, ctx: ActionContext, *args, **kwargs): + # get all files from the model + artifacts = self._entity_model.artifacts + # get root + output_folder = self.root / "output" / self._entity_model.stage + output_folder.mkdir(parents=True, exist_ok=True) + + output_files = [] + + # This is far from , but will be replaced by bundlemap mappings. + for file in artifacts: + output_file = output_folder / file.name + + if file.is_file(): + shutil.copy(file, output_file) + elif file.is_dir(): + output_file.mkdir(parents=True, exist_ok=True) + shutil.copytree(file, output_file, dirs_exist_ok=True) + + output_files.append(output_file) + + return output_files + + def action_deploy(self, action_ctx: ActionContext, *args, **kwargs): + # After adding bundle map- we should use it's mapping here + + query = self.action_get_deploy_sql(action_ctx, *args, **kwargs) + result = self._sql_executor.execute_query(query) + return result + + def action_drop(self, action_ctx: ActionContext, *args, **kwargs): + return self._sql_executor.execute_query(self.action_get_drop_sql(action_ctx)) + + def action_execute( + self, action_ctx: ActionContext, *args, **kwargs + ) -> SnowflakeCursor: + return self._sql_executor.execute_query(self.action_get_execute_sql(action_ctx)) + + def action_get_url( + self, action_ctx: ActionContext, *args, **kwargs + ): # maybe this should be a property + name = self._entity_model.fqn.using_connection(self._conn) + return make_snowsight_url( + self._conn, f"/#/streamlit-apps/{name.url_identifier}" + ) + + def action_get_deploy_sql(self, action_ctx: ActionContext, *args, **kwargs): + pass + + def action_share( + self, action_ctx: ActionContext, to_role: str, *args, **kwargs + ) -> SnowflakeCursor: + return self._sql_executor.execute_query(self.get_share_sql(action_ctx, to_role)) + + def action_get_drop_sql(self, action_ctx: ActionContext, *args, **kwargs): + return f"DROP STREAMLIT {self._entity_model.fqn}" + + def action_get_execute_sql(self, action_ctx: ActionContext, *args, **kwargs): + return f"EXECUTE STREAMLIT {self._entity_model.fqn}()" + + def get_share_sql( + self, action_ctx: ActionContext, to_role: str, *args, **kwargs + ) -> str: + return f"GRANT USAGE ON STREAMLIT {{self._entity_model.fqn}} to role {to_role}" diff --git a/tests/streamlit/__snapshots__/test_commands.ambr b/tests/streamlit/__snapshots__/test_commands.ambr index a849210505..0bcfb9c0e6 100644 --- a/tests/streamlit/__snapshots__/test_commands.ambr +++ b/tests/streamlit/__snapshots__/test_commands.ambr @@ -1,15 +1,12 @@ # serializer version: 1 -# name: test_artifacts_must_exists +# name: test_artifacts_must_exist ''' +- Error ----------------------------------------------------------------------+ | During evaluation of DefinitionV20 in project definition following errors | | were encountered: | - | For field entities.my_streamlit.streamlit you provided '{'artifacts': | - | ['streamlit_app.py', 'foo_bar.py', 'pages/', 'environment.yml'], | - | 'identifier': 'test_streamlit_deploy_snowcli', 'main_file': | - | 'streamlit_app.py', 'query_warehouse': 'xsmall', 'stage': 'streamlit', | - | 'title': 'My Fancy Streamlit', 'type': 'streamlit'}'. This caused: Value | - | error, Specified artifact foo_bar.py does not exist locally. | + | For field entities.my_streamlit you provided '{'artifacts': {'1': | + | 'foo_bar.py'}}'. This caused: Unable to extract tag using discriminator | + | 'type' | +------------------------------------------------------------------------------+ ''' diff --git a/tests/streamlit/test_actions.py b/tests/streamlit/test_actions.py new file mode 100644 index 0000000000..0c870fa162 --- /dev/null +++ b/tests/streamlit/test_actions.py @@ -0,0 +1,114 @@ +from pathlib import Path +from unittest import mock + +import pytest +import yaml +from snowflake.cli._plugins.streamlit.streamlit_entity import StreamlitEntity +from snowflake.cli._plugins.streamlit.streamlit_entity_model import StreamlitEntityModel +from snowflake.cli._plugins.workspace.context import ActionContext, WorkspaceContext + +CONNECTOR = "snowflake.connector.connect" +CONTEXT = "" +EXECUTE_QUERY = "snowflake.cli.api.sql_execution.BaseSqlExecutor.execute_query" + +GET_UI_PARAMETERS = "snowflake.cli._plugins.connection.util.get_ui_parameters" + + +@pytest.fixture +def example_streamlit_workspace(project_directory): + with project_directory("example_streamlit_v2") as pdir: + with Path(pdir / "snowflake.yml").open() as definition_file: + definition = yaml.safe_load(definition_file) + print(definition) + model = StreamlitEntityModel( + **definition.get("entities", {}).get("test_streamlit") + ) + + workspace_context = WorkspaceContext( + console=mock.MagicMock(), + project_root=pdir, + get_default_role=lambda: "test_role", + get_default_warehouse=lambda: "test_warehouse", + ) + + return ( + StreamlitEntity(workspace_ctx=workspace_context, entity_model=model), + ActionContext( + get_entity=lambda *args: None, + ), + ) + + +def test_bundle(example_streamlit_workspace): + + entity, action_ctx = example_streamlit_workspace + entity.action_bundle(action_ctx) + + output = entity.root / "output" / entity._entity_model.stage # noqa + assert output.exists() + assert (output / "streamlit_app.py").exists() + assert (output / "environment.yml").exists() + assert (output / "pages" / "my_page.py").exists() + + +@mock.patch(EXECUTE_QUERY) +def test_drop(mock_execute, example_streamlit_workspace): + entity, action_ctx = example_streamlit_workspace + entity.action_drop(action_ctx) + + mock_execute.assert_called_with("DROP STREAMLIT test_streamlit_deploy_snowcli") + + +@mock.patch(CONNECTOR) +@mock.patch( + GET_UI_PARAMETERS, + return_value={"UI_SNOWSIGHT_ENABLE_REGIONLESS_REDIRECT": "false"}, +) +@mock.patch("click.get_current_context") +def test_get_url( + mock_get_ctx, + mock_param, + mock_connect, + mock_cursor, + example_streamlit_workspace, + mock_ctx, +): + ctx = mock_ctx( + mock_cursor( + rows=[ + {"SYSTEM$GET_SNOWSIGHT_HOST()": "https://snowsight.domain"}, + {"SYSTEM$RETURN_CURRENT_ORG_NAME()": "FOOBARBAZ"}, + {"CURRENT_ACCOUNT_NAME()": "https://snowsight.domain"}, + ], + columns=["SYSTEM$GET_SNOWSIGHT_HOST()"], + ) + ) + mock_connect.return_value = ctx + mock_get_ctx.return_value = ctx + + entity, action_ctx = example_streamlit_workspace + result = entity.action_get_url(action_ctx) + + mock_connect.assert_called() + + +@mock.patch(EXECUTE_QUERY) +def test_execute(mock_execute, example_streamlit_workspace): + entity, action_ctx = example_streamlit_workspace + entity.action_execute(action_ctx) + + mock_execute.assert_called_with("EXECUTE STREAMLIT test_streamlit_deploy_snowcli()") + + +def test_get_execute_sql(example_streamlit_workspace): + entity, action_ctx = example_streamlit_workspace + execute_sql = entity.action_get_execute_sql(action_ctx) + + assert execute_sql == "EXECUTE STREAMLIT test_streamlit_deploy_snowcli()" + + +def test_get_drop_sql(example_streamlit_workspace): + entity, action_ctx = example_streamlit_workspace + drop_sql = entity.action_get_drop_sql(action_ctx) + + assert drop_sql == "DROP STREAMLIT test_streamlit_deploy_snowcli" diff --git a/tests/streamlit/test_commands.py b/tests/streamlit/test_commands.py index 68464053bd..deb6701695 100644 --- a/tests/streamlit/test_commands.py +++ b/tests/streamlit/test_commands.py @@ -281,7 +281,7 @@ def test_deploy_only_streamlit_file_replace( mock_typer.launch.assert_not_called() -def test_artifacts_must_exists( +def test_artifacts_must_exist( runner, mock_ctx, project_directory, alter_snowflake_yml, snapshot ): with project_directory("example_streamlit_v2") as pdir: diff --git a/tests/test_data/projects/example_streamlit_v2/snowflake.yml b/tests/test_data/projects/example_streamlit_v2/snowflake.yml index f10b3d71ee..607b94bd0a 100644 --- a/tests/test_data/projects/example_streamlit_v2/snowflake.yml +++ b/tests/test_data/projects/example_streamlit_v2/snowflake.yml @@ -1,6 +1,6 @@ definition_version: 2 entities: - my_streamlit: + test_streamlit: type: "streamlit" identifier: test_streamlit_deploy_snowcli title: "My Fancy Streamlit" diff --git a/tests/test_data/projects/example_streamlit_v2/utils/utils.py b/tests/test_data/projects/example_streamlit_v2/utils/utils.py new file mode 100644 index 0000000000..e69de29bb2