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

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

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

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__}"

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,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