Make proper put commands for additional_source_files that aren't in a…

… folder (#581)
snowflakedb · Dec 1, 2023 · 457e978 · 457e978
1 parent a2707d6
commit 457e978
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/src/snowcli/cli/streamlit/manager.py b/src/snowcli/cli/streamlit/manager.py
@@ -45,9 +45,14 @@ def _put_streamlit_files(
 
         if additional_source_files:
             for file in additional_source_files:
-                stage_manager.put(
-                    file, f"{root_location}/{str(Path(file).parent)}", 4, True
+                # If the file is in a folder, PUT it to the same folder in the stage
+                # If not, just PUT it to the root of the stage
+                destination = (
+                    f"{root_location}/{str(Path(file).parent)}"
+                    if "/" in file
+                    else root_location
                 )
+                stage_manager.put(file, destination, 4, True)
 
     def _create_streamlit(
         self,

diff --git a/tests/streamlit/test_commands.py b/tests/streamlit/test_commands.py
@@ -272,6 +272,7 @@ def test_deploy_all_streamlit_files(
         _put_query("environment.yml", root_path),
         _put_query("pages/*.py", f"{root_path}/pages"),
         _put_query("utils/utils.py", f"{root_path}/utils"),
+        _put_query("extra_file.py", root_path),
         dedent(
             f"""
             CREATE STREAMLIT {STREAMLIT_NAME}

diff --git a/tests/test_data/projects/streamlit_full_definition/snowflake.yml b/tests/test_data/projects/streamlit_full_definition/snowflake.yml
@@ -8,3 +8,4 @@ streamlit:
   pages_dir: pages
   additional_source_files:
     - utils/utils.py
+    - extra_file.py