Add support for multiple Streamlits (#1440)

RELEASE-NOTES.md

@@ -33,6 +33,7 @@
 * Added `snow app events` command to fetch logs and traces from local and customer app installations
 * Added support for project definition file defaults in templates
 * Added support for external access (api integrations and secrets) in Streamlit.
+* Support multiple Streamlit application in single snowflake.yml project definition file.
 
 ## Fixes and improvements
 * Fixed problem with whitespaces in `snow connection add` command

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

@@ -20,7 +20,7 @@
 
 import click
 import typer
-from click import ClickException
+from click import ClickException, UsageError
 from snowflake.cli._plugins.object.command_aliases import (
     add_object_command_aliases,
     scope_option,
@@ -33,6 +33,7 @@
 )
 from snowflake.cli.api.commands.flags import (
     ReplaceOption,
+    entity_argument,
     identifier_argument,
     like_option,
 )
@@ -120,13 +121,15 @@ def streamlit_deploy(
     replace: bool = ReplaceOption(
         help="Replace the Streamlit app if it already exists."
     ),
+    entity_id: str = entity_argument("streamlit"),
     open_: bool = OpenOption,
     **options,
 ) -> CommandResult:
     """
     Deploys a Streamlit app defined in the project definition file (snowflake.yml). By default, the command uploads
     environment.yml and any other pages or folders, if present. If you don’t specify a stage name, the `streamlit`
-    stage is used. If the specified stage does not exist, the command creates it.
+    stage is used. If the specified stage does not exist, the command creates it. If multiple Streamlits are defined
+    in snowflake.yml and no entity_id is provided then command will raise an error.
     """
 
     cli_context = get_cli_context()
@@ -147,14 +150,19 @@ def streamlit_deploy(
             project_type="streamlit", project_file=cli_context.project_root
         )
 
-    # TODO: fix in follow-up
-    if len(list(streamlits)) > 1:
-        raise ClickException(
-            "Currently only single streamlit entity per project is supported."
+    if entity_id and entity_id not in streamlits:
+        raise UsageError(f"No '{entity_id}' entity in project definition file.")
+
+    if len(streamlits.keys()) == 1:
+        entity_id = list(streamlits.keys())[0]
+
+    if entity_id is None:
+        raise UsageError(
+            "Multiple Streamlit apps found. Please provide entity id for the operation."
         )
 
     # Get first streamlit
-    streamlit: StreamlitEntityModel = streamlits[list(streamlits)[0]]
+    streamlit: StreamlitEntityModel = streamlits[entity_id]
     url = StreamlitManager().deploy(streamlit=streamlit, replace=replace)
 
     if open_:

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

@@ -27,6 +27,7 @@
 from snowflake.cli.api.commands.experimental_behaviour import (
     experimental_behaviour_enabled,
 )
+from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
@@ -50,6 +51,7 @@ def _put_streamlit_files(
         root_location: str,
         artifacts: Optional[List[Path]] = None,
     ):
+        cli_console.step(f"Deploying files to {root_location}")
         if not artifacts:
             return
         stage_manager = StageManager()
@@ -71,6 +73,7 @@ def _create_streamlit(
         from_stage_name: Optional[str] = None,
     ):
         streamlit_id = streamlit.fqn.using_connection(self._conn)
+        cli_console.step(f"Creating {streamlit_id} Streamlit")
         query = []
         if replace:
             query.append(f"CREATE OR REPLACE STREAMLIT {streamlit_id.sql_identifier}")
@@ -167,6 +170,7 @@ def deploy(self, streamlit: StreamlitEntityModel, replace: bool = False):
             stage_name = streamlit.stage or "streamlit"
             stage_name = FQN.from_string(stage_name).using_connection(self._conn)
 
+            cli_console.step(f"Creating {stage_name} stage")
             stage_manager.create(fqn=stage_name)
 
             root_location = stage_manager.get_standard_stage_prefix(

src/snowflake/cli/api/commands/flags.py

@@ -372,6 +372,10 @@ def _diag_log_path_callback(path: str):
 NoInteractiveOption = typer.Option(False, "--no-interactive", help="Disable prompting.")
 
 
+def entity_argument(entity_type: str) -> typer.Argument:
+    return typer.Argument(None, help=f"ID of {entity_type} entity.")
+
+
 def variables_option(description: str):
     return typer.Option(
         None,

tests/__snapshots__/test_help_messages.ambr

@@ -7410,14 +7410,18 @@
 # name: test_help_messages[streamlit.deploy]
   '''
 
-   Usage: default streamlit deploy [OPTIONS]                                      
+   Usage: default streamlit deploy [OPTIONS] [ENTITY_ID]                          
 
    Deploys a Streamlit app defined in the project definition file                 
    (snowflake.yml). By default, the command uploads environment.yml and any other 
    pages or folders, if present. If you don’t specify a stage name, the           
    `streamlit` stage is used. If the specified stage does not exist, the command  
-   creates it.                                                                    
+   creates it. If multiple Streamlits are defined in snowflake.yml and no         
+   entity_id is provided then command will raise an error.                        
 
+  +- Arguments ------------------------------------------------------------------+
+  |   entity_id      [ENTITY_ID]  ID of streamlit entity. [default: None]        |
+  +------------------------------------------------------------------------------+
   +- Options --------------------------------------------------------------------+
   | --replace                Replace the Streamlit app if it already exists.     |
   | --open                   Whether to open the Streamlit app in a browser.     |
@@ -7871,7 +7875,9 @@
   |            (snowflake.yml). By default, the command uploads environment.yml  |
   |            and any other pages or folders, if present. If you don’t specify  |
   |            a stage name, the `streamlit` stage is used. If the specified     |
-  |            stage does not exist, the command creates it.                     |
+  |            stage does not exist, the command creates it. If multiple         |
+  |            Streamlits are defined in snowflake.yml and no entity_id is       |
+  |            provided then command will raise an error.                        |
   | describe   Provides description of streamlit.                                |
   | drop       Drops streamlit with given name.                                  |
   | get-url    Returns a URL to the specified Streamlit app                      |
@@ -8308,7 +8314,9 @@
   |            (snowflake.yml). By default, the command uploads environment.yml  |
   |            and any other pages or folders, if present. If you don’t specify  |
   |            a stage name, the `streamlit` stage is used. If the specified     |
-  |            stage does not exist, the command creates it.                     |
+  |            stage does not exist, the command creates it. If multiple         |
+  |            Streamlits are defined in snowflake.yml and no entity_id is       |
+  |            provided then command will raise an error.                        |
   | describe   Provides description of streamlit.                                |
   | drop       Drops streamlit with given name.                                  |
   | get-url    Returns a URL to the specified Streamlit app                      |

tests/streamlit/__snapshots__/test_commands.ambr

@@ -29,3 +29,13 @@
 
   '''
 # ---
+# name: test_multiple_streamlit_raise_error_if_multiple_entities
+  '''
+  Usage: default streamlit deploy [OPTIONS] [ENTITY_ID]
+  Try 'default streamlit deploy --help' for help.
+  +- Error ----------------------------------------------------------------------+
+  | Multiple Streamlit apps found. Please provide entity id for the operation.   |
+  +------------------------------------------------------------------------------+
+
+  '''
+# ---

tests/streamlit/test_commands.py

@@ -807,3 +807,36 @@ def test_command_aliases(mock_connector, runner, mock_ctx, command, parameters):
 
     queries = ctx.get_queries()
     assert queries[0] == queries[1]
+
+
+@pytest.mark.parametrize("entity_id", ["app_1", "app_2"])
+@mock.patch("snowflake.cli._plugins.streamlit.commands.StreamlitManager")
+@mock.patch("snowflake.cli._plugins.streamlit.manager.StageManager")
+@mock.patch("snowflake.connector.connect")
+def test_selecting_streamlit_from_pdf(
+    _, __, mock_manager, project_directory, runner, entity_id
+):
+
+    with project_directory("example_streamlit_multiple_v2"):
+        result = runner.invoke(["streamlit", "deploy", entity_id])
+
+    assert result.exit_code == 0, result.output
+
+    calls = mock_manager().deploy
+    assert calls.call_count == 1
+
+    # Make sure the streamlit was called with proper app definition
+    st = calls.call_args.kwargs
+    assert st["streamlit"].entity_id == entity_id
+
+
+@mock.patch("snowflake.connector.connect")
+def test_multiple_streamlit_raise_error_if_multiple_entities(
+    _, runner, project_directory, os_agnostic_snapshot
+):
+
+    with project_directory("example_streamlit_multiple_v2"):
+        result = runner.invoke(["streamlit", "deploy"])
+
+    assert result.exit_code == 2, result.output
+    assert result.output == os_agnostic_snapshot

tests/test_data/projects/example_streamlit_multiple_v2/snowflake.yml

@@ -0,0 +1,12 @@
+definition_version: 2
+entities:
+  app_1:
+    type: "streamlit"
+    main_file: streamlit_app.py
+    artifacts:
+      - streamlit_app.py
+  app_2:
+    type: "streamlit"
+    main_file: streamlit_app.py
+    artifacts:
+      - streamlit_app.py

tests/test_data/projects/example_streamlit_multiple_v2/streamlit_app.py

@@ -0,0 +1,3 @@
+import streamlit as st
+
+st.title("Example streamlit app")