test

singlestore-labs · Dec 9, 2024 · 5b1fe51 · 5b1fe51
1 parent c89f411
commit 5b1fe51
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 25 deletions.
diff --git a/singlestoredb/magics/run_personal.py b/singlestoredb/magics/run_personal.py
@@ -1,5 +1,5 @@
 import os
-import time
+import tempfile
 from typing import Any
 
 from IPython.core.interactiveshell import InteractiveShell
@@ -32,6 +32,7 @@ def run_personal(self, line: str, local_ns: Any = None) -> Any:
           %run_personal {{ sample_notebook_name }}
 
         """
+
         template = Template(line.strip())
         personal_file = template.render(local_ns)
         if not personal_file:
@@ -42,16 +43,14 @@ def run_personal(self, line: str, local_ns: Any = None) -> Any:
             if not personal_file:
                 raise ValueError('No personal file specified.')
 
-        local_filename = (
-            f'{int(time.time() * 1_000_000)}_{personal_file}'.replace(' ', '_')
-        )
-        sql_command = f"DOWNLOAD PERSONAL FILE '{personal_file}' TO '{local_filename}'"
-
-        # Execute the SQL command
-        self.shell.run_line_magic('sql', sql_command)
-        # Run the downloaded file
-        self.shell.run_line_magic('run', local_filename)
-
-        # Delete the local file after running it
-        if os.path.exists(local_filename):
-            os.remove(local_filename)
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_file_path = os.path.join(temp_dir, personal_file)
+            sql_command = (
+                f"DOWNLOAD PERSONAL FILE '{personal_file}' "
+                f"TO '{temp_file_path}'"
+            )
+
+            # Execute the SQL command
+            self.shell.run_line_magic('sql', sql_command)
+            # Run the downloaded file
+            self.shell.run_line_magic('run', f'"{temp_file_path}"')
diff --git a/singlestoredb/magics/run_shared.py b/singlestoredb/magics/run_shared.py
@@ -1,5 +1,5 @@
 import os
-import time
+import tempfile
 from typing import Any
 
 from IPython.core.interactiveshell import InteractiveShell
@@ -32,6 +32,7 @@ def run_shared(self, line: str, local_ns: Any = None) -> Any:
           %run_shared {{ sample_notebook_name }}
 
         """
+
         template = Template(line.strip())
         shared_file = template.render(local_ns)
         if not shared_file:
@@ -42,14 +43,11 @@ def run_shared(self, line: str, local_ns: Any = None) -> Any:
             if not shared_file:
                 raise ValueError('No personal file specified.')
 
-        local_filename = f'{int(time.time() * 1_000_000)}_{shared_file}'.replace(' ', '_')
-        sql_command = f"DOWNLOAD SHARED FILE '{shared_file}' TO '{local_filename}'"
-
-        # Execute the SQL command
-        self.shell.run_line_magic('sql', sql_command)
-        # Run the downloaded file
-        self.shell.run_line_magic('run', local_filename)
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_file_path = os.path.join(temp_dir, shared_file)
+            sql_command = f"DOWNLOAD SHARED FILE '{shared_file}' TO '{temp_file_path}'"
 
-        # Delete the local file after running it
-        if os.path.exists(local_filename):
-            os.remove(local_filename)
+            # Execute the SQL command
+            self.shell.run_line_magic('sql', sql_command)
+            # Run the downloaded file
+            self.shell.run_line_magic('run', f'"{temp_file_path}"')