From e12d0281c4a35704615ddd1e5c9d56bc49a5b39a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nuno=20Gon=C3=A7alves?= Date: Mon, 9 Dec 2024 10:28:40 +0000 Subject: [PATCH 1/4] feat(magic): add run_personal and run_shared magic commands --- run_personal/__init__.py | 4 ++ run_personal/magic.py | 82 ++++++++++++++++++++++++++++++++++++++++ run_shared/__init__.py | 4 ++ run_shared/magic.py | 79 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 169 insertions(+) create mode 100644 run_personal/__init__.py create mode 100644 run_personal/magic.py create mode 100644 run_shared/__init__.py create mode 100644 run_shared/magic.py diff --git a/run_personal/__init__.py b/run_personal/__init__.py new file mode 100644 index 000000000..71629f1e3 --- /dev/null +++ b/run_personal/__init__.py @@ -0,0 +1,4 @@ +from run_personal.magic import load_ipython_extension + + +__all__ = ['load_ipython_extension'] diff --git a/run_personal/magic.py b/run_personal/magic.py new file mode 100644 index 000000000..e5dc6bfa0 --- /dev/null +++ b/run_personal/magic.py @@ -0,0 +1,82 @@ +import os +import time +from typing import Any + +from IPython.core.interactiveshell import InteractiveShell +from IPython.core.magic import line_magic +from IPython.core.magic import Magics +from IPython.core.magic import magics_class +from IPython.core.magic import needs_local_scope +from IPython.core.magic import no_var_expand + + +@magics_class +class RunPersonalMagic(Magics): + def __init__(self, shell: InteractiveShell): + Magics.__init__(self, shell=shell) + + @no_var_expand + @needs_local_scope + @line_magic('run_personal') + def run_personal(self, line: str, local_ns: Any = None) -> Any: + """ + Downloads a personal file using the %sql magic and then runs it using %run. + + Examples:: + + # Line usage + + %run_personal personal_file.ipynb + + """ + personal_file = line.strip() + if not personal_file: + raise ValueError('No personal file specified.') + if personal_file.startswith("'") and personal_file.endswith("'"): + personal_file = personal_file[1:-1] + 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) + + +# In order to actually use these magics, you must register them with a +# running IPython. + + +def load_ipython_extension(ip: InteractiveShell) -> None: + """ + Any module file that define a function named `load_ipython_extension` + can be loaded via `%load_ext module.path` or be configured to be + autoloaded by IPython at startup time. + """ + + # Load jupysql extension + # This is necessary for jupysql to initialize internal state + # required to render messages + assert ip.extension_manager is not None + result = ip.extension_manager.load_extension('sql') + if result == 'no load function': + raise RuntimeError('Could not load sql extension. Is jupysql installed?') + + # Check if %run magic command is defined + if ip.find_line_magic('run') is None: + raise RuntimeError( + '%run magic command is not defined. ' + 'Is it available in your IPython environment?', + ) + + # Register run_personal + ip.register_magics(RunPersonalMagic(ip)) diff --git a/run_shared/__init__.py b/run_shared/__init__.py new file mode 100644 index 000000000..a08fabd95 --- /dev/null +++ b/run_shared/__init__.py @@ -0,0 +1,4 @@ +from run_shared.magic import load_ipython_extension + + +__all__ = ['load_ipython_extension'] diff --git a/run_shared/magic.py b/run_shared/magic.py new file mode 100644 index 000000000..87949bde6 --- /dev/null +++ b/run_shared/magic.py @@ -0,0 +1,79 @@ +import os +import time +from typing import Any + +from IPython.core.interactiveshell import InteractiveShell +from IPython.core.magic import line_magic +from IPython.core.magic import Magics +from IPython.core.magic import magics_class +from IPython.core.magic import needs_local_scope +from IPython.core.magic import no_var_expand + + +@magics_class +class RunSharedMagic(Magics): + def __init__(self, shell: InteractiveShell): + Magics.__init__(self, shell=shell) + + @no_var_expand + @needs_local_scope + @line_magic('run_shared') + def run_shared(self, line: str, local_ns: Any = None) -> Any: + """ + Downloads a shared file using the %sql magic and then runs it using %run. + + Examples:: + + # Line usage + + %run_shared shared_file.ipynb + + """ + shared_file = line.strip() + if not shared_file: + raise ValueError('No shared file specified.') + if shared_file.startswith("'") and shared_file.endswith("'"): + shared_file = shared_file[1:-1] + 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) + + # Delete the local file after running it + if os.path.exists(local_filename): + os.remove(local_filename) + +# In order to actually use these magics, you must register them with a +# running IPython. + + +def load_ipython_extension(ip: InteractiveShell) -> None: + """ + Any module file that define a function named `load_ipython_extension` + can be loaded via `%load_ext module.path` or be configured to be + autoloaded by IPython at startup time. + """ + + # Load jupysql extension + # This is necessary for jupysql to initialize internal state + # required to render messages + assert ip.extension_manager is not None + result = ip.extension_manager.load_extension('sql') + if result == 'no load function': + raise RuntimeError('Could not load sql extension. Is jupysql installed?') + + # Check if %run magic command is defined + if ip.find_line_magic('run') is None: + raise RuntimeError( + '%run magic command is not defined. ' + 'Is it available in your IPython environment?', + ) + + # Register run_shared + ip.register_magics(RunSharedMagic(ip)) From 7b46bd910f80253a405f8c2bb27f51947fec4c7b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nuno=20Gon=C3=A7alves?= Date: Mon, 9 Dec 2024 14:23:08 +0000 Subject: [PATCH 2/4] feat(magic): add support for jinja2 templates --- run_personal/magic.py | 9 +++++++-- run_shared/magic.py | 9 +++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/run_personal/magic.py b/run_personal/magic.py index e5dc6bfa0..209999748 100644 --- a/run_personal/magic.py +++ b/run_personal/magic.py @@ -8,6 +8,7 @@ from IPython.core.magic import magics_class from IPython.core.magic import needs_local_scope from IPython.core.magic import no_var_expand +from jinja2 import Template @magics_class @@ -28,11 +29,15 @@ def run_personal(self, line: str, local_ns: Any = None) -> Any: %run_personal personal_file.ipynb + %run_personal {{ sample_notebook_name }} + """ - personal_file = line.strip() + template = Template(line.strip()) + personal_file = template.render(local_ns) if not personal_file: raise ValueError('No personal file specified.') - if personal_file.startswith("'") and personal_file.endswith("'"): + if (personal_file.startswith("'") and personal_file.endswith("'")) or \ + (personal_file.startswith('"') and personal_file.endswith('"')): personal_file = personal_file[1:-1] if not personal_file: raise ValueError('No personal file specified.') diff --git a/run_shared/magic.py b/run_shared/magic.py index 87949bde6..442578f38 100644 --- a/run_shared/magic.py +++ b/run_shared/magic.py @@ -8,6 +8,7 @@ from IPython.core.magic import magics_class from IPython.core.magic import needs_local_scope from IPython.core.magic import no_var_expand +from jinja2 import Template @magics_class @@ -28,11 +29,15 @@ def run_shared(self, line: str, local_ns: Any = None) -> Any: %run_shared shared_file.ipynb + %run_shared {{ sample_notebook_name }} + """ - shared_file = line.strip() + template = Template(line.strip()) + shared_file = template.render(local_ns) if not shared_file: raise ValueError('No shared file specified.') - if shared_file.startswith("'") and shared_file.endswith("'"): + if (shared_file.startswith("'") and shared_file.endswith("'")) or \ + (shared_file.startswith('"') and shared_file.endswith('"')): shared_file = shared_file[1:-1] if not shared_file: raise ValueError('No personal file specified.') From c89f4113b138362b97e55a24ebad4575f210afe1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nuno=20Gon=C3=A7alves?= Date: Mon, 9 Dec 2024 15:21:36 +0000 Subject: [PATCH 3/4] refactor(magic): create singlestoredb.magics module --- run_personal/__init__.py | 4 --- run_shared/__init__.py | 4 --- singlestoredb/magics/__init__.py | 34 +++++++++++++++++++ .../magics/run_personal.py | 30 ---------------- .../magics/run_shared.py | 29 ---------------- 5 files changed, 34 insertions(+), 67 deletions(-) delete mode 100644 run_personal/__init__.py delete mode 100644 run_shared/__init__.py create mode 100644 singlestoredb/magics/__init__.py rename run_personal/magic.py => singlestoredb/magics/run_personal.py (65%) rename run_shared/magic.py => singlestoredb/magics/run_shared.py (64%) diff --git a/run_personal/__init__.py b/run_personal/__init__.py deleted file mode 100644 index 71629f1e3..000000000 --- a/run_personal/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from run_personal.magic import load_ipython_extension - - -__all__ = ['load_ipython_extension'] diff --git a/run_shared/__init__.py b/run_shared/__init__.py deleted file mode 100644 index a08fabd95..000000000 --- a/run_shared/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from run_shared.magic import load_ipython_extension - - -__all__ = ['load_ipython_extension'] diff --git a/singlestoredb/magics/__init__.py b/singlestoredb/magics/__init__.py new file mode 100644 index 000000000..174e182ff --- /dev/null +++ b/singlestoredb/magics/__init__.py @@ -0,0 +1,34 @@ +from IPython.core.interactiveshell import InteractiveShell + +from .run_personal import RunPersonalMagic +from .run_shared import RunSharedMagic + +# In order to actually use these magics, we must register them with a +# running IPython. + + +def load_ipython_extension(ip: InteractiveShell) -> None: + """ + Any module file that define a function named `load_ipython_extension` + can be loaded via `%load_ext module.path` or be configured to be + autoloaded by IPython at startup time. + """ + + # Load jupysql extension + # This is necessary for jupysql to initialize internal state + # required to render messages + assert ip.extension_manager is not None + result = ip.extension_manager.load_extension('sql') + if result == 'no load function': + raise RuntimeError('Could not load sql extension. Is jupysql installed?') + + # Check if %run magic command is defined + if ip.find_line_magic('run') is None: + raise RuntimeError( + '%run magic command is not defined. ' + 'Is it available in your IPython environment?', + ) + + # Register run_personal and run_shared + ip.register_magics(RunPersonalMagic(ip)) + ip.register_magics(RunSharedMagic(ip)) diff --git a/run_personal/magic.py b/singlestoredb/magics/run_personal.py similarity index 65% rename from run_personal/magic.py rename to singlestoredb/magics/run_personal.py index 209999748..6c624b8f3 100644 --- a/run_personal/magic.py +++ b/singlestoredb/magics/run_personal.py @@ -55,33 +55,3 @@ def run_personal(self, line: str, local_ns: Any = None) -> Any: # Delete the local file after running it if os.path.exists(local_filename): os.remove(local_filename) - - -# In order to actually use these magics, you must register them with a -# running IPython. - - -def load_ipython_extension(ip: InteractiveShell) -> None: - """ - Any module file that define a function named `load_ipython_extension` - can be loaded via `%load_ext module.path` or be configured to be - autoloaded by IPython at startup time. - """ - - # Load jupysql extension - # This is necessary for jupysql to initialize internal state - # required to render messages - assert ip.extension_manager is not None - result = ip.extension_manager.load_extension('sql') - if result == 'no load function': - raise RuntimeError('Could not load sql extension. Is jupysql installed?') - - # Check if %run magic command is defined - if ip.find_line_magic('run') is None: - raise RuntimeError( - '%run magic command is not defined. ' - 'Is it available in your IPython environment?', - ) - - # Register run_personal - ip.register_magics(RunPersonalMagic(ip)) diff --git a/run_shared/magic.py b/singlestoredb/magics/run_shared.py similarity index 64% rename from run_shared/magic.py rename to singlestoredb/magics/run_shared.py index 442578f38..393ca4bd6 100644 --- a/run_shared/magic.py +++ b/singlestoredb/magics/run_shared.py @@ -53,32 +53,3 @@ def run_shared(self, line: str, local_ns: Any = None) -> Any: # Delete the local file after running it if os.path.exists(local_filename): os.remove(local_filename) - -# In order to actually use these magics, you must register them with a -# running IPython. - - -def load_ipython_extension(ip: InteractiveShell) -> None: - """ - Any module file that define a function named `load_ipython_extension` - can be loaded via `%load_ext module.path` or be configured to be - autoloaded by IPython at startup time. - """ - - # Load jupysql extension - # This is necessary for jupysql to initialize internal state - # required to render messages - assert ip.extension_manager is not None - result = ip.extension_manager.load_extension('sql') - if result == 'no load function': - raise RuntimeError('Could not load sql extension. Is jupysql installed?') - - # Check if %run magic command is defined - if ip.find_line_magic('run') is None: - raise RuntimeError( - '%run magic command is not defined. ' - 'Is it available in your IPython environment?', - ) - - # Register run_shared - ip.register_magics(RunSharedMagic(ip)) From cf613fff4eea0dcab5dff08c23bc1803a16f5235 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nuno=20Gon=C3=A7alves?= Date: Mon, 9 Dec 2024 21:27:13 +0000 Subject: [PATCH 4/4] 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 6c624b8f3..a83e7b1a2 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 393ca4bd6..c021e894c 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}"')