SNOW-2037785 Add runtime_environment_version for Notebook (#2478)

* SNOW-2037785 Add `runtime_environment_version` for Notebook

* verify `runtime_environment_version` via SQL

Author: sfc-gh-daniszewski

RELEASE-NOTES.md

@@ -20,6 +20,7 @@
 
 ## New additions
 * Add `snow connection remove` command
+* Added support for `runtime_environment_version` field in notebook entity configuration, allowing specification of runtime environment version for containerized notebooks.
 
 ## Fixes and improvements
 * Fixed DBT deploy command to properly handle fully qualified names

src/snowflake/cli/_plugins/notebook/notebook_entity.py

@@ -60,6 +60,8 @@ def get_create_sql(self, replace: bool) -> str:
             query += f"\nCOMPUTE_POOL = '{self.model.compute_pool}'"
         if self.model.runtime_name:
             query += f"\nRUNTIME_NAME = '{self.model.runtime_name}'"
+        if self.model.runtime_environment_version and not self.model.compute_pool:
+            query += f"\nRUNTIME_ENVIRONMENT_VERSION = '{self.model.runtime_environment_version}'"
 
         query += (
             ";\n// Cannot use IDENTIFIER(...)"

src/snowflake/cli/_plugins/notebook/notebook_entity_model.py

@@ -20,6 +20,9 @@ class NotebookEntityModel(EntityModelBaseWithArtifacts):
     )
     notebook_file: Path = Field(title="Notebook file")
     query_warehouse: str = Field(title="Snowflake warehouse to execute the notebook")
+    runtime_environment_version: Optional[str] = Field(
+        title="Runtime environment version", default=None
+    )
     compute_pool: Optional[str] = Field(
         title="Compute pool to run the notebook in", default=None
     )
@@ -37,6 +40,10 @@ def validate_notebook_file(self):
     def validate_container_setup(self):
         if self.compute_pool and not self.runtime_name:
             raise ValueError("compute_pool is specified without runtime_name")
-        if self.runtime_name and not self.compute_pool and not self:
+        if self.runtime_name and not self.compute_pool:
             raise ValueError("runtime_name is specified without compute_pool")
+        if self.compute_pool and self.runtime_environment_version:
+            raise ValueError(
+                "runtime_environment_version is only applicable for notebooks using warehouse, not compute pool"
+            )
         return self

tests/notebook/test_notebook_commands.py

@@ -225,3 +225,72 @@ def test_deploy_project_definition_version_error(
             "This command requires project definition of version at least 2."
             in result.output
         )
+
+
+@mock.patch("snowflake.connector.connect")
+@mock.patch("snowflake.cli._plugins.notebook.notebook_entity.make_snowsight_url")
+@mock.patch(f"{STAGE_MANAGER}.list_files")
+def test_deploy_notebook_with_runtime_environment_version_sql(
+    mock_list_files,
+    mock_make_url,
+    mock_connector,
+    mock_ctx,
+    runner,
+    project_directory,
+):
+    """Test that runtime_environment_version field generates correct SQL for warehouse-only notebooks."""
+    ctx = mock_ctx()
+    mock_connector.return_value = ctx
+    mock_make_url.return_value = "http://the.notebook.url.mock"
+
+    with project_directory("notebook_with_runtime_env"):
+        result = runner.invoke(
+            ["notebook", "deploy", "warehouse_notebook", "--replace"]
+        )
+        assert result.exit_code == 0, result.output
+
+        # Get the generated SQL and verify it contains RUNTIME_ENVIRONMENT_VERSION
+        query = ctx.get_query()
+        assert "RUNTIME_ENVIRONMENT_VERSION = 'WH-RUNTIME-2.0'" in query
+
+        # Verify it does NOT contain compute pool fields (should be warehouse only)
+        assert "COMPUTE_POOL" not in query
+        assert "RUNTIME_NAME" not in query
+
+        # Verify the full SQL structure is correct
+        lines = [line for line in query.split("\n") if not line.startswith("put")]
+        sql_content = "\n".join(lines)
+        assert "CREATE OR REPLACE NOTEBOOK" in sql_content
+        assert "QUERY_WAREHOUSE = 'xsmall'" in sql_content
+        assert "RUNTIME_ENVIRONMENT_VERSION = 'WH-RUNTIME-2.0'" in sql_content
+
+
+def test_deploy_notebook_validation_error_compute_pool_with_runtime_env(
+    runner, project_directory, alter_snowflake_yml
+):
+    """Test that using runtime_environment_version with compute_pool raises validation error."""
+    with project_directory("notebook_v2") as project_root:
+        # Add compute_pool, runtime_name, and runtime_environment_version to trigger validation error
+        alter_snowflake_yml(
+            project_root / "snowflake.yml",
+            parameter_path="entities.my_notebook.compute_pool",
+            value="test_compute_pool",
+        )
+        alter_snowflake_yml(
+            project_root / "snowflake.yml",
+            parameter_path="entities.my_notebook.runtime_name",
+            value="system$basic_runtime",
+        )
+        alter_snowflake_yml(
+            project_root / "snowflake.yml",
+            parameter_path="entities.my_notebook.runtime_environment_version",
+            value="WH-RUNTIME-2.0",
+        )
+
+        result = runner.invoke(["notebook", "deploy", "--replace"])
+        assert result.exit_code == 1, result.output
+        assert (
+            "runtime_environment_version is only applicable for notebooks"
+            in result.output
+        )
+        assert "using warehouse, not compute pool" in result.output

tests/test_data/projects/notebook_with_runtime_env/notebook.ipynb

@@ -0,0 +1,20 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "print(\"Hello world!\")\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 2
+}

tests/test_data/projects/notebook_with_runtime_env/snowflake.yml

@@ -0,0 +1,9 @@
+definition_version: 2
+entities:
+  warehouse_notebook:
+    type: "notebook"
+    query_warehouse: xsmall
+    notebook_file: notebook.ipynb
+    runtime_environment_version: "WH-RUNTIME-2.0"
+    artifacts:
+     - notebook.ipynb 

tests_integration/test_data/projects/notebook_with_runtime_env/notebook.ipynb

@@ -0,0 +1,20 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "print(\"Hello world!\")\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 2
+}

tests_integration/test_data/projects/notebook_with_runtime_env/snowflake.yml

@@ -0,0 +1,9 @@
+definition_version: 2
+entities:
+  warehouse_notebook:
+    type: "notebook"
+    query_warehouse: xsmall
+    notebook_file: notebook.ipynb
+    runtime_environment_version: "WH-RUNTIME-2.0"
+    artifacts:
+     - notebook.ipynb 

tests_integration/test_notebooks.py

@@ -96,3 +96,79 @@ def test_deploy_single_notebook(runner, project_directory, test_database):
                 "custom_stage/particular_notebook_path/notebook.ipynb",
             ],
         )
+
+
+@pytest.mark.integration
+def test_deploy_containerized_notebook_with_compute_pool(
+    runner, project_directory, test_database
+):
+    """Test deployment of containerized notebook with compute pool."""
+    with project_directory("notebook_containerized_v2"):
+        # Deploy the containerized notebook with compute pool
+        result = runner.invoke_with_connection(
+            ["notebook", "deploy", "containerized_notebook", "--replace"]
+        )
+        assert result.exit_code == 0, result.output
+        assert (
+            "Uploading artifacts to @notebooks/containerized_notebook" in result.output
+        )
+        assert "Notebook successfully deployed and available under" in result.output
+
+        # Verify the notebook was created with the correct artifacts
+        assert_stage_has_files(
+            runner,
+            "notebooks",
+            ["notebooks/containerized_notebook/notebook.ipynb"],
+        )
+
+        # Test that the notebook can be deployed again with --replace
+        result = runner.invoke_with_connection(
+            ["notebook", "deploy", "containerized_notebook", "--replace"]
+        )
+        assert result.exit_code == 0, result.output
+        assert "Notebook successfully deployed and available under" in result.output
+
+
+@pytest.mark.integration
+def test_deploy_notebook_with_runtime_environment_version(
+    runner, project_directory, test_database
+):
+    """Test deployment of notebook with runtime_environment_version field for warehouse (no compute pool)."""
+    with project_directory("notebook_with_runtime_env"):
+        # Deploy the notebook with runtime_environment_version
+        result = runner.invoke_with_connection(
+            ["notebook", "deploy", "warehouse_notebook", "--replace"]
+        )
+        assert result.exit_code == 0, result.output
+        assert "Uploading artifacts to @notebooks/warehouse_notebook" in result.output
+        assert "Notebook successfully deployed and available under" in result.output
+
+        # Verify the notebook was created with the correct artifacts
+        assert_stage_has_files(
+            runner,
+            "notebooks",
+            ["notebooks/warehouse_notebook/notebook.ipynb"],
+        )
+
+        # Test that the notebook can be deployed again with --replace
+        result = runner.invoke_with_connection(
+            ["notebook", "deploy", "warehouse_notebook", "--replace"]
+        )
+        assert result.exit_code == 0, result.output
+        assert "Notebook successfully deployed and available under" in result.output
+
+        # Verify describe notebook shows runtime_environment_version is set using direct SQL
+        result = runner.invoke_with_connection_json(
+            ["sql", "-q", "DESCRIBE NOTEBOOK warehouse_notebook"]
+        )
+        assert result.exit_code == 0, result.output
+        # Check that runtime_environment_version appears in the describe result
+        describe_output = (
+            str(result.json).lower() if result.json else result.output.lower()
+        )
+        assert (
+            "runtime_environment_version" in describe_output
+        ), f"runtime_environment_version not found in describe output: {result.output}"
+        assert (
+            "wh-runtime-2.0" in describe_output
+        ), f"Expected runtime version 'WH-RUNTIME-2.0' not found in describe output: {result.output}"