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

deepnote_toolkit/dataframe_storage.py

@@ -1,14 +1,24 @@
-"""Writes DataFrames into the project's filesystem as named Arrow files."""
+"""Writes DataFrames into the project's filesystem as named Arrow files.
+
+Besides the function behind `artifacts.store_dataframe`, this holds the hook that
+stores the DataFrame a block with a storage setting returns.
+"""
 
 import contextlib
+import errno
 import os
 import re
 import signal
 import time
 import uuid
+import warnings
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Iterator, Literal, Optional
 
+import requests
+from IPython import get_ipython
+from IPython.core.interactiveshell import ExecutionResult
 from packaging.version import Version
 
 from deepnote_toolkit import env
@@ -19,14 +29,25 @@
     read_manifest,
     write_manifest,
 )
+from deepnote_toolkit.get_webapp_url import (
+    get_absolute_userpod_api_url,
+    get_project_auth_headers,
+)
+from deepnote_toolkit.logging import get_logger
 from deepnote_toolkit.ocelots.utils import (
     is_pandas_dataframe,
     is_polars_eager_dataframe,
 )
 from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview
 
+ENABLED_ENV_VAR = "DEEPNOTE_DATAFRAME_STORAGE_ENABLED"
 FRAMES_DIR = Path("deepnote_dataframes")
 STALE_FRAME_AGE_SECONDS = 24 * 60 * 60
+# The webapp's toolkit/errors endpoint rejects any other type with a 400, and groups
+# runtime errors in Sentry by `code`.
+REPORT_TYPE = "TOOLKIT_RUNTIME_ERROR"
+REPORT_CODE = "DATAFRAME_STORAGE_WRITE_FAILED"
+REPORT_TIMEOUT_SECONDS = 2
 # Each record batch is one page of the read endpoint; pandas already writes 64k-row
 # batches, and polars only honours a batch size from 1.37.
 POLARS_BATCH_ROWS = 65_536
@@ -37,6 +58,18 @@
 _SAFE_ID = re.compile(r"[A-Za-z0-9_-]{1,128}")
 
 
+@dataclass(frozen=True)
+class StorageTarget:
+    """What the platform asked the current execution to store, and for which block."""
+
+    name: str
+    block_id: str
+
+
+# Failures already reported in this kernel session, keyed by errno or exception type.
+_reported: set[str] = set()
+
+
 def is_safe_id(value: object) -> bool:
     """Return whether `value` can be used as a single path segment.
 
@@ -102,6 +135,18 @@ def _sigint_handled_by(handler: Callable[..., Any]) -> Iterator[None]:
             signal.signal(signal.SIGINT, previous)
 
 
+@contextlib.contextmanager
+def _sigint_raises() -> Iterator[None]:
+    """Make SIGINT raise KeyboardInterrupt.
+
+    Ordinary cells already run post_run_cell under this handler. A cell with
+    top-level await runs it under one that only queues the interrupt, so without
+    this, Stop would wait for the whole write to finish.
+    """
+    with _sigint_handled_by(signal.default_int_handler):
+        yield
+
+
 @contextlib.contextmanager
 def _sigint_deferred() -> Iterator[None]:
     """Hold SIGINT back until the block ends, then raise it as KeyboardInterrupt."""
@@ -181,9 +226,17 @@ def _delete_replaced_frames(
                     (directory / entry.name).unlink()
 
 
-def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str:
+def store_frame(
+    frame: Any,
+    kind: FrameKind,
+    name: str,
+    root: Path,
+    block_id: Optional[str] = None,
+) -> str:
     """Store the frame as the current one under `name` and return its reference ID.
 
+    `block_id` records the block whose result this is; a call stores none.
+
     The order matters because S3 replaces the manifest in one step: the frame is
     complete before the manifest names it, and nothing is deleted before then.
     An interrupt stops everything but the manifest update, which finishes first:
@@ -200,7 +253,9 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str:
         _write_frame_file(frame, kind, directory / new_file, _select_compression())
         with _sigint_deferred():
             old = read_manifest(directory)
-            write_manifest(directory, new_file, old.file if old else None, len(frame))
+            write_manifest(
+                directory, new_file, old.file if old else None, len(frame), block_id
+            )
             committed = True
     except BaseException:
         if not committed:
@@ -209,3 +264,134 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str:
         raise
     _delete_replaced_frames(directory, old, new_file)
     return reference_id
+
+
+def register_dataframe_storage() -> None:
+    """Register the post_run_cell hook that stores the results of blocks with a setting.
+
+    Registration is unconditional: whether a given execution stores a frame is
+    decided on every execution from the environment and the request metadata.
+    """
+    get_ipython().events.register("post_run_cell", _on_post_run_cell)
+
+
+def _is_enabled_by_environment() -> bool:
+    """Return whether the platform switched the feature on for this kernel."""
+    value = env.get_env(ENABLED_ENV_VAR)
+    return value is not None and value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _read_storage_target() -> Optional[StorageTarget]:
+    """Return the storage target in the current execute_request, if it is usable.
+
+    The metadata is written by other services, so every level is validated and the
+    name must be safe to use as a path segment.
+    """
+    kernel = getattr(get_ipython(), "kernel", None)
+    if kernel is None:
+        return None
+
+    node: Any = kernel.get_parent()
+    for key in ("metadata", "deepnote", "dataframeStorage"):
+        if not isinstance(node, dict):
+            return None
+        node = node.get(key)
+    if not isinstance(node, dict):
+        return None
+
+    name, block_id = node.get("name"), node.get("blockId")
+    if not (is_safe_id(name) and is_safe_id(block_id)):
+        return None
+    return StorageTarget(name=name, block_id=block_id)
+
+
+def _post_error_report(context: dict[str, str]) -> None:
+    """Tell the webapp about a failure, with one attempt and a short timeout.
+
+    Not report_error_to_webapp(): that logs a warning to the root logger before
+    every attempt, and IPython prints root-logger warnings into the user's cell.
+    """
+    try:
+        response = requests.post(
+            get_absolute_userpod_api_url("toolkit/errors"),
+            json={
+                "type": REPORT_TYPE,
+                "message": "DataFrame storage write failed",
+                "code": REPORT_CODE,
+                "context": context,
+            },
+            headers=get_project_auth_headers(),
+            timeout=REPORT_TIMEOUT_SECONDS,
+        )
+        response.raise_for_status()
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        get_logger().warning(
+            "Failed to report DataFrame storage failure: %s", type(e).__name__
+        )
+
+
+def _report_failure(exc: Exception) -> None:
+    """Report a failed write once per kernel session per errno or exception type.
+
+    The context carries the exception's type, never its message: pyarrow quotes
+    the offending value in its messages, and frames often hold personal data.
+    """
+    exception_type = f"{type(exc).__module__}.{type(exc).__qualname__}"
+    if isinstance(exc, OSError) and exc.errno is not None:
+        code = errno.errorcode.get(exc.errno, str(exc.errno))
+        key = f"io:{code}"
+        context = {"kind": "io", "errno": code, "exception_type": exception_type}
+    else:
+        key = f"error:{exception_type}"
+        context = {"kind": "error", "exception_type": exception_type}
+
+    if key in _reported:
+        return
+    _reported.add(key)
+    _post_error_report(context)
+
+
+def _store_result_frame(result: ExecutionResult) -> None:
+    """Store the cell's DataFrame result, reporting any failure to the webapp.
+
+    Only KeyboardInterrupt propagates, so that Stop during the report is not lost.
+    """
+    try:
+        target = _read_storage_target()
+        # Helper executions (paging, export, the variable explorer) do not store
+        # history. IPython 9 gives them an execution_count anyway, so ask the info.
+        if target is None or not result.info.store_history:
+            return
+        kind = classify_frame(result.result)
+        if kind is None:
+            return
+        root = resolve_project_root()
+        # a read-only mount is a deliberate choice, not a failure
+        if is_read_only(root):
+            return
+        with _sigint_raises():
+            store_frame(result.result, kind, target.name, root, target.block_id)
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        if isinstance(exc, OSError) and exc.errno == errno.EROFS:
+            return
+        _report_failure(exc)
+
+
+def _on_post_run_cell(result: Optional[ExecutionResult]) -> None:
+    """Store the cell's DataFrame result when the platform asked for it.
+
+    Never raises and never emits warnings, because IPython prints both into the
+    user's cell. A failed write is reported to the webapp instead.
+    """
+    if result is None or not _is_enabled_by_environment() or not result.success:
+        return
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            _store_result_frame(result)
+        except KeyboardInterrupt as exc:
+            # The reply is already decided, so this is the only way to make the
+            # executor see the interrupt and stop its queue.
+            result.error_in_exec = exc
+            get_ipython().showtraceback()

deepnote_toolkit/dataframe_storage_manifest.py

@@ -62,13 +62,19 @@ def read_manifest(directory: Path) -> Optional[Manifest]:
 
 
 def write_manifest(
-    directory: Path, new_file: str, previous: Optional[str], rows: int
+    directory: Path,
+    new_file: str,
+    previous: Optional[str],
+    rows: int,
+    block_id: Optional[str] = None,
 ) -> None:
-    """Point the manifest at `new_file`."""
+    """Point the manifest at `new_file`, recording the block that wrote it, if any."""
     manifest: dict[str, Any] = {"version": MANIFEST_VERSION, "file": new_file}
     if previous is not None:
         manifest["previous"] = previous
     manifest["written_at"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
     manifest["rows"] = rows
     manifest["toolkit_version"] = __version__
+    if block_id is not None:
+        manifest["block_id"] = block_id
     (directory / MANIFEST_NAME).write_text(json.dumps(manifest), encoding="utf-8")

deepnote_toolkit/runtime_initialization.py

@@ -8,6 +8,7 @@
 
 from deepnote_toolkit.runtime_patches import apply_runtime_patches
 
+from .dataframe_storage import register_dataframe_storage
 from .dataframe_utils import add_formatters
 from .execute_post_start_hooks import execute_post_start_hooks
 from .logging import LoggerManager
@@ -52,6 +53,13 @@ def init_deepnote_runtime():
     except Exception as e:  # pylint: disable=broad-exception-caught
         logger.error("Failed to add output middleware with a error: %s", e)
 
+    # Store flagged blocks' full DataFrames
+    try:
+        logger.debug("Registering DataFrame storage hook.")
+        register_dataframe_storage()
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        logger.error("Failed to register DataFrame storage hook with error: %s", e)
+
     # Disable sqlparse grouping limits for large analytical queries
     try:
         logger.debug("Configuring sqlparse limits.")