diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 95efac1..a471b60 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -6,15 +6,31 @@ import re import time import uuid +import warnings +from dataclasses import dataclass from pathlib import Path from typing import Any, Optional, Union import pandas as pd +import requests +from IPython import get_ipython +from IPython.core.interactiveshell import ExecutionResult from packaging.version import Version -from typing_extensions import TypeAlias, assert_never +from typing_extensions import TypeAlias, TypeGuard, assert_never +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.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, @@ -23,8 +39,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") @@ -34,11 +56,29 @@ _SAFE_ID_REGEX = re.compile(r"[A-Za-z0-9_-]{1,128}") -def is_safe_id(value: object) -> bool: +@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) -> 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() @@ -114,7 +154,9 @@ def _delete_replaced_frames( (directory / entry.name).unlink() -def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]: +def store_frame( + frame: StorableFrame, 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 @@ -134,6 +176,121 @@ def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]: return None raise - 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: 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 + + 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 + if not is_storable_frame(result.result): + return + root = resolve_project_root() + store_frame(result.result, 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 bdf94c3..b6a7c98 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py @@ -41,6 +41,7 @@ class Manifest(BaseModel): written_at: StrictStr rows: StrictInt toolkit_version: StrictStr + block_id: Optional[StrictStr] = None def read_manifest(directory: Path) -> Optional[Manifest]: @@ -69,9 +70,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. @@ -83,6 +88,7 @@ def write_manifest( written_at=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), rows=rows, toolkit_version=__version__, + 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 diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index aa11fd8..612e89c 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 @@ -54,7 +50,7 @@ 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__}" diff --git a/deepnote_toolkit/runtime_initialization.py b/deepnote_toolkit/runtime_initialization.py index 63b7355..bc96455 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 acab15c..e98536c 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -1,13 +1,78 @@ """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 IO, Any, Dict, Iterator, List, Optional import pytest +import responses +from IPython.core.interactiveshell import InteractiveShell from pyspark.sql import SparkSession +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 9d1c28e..922e5ba 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -10,11 +10,12 @@ from pyarrow import feather from deepnote_toolkit import artifacts -from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview from tests.unit.helpers.dataframe_storage import ( FRAME_A, NAME, arrow_files, + cleared_query_preview, + query_preview, stored_table, ) @@ -55,17 +56,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"), @@ -75,9 +65,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..07a1b54 --- /dev/null +++ b/tests/unit/artifacts/test_dataframe_storage.py @@ -0,0 +1,835 @@ +"""Tests for the post_run_cell hook that stores a block's DataFrame result.""" + +import errno +import json +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 ( + BLOCK_ID, + NAME, + REPORT_URL, + 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(tmp_project_root: Path, reports: SimpleNamespace) -> None: + assert not (tmp_project_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, tmp_project_root: Path) -> None: + _run(shell, pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})) + + assert stored_table(tmp_project_root).to_pydict() == { + "a": [1, 2, 3], + "b": ["x", "y", "z"], + } + assert stored_manifest(tmp_project_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, + tmp_project_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(tmp_project_root, reports) + + +def test_calls_store_frames_with_the_flag_off( + tmp_project_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(tmp_project_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, + 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 stored_table(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) + + +@pytest.mark.parametrize( + "cell", ["df\n1/0", "df\nraise KeyboardInterrupt"], ids=["raised", "interrupted"] +) +def test_failed_cell_writes_nothing( + shell: InteractiveShell, tmp_project_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(tmp_project_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, + tmp_project_root: Path, + reports: SimpleNamespace, + value: Any, +) -> None: + _run(shell, value) + + _assert_nothing_stored(tmp_project_root, reports) + + +def test_cell_ending_in_semicolon_writes_nothing( + shell: InteractiveShell, tmp_project_root: Path, reports: SimpleNamespace +) -> None: + shell.user_ns["df"] = pd.DataFrame({"a": [1]}) + + shell.run_cell("df;", store_history=True) + + _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: Any, +) -> 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) + + +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 _qualified(exc_type: type) -> str: + return f"{exc_type.__module__}.{exc_type.__qualname__}" + + +def _failing_frame_runs( + shell: InteractiveShell, + tmp_project_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(tmp_project_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()), + } + + +def _mixed_types() -> pd.DataFrame: + return pd.DataFrame({"a": [1, "two"]}) + + +def test_unserializable_frame_is_not_stored_and_reported_once( + shell: InteractiveShell, + tmp_project_root: Path, + reports: SimpleNamespace, +) -> None: + probed = probe_failure(_mixed_types()) + assert probed is not None, "pyarrow is expected to refuse this frame" + + left = _failing_frame_runs(shell, tmp_project_root, reports, _mixed_types) + + assert all(result.success for result in left["results"]) + assert (frame_dir(tmp_project_root) / "manifest.json").read_bytes() == left[ + "manifest_bytes" + ] + assert left["listing"] == sorted( + ["manifest.json", stored_manifest(tmp_project_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), + } + + +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 [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, 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_polars_object_column_is_not_stored_and_reported( + shell: InteractiveShell, tmp_project_root: Path, reports: SimpleNamespace +) -> None: + _run(shell, pl.DataFrame({"ok": [1]})) + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() + + result = _run(shell, pl.DataFrame({"a": pl.Series([object()], dtype=pl.Object)})) + + assert result.success + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_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, + tmp_project_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, + 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( + _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, + tmp_project_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 _callback_errors(printed) == [] + assert printed.startswith(expected_output) + + +def test_handler_never_raises( + 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"]["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((tmp_project_root / "deepnote_dataframes").rglob("*.arrow")) + assert not list((tmp_project_root / "deepnote_dataframes").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: + for error in ( + OSError(errno.EIO, "Input/output error"), + OSError(errno.EIO, "Input/output error"), + OSError(errno.ENOSPC, "No space left on device"), + ): + close_upload_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, + tmp_project_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(tmp_project_root) == [] + assert len(reports.bodies) == 1 + + +def test_failure_after_the_manifest_step_keeps_the_new_frame( + shell: InteractiveShell, + tmp_project_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(tmp_project_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, + 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) + 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 (tmp_project_root / "deepnote_dataframes").exists() + assert reports.bodies == [] + assert _displayed(capsys) == expected_output + + +def test_erofs_from_the_write_is_silent( + 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 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(tmp_project_root) == [] + assert _displayed(capsys) == expected_output + + +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]})) + replaced = frame_dir(tmp_project_root) / stored_manifest(tmp_project_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(tmp_project_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, + tmp_project_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(tmp_project_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(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_root)) == 1 + + +def test_interrupt_while_reporting_a_failure_fails_the_execution( + shell: InteractiveShell, + tmp_project_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 _callback_errors(printed) == [] + + +def test_stores_the_result_under_the_requested_name( + shell: InteractiveShell, tmp_project_root: Path +) -> None: + result = _run(shell, pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})) + + assert result.success + assert stored_table(tmp_project_root).to_pydict() == { + "a": [1, 2, 3], + "b": ["x", "y", "z"], + } + manifest = stored_manifest(tmp_project_root) + written = json.loads((frame_dir(tmp_project_root) / "manifest.json").read_text()) + assert set(written) == { + "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, tmp_project_root: Path, parent: Dict[str, Any] +) -> None: + parent.update(storage_request(name="monthly-2026")) + + _run(shell, pd.DataFrame({"a": [1]})) + + assert ( + tmp_project_root / "deepnote_dataframes" / "monthly-2026" / "manifest.json" + ).exists() + assert not (tmp_project_root / "deepnote_dataframes" / NAME).exists() + + +def test_block_frames_and_calls_share_one_name( + shell: InteractiveShell, tmp_project_root: Path, parent: Dict[str, Any] +) -> None: + _run(shell, pd.DataFrame({"a": [1]})) + first = stored_manifest(tmp_project_root) + parent.update(storage_request(blockId="blk2")) + _run(shell, pd.DataFrame({"a": [2]})) + second = stored_manifest(tmp_project_root) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + third = stored_manifest(tmp_project_root) + + assert first.block_id == BLOCK_ID + assert second.block_id == "blk2" + assert second.previous == first.file + assert third.block_id is None + assert third.previous == second.file + assert stored_table(tmp_project_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, + 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 9a7ecdc..b675079 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -6,7 +6,7 @@ import os import time from pathlib import Path -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional import pyarrow as pa import pytest @@ -14,8 +14,11 @@ 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" +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" @@ -84,3 +87,25 @@ 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 + + +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..7e786f9 --- /dev/null +++ b/tests/unit/test_runtime_initialization.py @@ -0,0 +1,61 @@ +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_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", + 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_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"