Support project definition V2 in streamlit deploy command

src/snowflake/cli/api/project/schemas/entities/entities.py

@@ -22,8 +22,9 @@
 from snowflake.cli.api.project.schemas.entities.application_package_entity import (
     ApplicationPackageEntity,
 )
+from snowflake.cli.api.project.schemas.entities.streamlit_entity import StreamlitEntity
 
-Entity = Union[ApplicationEntity, ApplicationPackageEntity]
+Entity = Union[ApplicationEntity, ApplicationPackageEntity, StreamlitEntity]
 
 ALL_ENTITIES = [*get_args(Entity)]
 

src/snowflake/cli/api/project/schemas/entities/streamlit_entity.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List, Literal, Optional
+
+from pydantic import Field, model_validator
+from snowflake.cli.api.project.schemas.entities.common import EntityBase
+from snowflake.cli.api.project.schemas.identifier_model import ObjectIdentifierModel
+
+
+class StreamlitEntity(EntityBase, ObjectIdentifierModel(object_name="Streamlit")):  # type: ignore
+    type: Literal["streamlit"]  # noqa: A003
+    title: Optional[str] = Field(
+        title="Human-readable title for the Streamlit dashboard", default=None
+    )
+    query_warehouse: str = Field(
+        title="Snowflake warehouse to host the app", default=None
+    )
+    main_file: Optional[str] = Field(
+        title="Entrypoint file of the Streamlit app", default="streamlit_app.py"
+    )
+    pages_dir: Optional[str] = Field(title="Streamlit pages", default=None)
+    stage: Optional[str] = Field(
+        title="Stage in which the app’s artifacts will be stored", default="streamlit"
+    )
+    # Possibly can be PathMapping
+    artifacts: Optional[List[Path]] = Field(
+        title="List of additional files which should be included into deployment artifacts",
+        default=None,
+    )
+
+    @model_validator(mode="after")
+    def main_file_must_be_in_artifacts(self):
+        if not self.artifacts:
+            return self
+
+        if Path(self.main_file) not in self.artifacts:
+            raise ValueError(
+                f"Specified main file {self.main_file} is not included in artifacts."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def artifacts_must_exists(self):
+        if not self.artifacts:
+            return self
+
+        for artefact in self.artifacts:
+            if not artefact.exists():
+                raise ValueError(
+                    f"Specified artefact {artefact} does not exist locally."
+                )
+
+        return self

src/snowflake/cli/api/project/schemas/project_definition.py

@@ -189,6 +189,9 @@ def _convert_env(
             return env
         return ProjectEnvironment(default_env=(env or {}), override_env={})
 
+    def get_entities_by_type(self, entity_type: str):
+        return {i: e for i, e in self.entities.items() if e.get_type() == entity_type}
+
 
 def build_project_definition(**data):
     """

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

@@ -16,6 +16,7 @@
 
 import logging
 from pathlib import Path
+from typing import Dict
 
 import click
 import typer
@@ -29,14 +30,18 @@
 from snowflake.cli.api.commands.project_initialisation import add_init_command
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.exceptions import NoProjectDefinitionError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
     CommandResult,
     MessageResult,
     SingleQueryResult,
 )
-from snowflake.cli.api.project.project_verification import assert_project_type
-from snowflake.cli.api.project.schemas.streamlit.streamlit import Streamlit
+from snowflake.cli.api.project.schemas.entities.streamlit_entity import StreamlitEntity
+from snowflake.cli.api.project.schemas.project_definition import (
+    ProjectDefinition,
+    ProjectDefinitionV2,
+)
 from snowflake.cli.plugins.object.command_aliases import (
     add_object_command_aliases,
     scope_option,
@@ -134,37 +139,37 @@ def streamlit_deploy(
     stage is used. If the specified stage does not exist, the command creates it.
     """
 
-    assert_project_type("streamlit")
+    pd = cli_context.project_definition
+    if not pd.meets_version_requirement("2"):
+        pd = _migrate_v1_streamlit_to_v2(pd)
 
-    streamlit: Streamlit = cli_context.project_definition.streamlit
-    if not streamlit:
-        return MessageResult("No streamlit were specified in project definition.")
+    streamlits: Dict[str, StreamlitEntity] = pd.get_entities_by_type(
+        entity_type="streamlit"
+    )
 
-    environment_file = streamlit.env_file
-    if environment_file and not Path(environment_file).exists():
-        raise ClickException(f"Provided file {environment_file} does not exist")
-    elif environment_file is None:
-        environment_file = "environment.yml"
+    if not streamlits:
+        raise NoProjectDefinitionError(
+            project_type="streamlit", project_file=cli_context.project_root
+        )
 
-    pages_dir = streamlit.pages_dir
-    if pages_dir and not Path(pages_dir).exists():
-        raise ClickException(f"Provided file {pages_dir} does not exist")
-    elif pages_dir is None:
-        pages_dir = "pages"
+    # TODO: fix in follow-up
+    if len(list(streamlits)) > 1:
+        raise ClickException(
+            "Currently only single streamlit entity per project is supported."
+        )
 
+    # Get first streamlit
+    streamlit: StreamlitEntity = streamlits[list(streamlits)[0]]
     streamlit_id = FQN.from_identifier_model(streamlit).using_context()
 
     url = StreamlitManager().deploy(
         streamlit_id=streamlit_id,
-        environment_file=Path(environment_file),
-        pages_dir=Path(pages_dir),
+        artifacts=streamlit.artifacts,
         stage_name=streamlit.stage,
-        main_file=Path(streamlit.main_file),
+        main_file=streamlit.main_file,
         replace=replace,
         query_warehouse=streamlit.query_warehouse,
-        additional_source_files=streamlit.additional_source_files,
         title=streamlit.title,
-        **options,
     )
 
     if open_:
@@ -173,6 +178,58 @@ def streamlit_deploy(
     return MessageResult(f"Streamlit successfully deployed and available under {url}")
 
 
+def _migrate_v1_streamlit_to_v2(pd: ProjectDefinition):
+    if not pd.streamlit:
+        raise NoProjectDefinitionError(
+            project_type="streamlit", project_file=cli_context.project_root
+        )
+
+    default_env_file = "environment.yml"
+    default_pages_dir = "pages"
+
+    # Process env file
+    environment_file = pd.streamlit.env_file
+    if environment_file and not Path(environment_file).exists():
+        raise ClickException(f"Provided file {environment_file} does not exist")
+    elif environment_file is None and Path(default_env_file).exists():
+        environment_file = default_env_file
+    # Process pages dir
+    pages_dir = pd.streamlit.pages_dir
+    if pages_dir and not Path(pages_dir).exists():
+        raise ClickException(f"Provided file {pages_dir} does not exist")
+    elif pages_dir is None and Path(default_pages_dir).exists():
+        pages_dir = default_pages_dir
+
+    # Build V2 definition
+    artefacts = [
+        pd.streamlit.main_file,
+        environment_file,
+        pages_dir,
+    ]
+    artefacts = [a for a in artefacts if a is not None]
+    if pd.streamlit.additional_source_files:
+        artefacts.extend(pd.streamlit.additional_source_files)
+
+    data = {
+        "definition_version": "2",
+        "entities": {
+            "streamlit_app": {
+                "type": "streamlit",
+                "name": pd.streamlit.name,
+                "schema": pd.streamlit.schema_name,
+                "database": pd.streamlit.database,
+                "title": pd.streamlit.title,
+                "query_warehouse": pd.streamlit.query_warehouse,
+                "main_file": str(pd.streamlit.main_file),
+                "pages_dir": str(pd.streamlit.pages_dir),
+                "stage": pd.streamlit.stage,
+                "artifacts": artefacts,
+            }
+        },
+    }
+    return ProjectDefinitionV2(**data)
+
+
 @app.command("get-url", requires_connection=True)
 def get_url(
     name: FQN = StreamlitNameArgument,

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

@@ -15,7 +15,6 @@
 from __future__ import annotations
 
 import logging
-import os
 from pathlib import Path
 from typing import List, Optional
 
@@ -45,33 +44,25 @@ def share(self, streamlit_name: FQN, to_role: str) -> SnowflakeCursor:
     def _put_streamlit_files(
         self,
         root_location: str,
-        main_file: Path,
-        environment_file: Optional[Path],
-        pages_dir: Optional[Path],
-        additional_source_files: Optional[List[Path]],
+        artifacts: Optional[List[Path]] = None,
     ):
+        if not artifacts:
+            return
         stage_manager = StageManager()
-
-        stage_manager.put(main_file, root_location, 4, True)
-
-        if environment_file and environment_file.exists():
-            stage_manager.put(environment_file, root_location, 4, True)
-
-        if pages_dir and pages_dir.exists():
-            stage_manager.put(pages_dir / "*.py", f"{root_location}/pages", 4, True)
-
-        if additional_source_files:
-            for file in additional_source_files:
-                if os.sep in str(file):
-                    destination = f"{root_location}/{str(file.parent)}"
-                else:
-                    destination = root_location
-                stage_manager.put(file, destination, 4, True)
+        for file in artifacts:
+            if file.is_dir():
+                stage_manager.put(
+                    f"{file.joinpath('*')}", f"{root_location}/{file}", 4, True
+                )
+            elif len(file.parts) > 1:
+                stage_manager.put(file, f"{root_location}/{file.parent}", 4, True)
+            else:
+                stage_manager.put(file, root_location, 4, True)
 
     def _create_streamlit(
         self,
         streamlit_id: FQN,
-        main_file: Path,
+        main_file: str,
         replace: Optional[bool] = None,
         experimental: Optional[bool] = None,
         query_warehouse: Optional[str] = None,
@@ -95,7 +86,7 @@ def _create_streamlit(
         if from_stage_name:
             query.append(f"ROOT_LOCATION = '{from_stage_name}'")
 
-        query.append(f"MAIN_FILE = '{main_file.name}'")
+        query.append(f"MAIN_FILE = '{main_file}'")
 
         if query_warehouse:
             query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
@@ -107,15 +98,12 @@ def _create_streamlit(
     def deploy(
         self,
         streamlit_id: FQN,
-        main_file: Path,
-        environment_file: Optional[Path] = None,
-        pages_dir: Optional[Path] = None,
+        main_file: str,
+        artifacts: Optional[List[Path]] = None,
         stage_name: Optional[str] = None,
         query_warehouse: Optional[str] = None,
         replace: Optional[bool] = False,
-        additional_source_files: Optional[List[Path]] = None,
         title: Optional[str] = None,
-        **options,
     ):
         # for backwards compatibility - quoted stage path might be case-sensitive
         # https://docs.snowflake.com/en/sql-reference/identifiers-syntax#double-quoted-identifiers
@@ -156,10 +144,7 @@ def deploy(
 
             self._put_streamlit_files(
                 root_location,
-                main_file,
-                environment_file,
-                pages_dir,
-                additional_source_files,
+                artifacts,
             )
         else:
             """
@@ -178,13 +163,7 @@ def deploy(
                 f"{stage_name}/{streamlit_name_for_root_location}"
             )
 
-            self._put_streamlit_files(
-                root_location,
-                main_file,
-                environment_file,
-                pages_dir,
-                additional_source_files,
-            )
+            self._put_streamlit_files(root_location, artifacts)
 
             self._create_streamlit(
                 streamlit_id,

src/snowflake/cli/templates/default_streamlit/snowflake.yml

@@ -1,10 +1,15 @@
-definition_version: "1.1"
-streamlit:
-  name: streamlit_app
-  stage: my_streamlit_stage
-  query_warehouse: my_streamlit_warehouse
-  main_file: streamlit_app.py
-  env_file: environment.yml
-  pages_dir: pages/
-  additional_source_files:
-    - common/hello.py
+definition_version: "2"
+entities:
+  my_streamlit:
+    type: "streamlit"
+    name: "my_dashboard"
+    title: "My Streamlit Dashboard"
+    stage: "my_streamlit_stage"
+    query_warehouse: 'my_streamlit_warehouse'
+    main_file: streamlit_app.py
+    pages_dir: pages/
+    artifacts:
+      - streamlit_app.py
+      - environment.yml
+      - common/hello.py
+      - pages/

tests/app/test_telemetry.py

@@ -107,6 +107,9 @@ def test_executing_command_sends_telemetry_result_data(
 
 @mock.patch("snowflake.connector.connect")
 @mock.patch("snowflake.cli.plugins.streamlit.commands.StreamlitManager")
+@mock.patch.dict(
+    os.environ, {"SNOWFLAKE_CLI_FEATURES_ENABLE_PROJECT_DEFINITION_V2": "true"}
+)
 def test_executing_command_sends_project_definition_in_telemetry_data(
     _, mock_conn, project_directory, runner
 ):