From cf613fff4eea0dcab5dff08c23bc1803a16f5235 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nuno=20Gon=C3=A7alves?= <ngoncalves@singlestore.com>
Date: Mon, 9 Dec 2024 21:27:13 +0000
Subject: [PATCH] feat(magic): use tempfile.TemporaryDirectory for downloaded
 file

---
 singlestoredb/magics/run_personal.py | 27 +++++++++++++--------------
 singlestoredb/magics/run_shared.py   | 20 +++++++++-----------
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/singlestoredb/magics/run_personal.py b/singlestoredb/magics/run_personal.py
index 6c624b8f..a83e7b1a 100644
--- 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
index 393ca4bd..c021e894 100644
--- 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}"')