Add Comment Support for Streamlit App Creation (#1580)

* Add support for 'comment' field in Streamlit app deployment

* Added unit tests for 'comment' support in Streamlit app deployment

* Added test case for StreamlitEntityModel to include comment field

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

@@ -104,6 +104,9 @@ def _create_streamlit(
         if streamlit.title:
             query.append(f"TITLE = '{streamlit.title}'")
 
+        if streamlit.comment:
+            query.append(f"COMMENT = '{streamlit.comment}'")
+
         if streamlit.external_access_integrations:
             query.append(streamlit.get_external_access_integrations_sql())
 

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

@@ -31,6 +31,7 @@ class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel):
     title: Optional[str] = Field(
         title="Human-readable title for the Streamlit dashboard", default=None
     )
+    comment: Optional[str] = Field(title="Comment for the Streamlit app", default=None)
     query_warehouse: str = Field(
         title="Snowflake warehouse to host the app", default=None
     )

src/snowflake/cli/api/project/definition_conversion.py

@@ -157,6 +157,7 @@ def convert_streamlit_to_v2_data(streamlit: Streamlit) -> Dict[str, Any]:
                 "type": "streamlit",
                 "identifier": identifier,
                 "title": streamlit.title,
+                "comment": streamlit.comment,
                 "query_warehouse": streamlit.query_warehouse,
                 "main_file": str(streamlit.main_file),
                 "pages_dir": str(streamlit.pages_dir),

src/snowflake/cli/api/project/schemas/v1/streamlit/streamlit.py

@@ -44,3 +44,4 @@ class Streamlit(UpdatableModel, ObjectIdentifierModel(object_name="Streamlit")):
     title: Optional[str] = Field(
         title="Human-readable title for the Streamlit dashboard", default=None
     )
+    comment: Optional[str] = Field(title="Comment for the Streamlit app", default=None)

tests/project/__snapshots__/test_config.ambr

@@ -774,6 +774,7 @@
         'utils/utils.py',
         'extra_file.py',
       ]),
+      'comment': None,
       'database': None,
       'env_file': 'environment.yml',
       'main_file': 'streamlit_app.py',

tests/streamlit/test_commands.py

@@ -911,3 +911,46 @@ def test_multiple_streamlit_raise_error_if_multiple_entities(
 
     assert result.exit_code == 2, result.output
     assert result.output == os_agnostic_snapshot
+
+
+@mock.patch("snowflake.connector.connect")
+def test_deploy_streamlit_with_comment_v2(
+    mock_connector, mock_cursor, runner, mock_ctx, project_directory
+):
+    ctx = mock_ctx(
+        mock_cursor(
+            rows=[
+                {"SYSTEM$GET_SNOWSIGHT_HOST()": "https://snowsight.domain"},
+                {"REGIONLESS": "false"},
+                {"CURRENT_ACCOUNT_NAME()": "https://snowsight.domain"},
+            ],
+            columns=["SYSTEM$GET_SNOWSIGHT_HOST()"],
+        )
+    )
+    mock_connector.return_value = ctx
+
+    with project_directory("example_streamlit_with_comment_v2"):
+        result = runner.invoke(["streamlit", "deploy", "--replace"])
+
+    root_path = f"@MockDatabase.MockSchema.streamlit/test_streamlit_deploy_snowcli"
+    assert result.exit_code == 0, result.output
+    assert ctx.get_queries() == [
+        f"describe streamlit IDENTIFIER('MockDatabase.MockSchema.test_streamlit_deploy_snowcli')",
+        "create stage if not exists IDENTIFIER('MockDatabase.MockSchema.streamlit')",
+        _put_query("streamlit_app.py", root_path),
+        _put_query("pages/*", f"{root_path}/pages"),
+        _put_query("environment.yml", root_path),
+        dedent(
+            f"""
+            CREATE OR REPLACE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.test_streamlit_deploy_snowcli')
+            ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit/test_streamlit_deploy_snowcli'
+            MAIN_FILE = 'streamlit_app.py'
+            QUERY_WAREHOUSE = xsmall
+            TITLE = 'My Streamlit App with Comment'
+            COMMENT = 'This is a test comment'
+            """
+        ).strip(),
+        "select system$get_snowsight_host()",
+        REGIONLESS_QUERY,
+        "select current_account_name()",
+    ]

tests/streamlit/test_streamlit_manager.py

@@ -90,3 +90,42 @@ def test_deploy_streamlit_with_api_integrations(
         secrets=('my_secret'=SecretOfTheSecrets, 'other'=other_secret)"""
         )
     )
+
+
+@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_streamlit_exists
+def test_deploy_streamlit_with_comment(
+    mock_execute_query, _, 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",
+        query_warehouse="My_WH",
+        main_file=str(main_file),
+        artifacts=[main_file],
+        comment="This is a test comment",
+    )
+
+    StreamlitManager(MagicMock(database="DB", schema="SH")).deploy(
+        streamlit=st, replace=False
+    )
+
+    mock_execute_query.assert_called_once_with(
+        dedent(
+            f"""\
+            CREATE STREAMLIT IDENTIFIER('DB.SH.my_streamlit_app')
+            ROOT_LOCATION = 'stage_root'
+            MAIN_FILE = '{main_file}'
+            QUERY_WAREHOUSE = My_WH
+            TITLE = 'MyStreamlit'
+            COMMENT = 'This is a test comment'"""
+        )
+    )

tests/test_data/projects/example_streamlit_with_comment_v2/environment.yml

@@ -0,0 +1,5 @@
+name: sf_env
+channels:
+  - snowflake
+dependencies:
+  - pandas

tests/test_data/projects/example_streamlit_with_comment_v2/pages/my_page.py

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

tests/test_data/projects/example_streamlit_with_comment_v2/snowflake.yml

@@ -0,0 +1,14 @@
+definition_version: 2
+entities:
+  my_streamlit:
+    type: "streamlit"
+    identifier: test_streamlit_deploy_snowcli
+    title: "My Streamlit App with Comment"
+    stage: streamlit
+    query_warehouse: xsmall
+    main_file: streamlit_app.py
+    artifacts:
+      - streamlit_app.py
+      - pages/
+      - environment.yml
+    comment: "This is a test comment"

tests/test_data/projects/example_streamlit_with_comment_v2/streamlit_app.py

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