From cfb1e68e82623eec5ba6ff234ce224eac2324e3d Mon Sep 17 00:00:00 2001
From: Jan Sikorski <jan.sikorski@snowflake.com>
Date: Tue, 27 Aug 2024 13:23:58 +0200
Subject: [PATCH] fix

---
 .../cli/_plugins/snowpark/commands.py         |  3 ++
 .../cli/_plugins/workspace/commands.py        |  1 +
 .../snowpark_templated_v1/snowflake.yml       |  2 ++
 .../streamlit_templated_v1/snowflake.yml      |  2 ++
 .../workspace/__snapshots__/test_manager.ambr | 33 +++++++++++++++----
 5 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/src/snowflake/cli/_plugins/snowpark/commands.py b/src/snowflake/cli/_plugins/snowpark/commands.py
index c2778d1cc5..c16d348d50 100644
--- a/src/snowflake/cli/_plugins/snowpark/commands.py
+++ b/src/snowflake/cli/_plugins/snowpark/commands.py
@@ -497,6 +497,9 @@ def migrate_v1_snowpark_to_v2(pd: ProjectDefinition):
 
         data["entities"][entity_name] = v2_entity
 
+        if pd.env:
+            data["env"] = {k: v for k, v in pd.env.items()}
+
     return ProjectDefinitionV2(**data)
 
 
diff --git a/src/snowflake/cli/_plugins/workspace/commands.py b/src/snowflake/cli/_plugins/workspace/commands.py
index 275e55924e..3a225ef51f 100644
--- a/src/snowflake/cli/_plugins/workspace/commands.py
+++ b/src/snowflake/cli/_plugins/workspace/commands.py
@@ -18,6 +18,7 @@
 
 import typer
 import yaml
+from click import ClickException
 from snowflake.cli._plugins.nativeapp.artifacts import BundleMap
 from snowflake.cli._plugins.snowpark.commands import migrate_v1_snowpark_to_v2
 from snowflake.cli._plugins.streamlit.commands import migrate_v1_streamlit_to_v2
diff --git a/tests/test_data/projects/snowpark_templated_v1/snowflake.yml b/tests/test_data/projects/snowpark_templated_v1/snowflake.yml
index e9b534f269..25884d6304 100644
--- a/tests/test_data/projects/snowpark_templated_v1/snowflake.yml
+++ b/tests/test_data/projects/snowpark_templated_v1/snowflake.yml
@@ -21,3 +21,5 @@ snowpark:
       handler: "procedures.test_procedure"
       signature: ""
       returns: string
+env:
+  project_source: "app/"
diff --git a/tests/test_data/projects/streamlit_templated_v1/snowflake.yml b/tests/test_data/projects/streamlit_templated_v1/snowflake.yml
index f805dca950..0d3d1253a1 100644
--- a/tests/test_data/projects/streamlit_templated_v1/snowflake.yml
+++ b/tests/test_data/projects/streamlit_templated_v1/snowflake.yml
@@ -9,3 +9,5 @@ streamlit:
   title: <% ctx.env.streamlit_title %>
   additional_source_files:
     - common/hello.py
+env:
+  streamlit_title: "My Fancy Streamlit"
diff --git a/tests/workspace/__snapshots__/test_manager.ambr b/tests/workspace/__snapshots__/test_manager.ambr
index 721f9194fe..647786381e 100644
--- a/tests/workspace/__snapshots__/test_manager.ambr
+++ b/tests/workspace/__snapshots__/test_manager.ambr
@@ -1,19 +1,20 @@
 # serializer version: 1
 # name: test_if_template_is_not_rendered_during_migration_with_option_checked[snowpark_templated_v1]
   '''
-  defaults:
-    project_name: <! project_name | to_snowflake_identifier !>
-    stage: <! stage | to_snowflake_identifier !>
   definition_version: '2'
   entities:
     hello_function:
       artifacts:
-      - <% ctx.env.project_source %>
+      - dest: <! project_name | to_snowflake_identifier !>
+        src: <% ctx.env.project_source %>
       external_access_integrations: []
       handler: functions.hello_function
       identifier:
         name: hello_function
       imports: []
+      meta:
+        use_mixins:
+        - snowpark_shared
       returns: string
       secrets: {}
       signature:
@@ -23,13 +24,17 @@
       type: function
     hello_procedure:
       artifacts:
-      - <% ctx.env.project_source %>
+      - dest: <! project_name | to_snowflake_identifier !>
+        src: <% ctx.env.project_source %>
       execute_as_caller: false
       external_access_integrations: []
       handler: procedures.hello_procedure
       identifier:
         name: hello_procedure
       imports: []
+      meta:
+        use_mixins:
+        - snowpark_shared
       returns: string
       secrets: {}
       signature:
@@ -39,18 +44,30 @@
       type: procedure
     test_procedure:
       artifacts:
-      - <% ctx.env.project_source %>
+      - dest: <! project_name | to_snowflake_identifier !>
+        src: <% ctx.env.project_source %>
       execute_as_caller: false
       external_access_integrations: []
       handler: procedures.test_procedure
       identifier:
         name: test_procedure
       imports: []
+      meta:
+        use_mixins:
+        - snowpark_shared
       returns: string
       secrets: {}
       signature: ''
       stage: <! stage | to_snowflake_identifier !>
       type: procedure
+  env:
+    project_source: app/
+  mixins:
+    snowpark_shared:
+      artifacts:
+      - dest: <! project_name | to_snowflake_identifier !>
+        src: <% ctx.env.project_source %>
+      stage: <! stage | to_snowflake_identifier !>
   
   '''
 # ---
@@ -79,6 +96,8 @@
         handler: "procedures.test_procedure"
         signature: ""
         returns: string
+  env:
+    project_source: "app/"
   
   '''
 # ---
@@ -116,6 +135,8 @@
     title: <% ctx.env.streamlit_title %>
     additional_source_files:
       - common/hello.py
+  env:
+    streamlit_title: "My Fancy Streamlit"
   
   '''
 # ---