Reset session in session_state if lost

app_utils/shared_utils.py

@@ -997,6 +997,9 @@ def upload_yaml(file_name: str) -> None:
         with open(tmp_file_path, "w", encoding='utf-8') as temp_file:
             temp_file.write(yaml)
 
+        if 'session' not in st.session_state:
+            set_snowpark_session(get_snowflake_connection())
+
         st.session_state.session.file.put(
             tmp_file_path,
             f"@{st.session_state.snowflake_stage.stage_name}",
@@ -1050,6 +1053,9 @@ def download_yaml(file_name: str, stage_name: str) -> str:
     import os
     import tempfile
 
+    if 'session' not in st.session_state:
+        set_snowpark_session(get_snowflake_connection())
+
     with tempfile.TemporaryDirectory() as temp_dir:
         # Downloads the YAML to {temp_dir}/{file_name}.
         st.session_state.session.file.get(