Fix %%sql cells in %run_shared magic

Author: kesmit13

singlestoredb/magics/run_shared.py

@@ -1,13 +1,17 @@
 import os
 import tempfile
+from pathlib import Path
 from typing import Any
+from warnings import warn
 
 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 IPython.utils.contexts import preserve_keys
+from IPython.utils.syspathcontext import prepended_to_syspath
 from jinja2 import Template
 
 
@@ -50,4 +54,72 @@ def run_shared(self, line: str, local_ns: Any = None) -> Any:
             # 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}"')
+            with preserve_keys(self.shell.user_ns, '__file__'):
+                self.shell.user_ns['__file__'] = temp_file_path
+                self.shell.safe_execfile_ipy(temp_file_path, raise_exceptions=True)
+
+    def safe_execfile_ipy(
+        self,
+        fname: str,
+        shell_futures: bool = False,
+        raise_exceptions: bool = False,
+    ) -> None:
+        """Like safe_execfile, but for .ipy or .ipynb files with IPython syntax.
+
+        Parameters
+        ----------
+        fname : str
+            The name of the file to execute.  The filename must have a
+            .ipy or .ipynb extension.
+        shell_futures : bool (False)
+            If True, the code will share future statements with the interactive
+            shell. It will both be affected by previous __future__ imports, and
+            any __future__ imports in the code will affect the shell. If False,
+            __future__ imports are not shared in either direction.
+        raise_exceptions : bool (False)
+            If True raise exceptions everywhere.  Meant for testing.
+        """
+        fpath = Path(fname).expanduser().resolve()
+
+        # Make sure we can open the file
+        try:
+            with fpath.open('rb'):
+                pass
+        except Exception:
+            warn('Could not open file <%s> for safe execution.' % fpath)
+            return
+
+        # Find things also in current directory.  This is needed to mimic the
+        # behavior of running a script from the system command line, where
+        # Python inserts the script's directory into sys.path
+        dname = str(fpath.parent)
+
+        def get_cells() -> Any:
+            """generator for sequence of code blocks to run"""
+            if fpath.suffix == '.ipynb':
+                from nbformat import read
+                nb = read(fpath, as_version=4)
+                if not nb.cells:
+                    return
+                for cell in nb.cells:
+                    if cell.cell_type == 'code':
+                        if getattr(cell, 'metadata', {}).get('language', '') == 'sql':
+                            yield f'%%sql\n{cell.source}'
+                        else:
+                            yield cell.source
+            else:
+                yield fpath.read_text(encoding='utf-8')
+
+        with prepended_to_syspath(dname):
+            try:
+                for cell in get_cells():
+                    result = self.run_cell(cell, silent=True, shell_futures=shell_futures)
+                    if raise_exceptions:
+                        result.raise_error()
+                    elif not result.success:
+                        break
+            except Exception:
+                if raise_exceptions:
+                    raise
+                self.showtraceback()
+                warn('Unknown failure executing file: <%s>' % fpath)