Skip to content

Commit 9620a71

Browse files
kesmit13kesmit13
authored
Fix implicit %%sql commands when running shared and personal notebooks from %run_shared / %run_personal (#66)
* Fix %%sql cells in %run_shared magic * Add output redirection * Skip empty cells * Make changes to personal magic too * Fix action name * Only test for output redirect for SQL cells
1 parent 2ab764d commit 9620a71

File tree

2 files changed

+164
-2
lines changed

2 files changed

+164
-2
lines changed

‎singlestoredb/magics/run_personal.py‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
import os
22
import tempfile
3+
from pathlib import Path
34
from typing import Any
5+
from warnings import warn
46

57
from IPython.core.interactiveshell import InteractiveShell
68
from IPython.core.magic import line_magic
79
from IPython.core.magic import Magics
810
from IPython.core.magic import magics_class
911
from IPython.core.magic import needs_local_scope
1012
from IPython.core.magic import no_var_expand
13+
from IPython.utils.contexts import preserve_keys
14+
from IPython.utils.syspathcontext import prepended_to_syspath
1115
from jinja2 import Template
1216

1317

@@ -53,4 +57,81 @@ def run_personal(self, line: str, local_ns: Any = None) -> Any:
5357
# Execute the SQL command
5458
self.shell.run_line_magic('sql', sql_command)
5559
# Run the downloaded file
56-
self.shell.run_line_magic('run', f'"{temp_file_path}"')
60+
with preserve_keys(self.shell.user_ns, '__file__'):
61+
self.shell.user_ns['__file__'] = temp_file_path
62+
self.safe_execfile_ipy(temp_file_path, raise_exceptions=True)
63+
64+
def safe_execfile_ipy(
65+
self,
66+
fname: str,
67+
shell_futures: bool = False,
68+
raise_exceptions: bool = False,
69+
) -> None:
70+
"""Like safe_execfile, but for .ipy or .ipynb files with IPython syntax.
71+
72+
Parameters
73+
----------
74+
fname : str
75+
The name of the file to execute. The filename must have a
76+
.ipy or .ipynb extension.
77+
shell_futures : bool (False)
78+
If True, the code will share future statements with the interactive
79+
shell. It will both be affected by previous __future__ imports, and
80+
any __future__ imports in the code will affect the shell. If False,
81+
__future__ imports are not shared in either direction.
82+
raise_exceptions : bool (False)
83+
If True raise exceptions everywhere. Meant for testing.
84+
"""
85+
fpath = Path(fname).expanduser().resolve()
86+
87+
# Make sure we can open the file
88+
try:
89+
with fpath.open('rb'):
90+
pass
91+
except Exception:
92+
warn('Could not open file <%s> for safe execution.' % fpath)
93+
return
94+
95+
# Find things also in current directory. This is needed to mimic the
96+
# behavior of running a script from the system command line, where
97+
# Python inserts the script's directory into sys.path
98+
dname = str(fpath.parent)
99+
100+
def get_cells() -> Any:
101+
"""generator for sequence of code blocks to run"""
102+
if fpath.suffix == '.ipynb':
103+
from nbformat import read
104+
nb = read(fpath, as_version=4)
105+
if not nb.cells:
106+
return
107+
for cell in nb.cells:
108+
if cell.cell_type == 'code':
109+
if not cell.source.strip():
110+
continue
111+
if getattr(cell, 'metadata', {}).get('language', '') == 'sql':
112+
output_redirect = getattr(
113+
cell, 'metadata', {},
114+
).get('output_variable', '') or ''
115+
if output_redirect:
116+
output_redirect = f' {output_redirect} <<'
117+
yield f'%%sql{output_redirect}\n{cell.source}'
118+
else:
119+
yield cell.source
120+
else:
121+
yield fpath.read_text(encoding='utf-8')
122+
123+
with prepended_to_syspath(dname):
124+
try:
125+
for cell in get_cells():
126+
result = self.shell.run_cell(
127+
cell, silent=True, shell_futures=shell_futures,
128+
)
129+
if raise_exceptions:
130+
result.raise_error()
131+
elif not result.success:
132+
break
133+
except Exception:
134+
if raise_exceptions:
135+
raise
136+
self.shell.showtraceback()
137+
warn('Unknown failure executing file: <%s>' % fpath)

‎singlestoredb/magics/run_shared.py‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
import os
22
import tempfile
3+
from pathlib import Path
34
from typing import Any
5+
from warnings import warn
46

57
from IPython.core.interactiveshell import InteractiveShell
68
from IPython.core.magic import line_magic
79
from IPython.core.magic import Magics
810
from IPython.core.magic import magics_class
911
from IPython.core.magic import needs_local_scope
1012
from IPython.core.magic import no_var_expand
13+
from IPython.utils.contexts import preserve_keys
14+
from IPython.utils.syspathcontext import prepended_to_syspath
1115
from jinja2 import Template
1216

1317

@@ -50,4 +54,81 @@ def run_shared(self, line: str, local_ns: Any = None) -> Any:
5054
# Execute the SQL command
5155
self.shell.run_line_magic('sql', sql_command)
5256
# Run the downloaded file
53-
self.shell.run_line_magic('run', f'"{temp_file_path}"')
57+
with preserve_keys(self.shell.user_ns, '__file__'):
58+
self.shell.user_ns['__file__'] = temp_file_path
59+
self.safe_execfile_ipy(temp_file_path, raise_exceptions=True)
60+
61+
def safe_execfile_ipy(
62+
self,
63+
fname: str,
64+
shell_futures: bool = False,
65+
raise_exceptions: bool = False,
66+
) -> None:
67+
"""Like safe_execfile, but for .ipy or .ipynb files with IPython syntax.
68+
69+
Parameters
70+
----------
71+
fname : str
72+
The name of the file to execute. The filename must have a
73+
.ipy or .ipynb extension.
74+
shell_futures : bool (False)
75+
If True, the code will share future statements with the interactive
76+
shell. It will both be affected by previous __future__ imports, and
77+
any __future__ imports in the code will affect the shell. If False,
78+
__future__ imports are not shared in either direction.
79+
raise_exceptions : bool (False)
80+
If True raise exceptions everywhere. Meant for testing.
81+
"""
82+
fpath = Path(fname).expanduser().resolve()
83+
84+
# Make sure we can open the file
85+
try:
86+
with fpath.open('rb'):
87+
pass
88+
except Exception:
89+
warn('Could not open file <%s> for safe execution.' % fpath)
90+
return
91+
92+
# Find things also in current directory. This is needed to mimic the
93+
# behavior of running a script from the system command line, where
94+
# Python inserts the script's directory into sys.path
95+
dname = str(fpath.parent)
96+
97+
def get_cells() -> Any:
98+
"""generator for sequence of code blocks to run"""
99+
if fpath.suffix == '.ipynb':
100+
from nbformat import read
101+
nb = read(fpath, as_version=4)
102+
if not nb.cells:
103+
return
104+
for cell in nb.cells:
105+
if cell.cell_type == 'code':
106+
if not cell.source.strip():
107+
continue
108+
if getattr(cell, 'metadata', {}).get('language', '') == 'sql':
109+
output_redirect = getattr(
110+
cell, 'metadata', {},
111+
).get('output_variable', '') or ''
112+
if output_redirect:
113+
output_redirect = f' {output_redirect} <<'
114+
yield f'%%sql{output_redirect}\n{cell.source}'
115+
else:
116+
yield cell.source
117+
else:
118+
yield fpath.read_text(encoding='utf-8')
119+
120+
with prepended_to_syspath(dname):
121+
try:
122+
for cell in get_cells():
123+
result = self.shell.run_cell(
124+
cell, silent=True, shell_futures=shell_futures,
125+
)
126+
if raise_exceptions:
127+
result.raise_error()
128+
elif not result.success:
129+
break
130+
except Exception:
131+
if raise_exceptions:
132+
raise
133+
self.shell.showtraceback()
134+
warn('Unknown failure executing file: <%s>' % fpath)

0 commit comments

Comments
 (0)