feat(artifacts): store a block's DataFrame result automatically

deepnote_toolkit/artifacts/_block_frames.py

@@ -0,0 +1,164 @@
+"""Stores a block's DataFrame result when the platform asks for it."""
+
+import errno
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import requests
+from IPython import get_ipython
+from IPython.core.interactiveshell import ExecutionResult
+
+from deepnote_toolkit import env
+from deepnote_toolkit.get_webapp_url import (
+    get_absolute_userpod_api_url,
+    get_project_auth_headers,
+)
+from deepnote_toolkit.logging import get_logger
+
+from . import _dataframe_storage
+
+ENABLED_ENV_VAR = "DEEPNOTE_DATAFRAME_STORAGE_ENABLED"
+# The webapp's dataframe reader serves block frames from here.
+BLOCK_FRAMES_DIR = Path(".deepnote") / "dataframes"
+# 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
+
+
+@dataclass(frozen=True)
+class StorageTarget:
+    """The block whose result the platform asked the current execution to store."""
+
+    notebook_id: str
+    block_id: str
+
+
+_reported: set[str] = set()
+
+
+def register_block_frame_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
+
+    notebook_id, block_id = node.get("notebookId"), node.get("blockId")
+    if not (
+        _dataframe_storage.is_safe_id(notebook_id)
+        and _dataframe_storage.is_safe_id(block_id)
+    ):
+        return None
+    return StorageTarget(notebook_id=notebook_id, 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; report any failure but EROFS 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 _dataframe_storage.is_storable_frame(result.result):
+            return
+        directory = (
+            _dataframe_storage.resolve_project_root()
+            / BLOCK_FRAMES_DIR
+            / target.notebook_id
+            / target.block_id
+        )
+        _dataframe_storage.store_frame(result.result, directory)
+    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",
+    so failures go to the webapp instead. Warnings are ignored because they would
+    reach the cell too.
+    """
+    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, the interrupt would leave 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.py

@@ -10,10 +10,15 @@
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
 import pandas as pd
-from typing_extensions import TypeAlias, assert_never
+from typing_extensions import TypeAlias, TypeGuard, assert_never
 
 from deepnote_toolkit.config import resolve_notebook_root
 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,
@@ -44,11 +49,18 @@
 _SAFE_ID_REGEX = re.compile(r"[A-Za-z0-9_-]{1,128}")
 
 
-def is_safe_id(value: object) -> bool:
+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()
@@ -242,13 +254,12 @@ def _delete_replaced_frames(
                     (directory / entry.name).unlink()
 
 
-def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]:
-    """Store the frame as the current one under `name` and return its reference ID.
+def store_frame(frame: StorableFrame, directory: Path) -> Optional[str]:
+    """Store the frame as the current one in `directory` and return its reference ID.
 
     Returns None, having stored nothing, if the filesystem is read-only. EROFS from
     the manifest step on raises like any other error: the new frame may be current.
     """
-    directory = root / FRAMES_DIR / name
     reference_id = str(uuid.uuid4())
     new_file = f"{reference_id}.arrow"
     try:

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
@@ -62,14 +58,16 @@ 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__}"
         )
 
     root = _dataframe_storage.resolve_project_root()
-    return _dataframe_storage.store_frame(df, name, root)
+    return _dataframe_storage.store_frame(
+        df, root / _dataframe_storage.FRAMES_DIR / name
+    )
 
 
 def delete_dataframe(name: str) -> bool:

deepnote_toolkit/runtime_initialization.py

@@ -8,6 +8,7 @@
 
 from deepnote_toolkit.runtime_patches import apply_runtime_patches
 
+from .artifacts._block_frames import register_block_frame_storage
 from .dataframe_utils import add_formatters
 from .execute_post_start_hooks import execute_post_start_hooks
 from .logging import LoggerManager
@@ -52,6 +53,14 @@ 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 the block DataFrame storage hook.")
+        register_block_frame_storage()
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        logger.error(
+            "Failed to register the block DataFrame storage hook with error: %s", e
+        )
+
     # Disable sqlparse grouping limits for large analytical queries
     try:
         logger.debug("Configuring sqlparse limits.")

tests/unit/artifacts/test_arguments.py

@@ -13,6 +13,8 @@
     FRAME_A,
     NAME,
     arrow_files,
+    cleared_query_preview,
+    query_preview,
     stored_table,
 )
 
@@ -51,17 +53,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"),
@@ -71,9 +62,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"),
 ]