diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 4b66294..aaa8bcb 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -6,12 +6,23 @@ import re import time import uuid +import warnings +from dataclasses import dataclass from pathlib import Path from typing import Any, 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 from deepnote_toolkit.config import resolve_notebook_root +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, @@ -25,8 +36,14 @@ write_manifest, ) +ENABLED_ENV_VAR = "DEEPNOTE_DATAFRAME_STORAGE_ENABLED" FRAMES_DIR = Path("deepnote_dataframes") STALE_FRAME_AGE_SECONDS = 24 * 60 * 60 +# 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 POLARS_BATCH_ROWS = 65_536 POLARS_BATCH_MIN_VERSION = Version("1.37") POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") @@ -36,6 +53,17 @@ _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 + + +_reported: set[str] = set() + + def is_safe_id(value: object) -> bool: """Return whether `value` can be used as a single path segment.""" return isinstance(value, str) and _SAFE_ID.fullmatch(value) is not None @@ -123,7 +151,13 @@ def _delete_replaced_frames( (directory / entry.name).unlink() -def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> Optional[str]: +def store_frame( + frame: Any, + kind: FrameKind, + name: str, + root: Path, + block_id: Optional[str] = None, +) -> Optional[str]: """Store the frame as the current one under `name` and return its reference ID. Returns None, having stored nothing, if the filesystem is read-only. EROFS from @@ -144,6 +178,122 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> Optional[ raise # A failure from here on keeps the new frame: an s3fs rename can put the new # manifest in place and still fail. - 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) _delete_replaced_frames(directory, old, new_file) return reference_id + + +def register_dataframe_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: 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(): 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, reporting a failure 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 + kind = classify_frame(result.result) + if kind is None: + return + root = resolve_project_root() + 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. + + IPython prints a callback's exceptions into the cell as "Error in callback", and + warnings reach the cell too, so failures go 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: + # Re-raised, IPython would print "Error in callback" and keep 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_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index a0cf1ff..9bc613e 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -64,9 +64,13 @@ 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 `block_id` if given. On s3fs a rename is a copy and then a delete, so a failed rename may already have replaced the manifest. @@ -77,6 +81,8 @@ def write_manifest( 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 # delete=False: s3fs uploads at close(), so the file must be closed before the # rename, and deleting it only when the context ends needs Python 3.12's # delete_on_close=False. diff --git a/deepnote_toolkit/runtime_initialization.py b/deepnote_toolkit/runtime_initialization.py index 63b7355..d573286 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.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,12 @@ 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 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/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index b58da10..396a270 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -1,12 +1,77 @@ """Fixtures for the DataFrame storage tests.""" +import builtins +import json +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.artifacts import dataframe_storage +from deepnote_toolkit.artifacts.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.""" + 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 diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index c66005c..25e6143 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -10,10 +10,11 @@ from pyarrow import feather from deepnote_toolkit import artifacts -from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview from tests.unit.helpers.dataframe_storage import ( NAME, arrow_files, + cleared_query_preview, + query_preview, read_only, stored_table, ) @@ -62,16 +63,6 @@ def test_invalid_name_raises_on_a_read_only_mount( artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) -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 - - UNSUPPORTED_VALUES = [ pytest.param(lambda: 1, id="int"), pytest.param(lambda: None, id="none"), @@ -81,9 +72,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_dataframe_storage.py b/tests/unit/artifacts/test_dataframe_storage.py new file mode 100644 index 0000000..a5a161e --- /dev/null +++ b/tests/unit/artifacts/test_dataframe_storage.py @@ -0,0 +1,819 @@ +"""Tests for the post_run_cell hook that stores a block's DataFrame result.""" + +import errno +import re +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 pytest +import requests +import responses +from IPython.core.interactiveshell import ExecutionResult, InteractiveShell +from pyarrow import feather + +from deepnote_toolkit import artifacts +from deepnote_toolkit.artifacts import 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, + arrow_files, + cleared_query_preview, + frame_dir, + probe_failure, + query_preview, + read_only, + storage_request, + stored_manifest, + stored_table, +) + +REPORT_TYPE = "TOOLKIT_RUNTIME_ERROR" +REPORT_CODE = "DATAFRAME_STORAGE_WRITE_FAILED" + + +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 stored_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) + + +@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_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(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, "pyarrow 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", stored_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_pyarrow_may_refuse_follow_what_pyarrow_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", stored_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 + 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 = feather.write_feather + + def noisy_write_feather(frame: pd.DataFrame, *args: Any, **kwargs: Any) -> Any: + warnings.warn("noisy", UserWarning) + return original(frame, *args, **kwargs) + + monkeypatch.setattr(feather, "write_feather", noisy_write_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(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + sink.write(b"ARROW1 partial") + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(feather, "write_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) + read_only(monkeypatch) + + with mock.patch.object(feather, "write_feather", autospec=True) as write: + result = _run(shell, frame) + + assert result.success + write.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(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + raise OSError(errno.EROFS, "Read-only file system") + + monkeypatch.setattr(feather, "write_feather", refuse) + + result = _run(shell, frame) + + assert result.success + assert reports.bodies == [] + assert arrow_files(root) == [] + assert _displayed(capsys) == expected_output + + +def test_erofs_after_the_manifest_step_is_silent( + shell: InteractiveShell, + 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]})) + replaced = frame_dir(root) / stored_manifest(root)["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 stored_table(root).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, + root: Path, + reports: SimpleNamespace, + capsys: pytest.CaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def write_some_then_interrupt( + frame: pd.DataFrame, sink: Any, **kwargs: Any + ) -> None: + sink.write(b"ARROW1 partial") + raise KeyboardInterrupt + + _run(shell, pd.DataFrame({"a": [1]})) + manifest_bytes = (frame_dir(root) / "manifest.json").read_bytes() + monkeypatch.setattr(feather, "write_feather", write_some_then_interrupt) + 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 (frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert len(arrow_files(root)) == 1 + + +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_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 = stored_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 = stored_manifest(root) + parent.update(storage_request(blockId="blk2")) + _run(shell, pd.DataFrame({"a": [2]})) + second = stored_manifest(root) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + third = stored_manifest(root) + + assert first["block_id"] == BLOCK_ID + assert second["block_id"] == "blk2" + assert second["previous"] == first["file"] + 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 28926be..14e5b68 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -15,7 +15,11 @@ import pytest from pyarrow import feather +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview + NAME = "revenue" +BLOCK_ID = "blk1" +REPORT_URL = "http://userpod/toolkit/errors" DAY = 24 * 60 * 60 FRAME_A = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee.arrow" FRAME_B = "11111111-2222-4333-8444-555555555555.arrow" @@ -105,3 +109,24 @@ def dictionary_concat() -> pd.DataFrame: pytest.param(lambda: pd.DataFrame({"a": [uuid.UUID(int=1)]}), 1, id="uuid"), pytest.param(dictionary_concat, 4, id="dictionary-concat"), ] + + +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 + + +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": {"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..b5d1789 --- /dev/null +++ b/tests/unit/test_runtime_initialization.py @@ -0,0 +1,58 @@ +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.""" + 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"