1513: Feature/jupyter notebook support - #9
Conversation
…s and examples This PR implements all items from the checklist in issue apache#1398: ## Implementation Checklist - [x] Add example .ipynb notebooks to python/examples/ - getting_started.ipynb - Basic connection and queries - dataframe_api.ipynb - DataFrame transformations - distributed_queries.ipynb - Multi-stage query examples - [x] Document notebook support in Python README - Added comprehensive Jupyter section with examples - [x] Create ballista.jupyter module with magic commands - Full implementation with BallistaMagics class - [x] Add %ballista connect/status/tables/schema line magics - connect: Connect to Ballista cluster - status: Show connection status - tables: List registered tables - schema: Show table schema - disconnect: Disconnect from cluster - history: Show query history - [x] Add %%sql cell magic - Line magic for single-line queries - Cell magic for multi-line queries - Variable assignment support - --no-display and --limit options - [x] Add explain_visual() method for query plan rendering - Generates DOT/SVG visualization - Supports Jupyter _repr_html_ - Fallback when graphviz not installed - [x] Add progress indicator support for long-running queries - collect_with_progress() method - Callback support for custom progress handling - Jupyter-aware display - [x] Consider JupySQL integration - Documented as alternative in README ## Additional Features - ExecutionPlanVisualization class for plan rendering - tables() method on BallistaSessionContext - Optional jupyter dependency in pyproject.toml - Comprehensive test coverage (45 tests passing) Closes apache#1398
WalkthroughThis pull request introduces comprehensive Jupyter notebook support to PyBallista. It adds three example notebooks demonstrating DataFrame API usage, distributed queries, and getting started with Ballista. The implementation includes a new IPython extension with magic commands for SQL execution and connection management, visualization of execution plans as SVG/HTML, progress tracking for long-running queries, table/schema exploration methods, and an optional jupyter dependency group. Supporting documentation is added to the Python README, gitignore entries for notebook checkpoints are updated, and test coverage is provided for all new functionality. ✨ Finishing Touches
🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment Tip You can customize the tone of the review comments and chat replies.Configure the |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 7 potential issues.
Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
| else: | ||
| def decorator(func): | ||
| return func | ||
| return decorator |
There was a problem hiding this comment.
Missing line_cell_magic stub breaks import without IPython
High Severity
The fallback stubs defined when IPython is not installed include line_magic, cell_magic, and magics_class, but do not include a stub for line_cell_magic. The sql method is decorated with @line_cell_magic, so importing ballista.jupyter without IPython installed will raise a NameError at class definition time, making the module completely unusable.
Additional Locations (1)
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The decorator fallback is not implemented and trying to use the decorator when IPyhton is not available will lead to a runtime error. Prevents an early error when IPython is not available, making this notebook support partially unusable.
| elif file_type == "csv": | ||
| self._ctx.register_csv(table_name, file_name) | ||
| else: | ||
| raise NotImplemented("Currently not supporting the inserted file format") |
There was a problem hiding this comment.
raise NotImplemented instead of raise NotImplementedError
Medium Severity
raise NotImplemented(...) is incorrect — NotImplemented is a special singleton constant used by rich comparison methods, not an exception class. This will raise a TypeError: exceptions must derive from BaseException instead of the intended error. The correct form is raise NotImplementedError(...).
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The name of the exception is wrong and any attempt to use it as an exception will fail because it is not extending from BaseException. Prevents hiding the real problem by failing with another error due to the wrong exception type.
| if not tables: | ||
| return "No tables registered.\n\nUse ctx.register_parquet() or ctx.register_csv() to register tables." | ||
| schema_count = len(tables.keys()) | ||
| table_count = len(tables.values()) |
There was a problem hiding this comment.
table_count incorrectly counts schemas instead of tables
Medium Severity
len(tables.values()) returns the number of schemas (dict entries), not the total number of tables. Since tables is a dict[str, list[str]], tables.values() is a view of lists, and len() gives the count of those lists (identical to schema_count). The total table count needs sum(len(v) for v in tables.values()).
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The values of the Dictionary are lists, so the length is the sum of their lengths, not just the number of values which would be the same as the length of the dictionary itself.
| LIMIT 10 | ||
| """ | ||
| for row in help_info.split("\n"): | ||
| print(row) |
There was a problem hiding this comment.
_show_help prints but never returns a string
Medium Severity
_show_help has return type Optional[str] but only prints to stdout and implicitly returns None. The ballista line magic (at "help" and empty-command branches) returns the result of _show_help(), which is always None. This also causes test failures since the tests call result.lower() on a None value.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The method return type is Optional but it does not return anything, i.e. implicitly it always returns None. This would break some unit tests because they try to use it as Some without checking
| except Exception: | ||
| # Fallback if logical_plan() fails | ||
| plan_str = "Unable to retrieve execution plan" | ||
| return ExecutionPlanVisualization(plan_str, analyze=analyze) |
There was a problem hiding this comment.
explain_visual ignores the analyze parameter for plan retrieval
Medium Severity
The analyze parameter is documented to include runtime statistics from actual execution, but the method always calls self.logical_plan().display_indent() regardless of the analyze value. The analyze flag is only passed through to ExecutionPlanVisualization where it's used solely for a display label string. When users call df.explain_visual(analyze=True), they get the same logical plan without any runtime statistics.
| if IPYTHON_AVAILABLE: | ||
| display(HTML(html)) | ||
| else: | ||
| print("\n".join(status_lines)) |
There was a problem hiding this comment.
_status returns None when connected, breaking callers
Low Severity
When connected, _status either calls display(HTML(...)) (IPython available) or print(...) (no IPython) but neither path returns a value. The method implicitly returns None, which causes the test test_status_when_connected to fail when it asserts string content in the result. The disconnected path correctly uses return.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The method return type is Optional but it does not return anything, i.e. implicitly it always returns None. This would break some unit tests because they try to use it as Some without checking
|
|
||
| try: | ||
| # Query the table with LIMIT 0 to get schema without data | ||
| df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0") |
There was a problem hiding this comment.
SQL injection risk in _schema via unescaped table name
Medium Severity
The _schema method constructs a SQL query by directly interpolating the user-provided table_name into an f-string: f"SELECT * FROM {table_name} LIMIT 0". Since table_name comes from user input via %ballista schema <table>, a malicious or accidental input could inject arbitrary SQL.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The table_name is an input provided by the user and used without any sanitation, so it could lead to an SQL injection. The value of table_name should be checked for any malicious characters which could break the SQL query.
🤖 Augment PR SummarySummary: This PR adds Jupyter notebook support to PyBallista via IPython magics and notebook-friendly helpers. Changes:
dot to produce SVG and falls back to HTML when unavailable.
🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| else: | ||
| raise NotImplemented("Currently not supporting the inserted file format") | ||
|
|
||
| @line_cell_magic |
There was a problem hiding this comment.
If IPython isn't available, the fallback stubs define line_magic/cell_magic but not line_cell_magic, so importing this module will raise NameError at the @line_cell_magic decorator.
Severity: high
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The decorator fallback is not implemented and trying to use the decorator when IPyhton is not available will lead to a runtime error. Prevents an early error when IPython is not available, making this notebook support partially unusable.
| elif file_type == "csv": | ||
| self._ctx.register_csv(table_name, file_name) | ||
| else: | ||
| raise NotImplemented("Currently not supporting the inserted file format") |
There was a problem hiding this comment.
Raising NotImplemented(...) will itself raise TypeError because NotImplemented isn't an exception type, so unsupported formats won't fail cleanly.
Severity: high
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The name of the exception is wrong and any attempt to use it as an exception will fail because it is not extending from BaseException. Prevents hiding the real problem by failing with another error due to the wrong exception type.
| lines = [ | ||
| {"content": f"Total: {table_count} table(s) in {schema_count} schema(s)", "is_info": True}, | ||
| {"content": "Registered tables:", "is_info": True}, | ||
| *[{"content": f"Schema: {schema_name}. Tables: {", ".join(list(table_names))}", "is_info": False} |
There was a problem hiding this comment.
This f-string has nested quotes ({", ".join(...)}) and is a syntax error, so ballista.jupyter won't import.
Severity: high
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The string concatenation is broken here due to the mixed usage of quotes and double quotes. Prevents broken rendering of the notebook
| return f"<p><strong>{name}:</strong> {value.strip()}</p>" | ||
|
|
||
| html = "".join(_format_html_status_output(line) for line in status_lines) | ||
| if IPYTHON_AVAILABLE: |
There was a problem hiding this comment.
Several magic helpers (e.g., _status) print/display output but return None, while callers/tests treat the return value as a string; this will break the added unit tests and makes %ballista commands harder to use programmatically.
Severity: high
Other Locations
python/python/ballista/jupyter.py:272python/python/ballista/jupyter.py:293python/python/ballista/jupyter.py:346python/python/ballista/jupyter.py:430python/python/ballista/jupyter.py:477
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The method return type is Optional but it does not return anything, i.e. implicitly it always returns None. This would break some unit tests because they try to use it as Some without checking
| if not tables: | ||
| return "No tables registered.\n\nUse ctx.register_parquet() or ctx.register_csv() to register tables." | ||
| schema_count = len(tables.keys()) | ||
| table_count = len(tables.values()) |
There was a problem hiding this comment.
table_count = len(tables.values()) counts schemas (dict values) rather than total tables, so the reported "Total: X table(s)" will be incorrect when a schema contains multiple tables.
Severity: low
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The values of the Dictionary are lists, so the length is the sum of their lengths, not just the number of values which would be the same as the length of the dictionary itself.
| def _execute_sql( | ||
| self, | ||
| query: str, | ||
| display_results: bool = True, |
There was a problem hiding this comment.
display_results/limit are parsed from %%sql options but _execute_sql doesn't use them, so --no-display and --limit currently don't change notebook output.
Severity: medium
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Augment AI reviewer is correct! The method does not use its arguments and this will confuse the notebook user, e.g. when the user tries to set a limit it will be ignored. The arguments should either be removed or used.
| df = self._to_internal_df() | ||
| df.write_parquet(str(path), compression.value, compression_level) | ||
|
|
||
| def explain_visual(self, analyze: bool = False) -> "ExecutionPlanVisualization": |
There was a problem hiding this comment.
explain_visual(analyze=True) sets the flag on ExecutionPlanVisualization, but plan_str is always built from logical_plan() and doesn't incorporate analyzed/runtime statistics, so the visualization won't actually differ when analyze=True.
Severity: medium
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
| def __init__(self, address: str, config=None, runtime=None): | ||
| super().__init__(config, runtime) | ||
| self.address = address | ||
| self.session_id = self.session_id() |
There was a problem hiding this comment.
Assigning self.session_id = self.session_id() shadows the inherited session_id() method, so any later call to session_id() would fail with a 'str' object is not callable-style error.
Severity: medium
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
| self.address = address | ||
| self.session_id = self.session_id() | ||
|
|
||
| def get_tables(self) -> Optional[dict[str, str]]: |
There was a problem hiding this comment.
get_tables() is annotated as Optional[dict[str, str]], but it always returns a dict and the values are lists of table names; this mismatch can confuse callers and type checkers.
Severity: low
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the PyBallista user experience by integrating robust Jupyter Notebook support. It empowers users with interactive SQL capabilities, rich data visualization, and detailed insights into query execution, making distributed data analysis more accessible and efficient within a notebook environment. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces excellent support for Jupyter Notebooks, including magic commands, rich HTML display for DataFrames, execution plan visualization, and progress indicators. The implementation is well-structured and adds significant value for interactive use of PyBallista. I've included a few suggestions for improving robustness, fixing minor bugs in the magic commands, and enhancing error reporting.
| if not tables: | ||
| return "No tables registered.\n\nUse ctx.register_parquet() or ctx.register_csv() to register tables." | ||
| schema_count = len(tables.keys()) | ||
| table_count = len(tables.values()) |
There was a problem hiding this comment.
The calculation for table_count is incorrect. len(tables.values()) will return the number of schemas, not the total number of tables. To get the total count of tables across all schemas, you should sum the lengths of the lists of table names.
| table_count = len(tables.values()) | |
| table_count = sum(len(v) for v in tables.values()) |
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Gemini AI reviewer is correct! The values of the Dictionary are lists, so the length is the sum of their lengths, not just the number of values which would be the same as the length of the dictionary itself.
| except Exception: | ||
| pass # Ignore display errors |
There was a problem hiding this comment.
Silencing all exceptions with except Exception: pass can hide potential bugs or issues with the environment (e.g., problems with IPython.display). It would be better to at least log the exception as a warning so that issues can be debugged if the progress indicator fails to display correctly.
| except Exception: | |
| pass # Ignore display errors | |
| except Exception as e: | |
| import warnings | |
| warnings.warn(f"Failed to display progress in Jupyter: {e}") |
| except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired): | ||
| pass |
There was a problem hiding this comment.
The try...except block for running the dot command silences all errors. While the fallback to an HTML representation is good, it would be beneficial to log a warning when the dot command fails. This would help users diagnose why they are not seeing the SVG visualization (e.g., graphviz not installed or not in PATH).
| except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired): | |
| pass | |
| except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired) as e: | |
| import warnings | |
| warnings.warn(f"Graphviz visualization failed: {e}. Falling back to text plan.") | |
| pass |
| def get_tables(self) -> Optional[dict[str, str]]: | ||
| """Get tables and their respective schemas (in terms of database schema).""" |
There was a problem hiding this comment.
The type hint for the return value of get_tables is incorrect. The method returns a dictionary mapping schema names (strings) to lists of table names (lists of strings), so the type hint should be Dict[str, List[str]]. Currently, it's Optional[dict[str, str]]. Also, the docstring could be more precise about what is being returned.
| def get_tables(self) -> Optional[dict[str, str]]: | |
| """Get tables and their respective schemas (in terms of database schema).""" | |
| def get_tables(self) -> Dict[str, List[str]]: | |
| """Get a dictionary of schemas and their tables.""" |
| except (AttributeError, NotImplementedError): | ||
| pass |
There was a problem hiding this comment.
Silencing AttributeError and NotImplementedError can hide issues if the underlying catalog API changes. While returning an empty dictionary is a safe fallback, logging a warning would make it easier to detect and debug such problems in the future.
| except (AttributeError, NotImplementedError): | |
| pass | |
| except (AttributeError, NotImplementedError) as e: | |
| import warnings | |
| warnings.warn(f"Could not retrieve tables from catalog: {e}") | |
| pass |
There was a problem hiding this comment.
value:useful; category:bug; feedback: The Gemini AI reviewer is correct! The name of the exception is wrong and any attempt to use it as an exception will fail because it is not extending from BaseException. Prevents hiding the real problem by failing with another error due to the wrong exception type.
| "Use: %ballista connect df://host:port" | ||
| ) | ||
| else: | ||
| args = line.strip().split() |
There was a problem hiding this comment.
The current argument parsing using line.strip().split() is not robust and will fail if file paths or table names contain spaces. Using shlex.split() provides more reliable parsing of command-line style arguments, correctly handling quoted strings. You will need to add import shlex at the top of the file.
| args = line.strip().split() | |
| args = shlex.split(line.strip()) |
|
|
||
| def _format_html_tables_output(line: str, is_info: bool = False) -> str: | ||
| if is_info: | ||
| return f"<p><strong>{line}<strong></p>" |
There was a problem hiding this comment.
There's a typo in the closing HTML tag. <strong> should be closed with </strong>.
| return f"<p><strong>{line}<strong></p>" | |
| return f"<p><strong>{line}</strong></p>" |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (6)
python/python/tests/test_notebook_features.py (2)
30-35: Consider reusing a session-scoped cluster fixture.Starting a cluster in every test can make this suite unnecessarily slow and flaky. A
scope="session"fixture with teardown would improve reliability.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/python/tests/test_notebook_features.py` around lines 30 - 35, The current ctx fixture calls setup_test_cluster() for every test which is slow/flaky; create a session-scoped fixture (e.g., cluster or test_cluster) that calls setup_test_cluster() once and yields the (address, port) pair with proper teardown, then change the existing ctx fixture to consume that session-scoped fixture and return BallistaSessionContext(address=f"df://{address}:{port}") so tests reuse the single cluster instance; reference the setup_test_cluster function and the ctx fixture/BallistaSessionContext to locate where to wire the new session-scoped fixture.
206-217: These assertions are too weak to catch regressions.
isinstance(..., list)/is not Nonewill pass even if table listing or HTML rendering is broken. Please assert expected content (e.g., registered table name present, HTML contains<table).Also applies to: 228-233
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/python/tests/test_notebook_features.py` around lines 206 - 217, The current assertions in test_tables_after_register_parquet and test_tables_after_register_csv are too weak; update them to verify actual content: after calling ctx.register_parquet("test_parquet", ...)/ctx.register_csv("test_csv", ...), call tables = ctx.tables() and assert the registered table name appears in the returned listing (e.g., "test_parquet" / "test_csv" in tables) and, if tables can render HTML, assert that the HTML representation contains "<table". Apply the same stronger assertions to the other similar tests referenced (the tests around the second block with the same pattern) using the ctx.tables(), ctx.register_parquet, ctx.register_csv and test_* function names to locate and update the assertions.python/python/tests/test_jupyter.py (1)
139-143: Test assertion is too weak.The test only checks
result is not None, but_tables()can returnNonewhen using IPython (it displays HTML instead). Consider either asserting the result contains expected content or testing in a non-IPython context explicitly.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/python/tests/test_jupyter.py` around lines 139 - 143, The test test_tables_empty currently only asserts result is not None which is too weak because connected_magics._tables() may return None in IPython (it renders HTML); update the test to assert specific expected behavior: either (a) assert the returned result contains an expected string/structure when _tables() returns a value (e.g., check for "No tables" or an empty table marker in result), or (b) force a non-IPython path by mocking/stubbing the IPython display/rendering so _tables() returns its textual output and then assert the exact output; locate the call to connected_magics._tables() in test_tables_empty and implement one of these two fixes to make the assertion deterministic.python/python/ballista/jupyter.py (2)
417-431: Inconsistent return behavior.When history exists, the method prints and returns
None. When empty, it returns a string. This inconsistency makes the return value unreliable for callers.Suggested fix for consistency
Either always return a string:
def _show_history(self) -> Optional[str]: """Show query history.""" if not self._query_history: return "No queries executed yet." lines = ["Query History:", "-" * 60] for i, entry in enumerate(self._query_history[-10:], 1): # Last 10 queries query_preview = entry["query"][:50] + "..." if len(entry["query"]) > 50 else entry["query"] query_preview = query_preview.replace("\n", " ") lines.append(f"{i}. [{entry['timestamp']}] ({entry['elapsed_seconds']:.2f}s)") lines.append(f" {query_preview}") lines.append("-" * 60) - for row in lines: - print(row) + return "\n".join(lines)Or always print:
def _show_history(self) -> None: """Show query history.""" if not self._query_history: - return "No queries executed yet." + print("No queries executed yet.") + return🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/python/ballista/jupyter.py` around lines 417 - 431, _show_history currently prints the history and returns None when history exists, but returns a string when empty; make its return type consistent by always returning a string: build the output by joining the lines (use the existing lines list), print the joined string for backward compatibility, and then return the same joined string; update usage of self._query_history and keep the method signature (_show_history) as Optional[str] or stricter if desired.
238-243: Silently ignoring invalid--limitvalue may confuse users.When
--limitis followed by a non-integer, the error is silently ignored and the default limit is used. Consider providing feedback.Suggested improvement
elif args[i] == "--limit" and i + 1 < len(args): try: limit = int(args[i + 1]) i += 1 except ValueError: - pass + print(f"Warning: Invalid --limit value '{args[i + 1]}', using default {limit}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/python/ballista/jupyter.py` around lines 238 - 243, The argument parsing block that handles "--limit" (the snippet using args, i, and variable limit in python/ballista/jupyter.py) currently swallows ValueError silently; update the except ValueError branch to notify the user about the invalid value instead of silently ignoring it (e.g., print a warning to stderr that includes the invalid token args[i+1] and that the default limit will be used, or use logging). Ensure you reference the same parsing logic (args, i, limit) and import sys or logging if necessary, then preserve the existing behavior of keeping the default limit while advancing i appropriately.python/python/ballista/extension.py (1)
484-492: Consider handling file write errors.The
savemethod doesn't handle potential file operation errors (permissions, disk full, invalid path). Users callingsave()would get an unhandled exception.Proposed enhancement
def save(self, path: str) -> None: """Save the visualization to a file (SVG or DOT format).""" if path.endswith(".dot"): content = self.to_dot() else: content = self.to_svg() - with open(path, "w") as f: - f.write(content) + with open(path, "w", encoding="utf-8") as f: + f.write(content)At minimum, adding
encoding="utf-8"ensures consistent behavior across platforms.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/python/ballista/extension.py` around lines 484 - 492, The save method lacks robust file error handling and explicit encoding; wrap the file write in a try/except that catches OSError (and optionally Exception) around the open/write, include encoding="utf-8" when opening the file, and re-raise a clearer exception or raise a custom error with context (include the path and original exception) so callers get a descriptive error instead of an unhandled traceback; update the save function (which calls to_dot() or to_svg()) to perform this guarded write and surface errors cleanly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@python/examples/getting_started.ipynb`:
- Around line 79-82: Replace the hard-coded scheduler host/port assignment with
the dynamic test-cluster setup: call setup_test_cluster() to obtain host and
port and use those values when constructing the BallistaSessionContext (the
current code assigns host, port = "localhost", "39431" and then calls
BallistaSessionContext(f"df://{host}:{port}")). Also update the other occurrence
that sets host/port manually (the later block around the other
BallistaSessionContext usage) to use setup_test_cluster() so the notebook uses a
reproducible, dynamic test cluster instead of a fixed port.
In `@python/python/ballista/extension.py`:
- Around line 536-548: The get_tables method has an incorrect return type: it is
annotated as Optional[dict[str, str]] but builds tables_info where each value is
a list of table names (list(catalog.schema(name=schema_name).table_names()));
update the signature of get_tables to Optional[dict[str, list[str]]] (or
dict[str, list[str]] if you never return None) and adjust any callers if
necessary to reflect that values are lists, keeping references to get_tables,
tables_info, and catalog.schema to locate the change.
- Around line 263-267: The parameter type for collect_with_progress currently
uses the built-in lowercase callable; change it to use typing.Callable by
importing Callable from typing and updating the signature of
collect_with_progress(callback: Optional[Callable] = None, ...). Locate the
function collect_with_progress in extension.py and add the Callable import
(alongside Optional if needed) to the module imports so the type hint is
correct.
In `@python/python/ballista/jupyter.py`:
- Around line 198-204: The code currently raises the singleton NotImplemented
(which is not an Exception) in the branch handling unsupported file types;
replace that with raising NotImplementedError (e.g., raise
NotImplementedError("Currently not supporting the inserted file format")) so
callers get the proper exception. Update the branch in the same block that calls
self._ctx.register_parquet and self._ctx.register_csv to raise
NotImplementedError with a clear message.
- Around line 330-331: schema_count is computed correctly as len(tables.keys())
but table_count is wrong because len(tables.values()) counts schemas, not
tables; change table_count to sum(len(v) for v in tables.values()) so it sums
the number of table names per schema (look for the variables schema_count,
table_count and the tables dict in the jupyter.py block to update this
computation).
- Around line 57-94: The stub fallback is missing line_cell_magic which causes a
NameError when IPython is absent; add a stub function named line_cell_magic that
mirrors the behavior of the existing line_magic and cell_magic stubs (accepts
either a callable or a name string and returns either the function or a
decorator wrapping the function), so code that uses `@line_cell_magic` or
`@line_cell_magic`("name") (and classes/functions like Magics) will work when
IPython is not installed.
- Around line 433-478: The _show_help method currently prints help_info
line-by-line and returns None, causing tests to fail; modify _show_help (the
function named _show_help and the local help_info variable) to return the full
help string (help_info) instead of (or in addition to) printing it so
callers/tests can inspect the text — e.g., remove the for loop printing or keep
printing but ensure the function ends with `return help_info` so a string is
returned.
- Around line 360-363: The SELECT uses an unvalidated table_name when calling
self.ctx.sql in the method (where df = self.ctx.sql(f"SELECT * FROM {table_name}
LIMIT 0") and schema = df.schema()), which allows SQL injection; validate
table_name before interpolating by enforcing a strict identifier pattern (e.g.
use a regex like ^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?$ to allow
optional qualified names) and reject or raise ValueError for anything that fails
the check, or alternatively normalize/quote the identifier using the engine’s
safe identifier quoting if available, before constructing the query.
In `@python/python/tests/test_jupyter.py`:
- Around line 87-93: The test fails because _show_help() currently returns None;
update the _show_help function in jupyter.py so it returns the help text string
(instead of printing or doing nothing) so the test can call result.lower();
ensure the returned string contains the commands "connect", "status", "tables",
and "sql" so the assertions in test_show_help pass.
- Around line 127-131: The _status() method in jupyter.py currently
prints/displays the connected status (using display/HTML or print) but returns
None, causing the test test_status_when_connected to fail; update the _status()
implementation so that after building the status string (the same content passed
to display(HTML(...)) or print), it returns that string (ensure both the
connected and not-connected branches return the status string) so callers/tests
can assert on its contents; refer to the _status() method and the places where
it uses display/HTML and print to locate and change the behavior.
In `@python/README.md`:
- Around line 84-105: Add an explicit install step for notebook users: update
the Jupyter section to instruct users to install the optional Jupyter extras
before loading the extension (e.g., run pip install "ballista[jupyter]"), and
place this note immediately before the example that imports
BallistaSessionContext and before the %load_ext ballista.jupyter line so users
know to install the dependency first.
---
Nitpick comments:
In `@python/python/ballista/extension.py`:
- Around line 484-492: The save method lacks robust file error handling and
explicit encoding; wrap the file write in a try/except that catches OSError (and
optionally Exception) around the open/write, include encoding="utf-8" when
opening the file, and re-raise a clearer exception or raise a custom error with
context (include the path and original exception) so callers get a descriptive
error instead of an unhandled traceback; update the save function (which calls
to_dot() or to_svg()) to perform this guarded write and surface errors cleanly.
In `@python/python/ballista/jupyter.py`:
- Around line 417-431: _show_history currently prints the history and returns
None when history exists, but returns a string when empty; make its return type
consistent by always returning a string: build the output by joining the lines
(use the existing lines list), print the joined string for backward
compatibility, and then return the same joined string; update usage of
self._query_history and keep the method signature (_show_history) as
Optional[str] or stricter if desired.
- Around line 238-243: The argument parsing block that handles "--limit" (the
snippet using args, i, and variable limit in python/ballista/jupyter.py)
currently swallows ValueError silently; update the except ValueError branch to
notify the user about the invalid value instead of silently ignoring it (e.g.,
print a warning to stderr that includes the invalid token args[i+1] and that the
default limit will be used, or use logging). Ensure you reference the same
parsing logic (args, i, limit) and import sys or logging if necessary, then
preserve the existing behavior of keeping the default limit while advancing i
appropriately.
In `@python/python/tests/test_jupyter.py`:
- Around line 139-143: The test test_tables_empty currently only asserts result
is not None which is too weak because connected_magics._tables() may return None
in IPython (it renders HTML); update the test to assert specific expected
behavior: either (a) assert the returned result contains an expected
string/structure when _tables() returns a value (e.g., check for "No tables" or
an empty table marker in result), or (b) force a non-IPython path by
mocking/stubbing the IPython display/rendering so _tables() returns its textual
output and then assert the exact output; locate the call to
connected_magics._tables() in test_tables_empty and implement one of these two
fixes to make the assertion deterministic.
In `@python/python/tests/test_notebook_features.py`:
- Around line 30-35: The current ctx fixture calls setup_test_cluster() for
every test which is slow/flaky; create a session-scoped fixture (e.g., cluster
or test_cluster) that calls setup_test_cluster() once and yields the (address,
port) pair with proper teardown, then change the existing ctx fixture to consume
that session-scoped fixture and return
BallistaSessionContext(address=f"df://{address}:{port}") so tests reuse the
single cluster instance; reference the setup_test_cluster function and the ctx
fixture/BallistaSessionContext to locate where to wire the new session-scoped
fixture.
- Around line 206-217: The current assertions in
test_tables_after_register_parquet and test_tables_after_register_csv are too
weak; update them to verify actual content: after calling
ctx.register_parquet("test_parquet", ...)/ctx.register_csv("test_csv", ...),
call tables = ctx.tables() and assert the registered table name appears in the
returned listing (e.g., "test_parquet" / "test_csv" in tables) and, if tables
can render HTML, assert that the HTML representation contains "<table". Apply
the same stronger assertions to the other similar tests referenced (the tests
around the second block with the same pattern) using the ctx.tables(),
ctx.register_parquet, ctx.register_csv and test_* function names to locate and
update the assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 7a2cc295-ab35-492a-91e2-6ccacb6e7f24
⛔ Files ignored due to path filters (1)
python/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (11)
.gitignorepython/README.mdpython/examples/dataframe_api.ipynbpython/examples/distributed_queries.ipynbpython/examples/getting_started.ipynbpython/pyproject.tomlpython/python/ballista/__init__.pypython/python/ballista/extension.pypython/python/ballista/jupyter.pypython/python/tests/test_jupyter.pypython/python/tests/test_notebook_features.py
| "# host, port = setup_test_cluster()\n", | ||
| "host, port = \"localhost\", \"39431\"\n", | ||
| "ctx = BallistaSessionContext(f\"df://{host}:{port}\")\n", | ||
| "\n", |
There was a problem hiding this comment.
Avoid hard-coded scheduler port in a getting-started notebook.
This makes the happy path fail for most users unless they happen to run a cluster on 39431. Prefer the dynamic test-cluster setup for reproducibility.
Also applies to: 167-169
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/examples/getting_started.ipynb` around lines 79 - 82, Replace the
hard-coded scheduler host/port assignment with the dynamic test-cluster setup:
call setup_test_cluster() to obtain host and port and use those values when
constructing the BallistaSessionContext (the current code assigns host, port =
"localhost", "39431" and then calls
BallistaSessionContext(f"df://{host}:{port}")). Also update the other occurrence
that sets host/port manually (the later block around the other
BallistaSessionContext usage) to use setup_test_cluster() so the notebook uses a
reproducible, dynamic test cluster instead of a fixed port.
| def collect_with_progress( | ||
| self, | ||
| callback: Optional[callable] = None, | ||
| poll_interval: float = 0.5, | ||
| ): |
There was a problem hiding this comment.
Incorrect type hint: callable should be Callable.
The type hint uses lowercase callable which is the built-in function, not the typing annotation. Use Callable from the typing module instead.
Proposed fix
Add to imports:
from typing import (
Union,
Optional,
List,
+ Callable,
)Then update the parameter type:
def collect_with_progress(
self,
- callback: Optional[callable] = None,
+ callback: Optional[Callable[[str, float], None]] = None,
poll_interval: float = 0.5,
):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/ballista/extension.py` around lines 263 - 267, The parameter
type for collect_with_progress currently uses the built-in lowercase callable;
change it to use typing.Callable by importing Callable from typing and updating
the signature of collect_with_progress(callback: Optional[Callable] = None,
...). Locate the function collect_with_progress in extension.py and add the
Callable import (alongside Optional if needed) to the module imports so the type
hint is correct.
| def get_tables(self) -> Optional[dict[str, str]]: | ||
| """Get tables and their respective schemas (in terms of database schema).""" | ||
| try: | ||
| catalog = self.catalog() | ||
| schema_names = list(catalog.schema_names()) | ||
| if schema_names is not None: | ||
| tables_info = {} | ||
| for schema_name in schema_names: | ||
| tables_info[schema_name] = list(catalog.schema(name=schema_name).table_names()) | ||
| return tables_info | ||
| except (AttributeError, NotImplementedError): | ||
| pass | ||
| return {} |
There was a problem hiding this comment.
Incorrect return type annotation.
The return type is annotated as Optional[dict[str, str]], but the method returns a dict where values are lists of table names (line 544 assigns list(catalog.schema(name=schema_name).table_names())).
Proposed fix
- def get_tables(self) -> Optional[dict[str, str]]:
- """Get tables and their respective schemas (in terms of database schema)."""
+ def get_tables(self) -> Optional[dict[str, list[str]]]:
+ """Get a mapping of schema names to their table names."""🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/ballista/extension.py` around lines 536 - 548, The get_tables
method has an incorrect return type: it is annotated as Optional[dict[str, str]]
but builds tables_info where each value is a list of table names
(list(catalog.schema(name=schema_name).table_names())); update the signature of
get_tables to Optional[dict[str, list[str]]] (or dict[str, list[str]] if you
never return None) and adjust any callers if necessary to reflect that values
are lists, keeping references to get_tables, tables_info, and catalog.schema to
locate the change.
| try: | ||
| from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic | ||
| from IPython.display import display, HTML | ||
|
|
||
| IPYTHON_AVAILABLE = True | ||
| except ImportError: | ||
| IPYTHON_AVAILABLE = False | ||
|
|
||
| # Provide stub classes for when IPython is not available | ||
| class Magics: | ||
| def __init__(self, shell=None): | ||
| self.shell = shell | ||
|
|
||
| def magics_class(cls): | ||
| return cls | ||
|
|
||
| def line_magic(name_or_func): | ||
| """Stub line_magic decorator for when IPython is not available.""" | ||
| # Handle both @line_magic and @line_magic("name") usage | ||
| if callable(name_or_func): | ||
| # Used as @line_magic without arguments | ||
| return name_or_func | ||
| else: | ||
| # Used as @line_magic("name") with arguments | ||
| def decorator(func): | ||
| return func | ||
| return decorator | ||
|
|
||
| def cell_magic(name_or_func): | ||
| """Stub cell_magic decorator for when IPython is not available.""" | ||
| # Handle both @cell_magic and @cell_magic("name") usage | ||
| if callable(name_or_func): | ||
| return name_or_func | ||
| else: | ||
| def decorator(func): | ||
| return func | ||
| return decorator | ||
|
|
There was a problem hiding this comment.
Missing line_cell_magic stub will cause NameError.
The code imports line_cell_magic from IPython (line 58) and uses it as a decorator (line 205), but the fallback stubs only define line_magic and cell_magic. When IPython is unavailable, this will raise a NameError.
Proposed fix
def cell_magic(name_or_func):
"""Stub cell_magic decorator for when IPython is not available."""
# Handle both `@cell_magic` and `@cell_magic`("name") usage
if callable(name_or_func):
return name_or_func
else:
def decorator(func):
return func
return decorator
+
+ def line_cell_magic(name_or_func):
+ """Stub line_cell_magic decorator for when IPython is not available."""
+ if callable(name_or_func):
+ return name_or_func
+ else:
+ def decorator(func):
+ return func
+ return decorator📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic | |
| from IPython.display import display, HTML | |
| IPYTHON_AVAILABLE = True | |
| except ImportError: | |
| IPYTHON_AVAILABLE = False | |
| # Provide stub classes for when IPython is not available | |
| class Magics: | |
| def __init__(self, shell=None): | |
| self.shell = shell | |
| def magics_class(cls): | |
| return cls | |
| def line_magic(name_or_func): | |
| """Stub line_magic decorator for when IPython is not available.""" | |
| # Handle both @line_magic and @line_magic("name") usage | |
| if callable(name_or_func): | |
| # Used as @line_magic without arguments | |
| return name_or_func | |
| else: | |
| # Used as @line_magic("name") with arguments | |
| def decorator(func): | |
| return func | |
| return decorator | |
| def cell_magic(name_or_func): | |
| """Stub cell_magic decorator for when IPython is not available.""" | |
| # Handle both @cell_magic and @cell_magic("name") usage | |
| if callable(name_or_func): | |
| return name_or_func | |
| else: | |
| def decorator(func): | |
| return func | |
| return decorator | |
| try: | |
| from IPython.core.magic import Magics, magics_class, line_magic, cell_magic, line_cell_magic | |
| from IPython.display import display, HTML | |
| IPYTHON_AVAILABLE = True | |
| except ImportError: | |
| IPYTHON_AVAILABLE = False | |
| # Provide stub classes for when IPython is not available | |
| class Magics: | |
| def __init__(self, shell=None): | |
| self.shell = shell | |
| def magics_class(cls): | |
| return cls | |
| def line_magic(name_or_func): | |
| """Stub line_magic decorator for when IPython is not available.""" | |
| # Handle both `@line_magic` and `@line_magic`("name") usage | |
| if callable(name_or_func): | |
| # Used as `@line_magic` without arguments | |
| return name_or_func | |
| else: | |
| # Used as `@line_magic`("name") with arguments | |
| def decorator(func): | |
| return func | |
| return decorator | |
| def cell_magic(name_or_func): | |
| """Stub cell_magic decorator for when IPython is not available.""" | |
| # Handle both `@cell_magic` and `@cell_magic`("name") usage | |
| if callable(name_or_func): | |
| return name_or_func | |
| else: | |
| def decorator(func): | |
| return func | |
| return decorator | |
| def line_cell_magic(name_or_func): | |
| """Stub line_cell_magic decorator for when IPython is not available.""" | |
| if callable(name_or_func): | |
| return name_or_func | |
| else: | |
| def decorator(func): | |
| return func | |
| return decorator |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/ballista/jupyter.py` around lines 57 - 94, The stub fallback is
missing line_cell_magic which causes a NameError when IPython is absent; add a
stub function named line_cell_magic that mirrors the behavior of the existing
line_magic and cell_magic stubs (accepts either a callable or a name string and
returns either the function or a decorator wrapping the function), so code that
uses `@line_cell_magic` or `@line_cell_magic`("name") (and classes/functions like
Magics) will work when IPython is not installed.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The decorator fallback is not implemented and trying to use the decorator when IPyhton is not available will lead to a runtime error. Prevents an early error when IPython is not available, making this notebook support partially unusable.
| if file_type == "parquet": | ||
| self._ctx.register_parquet(table_name, file_name) | ||
| elif file_type == "csv": | ||
| self._ctx.register_csv(table_name, file_name) | ||
| else: | ||
| raise NotImplemented("Currently not supporting the inserted file format") | ||
|
|
There was a problem hiding this comment.
Raise NotImplementedError, not NotImplemented.
NotImplemented is a special singleton for rich comparison methods, not an exception. This line will raise a TypeError instead of the intended error.
Proposed fix
if file_type == "parquet":
self._ctx.register_parquet(table_name, file_name)
elif file_type == "csv":
self._ctx.register_csv(table_name, file_name)
else:
- raise NotImplemented("Currently not supporting the inserted file format")
+ raise NotImplementedError("Currently not supporting the inserted file format")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/ballista/jupyter.py` around lines 198 - 204, The code currently
raises the singleton NotImplemented (which is not an Exception) in the branch
handling unsupported file types; replace that with raising NotImplementedError
(e.g., raise NotImplementedError("Currently not supporting the inserted file
format")) so callers get the proper exception. Update the branch in the same
block that calls self._ctx.register_parquet and self._ctx.register_csv to raise
NotImplementedError with a clear message.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The name of the exception is wrong and any attempt to use it as an exception will fail because it is not extending from BaseException. Prevents hiding the real problem by failing with another error due to the wrong exception type.
| try: | ||
| # Query the table with LIMIT 0 to get schema without data | ||
| df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0") | ||
| schema = df.schema() |
There was a problem hiding this comment.
SQL injection risk with unvalidated table name.
The table_name is directly interpolated into the SQL query without validation. A malicious input like users; DROP TABLE users; -- could cause unintended behavior.
Suggested mitigation
Add basic validation to ensure the table name contains only valid identifier characters:
def _schema(self, table_name: str) -> Optional[str]:
"""Show schema for a table."""
if not table_name:
return "Usage: %ballista schema <table_name>"
+ # Basic validation: allow only alphanumeric, underscore, dot (for schema.table)
+ import re
+ if not re.match(r'^[\w.]+$', table_name):
+ return f"Invalid table name: '{table_name}'"
+
try:
# Query the table with LIMIT 0 to get schema without data
df = self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/ballista/jupyter.py` around lines 360 - 363, The SELECT uses an
unvalidated table_name when calling self.ctx.sql in the method (where df =
self.ctx.sql(f"SELECT * FROM {table_name} LIMIT 0") and schema = df.schema()),
which allows SQL injection; validate table_name before interpolating by
enforcing a strict identifier pattern (e.g. use a regex like
^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?$ to allow optional qualified
names) and reject or raise ValueError for anything that fails the check, or
alternatively normalize/quote the identifier using the engine’s safe identifier
quoting if available, before constructing the query.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The table_name is an input provided by the user and used without any sanitation, so it could lead to an SQL injection. The value of table_name should be checked for any malicious characters which could break the SQL query.
| def _show_help(self) -> Optional[str]: | ||
| """Show help for Ballista magic commands.""" | ||
| help_info = """ | ||
| Ballista Jupyter Magic Commands | ||
| ================================ | ||
|
|
||
| Connection: | ||
| %ballista connect <url> - Connect to Ballista cluster | ||
| %ballista disconnect - Disconnect from cluster | ||
| %ballista status - Show connection status | ||
|
|
||
| Exploration: | ||
| %ballista tables - List registered tables | ||
| %ballista schema <table> - Show table schema | ||
|
|
||
| Table-register: | ||
| %register [format] [schema.table_name] [file_path] - Register a new table in the current Ballista Context | ||
|
|
||
| Query: | ||
| %sql <query> - Execute single-line SQL query | ||
|
|
||
| %%sql [options] [var] - Execute multi-line SQL query | ||
| Options: | ||
| --no-display - Don't display results | ||
| --limit N - Limit displayed rows (default: 50) | ||
| var - Store result in variable | ||
|
|
||
| History: | ||
| %ballista history - Show recent query history | ||
|
|
||
| Examples: | ||
| %ballista connect df://localhost:50050 | ||
| %ballista tables | ||
| %ballista schema orders | ||
|
|
||
| %sql SELECT COUNT(*) FROM orders | ||
|
|
||
| %%sql my_result | ||
| SELECT customer_id, SUM(amount) as total | ||
| FROM orders | ||
| GROUP BY customer_id | ||
| ORDER BY total DESC | ||
| LIMIT 10 | ||
| """ | ||
| for row in help_info.split("\n"): | ||
| print(row) |
There was a problem hiding this comment.
Method prints but doesn't return, causing test failures.
The _show_help method prints the help text but returns None. However, tests (lines 89-99 in test_jupyter.py) expect it to return a string containing "connect", "status", etc. The tests will fail with AttributeError: 'NoneType' object has no attribute 'lower'.
Proposed fix
def _show_help(self) -> Optional[str]:
"""Show help for Ballista magic commands."""
help_info = """
Ballista Jupyter Magic Commands
...
"""
- for row in help_info.split("\n"):
- print(row)
+ print(help_info)
+ return help_info🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/ballista/jupyter.py` around lines 433 - 478, The _show_help
method currently prints help_info line-by-line and returns None, causing tests
to fail; modify _show_help (the function named _show_help and the local
help_info variable) to return the full help string (help_info) instead of (or in
addition to) printing it so callers/tests can inspect the text — e.g., remove
the for loop printing or keep printing but ensure the function ends with `return
help_info` so a string is returned.
| def test_show_help(self, magics): | ||
| """Test that help is shown correctly.""" | ||
| result = magics._show_help() | ||
| assert "connect" in result.lower() | ||
| assert "status" in result.lower() | ||
| assert "tables" in result.lower() | ||
| assert "sql" in result.lower() |
There was a problem hiding this comment.
Test will fail due to _show_help returning None.
As noted in the review of jupyter.py, _show_help() returns None but this test expects a string. The result.lower() call will raise AttributeError.
This will be fixed when _show_help() is updated to return the help string.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/tests/test_jupyter.py` around lines 87 - 93, The test fails
because _show_help() currently returns None; update the _show_help function in
jupyter.py so it returns the help text string (instead of printing or doing
nothing) so the test can call result.lower(); ensure the returned string
contains the commands "connect", "status", "tables", and "sql" so the assertions
in test_show_help pass.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The method return type is Optional but it does not return anything, i.e. implicitly it always returns None. This would break some unit tests because they try to use it as Some without checking
| def test_status_when_connected(self, connected_magics): | ||
| """Test status shows connected state.""" | ||
| result = connected_magics._status() | ||
| assert "Connected" in result | ||
| assert "Session ID" in result |
There was a problem hiding this comment.
Test will fail: _status() returns None when connected.
Looking at _status() in jupyter.py (lines 298-321), when connected, the method either displays HTML or prints, but doesn't return a value. The test expects result to contain "Connected" and "Session ID", but result will be None.
This is related to the inconsistent return behavior in _status(). Either the method needs to return the status string, or the test needs to capture stdout.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/python/tests/test_jupyter.py` around lines 127 - 131, The _status()
method in jupyter.py currently prints/displays the connected status (using
display/HTML or print) but returns None, causing the test
test_status_when_connected to fail; update the _status() implementation so that
after building the status string (the same content passed to display(HTML(...))
or print), it returns that string (ensure both the connected and not-connected
branches return the status string) so callers/tests can assert on its contents;
refer to the _status() method and the places where it uses display/HTML and
print to locate and change the behavior.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The method return type is Optional but it does not return anything, i.e. implicitly it always returns None. This would break some unit tests because they try to use it as Some without checking
| PyBallista provides first-class Jupyter notebook support with SQL magic commands and rich HTML rendering. | ||
|
|
||
| ### HTML Table Rendering | ||
|
|
||
| DataFrames automatically render as styled HTML tables in Jupyter notebooks: | ||
|
|
||
| ```python | ||
| from ballista import BallistaSessionContext | ||
|
|
||
| ctx = BallistaSessionContext("df://localhost:50050") | ||
| df = ctx.sql("SELECT * FROM my_table LIMIT 10") | ||
| df # Renders as HTML table via _repr_html_() | ||
| ``` | ||
|
|
||
| ### SQL Magic Commands | ||
|
|
||
| For a more interactive SQL experience, load the Ballista Jupyter extension: | ||
|
|
||
| ```python | ||
| # Load the extension | ||
| %load_ext ballista.jupyter | ||
|
|
There was a problem hiding this comment.
Add explicit install step for notebook users.
The new section documents magics usage but doesn’t tell users to install the optional extra first, which can lead to immediate import/load failures.
Suggested doc patch
## Jupyter Notebook Support
PyBallista provides first-class Jupyter notebook support with SQL magic commands and rich HTML rendering.
+
+Install Jupyter extras first:
+
+```bash
+pip install "ballista[jupyter]"
+```🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/README.md` around lines 84 - 105, Add an explicit install step for
notebook users: update the Jupyter section to instruct users to install the
optional Jupyter extras before loading the extension (e.g., run pip install
"ballista[jupyter]"), and place this note immediately before the example that
imports BallistaSessionContext and before the %load_ext ballista.jupyter line so
users know to install the dependency first.
PR Review: Jupyter Notebook SupportOverall this is a useful feature addition with good documentation and a reasonable test suite. There are several bugs to address before merging. Critical Bugs1.
2. Missing The 3. F-string nested-quote SyntaxError (jupyter.py:336) Reusing the same quote character inside an f-string expression is a 4.
5.
Medium Issues6. Both methods print to stdout and return 7. Unclosed The closing tag reads 8. Wrong return type on Annotated as 9.
10. SQL injection risk in
Minor Issues
|
value:useful; category:bug; feedback: The Claude AI reviewer is correct! The decorator fallback is not implemented and trying to use the decorator when IPyhton is not available will lead to a runtime error. Prevents an early error when IPython is not available, making this notebook support partially unusable. |
value:useful; category:bug; feedback: The Claude AI reviewer is correct! The name of the exception is wrong and any attempt to use it as an exception will fail because it is not extending from BaseException. Prevents hiding the real problem by failing with another error due to the wrong exception type. |
value:useful; category:bug; feedback: The Claude AI reviewer is correct! The values of the Dictionary are lists, so the length is the sum of their lengths, not just the number of values which would be the same as the length of the dictionary itself. |
value:useful; category:bug; feedback: The Claude AI reviewer is correct! The method does not use its arguments and this will confuse the notebook user, e.g. when the user tries to set a limit it will be ignored. The arguments should either be removed or used. |
3 participants
1513: To review by AI