feat(dataframe-storage): store a block's result automatically

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()

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.

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.")

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

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"),
 ]