Add SiS on SPCS vNext support in Snow Cli (#2531)

Author: sfc-gh-svishnu

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

@@ -8,6 +8,7 @@
 from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli._plugins.streamlit.manager import StreamlitManager
 from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
+    SPCS_RUNTIME_V2_NAME,
     StreamlitEntityModel,
 )
 from snowflake.cli._plugins.workspace.context import ActionContext
@@ -65,6 +66,14 @@ def action_get_url(
             self._conn, f"/#/streamlit-apps/{name.url_identifier}"
         )
 
+    def _is_spcs_runtime_v2_mode(self, experimental: bool = False) -> bool:
+        """Check if SPCS runtime v2 mode is enabled."""
+        return (
+            experimental
+            and self.model.runtime_name == SPCS_RUNTIME_V2_NAME
+            and self.model.compute_pool
+        )
+
     def bundle(self, output_dir: Optional[Path] = None) -> BundleMap:
         return build_bundle(
             self.root,
@@ -84,7 +93,7 @@ def deploy(
         replace: bool,
         prune: bool = False,
         bundle_map: Optional[BundleMap] = None,
-        experimental: Optional[bool] = False,
+        experimental: bool = False,
         *args,
         **kwargs,
     ):
@@ -130,7 +139,11 @@ def deploy(
             console.step(f"Creating Streamlit object {self.model.fqn.sql_identifier}")
 
             self._execute_query(
-                self.get_deploy_sql(replace=replace, from_stage_name=stage_root)
+                self.get_deploy_sql(
+                    replace=replace,
+                    from_stage_name=stage_root,
+                    experimental=False,
+                )
             )
 
             StreamlitManager(connection=self._conn).grant_privileges(self.model)
@@ -159,6 +172,7 @@ def get_deploy_sql(
         artifacts_dir: Optional[Path] = None,
         schema: Optional[str] = None,
         database: Optional[str] = None,
+        experimental: bool = False,
         *args,
         **kwargs,
     ) -> str:
@@ -202,6 +216,12 @@ def get_deploy_sql(
         if self.model.secrets:
             query += "\n" + self.model.get_secrets_sql()
 
+        # SPCS runtime fields are only supported for FBE/versioned streamlits (FROM syntax)
+        # Never add these fields for stage-based deployments (ROOT_LOCATION syntax)
+        if not from_stage_name and self._is_spcs_runtime_v2_mode(experimental):
+            query += f"\nRUNTIME_NAME = '{self.model.runtime_name}'"
+            query += f"\nCOMPUTE_POOL = '{self.model.compute_pool}'"
+
         return query + ";"
 
     def get_describe_sql(self) -> str:
@@ -236,6 +256,7 @@ def _deploy_experimental(
             self.get_deploy_sql(
                 if_not_exists=True,
                 replace=replace,
+                experimental=True,
             )
         )
         try:

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

@@ -15,17 +15,18 @@
 
 from typing import Literal, Optional
 
-from pydantic import Field
+from pydantic import Field, model_validator
 from snowflake.cli.api.project.schemas.entities.common import (
     Artifacts,
     EntityModelBaseWithArtifacts,
     ExternalAccessBaseModel,
     GrantBaseModel,
     ImportsBaseModel,
 )
-from snowflake.cli.api.project.schemas.updatable_model import (
-    DiscriminatorField,
-)
+from snowflake.cli.api.project.schemas.updatable_model import DiscriminatorField
+
+# SPCS Runtime v2 constants
+SPCS_RUNTIME_V2_NAME = "SYSTEM$ST_CONTAINER_RUNTIME_PY3_11"
 
 
 class StreamlitEntityModel(
@@ -54,3 +55,22 @@ class StreamlitEntityModel(
         title="List of paths or file source/destination pairs to add to the deploy root",
         default=None,
     )
+    runtime_name: Optional[str] = Field(
+        title="The runtime name to run the streamlit app on", default=None
+    )
+    compute_pool: Optional[str] = Field(
+        title="The compute pool name of the snowservices running the streamlit app",
+        default=None,
+    )
+
+    @model_validator(mode="after")
+    def validate_spcs_runtime_fields(self):
+        """Validate that runtime_name and compute_pool are provided together for SPCS container runtime."""
+        # Only validate for SPCS container runtime, not warehouse runtime
+        if self.compute_pool and not self.runtime_name:
+            raise ValueError("compute_pool is specified without runtime_name")
+        if self.runtime_name == SPCS_RUNTIME_V2_NAME and not self.compute_pool:
+            raise ValueError(
+                f"compute_pool is required when using {SPCS_RUNTIME_V2_NAME}"
+            )
+        return self

src/snowflake/cli/api/feature_flags.py

@@ -73,3 +73,6 @@ class FeatureFlag(FeatureFlagMixin):
     # TODO 4.0: remove ENABLE_RELEASE_CHANNELS
     ENABLE_RELEASE_CHANNELS = BooleanFlag("ENABLE_RELEASE_CHANNELS", None)
     ENABLE_SNOWFLAKE_PROJECTS = BooleanFlag("ENABLE_SNOWFLAKE_PROJECTS", False)
+    ENABLE_STREAMLIT_SPCS_RUNTIME_V2 = BooleanFlag(
+        "ENABLE_STREAMLIT_SPCS_RUNTIME_V2", False
+    )

tests/streamlit/test_streamlit_entity.py

@@ -1,6 +1,12 @@
+from pathlib import Path
 from unittest import mock
 
 import pytest
+from snowflake.cli._plugins.streamlit.streamlit_entity import StreamlitEntity
+from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
+    SPCS_RUNTIME_V2_NAME,
+    StreamlitEntityModel,
+)
 
 from tests.streamlit.streamlit_test_class import STREAMLIT_NAME, StreamlitTestClass
 
@@ -112,3 +118,212 @@ def test_if_attribute_is_not_set_correct_error_is_raised(
         with pytest.raises(ValueError) as e:
             result = getattr(example_entity, attribute)
         assert str(e.value) == f"Could not determine {attribute} for {STREAMLIT_NAME}"
+
+    def test_spcs_runtime_v2_model_fields(self, workspace_context):
+        """Test that StreamlitEntityModel accepts runtime_name and compute_pool fields"""
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name=SPCS_RUNTIME_V2_NAME,
+            compute_pool="MYPOOL",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        model.set_entity_id("test_streamlit")
+
+        assert model.runtime_name == SPCS_RUNTIME_V2_NAME
+        assert model.compute_pool == "MYPOOL"
+
+    def test_get_deploy_sql_with_spcs_runtime_v2(self, workspace_context):
+        """Test that get_deploy_sql includes RUNTIME_NAME and COMPUTE_POOL when experimental is True"""
+
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name=SPCS_RUNTIME_V2_NAME,
+            compute_pool="MYPOOL",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        model.set_entity_id("test_streamlit")
+
+        entity = StreamlitEntity(workspace_ctx=workspace_context, entity_model=model)
+
+        # Test with FROM syntax (artifacts_dir provided)
+        sql = entity.get_deploy_sql(
+            artifacts_dir=Path("/tmp/artifacts"), experimental=True
+        )
+
+        assert f"RUNTIME_NAME = '{SPCS_RUNTIME_V2_NAME}'" in sql
+        assert "COMPUTE_POOL = 'MYPOOL'" in sql
+
+    def test_get_deploy_sql_spcs_runtime_v2_with_stage(self, workspace_context):
+        """Test that SPCS runtime v2 clauses are NOT added with stage-based deployment (old-style streamlits)"""
+
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name=SPCS_RUNTIME_V2_NAME,
+            compute_pool="MYPOOL",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        model.set_entity_id("test_streamlit")
+
+        entity = StreamlitEntity(workspace_ctx=workspace_context, entity_model=model)
+
+        # Test with stage-based deployment - should NOT include SPCS runtime fields
+        # even when experimental=True, as stage-based deployments are old-style
+        sql = entity.get_deploy_sql(from_stage_name="@stage/path", experimental=True)
+
+        assert "ROOT_LOCATION = '@stage/path'" in sql
+        assert "RUNTIME_NAME" not in sql
+        assert "COMPUTE_POOL" not in sql
+
+    def test_get_deploy_sql_without_spcs_runtime_v2(self, workspace_context):
+        """Test that get_deploy_sql works normally when experimental is False"""
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name=SPCS_RUNTIME_V2_NAME,
+            compute_pool="MYPOOL",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        model.set_entity_id("test_streamlit")
+
+        entity = StreamlitEntity(workspace_ctx=workspace_context, entity_model=model)
+
+        # Test without experimental flag enabled
+        sql = entity.get_deploy_sql(
+            artifacts_dir=Path("/tmp/artifacts"), experimental=False
+        )
+
+        assert "RUNTIME_NAME" not in sql
+        assert "COMPUTE_POOL" not in sql
+
+    def test_spcs_runtime_v2_requires_correct_runtime_name(self, workspace_context):
+        """Test that SPCS runtime v2 requires correct runtime name to be enabled"""
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name=SPCS_RUNTIME_V2_NAME,
+            compute_pool="MYPOOL",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        model.set_entity_id("test_streamlit")
+
+        entity = StreamlitEntity(workspace_ctx=workspace_context, entity_model=model)
+
+        # Test with experimental=True and correct runtime_name
+        sql = entity.get_deploy_sql(
+            artifacts_dir=Path("/tmp/artifacts"), experimental=True
+        )
+        assert f"RUNTIME_NAME = '{SPCS_RUNTIME_V2_NAME}'" in sql
+        assert "COMPUTE_POOL = 'MYPOOL'" in sql
+
+        # Test with experimental=False, should not add SPCS fields
+        sql = entity.get_deploy_sql(
+            artifacts_dir=Path("/tmp/artifacts"), experimental=False
+        )
+        assert "RUNTIME_NAME" not in sql
+        assert "COMPUTE_POOL" not in sql
+
+        # Test with wrong runtime_name
+        model.runtime_name = "SOME_OTHER_RUNTIME"
+        entity = StreamlitEntity(workspace_ctx=workspace_context, entity_model=model)
+        sql = entity.get_deploy_sql(
+            artifacts_dir=Path("/tmp/artifacts"), experimental=True
+        )
+        assert "RUNTIME_NAME" not in sql
+        assert "COMPUTE_POOL" not in sql
+
+    def test_spcs_runtime_v2_requires_runtime_and_pool(self, workspace_context):
+        """Test that SPCS runtime v2 SQL generation works with valid models"""
+
+        # Test with valid container runtime and compute_pool
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name=SPCS_RUNTIME_V2_NAME,
+            compute_pool="MYPOOL",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        model.set_entity_id("test_streamlit")
+        entity = StreamlitEntity(workspace_ctx=workspace_context, entity_model=model)
+        sql = entity.get_deploy_sql(
+            artifacts_dir=Path("/tmp/artifacts"), experimental=True
+        )
+        assert f"RUNTIME_NAME = '{SPCS_RUNTIME_V2_NAME}'" in sql
+        assert "COMPUTE_POOL = 'MYPOOL'" in sql
+
+        # Test with warehouse runtime (no compute_pool needed)
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name="SYSTEM$WAREHOUSE_RUNTIME",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        model.set_entity_id("test_streamlit")
+        entity = StreamlitEntity(workspace_ctx=workspace_context, entity_model=model)
+        sql = entity.get_deploy_sql(
+            artifacts_dir=Path("/tmp/artifacts"), experimental=True
+        )
+        # Warehouse runtime should not trigger SPCS runtime v2 mode
+        assert "RUNTIME_NAME" not in sql
+        assert "COMPUTE_POOL" not in sql
+
+    def test_spcs_runtime_validation(self, workspace_context):
+        """Test validation for SPCS runtime configuration"""
+
+        # Test: SYSTEM$ST_CONTAINER_RUNTIME_PY3_11 requires compute_pool
+        escaped_runtime_name = SPCS_RUNTIME_V2_NAME.replace("$", r"\$")
+        with pytest.raises(
+            ValueError,
+            match=rf"compute_pool is required when using {escaped_runtime_name}",
+        ):
+            StreamlitEntityModel(
+                type="streamlit",
+                identifier="test_streamlit",
+                runtime_name=SPCS_RUNTIME_V2_NAME,
+                main_file="streamlit_app.py",
+                artifacts=["streamlit_app.py"],
+            )
+
+        # Test: compute_pool without runtime_name is invalid
+        with pytest.raises(
+            ValueError, match="compute_pool is specified without runtime_name"
+        ):
+            StreamlitEntityModel(
+                type="streamlit",
+                identifier="test_streamlit",
+                compute_pool="MYPOOL",
+                main_file="streamlit_app.py",
+                artifacts=["streamlit_app.py"],
+            )
+
+        # Test: warehouse runtime without compute_pool is valid
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name="SYSTEM$WAREHOUSE_RUNTIME",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        assert model.runtime_name == "SYSTEM$WAREHOUSE_RUNTIME"
+        assert model.compute_pool is None
+
+        # Test: container runtime with compute_pool is valid
+        model = StreamlitEntityModel(
+            type="streamlit",
+            identifier="test_streamlit",
+            runtime_name=SPCS_RUNTIME_V2_NAME,
+            compute_pool="MYPOOL",
+            main_file="streamlit_app.py",
+            artifacts=["streamlit_app.py"],
+        )
+        assert model.runtime_name == SPCS_RUNTIME_V2_NAME
+        assert model.compute_pool == "MYPOOL"

tests_integration/test_data/projects/streamlit_spcs_v2/pages/my_page.py

@@ -0,0 +1,30 @@
+import streamlit as st
+from utils.utils import hello_world
+
+st.title("Simple Page")
+st.write("This page demonstrates basic Streamlit functionality in SPCS runtime v2.")
+
+# Simple content
+st.header("Page Content")
+st.write(hello_world())
+
+# Simple interactive elements
+st.header("Interactive Elements")
+name = st.text_input("Enter your name:", "World")
+st.write(f"Hello, {name}!")
+
+# Simple columns layout
+col1, col2 = st.columns(2)
+
+with col1:
+    st.subheader("Column 1")
+    st.write("This is the first column")
+
+with col2:
+    st.subheader("Column 2")
+    st.write("This is the second column")
+
+# Simple metrics
+st.header("Simple Metrics")
+st.metric("Test Metric", 42)
+st.metric("Another Metric", 100, delta=10)

tests_integration/test_data/projects/streamlit_spcs_v2/requirements.txt

@@ -0,0 +1 @@
+# No additional requirements needed for basic Streamlit testing