Conversation
Proposes a pluto.query module wrapping existing server REST endpoints to support programmatic run reads, metric fetching, file downloads, and run resumption for the Neptune-to-Pluto migration use case. https://claude.ai/code/session_01VUDsz1aDmSAeuThR5yJraP
Implements a thin Python client wrapping the existing Pluto server REST API (/api/runs/*) with API-key authentication. This enables the Neptune-to-Pluto migration use case where users need to programmatically read run data (configs, metrics, files) for async evaluation pipelines. API surface: - Client class with explicit auth/host config - list_projects, list_runs, get_run (by numeric or display ID) - get_metrics (returns pandas DataFrame), get_metric_names - get_files, download_file - get_logs, get_statistics, compare_runs, leaderboard - Module-level convenience functions (use PLUTO_API_TOKEN env var) Includes 44 unit tests with mocked HTTP responses. https://claude.ai/code/session_01VUDsz1aDmSAeuThR5yJraP
Summary of ChangesHello @asaiacai, 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 Pluto Python SDK by introducing a dedicated 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
|
There was a problem hiding this comment.
Code Review
This pull request introduces a much-needed read/query API to the Pluto SDK, which is a great step towards feature parity with Neptune and enabling customer migration workflows. However, a high-severity path traversal vulnerability was identified in the download_file method. This vulnerability arises from using untrusted filenames from the server to construct local file paths without sanitization, and it must be addressed by ensuring only the base filename is used for destination paths. While the implementation is well-structured with clear separation between the Client class and convenience functions, and includes a comprehensive test suite and design document, my review also provides suggestions to improve design consistency, code clarity, and resilience.
|
|
||
| dest = Path(destination) | ||
| if dest.is_dir(): | ||
| dest = dest / file_info.get('fileName', file_name) |
There was a problem hiding this comment.
The download_file method is vulnerable to path traversal. It uses the fileName returned by the Pluto server (or the user-supplied file_name as a fallback) to construct the destination path without any sanitization. If a malicious or compromised server returns a fileName containing path traversal sequences (e.g., ../../../../etc/passwd), the client will write the downloaded content to an arbitrary location on the user's file system, potentially overwriting sensitive files. To remediate this, sanitize the filename using os.path.basename() to extract only the filename part.
| dest = dest / file_info.get('fileName', file_name) | |
| if dest.is_dir(): | |
| dest = dest / os.path.basename(file_info.get('fileName', file_name)) |
|
|
||
| ### Module: `pluto.query` | ||
|
|
||
| New module at `pluto/query.py`. All functions are standalone (no global state). Users pass project name and authenticate via the same `PLUTO_API_TOKEN` used for logging. |
There was a problem hiding this comment.
The statement "All functions are standalone (no global state)" is not entirely accurate. The implementation in pluto/query.py uses a shared, module-level _default_client instance for the convenience functions (e.g., pq.list_runs). This is a form of global state. While this is a common and often acceptable pattern for client libraries, the design document should reflect this to avoid confusion. Consider rephrasing to mention that the module-level functions use a shared client instance for convenience.
| client.get_metric_names( | ||
| project: str, | ||
| run_id: int | None = None, # None = all metrics in project | ||
| search: str | None = None, # Filter by name substring | ||
| ) -> list[str] |
There was a problem hiding this comment.
The parameter run_id: int | None = None seems inconsistent with the implementation and the corresponding API endpoint. The implementation uses run_ids: Optional[List[int]] = None, which aligns with the GET /api/runs/metric-names endpoint that accepts a runIds query parameter (a comma-separated list). To maintain consistency between the design and implementation, consider updating the signature here to run_ids: list[int] | None = None.
| """Fetch time-series metric data for a run. | ||
|
|
||
| The server returns up to ``limit`` data points per metric, sampled | ||
| via reservoir sampling when the full series exceeds the limit. | ||
|
|
||
| When *pandas* is installed the return value is a | ||
| :class:`~pandas.DataFrame` with columns ``metric``, ``step``, | ||
| ``value``, ``time``. Otherwise a list of dicts is returned. | ||
|
|
||
| Args: | ||
| project: Project name. | ||
| run_id: Numeric server ID. | ||
| metric_names: Metric names to fetch. ``None`` fetches all. | ||
| limit: Max data points per metric (max 10 000). | ||
|
|
||
| Returns: | ||
| ``pandas.DataFrame`` or ``list[dict]``. | ||
| """ |
There was a problem hiding this comment.
The implementation for fetching multiple metrics makes a separate API call for each metric name. This is an important performance consideration that users should be aware of. Please add a note to the docstring to clarify this behavior.
"""Fetch time-series metric data for a run.
The server returns up to ``limit`` data points per metric, sampled
via reservoir sampling when the full series exceeds the limit.
When *pandas* is installed the return value is a
:class:`~pandas.DataFrame` with columns ``metric``, ``step``,
``value``, ``time``. Otherwise a list of dicts is returned.
.. note::
When multiple metric names are provided, a separate API request is made
for each metric. This may have performance implications if a large
number of metrics are requested.
Args:
project: Project name.
run_id: Numeric server ID.
metric_names: Metric names to fetch. ``None`` fetches all.
limit: Max data points per metric (max 10 000).
Returns:
``pandas.DataFrame`` or ``list[dict]``.
"""| ) -> Any: | ||
| if retry >= _RETRY_MAX: | ||
| raise QueryError(f'Failed after {retry} retries: {error_info}') | ||
| wait = min(_RETRY_WAIT_MIN * (2 ** (retry + 1)), _RETRY_WAIT_MAX) |
There was a problem hiding this comment.
The current retry mechanism uses exponential backoff, which is great. To further improve resilience, especially in scenarios with multiple clients retrying simultaneously, consider adding jitter to the backoff delay. This helps to spread out retry attempts and avoid thundering herd effects. A common approach is to add a small, random amount of time to the delay, for example: import random; ...; wait_with_jitter = wait + random.uniform(0, wait * 0.5).
| assert __import__('sys').platform == 'darwin' | ||
| token = keyring.get_password('pluto', 'pluto') |
There was a problem hiding this comment.
Using assert for control flow is unconventional and can be confusing. It's better to use a standard if statement to check the platform and explicitly raise an exception to trigger the fallback logic. This improves readability.
| assert __import__('sys').platform == 'darwin' | |
| token = keyring.get_password('pluto', 'pluto') | |
| if __import__('sys').platform != 'darwin': | |
| raise AssertionError | |
| token = keyring.get_password('pluto', 'pluto') |
Tested (run the relevant ones):
bash format.sh