From 0a9cde5ab31c1d4b7cacb9df42b38c4746c4acf5 Mon Sep 17 00:00:00 2001 From: Vida Maleki Date: Thu, 19 Sep 2024 00:29:56 -0700 Subject: [PATCH] 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 --- .../cli/_plugins/streamlit/manager.py | 3 ++ .../streamlit/streamlit_entity_model.py | 1 + .../cli/api/project/definition_conversion.py | 1 + .../project/schemas/v1/streamlit/streamlit.py | 1 + tests/project/__snapshots__/test_config.ambr | 1 + tests/streamlit/test_commands.py | 43 +++++++++++++++++++ tests/streamlit/test_streamlit_manager.py | 39 +++++++++++++++++ .../environment.yml | 5 +++ .../pages/my_page.py | 3 ++ .../snowflake.yml | 14 ++++++ .../streamlit_app.py | 3 ++ 11 files changed, 114 insertions(+) create mode 100644 tests/test_data/projects/example_streamlit_with_comment_v2/environment.yml create mode 100644 tests/test_data/projects/example_streamlit_with_comment_v2/pages/my_page.py create mode 100644 tests/test_data/projects/example_streamlit_with_comment_v2/snowflake.yml create mode 100644 tests/test_data/projects/example_streamlit_with_comment_v2/streamlit_app.py diff --git a/src/snowflake/cli/_plugins/streamlit/manager.py b/src/snowflake/cli/_plugins/streamlit/manager.py index 62df6329d3..cb7e93dc0e 100644 --- a/src/snowflake/cli/_plugins/streamlit/manager.py +++ b/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()) diff --git a/src/snowflake/cli/_plugins/streamlit/streamlit_entity_model.py b/src/snowflake/cli/_plugins/streamlit/streamlit_entity_model.py index c4785292d8..afe6fe0404 100644 --- a/src/snowflake/cli/_plugins/streamlit/streamlit_entity_model.py +++ b/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 ) diff --git a/src/snowflake/cli/api/project/definition_conversion.py b/src/snowflake/cli/api/project/definition_conversion.py index d3c93ef77a..6abcfdf804 100644 --- a/src/snowflake/cli/api/project/definition_conversion.py +++ b/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), diff --git a/src/snowflake/cli/api/project/schemas/v1/streamlit/streamlit.py b/src/snowflake/cli/api/project/schemas/v1/streamlit/streamlit.py index ca98e5a9b1..c7a454d077 100644 --- a/src/snowflake/cli/api/project/schemas/v1/streamlit/streamlit.py +++ b/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) diff --git a/tests/project/__snapshots__/test_config.ambr b/tests/project/__snapshots__/test_config.ambr index 85e1066a44..4068fc629a 100644 --- a/tests/project/__snapshots__/test_config.ambr +++ b/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', diff --git a/tests/streamlit/test_commands.py b/tests/streamlit/test_commands.py index 02720eea84..35d69f6310 100644 --- a/tests/streamlit/test_commands.py +++ b/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()", + ] diff --git a/tests/streamlit/test_streamlit_manager.py b/tests/streamlit/test_streamlit_manager.py index 2d6f82e9d0..bb50592f0f 100644 --- a/tests/streamlit/test_streamlit_manager.py +++ b/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'""" + ) + ) diff --git a/tests/test_data/projects/example_streamlit_with_comment_v2/environment.yml b/tests/test_data/projects/example_streamlit_with_comment_v2/environment.yml new file mode 100644 index 0000000000..ac8feac3e8 --- /dev/null +++ b/tests/test_data/projects/example_streamlit_with_comment_v2/environment.yml @@ -0,0 +1,5 @@ +name: sf_env +channels: + - snowflake +dependencies: + - pandas diff --git a/tests/test_data/projects/example_streamlit_with_comment_v2/pages/my_page.py b/tests/test_data/projects/example_streamlit_with_comment_v2/pages/my_page.py new file mode 100644 index 0000000000..bc3ecbccba --- /dev/null +++ b/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") diff --git a/tests/test_data/projects/example_streamlit_with_comment_v2/snowflake.yml b/tests/test_data/projects/example_streamlit_with_comment_v2/snowflake.yml new file mode 100644 index 0000000000..682b2c6159 --- /dev/null +++ b/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" diff --git a/tests/test_data/projects/example_streamlit_with_comment_v2/streamlit_app.py b/tests/test_data/projects/example_streamlit_with_comment_v2/streamlit_app.py new file mode 100644 index 0000000000..b699538d7e --- /dev/null +++ b/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")