diff --git a/deepnote_toolkit/artifacts/_block_frames.py b/deepnote_toolkit/artifacts/_block_frames.py new file mode 100644 index 0000000..383011b --- /dev/null +++ b/deepnote_toolkit/artifacts/_block_frames.py @@ -0,0 +1,164 @@ +"""Stores a block's DataFrame result when the platform asks for it.""" + +import errno +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import requests +from IPython import get_ipython +from IPython.core.interactiveshell import ExecutionResult + +from deepnote_toolkit import env +from deepnote_toolkit.get_webapp_url import ( + get_absolute_userpod_api_url, + get_project_auth_headers, +) +from deepnote_toolkit.logging import get_logger + +from . import _dataframe_storage + +ENABLED_ENV_VAR = "DEEPNOTE_DATAFRAME_STORAGE_ENABLED" +# The webapp's dataframe reader serves block frames from here. +BLOCK_FRAMES_DIR = Path(".deepnote") / "dataframes" +# toolkit/errors accepts only TOOLKIT_IMPORT_ERROR and TOOLKIT_RUNTIME_ERROR, and +# sends the latter to Sentry, grouped by `code`. +REPORT_TYPE = "TOOLKIT_RUNTIME_ERROR" +REPORT_CODE = "DATAFRAME_STORAGE_WRITE_FAILED" +REPORT_TIMEOUT_SECONDS = 2 + + +@dataclass(frozen=True) +class StorageTarget: + """The block whose result the platform asked the current execution to store.""" + + notebook_id: str + block_id: str + + +_reported: set[str] = set() + + +def register_block_frame_storage() -> None: + """Register the post_run_cell hook that stores a block's DataFrame result.""" + get_ipython().events.register("post_run_cell", _on_post_run_cell) + + +def _is_enabled_by_environment() -> bool: + """Return whether the platform switched the feature on for this kernel.""" + value = env.get_env(ENABLED_ENV_VAR) + return value is not None and value.strip().lower() in {"1", "true", "yes", "on"} + + +def _read_storage_target() -> Optional[StorageTarget]: + """Return the storage target in the current execute_request, if it is usable.""" + kernel = getattr(get_ipython(), "kernel", None) + if kernel is None: + return None + + node: object = kernel.get_parent() + for key in ("metadata", "deepnote", "dataframeStorage"): + if not isinstance(node, dict): + return None + node = node.get(key) + if not isinstance(node, dict): + return None + + notebook_id, block_id = node.get("notebookId"), node.get("blockId") + if not ( + _dataframe_storage.is_safe_id(notebook_id) + and _dataframe_storage.is_safe_id(block_id) + ): + return None + return StorageTarget(notebook_id=notebook_id, block_id=block_id) + + +def _post_error_report(context: dict[str, str]) -> None: + """Tell the webapp about a failure, with one attempt and a short timeout. + + Not report_error_to_webapp(): it sends no `code`, retries with sleeps, and logs + a root-logger warning, which reaches the user's cell. + """ + try: + response = requests.post( + get_absolute_userpod_api_url("toolkit/errors"), + json={ + "type": REPORT_TYPE, + "message": "DataFrame storage write failed", + "code": REPORT_CODE, + "context": context, + }, + headers=get_project_auth_headers(), + timeout=REPORT_TIMEOUT_SECONDS, + ) + response.raise_for_status() + except Exception as e: # pylint: disable=broad-exception-caught + get_logger().warning( + "Failed to report DataFrame storage failure: %s", type(e).__name__ + ) + + +def _report_failure(exc: Exception) -> None: + """Report a failed write once per kernel session per errno or exception type. + + The context carries the exception's type, never its message: pyarrow quotes + the offending value in its messages, and frames often hold personal data. + """ + exception_type = f"{type(exc).__module__}.{type(exc).__qualname__}" + if isinstance(exc, OSError) and exc.errno is not None: + code = errno.errorcode.get(exc.errno, str(exc.errno)) + key = f"io:{code}" + context = {"kind": "io", "errno": code, "exception_type": exception_type} + else: + key = f"error:{exception_type}" + context = {"kind": "error", "exception_type": exception_type} + + if key in _reported: + return + _reported.add(key) + _post_error_report(context) + + +def _store_result_frame(result: ExecutionResult) -> None: + """Store the cell's DataFrame result; report any failure but EROFS to the webapp.""" + try: + target = _read_storage_target() + # Helper executions don't store history; IPython 9 may give them an + # execution_count anyway. + if target is None or not result.info.store_history: + return + if not _dataframe_storage.is_storable_frame(result.result): + return + directory = ( + _dataframe_storage.resolve_project_root() + / BLOCK_FRAMES_DIR + / target.notebook_id + / target.block_id + ) + _dataframe_storage.store_frame(result.result, directory) + except Exception as exc: # pylint: disable=broad-exception-caught + if isinstance(exc, OSError) and exc.errno == errno.EROFS: + return + _report_failure(exc) + + +def _on_post_run_cell(result: Optional[ExecutionResult]) -> None: + """Store the cell's DataFrame result when the platform asked for it. + + IPython prints a callback's exceptions into the cell as "Error in callback", + so failures go to the webapp instead. Warnings are ignored because they would + reach the cell too. + """ + if result is None or not _is_enabled_by_environment() or not result.success: + return + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + _store_result_frame(result) + except KeyboardInterrupt as exc: + # Re-raised, the interrupt would leave the cell successful; ipykernel + # builds the reply from `result` after this hook. + result.error_in_exec = exc + get_ipython().showtraceback() diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index c8c6b98..9402dad 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -10,10 +10,15 @@ from typing import TYPE_CHECKING, Callable, Optional, Union import pandas as pd -from typing_extensions import TypeAlias, assert_never +from typing_extensions import TypeAlias, TypeGuard, assert_never from deepnote_toolkit.config import resolve_notebook_root from deepnote_toolkit.ocelots.types import PandasDF, PolarsEagerDF +from deepnote_toolkit.ocelots.utils import ( + is_pandas_dataframe, + is_polars_eager_dataframe, +) +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview from ._dataframe_storage_manifest import ( Manifest, @@ -44,11 +49,18 @@ _SAFE_ID_REGEX = re.compile(r"[A-Za-z0-9_-]{1,128}") -def is_safe_id(value: object) -> bool: +def is_safe_id(value: object) -> TypeGuard[str]: """Return whether `value` can be used as a single path segment.""" return isinstance(value, str) and _SAFE_ID_REGEX.fullmatch(value) is not None +def is_storable_frame(value: object) -> TypeGuard[StorableFrame]: + """Return whether `value` is a pandas or polars DataFrame, not a query preview.""" + if isinstance(value, DeepnoteQueryPreview): + return False + return is_pandas_dataframe(value) or is_polars_eager_dataframe(value) + + def resolve_project_root() -> Path: """Return the project's root directory, the one set_notebook_path() uses.""" return resolve_notebook_root() @@ -242,13 +254,12 @@ def _delete_replaced_frames( (directory / entry.name).unlink() -def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]: - """Store the frame as the current one under `name` and return its reference ID. +def store_frame(frame: StorableFrame, directory: Path) -> Optional[str]: + """Store the frame as the current one in `directory` and return its reference ID. Returns None, having stored nothing, if the filesystem is read-only. EROFS from the manifest step on raises like any other error: the new frame may be current. """ - directory = root / FRAMES_DIR / name reference_id = str(uuid.uuid4()) new_file = f"{reference_id}.arrow" try: diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index 17492e4..92c6ac7 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -4,10 +4,6 @@ import shutil from typing import Any, Optional -from deepnote_toolkit.ocelots.utils import ( - is_pandas_dataframe, - is_polars_eager_dataframe, -) from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview from . import _dataframe_storage @@ -62,14 +58,16 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: "block's output mode to DataFrame, or pass pd.DataFrame(df) to store " "these rows" ) - if not (is_pandas_dataframe(df) or is_polars_eager_dataframe(df)): + if not _dataframe_storage.is_storable_frame(df): raise TypeError( "store_dataframe needs a pandas or polars DataFrame, " f"got {type(df).__qualname__}" ) root = _dataframe_storage.resolve_project_root() - return _dataframe_storage.store_frame(df, name, root) + return _dataframe_storage.store_frame( + df, root / _dataframe_storage.FRAMES_DIR / name + ) def delete_dataframe(name: str) -> bool: diff --git a/deepnote_toolkit/runtime_initialization.py b/deepnote_toolkit/runtime_initialization.py index 63b7355..a94600a 100644 --- a/deepnote_toolkit/runtime_initialization.py +++ b/deepnote_toolkit/runtime_initialization.py @@ -8,6 +8,7 @@ from deepnote_toolkit.runtime_patches import apply_runtime_patches +from .artifacts._block_frames import register_block_frame_storage from .dataframe_utils import add_formatters from .execute_post_start_hooks import execute_post_start_hooks from .logging import LoggerManager @@ -52,6 +53,14 @@ def init_deepnote_runtime(): except Exception as e: # pylint: disable=broad-exception-caught logger.error("Failed to add output middleware with a error: %s", e) + try: + logger.debug("Registering the block DataFrame storage hook.") + register_block_frame_storage() + except Exception as e: # pylint: disable=broad-exception-caught + logger.error( + "Failed to register the block DataFrame storage hook with error: %s", e + ) + # Disable sqlparse grouping limits for large analytical queries try: logger.debug("Configuring sqlparse limits.") diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index f1bad71..f81a83a 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -13,6 +13,8 @@ FRAME_A, NAME, arrow_files, + cleared_query_preview, + query_preview, stored_table, ) @@ -51,17 +53,6 @@ def test_invalid_name_raises_before_anything_is_written( assert list(tmp_project_root.iterdir()) == [] -def _query_preview() -> DeepnoteQueryPreview: - return DeepnoteQueryPreview({"a": [1]}, deepnote_query="select 1") - - -def _cleared_query_preview() -> DeepnoteQueryPreview: - preview = _query_preview() - # Mutating the frame clears its attached SQL query metadata. - preview["b"] = 2 - return preview - - UNSUPPORTED_VALUES = [ pytest.param(lambda: 1, id="int"), pytest.param(lambda: None, id="none"), @@ -71,9 +62,9 @@ def _cleared_query_preview() -> DeepnoteQueryPreview: ] QUERY_PREVIEWS = [ - pytest.param(_query_preview, id="query-preview"), - pytest.param(_cleared_query_preview, id="cleared-query-preview"), - pytest.param(lambda: _query_preview().copy(), id="copied-query-preview"), + pytest.param(query_preview, id="query-preview"), + pytest.param(cleared_query_preview, id="cleared-query-preview"), + pytest.param(lambda: query_preview().copy(), id="copied-query-preview"), ] diff --git a/tests/unit/artifacts/test_block_frames.py b/tests/unit/artifacts/test_block_frames.py new file mode 100644 index 0000000..c405869 --- /dev/null +++ b/tests/unit/artifacts/test_block_frames.py @@ -0,0 +1,702 @@ +"""Tests for the post_run_cell hook that stores a block's DataFrame result.""" + +import builtins +import errno +import json +import re +import sys +import warnings +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, Iterator, List, Optional + +import pandas as pd +import polars as pl +import pyarrow as pa +import pytest +import requests +import responses +from IPython.core.interactiveshell import ExecutionResult, InteractiveShell +from traitlets.config import Config + +from deepnote_toolkit.artifacts import _block_frames, _dataframe_storage +from deepnote_toolkit.artifacts._block_frames import register_block_frame_storage +from deepnote_toolkit.ocelots.data_preview import DeepnoteDataFrameWithDataPreview +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview +from tests.unit.helpers.dataframe_storage import ( + cleared_query_preview, + manifest_in, + query_preview, + table_in, +) + +NOTEBOOK_ID = "nb1" +BLOCK_ID = "blk1" +REPORT_URL = "http://userpod/toolkit/errors" +REPORT_TYPE = "TOOLKIT_RUNTIME_ERROR" +REPORT_CODE = "DATAFRAME_STORAGE_WRITE_FAILED" + + +def storage_request(**overrides: Any) -> Dict[str, Any]: + """Return the request message that asks the kernel to store a block's result.""" + return { + "metadata": { + "deepnote": { + "dataframeStorage": { + "notebookId": NOTEBOOK_ID, + "blockId": BLOCK_ID, + **overrides, + } + } + } + } + + +def block_dir(root: Path) -> Path: + return root / ".deepnote" / "dataframes" / NOTEBOOK_ID / BLOCK_ID + + +def block_arrow_files(root: Path) -> List[Path]: + return sorted(block_dir(root).glob("*.arrow")) + + +@pytest.fixture +def parent() -> Dict[str, Any]: + """The execute_request the fake kernel reports; tests mutate it in place.""" + return storage_request() + + +@pytest.fixture +def reports(monkeypatch: pytest.MonkeyPatch) -> Iterator[SimpleNamespace]: + """Replace the webapp transport; `bodies` holds each posted JSON body.""" + transport = SimpleNamespace(bodies=[], calls=[], mock=None) + + def record(request: Any) -> Any: + transport.calls.append(request) + transport.bodies.append(json.loads(request.body)) + return 200, {}, "" + + monkeypatch.setattr( + _block_frames, + "get_absolute_userpod_api_url", + lambda path: f"http://userpod/{path}", + ) + _block_frames._reported.clear() + with responses.RequestsMock(assert_all_requests_are_fired=False) as mock: + mock.add_callback(responses.POST, REPORT_URL, callback=record) + transport.mock = mock + yield transport + _block_frames._reported.clear() + + +@pytest.fixture +def shell( + parent: Dict[str, Any], + reports: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> Iterator[InteractiveShell]: + """A shell whose kernel reports `parent`, with the hook registered and enabled.""" + monkeypatch.setenv(_block_frames.ENABLED_ENV_VAR, "true") + # A shell swaps in its own __main__ and adds builtins; leave both as found + monkeypatch.setitem(sys.modules, "__main__", sys.modules["__main__"]) + builtins_before = set(vars(builtins)) + # No history: it would write a sqlite file into the user's profile, and its + # sqlite3 usage raises a DeprecationWarning on Python 3.12+ + config = Config() + config.HistoryManager.enabled = False + InteractiveShell.clear_instance() + shell = InteractiveShell.instance(config=config) + shell.ast_node_interactivity = "all" + shell.kernel = SimpleNamespace(get_parent=lambda channel=None: parent) + register_block_frame_storage() + yield shell + shell.events.unregister("post_run_cell", _block_frames._on_post_run_cell) + InteractiveShell.clear_instance() + for name in set(vars(builtins)) - builtins_before: + delattr(builtins, name) + + +def _assert_nothing_stored(tmp_project_root: Path, reports: SimpleNamespace) -> None: + assert list(tmp_project_root.iterdir()) == [] + assert reports.bodies == [] + + +def _run( + shell: InteractiveShell, value: Any, store_history: bool = True +) -> ExecutionResult: + shell.user_ns["df"] = value + return shell.run_cell("df", store_history=store_history) + + +def test_stores_the_result_in_the_blocks_directory( + shell: InteractiveShell, tmp_project_root: Path +) -> None: + result = _run(shell, pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})) + + assert result.success + directory = tmp_project_root / ".deepnote" / "dataframes" / "nb1" / "blk1" + assert table_in(directory).to_pydict() == { + "a": [1, 2, 3], + "b": ["x", "y", "z"], + } + assert manifest_in(directory).rows == 3 + + +def test_stores_polars_frame(shell: InteractiveShell, tmp_project_root: Path) -> None: + _run(shell, pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})) + + assert table_in(block_dir(tmp_project_root)).to_pydict() == { + "a": [1, 2, 3], + "b": ["x", "y", "z"], + } + + +def test_env_var_unset_writes_nothing( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED") + + result = _run(shell, pd.DataFrame({"a": [1]})) + + assert result.success + assert _callback_errors(_displayed(capsys)) == [] + _assert_nothing_stored(tmp_project_root, reports) + + +def test_env_var_set_to_another_value_writes_nothing( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED", "2") + + result = _run(shell, pd.DataFrame({"a": [1]})) + + assert result.success + _assert_nothing_stored(tmp_project_root, reports) + + +@pytest.mark.parametrize("env_value", ["1", "TRUE", " yes ", "on"]) +def test_env_var_on_writes_the_frame( + shell: InteractiveShell, + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + env_value: str, +) -> None: + monkeypatch.setenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED", env_value) + + _run(shell, pd.DataFrame({"a": [1]})) + + assert table_in(block_dir(tmp_project_root)).to_pydict() == {"a": [1]} + + +def test_shell_without_kernel_writes_nothing( + shell: InteractiveShell, tmp_project_root: Path, reports: SimpleNamespace +) -> None: + del shell.kernel + + result = _run(shell, pd.DataFrame({"a": [1]})) + + assert result.success + _assert_nothing_stored(tmp_project_root, reports) + + +def test_missing_result_is_ignored( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, +) -> None: + """ipykernel can trigger post_run_cell without a result.""" + shell.events.trigger("post_run_cell", None) + + assert _callback_errors(_displayed(capsys)) == [] + _assert_nothing_stored(tmp_project_root, reports) + + +def test_helper_execution_writes_nothing( + shell: InteractiveShell, tmp_project_root: Path, reports: SimpleNamespace +) -> None: + result = _run(shell, pd.DataFrame({"a": [1]}), store_history=False) + + assert result.success + _assert_nothing_stored(tmp_project_root, reports) + + +def test_interrupted_cell_writes_nothing( + shell: InteractiveShell, tmp_project_root: Path, reports: SimpleNamespace +) -> None: + shell.user_ns["df"] = pd.DataFrame({"a": [1]}) + + result = shell.run_cell("df\nraise KeyboardInterrupt", store_history=True) + + assert not result.success + _assert_nothing_stored(tmp_project_root, reports) + + +@pytest.mark.parametrize( + "value", + [pd.Series([1, 2]), pl.DataFrame({"a": [1]}).lazy()], + ids=["series", "lazyframe"], +) +def test_non_dataframe_result_writes_nothing( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + value: Any, +) -> None: + _run(shell, value) + + _assert_nothing_stored(tmp_project_root, reports) + + +@pytest.mark.parametrize( + "make_preview", + [ + query_preview, + cleared_query_preview, + lambda: query_preview().copy(), + ], + ids=["preview", "cleared-query", "copy"], +) +def test_sql_query_preview_writes_nothing( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + make_preview: Callable[[], DeepnoteQueryPreview], +) -> None: + preview = make_preview() + assert isinstance(preview, DeepnoteQueryPreview) + + _run(shell, preview) + + _assert_nothing_stored(tmp_project_root, reports) + + +def test_pyspark_result_writes_nothing( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + spark: Any, +) -> None: + wrapped = DeepnoteDataFrameWithDataPreview(spark.range(3)) + + _run(shell, wrapped) + + _assert_nothing_stored(tmp_project_root, reports) + + +# Displaying the result runs pyspark 3.5 code that uses deprecated distutils and +# numpy APIs; Python < 3.13 gets pyspark 3.5. +@pytest.mark.filterwarnings("ignore::DeprecationWarning:pyspark") +def test_pandas_on_spark_result_writes_nothing( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + spark_without_ansi_mode: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # pyspark.pandas warns on import unless told that pyarrow may ignore time zones + monkeypatch.setenv("PYARROW_IGNORE_TIMEZONE", "1") + # pyspark 3.5, which Python < 3.13 gets, imports distutils, gone since Python 3.12 + pytest.importorskip("pyspark.pandas") + wrapped = DeepnoteDataFrameWithDataPreview( + spark_without_ansi_mode.range(3).pandas_api() + ) + + _run(shell, wrapped) + + _assert_nothing_stored(tmp_project_root, reports) + + +def _displayed(capsys: pytest.CaptureFixture) -> str: + """Everything printed since the last read, with execution counts masked.""" + captured = capsys.readouterr() + return re.sub(r"Out\[\d+\]", "Out[N]", captured.out + captured.err) + + +def _callback_errors(printed: str) -> List[str]: + """IPython's "Error in callback" lines, skipping traceback lines that quote it.""" + return re.findall(r"^Error in callback .*", printed, flags=re.MULTILINE) + + +def _output_without_storage( + shell: InteractiveShell, + value: Any, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> str: + """What the cell prints when the writer does nothing.""" + monkeypatch.setenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED", "false") + _run(shell, value) + output = _displayed(capsys) + monkeypatch.setenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED", "true") + return output + + +def test_mixed_type_column_failure_names_the_exception_type_only( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + frame = pd.DataFrame({"a": [1, "two"]}) + expected_output = _output_without_storage(shell, frame, capsys, monkeypatch) + + result = _run(shell, frame) + + assert result.success + assert _displayed(capsys) == expected_output + assert reports.bodies == [ + { + "type": REPORT_TYPE, + "message": "DataFrame storage write failed", + "code": REPORT_CODE, + "context": {"kind": "error", "exception_type": "pyarrow.lib.ArrowInvalid"}, + } + ] + + +def test_each_exception_type_is_reported_separately( + shell: InteractiveShell, tmp_project_root: Path, reports: SimpleNamespace +) -> None: + _run(shell, pd.DataFrame({"a": [1, "two"]})) + _run(shell, pd.DataFrame([[1, 2]], columns=["a", "a"])) + _run(shell, pd.DataFrame({"a": [1, "two"]})) + + assert [body["context"]["exception_type"] for body in reports.bodies] == [ + "pyarrow.lib.ArrowInvalid", + "builtins.ValueError", + ] + + +def test_report_is_one_short_authenticated_attempt( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + _block_frames, + "get_project_auth_headers", + lambda: {"Authorization": "Bearer t"}, + ) + + _run(shell, pd.DataFrame({"a": [1, "two"]})) + + (request,) = reports.calls + assert request.url == REPORT_URL + assert request.method == "POST" + assert request.headers["Content-Type"] == "application/json" + assert request.headers["Authorization"] == "Bearer t" + assert request.req_kwargs["timeout"] == 2 + + +def _refuse_url(path: str) -> str: + raise ValueError("no webapp url") + + +@pytest.mark.parametrize( + "break_report, logged_type", + [ + pytest.param( + lambda reports, monkeypatch: reports.mock.replace( + responses.POST, REPORT_URL, body=requests.ConnectionError("down") + ), + "ConnectionError", + id="connection", + ), + pytest.param( + lambda reports, monkeypatch: reports.mock.replace( + responses.POST, REPORT_URL, body=requests.Timeout("slow") + ), + "Timeout", + id="timeout", + ), + pytest.param( + lambda reports, monkeypatch: reports.mock.replace( + responses.POST, REPORT_URL, status=400 + ), + "HTTPError", + id="rejected-by-the-webapp", + ), + pytest.param( + lambda reports, monkeypatch: monkeypatch.setattr( + _block_frames, "get_absolute_userpod_api_url", _refuse_url + ), + "ValueError", + id="no-url", + ), + ], +) +def test_report_that_cannot_be_delivered_is_silent( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + break_report: Callable[[SimpleNamespace, pytest.MonkeyPatch], Any], + logged_type: str, +) -> None: + frame = pd.DataFrame({"a": [1, "two"]}) + expected_output = _output_without_storage(shell, frame, capsys, monkeypatch) + logged: List[Any] = [] + break_report(reports, monkeypatch) + monkeypatch.setattr( + _block_frames, + "get_logger", + lambda: SimpleNamespace(warning=lambda *args: logged.append(args)), + ) + + result = _run(shell, frame) + + assert result.success + assert _displayed(capsys) == expected_output + assert caplog.records == [] + assert [args[1:] for args in logged] == [(logged_type,)] + + +def test_writer_adds_nothing_to_the_cell( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + unserializable = pd.DataFrame({"a": [1, "two"]}) + storable = pd.DataFrame({"a": [1]}) + expected_output = _output_without_storage( + shell, unserializable, capsys, monkeypatch + ) + _output_without_storage(shell, storable, capsys, monkeypatch) + real_to_arrow_table = _dataframe_storage._to_arrow_table + + def noisy_to_arrow_table(frame: Any) -> pa.Table: + warnings.warn("noisy", UserWarning) + return real_to_arrow_table(frame) + + monkeypatch.setattr(_dataframe_storage, "_to_arrow_table", noisy_to_arrow_table) + + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + _run(shell, unserializable) + _run(shell, storable) + + assert "noisy" not in [str(w.message) for w in recorded] + assert caplog.records == [] + assert len(reports.bodies) == 1 + assert _displayed(capsys) == expected_output + + +def test_unexpected_error_is_reported_not_raised( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + frame = pd.DataFrame({"a": [1]}) + expected_output = _output_without_storage(shell, frame, capsys, monkeypatch) + + def broken_root() -> Path: + raise RuntimeError("no root") + + monkeypatch.setattr(_dataframe_storage, "resolve_project_root", broken_root) + + result = _run(shell, frame) + + assert result.success + assert _displayed(capsys) == expected_output + assert [body["context"] for body in reports.bodies] == [ + {"kind": "error", "exception_type": "builtins.RuntimeError"} + ] + + +def test_eperm_from_close_is_reported_once( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + parent: Dict[str, Any], + close_upload_error: Dict[str, Optional[OSError]], +) -> None: + close_upload_error["error"] = PermissionError( + errno.EPERM, "Operation not permitted" + ) + + first = _run(shell, pd.DataFrame({"a": [1]})) + parent["metadata"]["deepnote"]["dataframeStorage"]["blockId"] = "blk2" + second = _run(shell, pd.DataFrame({"a": [2]})) + + assert first.success and second.success + assert [body["context"] for body in reports.bodies] == [ + {"kind": "io", "errno": "EPERM", "exception_type": "builtins.PermissionError"} + ] + assert not list(tmp_project_root.rglob("*.arrow")) + assert not list(tmp_project_root.rglob("manifest.json")) + + +def test_each_io_error_code_is_reported_once_and_not_swallowed( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + close_upload_error: Dict[str, Optional[OSError]], +) -> None: + close_upload_error["error"] = OSError(errno.EIO, "Input/output error") + first = _run(shell, pd.DataFrame({"a": [1]})) + second = _run(shell, pd.DataFrame({"a": [1]})) + close_upload_error["error"] = OSError(errno.ENOSPC, "No space left on device") + third = _run(shell, pd.DataFrame({"a": [1]})) + + assert (first.success, second.success, third.success) == (True, True, True) + assert [body["context"] for body in reports.bodies] == [ + {"kind": "io", "errno": "EIO", "exception_type": "builtins.OSError"}, + {"kind": "io", "errno": "ENOSPC", "exception_type": "builtins.OSError"}, + ] + + +def _write_a_batch_then_raise( + monkeypatch: pytest.MonkeyPatch, exc: BaseException +) -> None: + real_write_batch = pa.ipc.RecordBatchFileWriter.write_batch + + def write_a_batch_then_raise( + self: pa.ipc.RecordBatchFileWriter, batch: pa.RecordBatch, **kwargs: Any + ) -> None: + real_write_batch(self, batch, **kwargs) + raise exc + + monkeypatch.setattr( + pa.ipc.RecordBatchFileWriter, "write_batch", write_a_batch_then_raise + ) + + +def test_erofs_after_the_manifest_step_is_silent( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + frame = pd.DataFrame({"a": [3]}) + expected_output = _output_without_storage(shell, frame, capsys, monkeypatch) + _run(shell, pd.DataFrame({"a": [1]})) + _run(shell, pd.DataFrame({"a": [2]})) + block = block_dir(tmp_project_root) + previous = manifest_in(block).previous + assert previous is not None + replaced = block / previous + real_unlink = Path.unlink + + def unlink_refused_for_the_replaced_frame( + self: Path, *args: Any, **kwargs: Any + ) -> None: + if self == replaced: + raise OSError(errno.EROFS, "Read-only file system") + real_unlink(self, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", unlink_refused_for_the_replaced_frame) + capsys.readouterr() + + result = _run(shell, frame) + + assert result.success + assert table_in(block).to_pydict() == {"a": [3]} + assert replaced.exists() + assert reports.bodies == [] + assert _displayed(capsys) == expected_output + + +def test_interrupt_during_the_write_fails_the_execution( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run(shell, pd.DataFrame({"a": [1]})) + block = block_dir(tmp_project_root) + manifest_bytes = (block / "manifest.json").read_bytes() + _write_a_batch_then_raise(monkeypatch, KeyboardInterrupt()) + capsys.readouterr() + + result = _run(shell, pd.DataFrame({"a": [2]})) + + assert isinstance(result.error_in_exec, KeyboardInterrupt) + assert not result.success + assert "KeyboardInterrupt" in _displayed(capsys) + assert reports.bodies == [] + assert (block / "manifest.json").read_bytes() == manifest_bytes + assert len(block_arrow_files(tmp_project_root)) == 1 + + +def test_interrupt_while_reporting_a_failure_fails_the_execution( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, +) -> None: + def interrupted(request: Any) -> Any: + raise KeyboardInterrupt + + reports.mock.replace( + responses.CallbackResponse(responses.POST, REPORT_URL, callback=interrupted) + ) + capsys.readouterr() + + result = _run(shell, pd.DataFrame({"a": [1, "two"]})) + + assert isinstance(result.error_in_exec, KeyboardInterrupt) + assert not result.success + printed = _displayed(capsys) + assert "KeyboardInterrupt" in printed + assert _callback_errors(printed) == [] + + +@pytest.mark.parametrize( + "request_message", + [ + pytest.param({}, id="no-metadata"), + pytest.param({"metadata": {}}, id="no-deepnote"), + pytest.param({"metadata": {"deepnote": {}}}, id="no-storage"), + pytest.param({"metadata": {"deepnote": "x"}}, id="deepnote-not-a-dict"), + pytest.param( + {"metadata": {"deepnote": {"dataframeStorage": []}}}, + id="storage-not-a-dict", + ), + pytest.param( + {"metadata": {"deepnote": {"dataframeStorage": {"blockId": BLOCK_ID}}}}, + id="no-notebook-id", + ), + pytest.param( + { + "metadata": { + "deepnote": {"dataframeStorage": {"notebookId": NOTEBOOK_ID}} + } + }, + id="no-block-id", + ), + pytest.param(storage_request(notebookId="../x"), id="notebook-traversal"), + pytest.param(storage_request(blockId="../x"), id="block-traversal"), + ], +) +def test_unusable_request_metadata_writes_nothing( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, + parent: Dict[str, Any], + request_message: Dict[str, Any], +) -> None: + parent.clear() + parent.update(request_message) + + result = _run(shell, pd.DataFrame({"a": [1]})) + + assert result.success + _assert_nothing_stored(tmp_project_root, reports) diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 2d0c4eb..73f2ae9 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -12,6 +12,7 @@ from deepnote_core.pydantic_compat_helpers import model_validate_compat from deepnote_toolkit.artifacts._dataframe_storage_manifest import Manifest +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview NAME = "revenue" DAY = 24 * 60 * 60 @@ -33,17 +34,25 @@ def non_arrow_files(root: Path) -> List[str]: ) -def stored_manifest(root: Path, name: str = NAME) -> Manifest: - raw = json.loads((frame_dir(root, name) / "manifest.json").read_text()) +def manifest_in(directory: Path) -> Manifest: + raw = json.loads((directory / "manifest.json").read_text()) return model_validate_compat(Manifest, raw) -def stored_table(root: Path, name: str = NAME) -> pa.Table: - """The frame the name's manifest currently names.""" - path = frame_dir(root, name) / stored_manifest(root, name).file +def table_in(directory: Path) -> pa.Table: + """The frame the directory's manifest currently names.""" + path = directory / manifest_in(directory).file return pa.ipc.open_file(path).read_all() +def stored_manifest(root: Path, name: str = NAME) -> Manifest: + return manifest_in(frame_dir(root, name)) + + +def stored_table(root: Path, name: str = NAME) -> pa.Table: + return table_in(frame_dir(root, name)) + + def backdate(path: Path, age_seconds: float) -> None: mtime = time.time() - age_seconds os.utime(path, (mtime, mtime), follow_symlinks=False) @@ -82,3 +91,14 @@ def probe_failure(frame: Any) -> Optional[type]: except Exception as exc: # pylint: disable=broad-exception-caught return type(exc) return None + + +def query_preview() -> DeepnoteQueryPreview: + return DeepnoteQueryPreview({"a": [1]}, deepnote_query="select 1") + + +def cleared_query_preview() -> DeepnoteQueryPreview: + preview = query_preview() + # Mutating the frame clears its attached SQL query metadata. + preview["b"] = 2 + return preview diff --git a/tests/unit/test_runtime_initialization.py b/tests/unit/test_runtime_initialization.py new file mode 100644 index 0000000..5ef5e4c --- /dev/null +++ b/tests/unit/test_runtime_initialization.py @@ -0,0 +1,63 @@ +import logging +from types import SimpleNamespace +from typing import Iterator +from unittest import mock + +import pytest + +from deepnote_toolkit.runtime_initialization import init_deepnote_runtime + +STEPS = ( + "apply_runtime_patches", + "add_formatters", + "add_output_middleware", + "configure_sqlparse_limits", + "set_ipython_page_printer", + "set_notebook_path", + "set_integration_env", + "execute_post_start_hooks", + "register_block_frame_storage", +) + + +@pytest.fixture +def steps() -> Iterator[SimpleNamespace]: + """Replace every startup step so the test sees only the wiring.""" + with ( + mock.patch.multiple( + "deepnote_toolkit.runtime_initialization", + autospec=True, + LoggerManager=mock.DEFAULT, + **{name: mock.DEFAULT for name in STEPS}, + ) as mocks, + mock.patch("IPython.get_ipython"), + mock.patch("psycopg2.extensions.set_wait_callback"), + ): + logger = mock.create_autospec(logging.Logger, instance=True) + mocks["LoggerManager"].return_value.get_logger.return_value = logger + yield SimpleNamespace(logger=logger, **mocks) + + +def test_registers_the_dataframe_storage_hook(steps: SimpleNamespace) -> None: + init_deepnote_runtime() + + steps.register_block_frame_storage.assert_called_once_with() + steps.logger.error.assert_not_called() + + +def test_failing_to_register_the_hook_does_not_stop_later_steps( + steps: SimpleNamespace, +) -> None: + steps.register_block_frame_storage.side_effect = RuntimeError("boom") + + init_deepnote_runtime() + + steps.set_notebook_path.assert_called_once_with() + steps.set_integration_env.assert_called_once_with() + steps.execute_post_start_hooks.assert_called_once_with() + steps.logger.error.assert_called_once() + message, error = steps.logger.error.call_args.args + assert ( + message == "Failed to register the block DataFrame storage hook with error: %s" + ) + assert str(error) == "boom"