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/singlestoredb/magics/run_personal.py b/singlestoredb/magics/run_personal.py new file mode 100644 index 000000000..a83e7b1a2 --- /dev/null +++ b/singlestoredb/magics/run_personal.py @@ -0,0 +1,56 @@ +import os +import tempfile +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 +from jinja2 import Template + + +@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 + + %run_personal {{ sample_notebook_name }} + + """ + + 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("'")) or \ + (personal_file.startswith('"') and personal_file.endswith('"')): + personal_file = personal_file[1:-1] + if not personal_file: + raise ValueError('No personal file specified.') + + 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 new file mode 100644 index 000000000..c021e894c --- /dev/null +++ b/singlestoredb/magics/run_shared.py @@ -0,0 +1,53 @@ +import os +import tempfile +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 +from jinja2 import Template + + +@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 + + %run_shared {{ sample_notebook_name }} + + """ + + 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("'")) or \ + (shared_file.startswith('"') and shared_file.endswith('"')): + shared_file = shared_file[1:-1] + if not shared_file: + raise ValueError('No personal file specified.') + + 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}'" + + # 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}"')