diff --git a/deepnote_toolkit/dataframe_storage.py b/deepnote_toolkit/dataframe_storage.py index df2e432..9c088a5 100644 --- a/deepnote_toolkit/dataframe_storage.py +++ b/deepnote_toolkit/dataframe_storage.py @@ -1,14 +1,24 @@ -"""Writes DataFrames into the project's filesystem as named Arrow files.""" +"""Writes DataFrames into the project's filesystem as named Arrow files. + +Besides the function behind `artifacts.store_dataframe`, this holds the hook that +stores the DataFrame a block with a storage setting returns. +""" import contextlib +import errno import os import re import signal import time import uuid +import warnings +from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, Iterator, Literal, Optional +import requests +from IPython import get_ipython +from IPython.core.interactiveshell import ExecutionResult from packaging.version import Version from deepnote_toolkit import env @@ -19,14 +29,25 @@ read_manifest, write_manifest, ) +from deepnote_toolkit.get_webapp_url import ( + get_absolute_userpod_api_url, + get_project_auth_headers, +) +from deepnote_toolkit.logging import get_logger from deepnote_toolkit.ocelots.utils import ( is_pandas_dataframe, is_polars_eager_dataframe, ) from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview +ENABLED_ENV_VAR = "DEEPNOTE_DATAFRAME_STORAGE_ENABLED" FRAMES_DIR = Path("deepnote_dataframes") STALE_FRAME_AGE_SECONDS = 24 * 60 * 60 +# The webapp's toolkit/errors endpoint rejects any other type with a 400, and groups +# runtime errors in Sentry by `code`. +REPORT_TYPE = "TOOLKIT_RUNTIME_ERROR" +REPORT_CODE = "DATAFRAME_STORAGE_WRITE_FAILED" +REPORT_TIMEOUT_SECONDS = 2 # Each record batch is one page of the read endpoint; pandas already writes 64k-row # batches, and polars only honours a batch size from 1.37. POLARS_BATCH_ROWS = 65_536 @@ -37,6 +58,18 @@ _SAFE_ID = re.compile(r"[A-Za-z0-9_-]{1,128}") +@dataclass(frozen=True) +class StorageTarget: + """What the platform asked the current execution to store, and for which block.""" + + name: str + block_id: str + + +# Failures already reported in this kernel session, keyed by errno or exception type. +_reported: set[str] = set() + + def is_safe_id(value: object) -> bool: """Return whether `value` can be used as a single path segment. @@ -102,6 +135,18 @@ def _sigint_handled_by(handler: Callable[..., Any]) -> Iterator[None]: signal.signal(signal.SIGINT, previous) +@contextlib.contextmanager +def _sigint_raises() -> Iterator[None]: + """Make SIGINT raise KeyboardInterrupt. + + Ordinary cells already run post_run_cell under this handler. A cell with + top-level await runs it under one that only queues the interrupt, so without + this, Stop would wait for the whole write to finish. + """ + with _sigint_handled_by(signal.default_int_handler): + yield + + @contextlib.contextmanager def _sigint_deferred() -> Iterator[None]: """Hold SIGINT back until the block ends, then raise it as KeyboardInterrupt.""" @@ -181,9 +226,17 @@ def _delete_replaced_frames( (directory / entry.name).unlink() -def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str: +def store_frame( + frame: Any, + kind: FrameKind, + name: str, + root: Path, + block_id: Optional[str] = None, +) -> str: """Store the frame as the current one under `name` and return its reference ID. + `block_id` records the block whose result this is; a call stores none. + The order matters because S3 replaces the manifest in one step: the frame is complete before the manifest names it, and nothing is deleted before then. An interrupt stops everything but the manifest update, which finishes first: @@ -200,7 +253,9 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str: _write_frame_file(frame, kind, directory / new_file, _select_compression()) with _sigint_deferred(): old = read_manifest(directory) - write_manifest(directory, new_file, old.file if old else None, len(frame)) + write_manifest( + directory, new_file, old.file if old else None, len(frame), block_id + ) committed = True except BaseException: if not committed: @@ -209,3 +264,134 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str: raise _delete_replaced_frames(directory, old, new_file) return reference_id + + +def register_dataframe_storage() -> None: + """Register the post_run_cell hook that stores the results of blocks with a setting. + + Registration is unconditional: whether a given execution stores a frame is + decided on every execution from the environment and the request metadata. + """ + 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. + + The metadata is written by other services, so every level is validated and the + name must be safe to use as a path segment. + """ + kernel = getattr(get_ipython(), "kernel", None) + if kernel is None: + return None + + node: Any = 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 + + name, block_id = node.get("name"), node.get("blockId") + if not (is_safe_id(name) and is_safe_id(block_id)): + return None + return StorageTarget(name=name, 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(): that logs a warning to the root logger before + every attempt, and IPython prints root-logger warnings into 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, reporting any failure to the webapp. + + Only KeyboardInterrupt propagates, so that Stop during the report is not lost. + """ + try: + target = _read_storage_target() + # Helper executions (paging, export, the variable explorer) do not store + # history. IPython 9 gives them an execution_count anyway, so ask the info. + if target is None or not result.info.store_history: + return + kind = classify_frame(result.result) + if kind is None: + return + root = resolve_project_root() + # a read-only mount is a deliberate choice, not a failure + if is_read_only(root): + return + with _sigint_raises(): + store_frame(result.result, kind, target.name, root, target.block_id) + 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. + + Never raises and never emits warnings, because IPython prints both into the + user's cell. A failed write is reported to the webapp instead. + """ + 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: + # The reply is already decided, so this is the only way to make the + # executor see the interrupt and stop its queue. + result.error_in_exec = exc + get_ipython().showtraceback() diff --git a/deepnote_toolkit/dataframe_storage_manifest.py b/deepnote_toolkit/dataframe_storage_manifest.py index 0bd535b..657de70 100644 --- a/deepnote_toolkit/dataframe_storage_manifest.py +++ b/deepnote_toolkit/dataframe_storage_manifest.py @@ -62,13 +62,19 @@ def read_manifest(directory: Path) -> Optional[Manifest]: def write_manifest( - directory: Path, new_file: str, previous: Optional[str], rows: int + directory: Path, + new_file: str, + previous: Optional[str], + rows: int, + block_id: Optional[str] = None, ) -> None: - """Point the manifest at `new_file`.""" + """Point the manifest at `new_file`, recording the block that wrote it, if any.""" manifest: dict[str, Any] = {"version": MANIFEST_VERSION, "file": new_file} if previous is not None: manifest["previous"] = previous manifest["written_at"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") manifest["rows"] = rows manifest["toolkit_version"] = __version__ + if block_id is not None: + manifest["block_id"] = block_id (directory / MANIFEST_NAME).write_text(json.dumps(manifest), encoding="utf-8") diff --git a/deepnote_toolkit/runtime_initialization.py b/deepnote_toolkit/runtime_initialization.py index 63b7355..6753185 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 .dataframe_storage import register_dataframe_storage from .dataframe_utils import add_formatters from .execute_post_start_hooks import execute_post_start_hooks from .logging import LoggerManager @@ -52,6 +53,13 @@ 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) + # Store flagged blocks' full DataFrames + try: + logger.debug("Registering DataFrame storage hook.") + register_dataframe_storage() + except Exception as e: # pylint: disable=broad-exception-caught + logger.error("Failed to register 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/dataframe_storage/conftest.py b/tests/unit/dataframe_storage/conftest.py index 83879d8..69e1def 100644 --- a/tests/unit/dataframe_storage/conftest.py +++ b/tests/unit/dataframe_storage/conftest.py @@ -3,13 +3,83 @@ Kept out of tests/unit/conftest.py because the names are generic. """ +import builtins +import json +import signal +import sys import time from pathlib import Path +from types import SimpleNamespace from typing import Any, Dict, Iterator, List, Optional import pytest +import responses +from IPython.core.interactiveshell import InteractiveShell +from traitlets.config import Config from deepnote_toolkit import dataframe_storage +from deepnote_toolkit.dataframe_storage import register_dataframe_storage +from tests.unit.helpers.dataframe_storage import REPORT_URL, storage_request + + +@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. + + The writer remembers what it already reported in module state, so it is + cleared around every test. + """ + 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( + dataframe_storage, + "get_absolute_userpod_api_url", + lambda path: f"http://userpod/{path}", + ) + dataframe_storage._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 + dataframe_storage._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(dataframe_storage.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_dataframe_storage() + yield shell + shell.events.unregister("post_run_cell", dataframe_storage._on_post_run_cell) + InteractiveShell.clear_instance() + for name in set(vars(builtins)) - builtins_before: + delattr(builtins, name) @pytest.fixture @@ -102,6 +172,39 @@ def open_failing_on_close( return state +@pytest.fixture +def queued_sigints() -> Iterator[List[int]]: + """Stand in for the handler ipykernel installs while it runs a cell. + + It only queues the interrupt. Under pytest SIGINT already raises + KeyboardInterrupt, so without this the interrupt tests would pass vacuously. + """ + queued: List[int] = [] + original = signal.signal(signal.SIGINT, lambda *args: queued.append(1)) + yield queued + signal.signal(signal.SIGINT, original) + + +@pytest.fixture(params=["ordinary-cell", "top-level-await-cell"]) +def sigint_handler(request: pytest.FixtureRequest) -> Iterator[List[int]]: + """The handler a cell's post_run_cell runs under, and what it has queued. + + ipykernel keeps default_int_handler for ordinary cells, so SIGINT raises + anywhere. Only a cell with top-level await runs under a handler that queues. + """ + queued: List[int] = [] + + def only_queue(*args: Any) -> None: + queued.append(1) + + handler: Any = only_queue + if request.param == "ordinary-cell": + handler = signal.default_int_handler + original = signal.signal(signal.SIGINT, handler) + yield queued + signal.signal(signal.SIGINT, original) + + @pytest.fixture def no_root_settings(monkeypatch: pytest.MonkeyPatch) -> None: for name in ( diff --git a/tests/unit/dataframe_storage/test_dataframe_storage.py b/tests/unit/dataframe_storage/test_dataframe_storage.py new file mode 100644 index 0000000..d3cda02 --- /dev/null +++ b/tests/unit/dataframe_storage/test_dataframe_storage.py @@ -0,0 +1,876 @@ +"""Tests for the hook that stores the DataFrame a block with a storage setting returns.""" + +import errno +import json +import os +import re +import signal +import warnings +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, List, Optional +from unittest import mock + +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 deepnote_toolkit import artifacts, dataframe_storage +from deepnote_toolkit.ocelots.data_preview import DeepnoteDataFrameWithDataPreview +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview +from tests.unit.helpers.dataframe_storage import ( + ALWAYS_FAILING, + BLOCK_ID, + NAME, + REPORT_URL, + VERSION_DEPENDENT, + interrupt_at, + probe_failure, + storage_request, +) + +# the only runtime type the webapp's toolkit/errors endpoint accepts; its Sentry +# fingerprint comes from `code` +REPORT_TYPE = "TOOLKIT_RUNTIME_ERROR" +REPORT_CODE = "DATAFRAME_STORAGE_WRITE_FAILED" + + +def _frame_dir(root: Path) -> Path: + return root / "deepnote_dataframes" / NAME + + +def _arrow_files(root: Path) -> List[Path]: + return sorted(_frame_dir(root).glob("*.arrow")) + + +def _manifest(root: Path) -> Dict[str, Any]: + return json.loads((_frame_dir(root) / "manifest.json").read_text()) + + +def _stored_table(root: Path) -> pa.Table: + """The frame the name's manifest currently names.""" + return pa.ipc.open_file(_frame_dir(root) / _manifest(root)["file"]).read_all() + + +def _assert_nothing_stored(root: Path, reports: SimpleNamespace) -> None: + assert not (root / "deepnote_dataframes").exists() + 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_polars_frame(shell: InteractiveShell, root: Path) -> None: + _run(shell, pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})) + + assert _stored_table(root).to_pydict() == {"a": [1, 2, 3], "b": ["x", "y", "z"]} + assert _manifest(root)["rows"] == 3 + + +@pytest.mark.parametrize( + "env_value", + [None, "false", "", "0", "off", "no", "2"], + ids=["unset", "false", "empty", "zero", "off", "no", "two"], +) +def test_env_var_off_writes_nothing( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, + env_value: Optional[str], +) -> None: + if env_value is None: + monkeypatch.delenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED") + else: + monkeypatch.setenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED", env_value) + + result = _run(shell, pd.DataFrame({"a": [1]})) + + assert result.success + _assert_nothing_stored(root, reports) + + +def test_calls_store_frames_with_the_flag_off( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED", raising=False) + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert _stored_table(root).to_pydict() == {"a": [1]} + + +@pytest.mark.parametrize("env_value", ["1", "true", "TRUE", " yes ", "on"]) +def test_env_var_on_writes_the_frame( + shell: InteractiveShell, + root: Path, + monkeypatch: pytest.MonkeyPatch, + env_value: str, +) -> None: + monkeypatch.setenv("DEEPNOTE_DATAFRAME_STORAGE_ENABLED", env_value) + + _run(shell, pd.DataFrame({"a": [1]})) + + assert _stored_table(root).to_pydict() == {"a": [1]} + + +def test_shell_without_kernel_writes_nothing( + shell: InteractiveShell, root: Path, reports: SimpleNamespace +) -> None: + del shell.kernel + + result = _run(shell, pd.DataFrame({"a": [1]})) + + assert result.success + _assert_nothing_stored(root, reports) + + +def test_missing_result_is_ignored(root: Path, reports: SimpleNamespace) -> None: + """ipykernel can trigger post_run_cell without a result.""" + dataframe_storage._on_post_run_cell(None) + + _assert_nothing_stored(root, reports) + + +def test_helper_execution_writes_nothing( + shell: InteractiveShell, root: Path, reports: SimpleNamespace +) -> None: + result = _run(shell, pd.DataFrame({"a": [1]}), store_history=False) + + assert result.success + _assert_nothing_stored(root, reports) + + +@pytest.mark.parametrize( + "cell", ["df\n1/0", "df\nraise KeyboardInterrupt"], ids=["raised", "interrupted"] +) +def test_failed_cell_writes_nothing( + shell: InteractiveShell, root: Path, reports: SimpleNamespace, cell: str +) -> None: + shell.user_ns["df"] = pd.DataFrame({"a": [1]}) + + result = shell.run_cell(cell, store_history=True) + + assert not result.success + _assert_nothing_stored(root, reports) + + +@pytest.mark.parametrize( + "value", + [1, pd.Series([1, 2]), pl.DataFrame({"a": [1]}).lazy(), None], + ids=["int", "series", "lazyframe", "none"], +) +def test_non_dataframe_result_writes_nothing( + shell: InteractiveShell, root: Path, reports: SimpleNamespace, value: Any +) -> None: + _run(shell, value) + + _assert_nothing_stored(root, reports) + + +def test_cell_ending_in_semicolon_writes_nothing( + shell: InteractiveShell, root: Path, reports: SimpleNamespace +) -> None: + shell.user_ns["df"] = pd.DataFrame({"a": [1]}) + + shell.run_cell("df;", store_history=True) + + _assert_nothing_stored(root, reports) + + +def _query_preview() -> DeepnoteQueryPreview: + return DeepnoteQueryPreview({"a": [1]}, deepnote_query="select 1") + + +def _cleared_query_preview() -> DeepnoteQueryPreview: + preview = _query_preview() + preview["b"] = 2 + return preview + + +@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, root: Path, reports: SimpleNamespace, make_preview: Any +) -> None: + preview = make_preview() + assert isinstance(preview, DeepnoteQueryPreview) + + _run(shell, preview) + + _assert_nothing_stored(root, reports) + + +def test_pyspark_result_writes_nothing( + shell: InteractiveShell, root: Path, reports: SimpleNamespace, spark: Any +) -> None: + wrapped = DeepnoteDataFrameWithDataPreview(spark.range(3)) + + _run(shell, wrapped) + + _assert_nothing_stored(root, reports) + + +def test_pandas_on_spark_result_writes_nothing( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + spark: 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") + # pandas-on-Spark refuses to run under ANSI mode, which Spark 4 enables by default + ansi_mode = spark.conf.get("spark.sql.ansi.enabled") + spark.conf.set("spark.sql.ansi.enabled", "false") + try: + wrapped = DeepnoteDataFrameWithDataPreview(spark.range(3).pandas_api()) + finally: + spark.conf.set("spark.sql.ansi.enabled", ansi_mode) + + _run(shell, wrapped) + + _assert_nothing_stored(root, reports) + + +def _displayed(capsys: pytest.CaptureFixture) -> str: + """Everything the last cells printed, with execution counts masked.""" + captured = capsys.readouterr() + return re.sub(r"Out\[\d+\]", "Out[N]", captured.out + captured.err) + + +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 _qualified(exc_type: type) -> str: + return f"{exc_type.__module__}.{exc_type.__qualname__}" + + +def _failing_frame_runs( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + make_frame: Callable[[], Any], +) -> Dict[str, Any]: + """Store a good frame, then run a bad one twice; return what was left behind.""" + _run(shell, pd.DataFrame({"ok": [1]})) + block = _frame_dir(root) + manifest_bytes = (block / "manifest.json").read_bytes() + first = _run(shell, make_frame()) + reported_after_first = len(reports.bodies) + second = _run(shell, make_frame()) + return { + "results": (first, second), + "manifest_bytes": manifest_bytes, + "reported_after_first": reported_after_first, + "listing": sorted(path.name for path in block.iterdir()), + } + + +@pytest.mark.parametrize("make_frame", ALWAYS_FAILING) +def test_unserializable_frame_is_not_stored_and_reported_once( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + make_frame: Callable[[], Any], +) -> None: + probed = probe_failure(make_frame()) + assert probed is not None, "pandas is expected to refuse this frame" + + left = _failing_frame_runs(shell, root, reports, make_frame) + + assert all(result.success for result in left["results"]) + assert (_frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] + assert left["listing"] == sorted(["manifest.json", _manifest(root)["file"]]) + assert left["reported_after_first"] == 1 + assert len(reports.bodies) == 1 + assert reports.bodies[0]["type"] == REPORT_TYPE + assert reports.bodies[0]["code"] == REPORT_CODE + assert reports.bodies[0]["context"] == { + "kind": "error", + "exception_type": _qualified(probed), + } + + +@pytest.mark.parametrize("make_frame, rows", VERSION_DEPENDENT) +def test_frames_pandas_may_refuse_follow_what_pandas_does( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + make_frame: Callable[[], Any], + rows: int, +) -> None: + probed = probe_failure(make_frame()) + + left = _failing_frame_runs(shell, root, reports, make_frame) + + assert all(result.success for result in left["results"]) + if probed is None: + assert reports.bodies == [] + assert _stored_table(root).num_rows == rows + return + assert (_frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] + assert left["listing"] == sorted(["manifest.json", _manifest(root)["file"]]) + assert [body["context"]["exception_type"] for body in reports.bodies] == [ + _qualified(probed) + ] + + +def test_mixed_type_column_failure_names_the_exception_type_only( + shell: InteractiveShell, + 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 [body["context"] for body in reports.bodies] == [ + {"kind": "error", "exception_type": "pyarrow.lib.ArrowInvalid"} + ] + posted = reports.calls[0].body.decode() + assert "two" not in posted + assert "tried to convert" not in posted + assert set(reports.bodies[0]) == {"type", "message", "code", "context"} + assert reports.bodies[0]["message"] == "DataFrame storage write failed" + + +def test_each_exception_type_is_reported_separately( + shell: InteractiveShell, 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_polars_object_column_is_not_stored_and_reported( + shell: InteractiveShell, root: Path, reports: SimpleNamespace +) -> None: + _run(shell, pl.DataFrame({"ok": [1]})) + manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() + + result = _run(shell, pl.DataFrame({"a": pl.Series([object()], dtype=pl.Object)})) + + assert result.success + assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert len(_arrow_files(root)) == 1 + assert [body["context"] for body in reports.bodies] == [ + {"kind": "error", "exception_type": "polars.exceptions.ComputeError"} + ] + + +def test_report_is_one_short_authenticated_attempt( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + dataframe_storage, + "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( + dataframe_storage, "get_absolute_userpod_api_url", _refuse_url + ), + "ValueError", + id="no-url", + ), + ], +) +def test_report_that_cannot_be_delivered_is_silent( + shell: InteractiveShell, + 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( + dataframe_storage, + "get_logger", + lambda: SimpleNamespace(warning=lambda *args: logged.append(args)), + ) + + result = _run(shell, frame) + + assert result.success + assert _displayed(capsys) == expected_output + # IPython prints root-logger warnings into the cell, so the failure must only + # go to the file logger, and as the exception type alone + assert caplog.records == [] + assert [args[1:] for args in logged] == [(logged_type,)] + + +def test_writer_adds_nothing_to_the_cell( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + frame = pd.DataFrame({"a": [1, "two"]}) + expected_output = _output_without_storage(shell, frame, capsys, monkeypatch) + original = pd.DataFrame.to_feather + + def noisy_to_feather(self: pd.DataFrame, *args: Any, **kwargs: Any) -> Any: + warnings.warn("noisy", UserWarning) + return original(self, *args, **kwargs) + + monkeypatch.setattr(pd.DataFrame, "to_feather", noisy_to_feather) + + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + for value in (frame, pd.DataFrame({"a": [1]})): + _run(shell, value) + + assert [str(w.message) for w in recorded if "noisy" in str(w.message)] == [] + assert caplog.records == [] + assert len(reports.bodies) == 1 + printed = _displayed(capsys) + assert printed.count("Error in callback") == 0 + assert printed.startswith(expected_output) + + +def test_handler_never_raises( + shell: InteractiveShell, + 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, + root: Path, + reports: SimpleNamespace, + parent: Dict[str, Any], + close_error: Dict[str, Optional[OSError]], +) -> None: + close_error["error"] = PermissionError(errno.EPERM, "Operation not permitted") + + first = _run(shell, pd.DataFrame({"a": [1]})) + parent["metadata"]["deepnote"]["dataframeStorage"]["name"] = "other" + 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((root / "deepnote_dataframes").rglob("*.arrow")) + assert not list((root / "deepnote_dataframes").rglob("manifest.json")) + + +def test_each_io_error_code_is_reported_once_and_not_swallowed( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + close_error: Dict[str, Optional[OSError]], +) -> None: + for error in ( + OSError(errno.EIO, "Input/output error"), + OSError(errno.EIO, "Input/output error"), + OSError(errno.ENOSPC, "No space left on device"), + ): + close_error["error"] = error + assert _run(shell, pd.DataFrame({"a": [1]})).success + + 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 test_partial_file_is_deleted_when_the_write_fails( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def write_some_then_fail(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + sink.write(b"ARROW1 partial") + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(pd.DataFrame, "to_feather", write_some_then_fail) + + result = _run(shell, pd.DataFrame({"a": [1]})) + + assert result.success + assert _arrow_files(root) == [] + assert len(reports.bodies) == 1 + + +def test_failure_after_the_manifest_step_keeps_the_new_frame( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _run(shell, pd.DataFrame({"a": [1]})) + + def broken_cleanup(*args: Any) -> Any: + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(dataframe_storage, "_delete_replaced_frames", broken_cleanup) + + result = _run(shell, pd.DataFrame({"a": [2]})) + + assert result.success + assert _stored_table(root).to_pydict() == {"a": [2]} + assert [body["context"] for body in reports.bodies] == [ + {"kind": "io", "errno": "EIO", "exception_type": "builtins.OSError"} + ] + + +def test_read_only_mount_is_skipped_before_converting( + shell: InteractiveShell, + 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) + monkeypatch.setattr( + dataframe_storage.os, + "statvfs", + lambda path: SimpleNamespace(f_flag=os.ST_RDONLY), + ) + + with mock.patch.object(pd.DataFrame, "to_feather", autospec=True) as to_feather: + result = _run(shell, frame) + + assert result.success + to_feather.assert_not_called() + assert not (root / "deepnote_dataframes").exists() + assert reports.bodies == [] + assert _displayed(capsys) == expected_output + + +def test_erofs_from_the_write_is_silent( + shell: InteractiveShell, + 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 refuse(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + raise OSError(errno.EROFS, "Read-only file system") + + monkeypatch.setattr(pd.DataFrame, "to_feather", refuse) + + result = _run(shell, frame) + + assert result.success + assert reports.bodies == [] + assert _arrow_files(root) == [] + assert _displayed(capsys) == expected_output + + +def _write_partial_file_then(interrupt: Callable[[], None]) -> Callable[..., None]: + def to_feather(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + sink.write(b"ARROW1 partial") + interrupt() + + return to_feather + + +def _raise_keyboard_interrupt() -> None: + raise KeyboardInterrupt + + +@pytest.mark.parametrize( + "interrupt", + [lambda: signal.raise_signal(signal.SIGINT), _raise_keyboard_interrupt], + ids=["sigint", "keyboard-interrupt"], +) +def test_interrupt_during_the_write_fails_the_execution( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + queued_sigints: List[int], + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, + interrupt: Callable[[], None], +) -> None: + _run(shell, pd.DataFrame({"a": [1]})) + manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() + handler_before = signal.getsignal(signal.SIGINT) + monkeypatch.setattr(pd.DataFrame, "to_feather", _write_partial_file_then(interrupt)) + capsys.readouterr() + + result = _run(shell, pd.DataFrame({"a": [2]})) + + assert isinstance(result.error_in_exec, KeyboardInterrupt) + assert not result.success + assert queued_sigints == [] + assert "KeyboardInterrupt" in _displayed(capsys) + assert reports.bodies == [] + assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert len(_arrow_files(root)) == 1 + assert signal.getsignal(signal.SIGINT) is handler_before + + +@pytest.mark.parametrize( + "point", ["manifest-truncated", "manifest-written", "old-frames-deleted"] +) +def test_interrupt_after_the_frame_is_written_leaves_a_usable_manifest( + shell: InteractiveShell, + root: Path, + reports: SimpleNamespace, + sigint_handler: List[int], + monkeypatch: pytest.MonkeyPatch, + point: str, +) -> None: + _run(shell, pd.DataFrame({"a": [1]})) + _run(shell, pd.DataFrame({"a": [2]})) + handler_before = signal.getsignal(signal.SIGINT) + interrupt_at(point, monkeypatch) + + result = _run(shell, pd.DataFrame({"a": [3]})) + + assert isinstance(result.error_in_exec, KeyboardInterrupt) + assert sigint_handler == [] + assert reports.bodies == [] + # parses, and names a frame that is there and complete + assert _stored_table(root).num_rows == 1 + assert signal.getsignal(signal.SIGINT) is handler_before + + +def test_interrupt_while_reporting_a_failure_fails_the_execution( + shell: InteractiveShell, + root: Path, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def interrupted(context: Dict[str, str]) -> None: + raise KeyboardInterrupt + + monkeypatch.setattr(dataframe_storage, "_post_error_report", 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 "Error in callback" not in printed + + +def test_sigint_handler_is_restored_after_a_failed_write( + shell: InteractiveShell, + root: Path, + queued_sigints: List[int], +) -> None: + handler_before = signal.getsignal(signal.SIGINT) + + _run(shell, pd.DataFrame({"a": [1, "two"]})) + + assert signal.getsignal(signal.SIGINT) is handler_before + + +def test_stores_the_result_under_the_requested_name( + shell: InteractiveShell, root: Path +) -> None: + result = _run(shell, pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})) + + assert result.success + assert _stored_table(root).to_pydict() == {"a": [1, 2, 3], "b": ["x", "y", "z"]} + manifest = _manifest(root) + assert set(manifest) == { + "version", + "file", + "written_at", + "rows", + "toolkit_version", + "block_id", + } + assert manifest["block_id"] == BLOCK_ID + assert manifest["rows"] == 3 + + +def test_the_name_in_the_request_chooses_the_directory( + shell: InteractiveShell, root: Path, parent: Dict[str, Any] +) -> None: + parent.update(storage_request(name="monthly-2026")) + + _run(shell, pd.DataFrame({"a": [1]})) + + assert (root / "deepnote_dataframes" / "monthly-2026" / "manifest.json").exists() + assert not (root / "deepnote_dataframes" / NAME).exists() + + +def test_block_frames_and_calls_share_one_name( + shell: InteractiveShell, root: Path, parent: Dict[str, Any] +) -> None: + _run(shell, pd.DataFrame({"a": [1]})) + first = _manifest(root) + parent.update(storage_request(blockId="blk2")) + _run(shell, pd.DataFrame({"a": [2]})) + second = _manifest(root) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + third = _manifest(root) + + assert first["block_id"] == BLOCK_ID + assert second["block_id"] == "blk2" + assert second["previous"] == first["file"] + # a frame that code stored last belongs to no block + assert "block_id" not in third + assert third["previous"] == second["file"] + assert _stored_table(root).to_pydict() == {"a": [3]} + + +@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": "blk1"}}}}, + id="no-name", + ), + pytest.param( + {"metadata": {"deepnote": {"dataframeStorage": {"name": "revenue"}}}}, + id="no-block-id", + ), + pytest.param(storage_request(name="../x"), id="name-traversal"), + pytest.param(storage_request(name=""), id="name-empty"), + pytest.param(storage_request(name="a/b"), id="name-slash"), + pytest.param(storage_request(name="."), id="name-dot"), + pytest.param(storage_request(name=".."), id="name-dotdot"), + pytest.param(storage_request(name=5), id="name-int"), + pytest.param(storage_request(name=None), id="name-null"), + pytest.param(storage_request(name="x" * 129), id="name-too-long"), + pytest.param(storage_request(blockId="../x"), id="block-traversal"), + pytest.param(storage_request(blockId=""), id="block-empty"), + pytest.param(storage_request(blockId="a/b"), id="block-slash"), + pytest.param(storage_request(blockId=".."), id="block-dotdot"), + pytest.param(storage_request(blockId=None), id="block-null"), + pytest.param(storage_request(blockId="x" * 129), id="block-too-long"), + ], +) +def test_unusable_request_metadata_writes_nothing( + shell: InteractiveShell, + 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(root, reports) diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index f7fd1fb..e33ce50 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -1,10 +1,10 @@ -"""Frames and interrupts shared by the DataFrame storage tests.""" +"""Frames, interrupts and execute requests shared by the DataFrame storage tests.""" import io import signal import uuid from pathlib import Path -from typing import Any, Optional +from typing import Any, Dict, Optional import numpy as np import pandas as pd @@ -13,6 +13,10 @@ from deepnote_toolkit import dataframe_storage +NAME = "revenue" +BLOCK_ID = "blk1" +REPORT_URL = "http://userpod/toolkit/errors" + def probe_failure(frame: Any) -> Optional[type]: """The exception type pandas itself raises for this frame, if any.""" @@ -102,3 +106,17 @@ def interrupt_then_delete(*args: Any, **kwargs: Any) -> None: monkeypatch.setattr( dataframe_storage, "_delete_replaced_frames", interrupt_then_delete ) + + +def storage_request(**overrides: Any) -> Dict[str, Any]: + """Return the request message that asks the kernel to store a block's result. + + `overrides` replace fields of `deepnote.dataframeStorage`. + """ + return { + "metadata": { + "deepnote": { + "dataframeStorage": {"name": NAME, "blockId": BLOCK_ID, **overrides} + } + } + } diff --git a/tests/unit/test_runtime_initialization.py b/tests/unit/test_runtime_initialization.py new file mode 100644 index 0000000..facf01f --- /dev/null +++ b/tests/unit/test_runtime_initialization.py @@ -0,0 +1,62 @@ +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_dataframe_storage", +) + + +@pytest.fixture +def steps() -> Iterator[SimpleNamespace]: + """Replace every startup step so the test sees only the wiring. + + The logger is replaced too: a real `logger.error` reports to the webapp over + HTTP, with retries. + """ + with ( + mock.patch.multiple( + "deepnote_toolkit.runtime_initialization", + 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 = mocks["LoggerManager"].return_value.get_logger.return_value + yield SimpleNamespace(logger=logger, **mocks) + + +def test_registers_the_dataframe_storage_hook(steps: SimpleNamespace) -> None: + init_deepnote_runtime() + + steps.register_dataframe_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_dataframe_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 DataFrame storage hook with error: %s" + assert str(error) == "boom"