Actions

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

@@ -1,12 +1,97 @@
+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]):
     """
     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}"

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'                                                                       |
   +------------------------------------------------------------------------------+
 
   '''

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"

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:

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"