feat(artifacts): add store_dataframe and delete_dataframe

deepnote_toolkit/__init__.py

@@ -38,6 +38,7 @@
     (".variable_explorer", "deepnote_get_data_preview_json"),
     (".variable_explorer", "get_var_list"),
     (".ocelots", None),
+    (".artifacts", None),
 ]
 
 # Perform safe imports

deepnote_toolkit/artifacts/__init__.py

@@ -0,0 +1,8 @@
+"""Stores and deletes named DataFrames in the project's filesystem.
+
+Only the names in __all__ are supported public API.
+"""
+
+from .dataframes import delete_dataframe, store_dataframe
+
+__all__ = ["delete_dataframe", "store_dataframe"]

deepnote_toolkit/artifacts/dataframe_storage.py

@@ -0,0 +1,200 @@
+"""Writes DataFrames into the project's filesystem as named Arrow files."""
+
+import contextlib
+import os
+import re
+import signal
+import time
+import uuid
+from pathlib import Path
+from typing import Any, Callable, Iterator, Literal, Optional
+
+from packaging.version import Version
+
+from deepnote_toolkit.config import resolve_notebook_root
+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,
+    is_frame_name,
+    read_manifest,
+    write_manifest,
+)
+
+FRAMES_DIR = Path("deepnote_dataframes")
+STALE_FRAME_AGE_SECONDS = 24 * 60 * 60
+# 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
+POLARS_BATCH_MIN_VERSION = Version("1.37")
+
+FrameKind = Literal["pandas", "polars"]
+
+_SAFE_ID = re.compile(r"[A-Za-z0-9_-]{1,128}")
+
+
+def is_safe_id(value: object) -> bool:
+    """Return whether `value` can be used as a single path segment.
+
+    Frame names become directory names, so they must not be able to leave
+    `FRAMES_DIR`.
+    """
+    return isinstance(value, str) and _SAFE_ID.fullmatch(value) is not None
+
+
+def classify_frame(value: object) -> Optional[FrameKind]:
+    """Return "pandas" or "polars" for a frame worth storing, else None.
+
+    PySpark and pandas-on-Spark frames are neither, so they are skipped.
+    """
+    # A query preview holds at most 100 rows, so storing it would pass off a
+    # truncated frame as the result. isinstance, not the cleared-query marker,
+    # because mutating a preview clears the marker but leaves it a preview.
+    if isinstance(value, DeepnoteQueryPreview):
+        return None
+    if is_pandas_dataframe(value):
+        return "pandas"
+    if is_polars_eager_dataframe(value):
+        return "polars"
+    return None
+
+
+def resolve_project_root() -> Path:
+    """Return the project's root directory, the one set_notebook_path() uses.
+
+    The working directory cannot be used: set_notebook_path() moves it into the
+    notebook's directory.
+    """
+    return resolve_notebook_root()
+
+
+@contextlib.contextmanager
+def _sigint_handled_by(handler: Callable[..., Any]) -> Iterator[None]:
+    """Install `handler` for SIGINT, then restore the previous one.
+
+    Off the main thread signal handlers cannot be set, and nothing needs to be.
+    """
+    try:
+        previous = signal.signal(signal.SIGINT, handler)
+    except ValueError:
+        yield
+        return
+    try:
+        yield
+    finally:
+        if previous is not None:
+            signal.signal(signal.SIGINT, previous)
+
+
+@contextlib.contextmanager
+def _sigint_deferred() -> Iterator[None]:
+    """Hold SIGINT back until the block ends, then raise it as KeyboardInterrupt.
+
+    It is raised even when the block fails, chained to that failure, so that an
+    error never swallows Stop.
+    """
+    interrupted: list[bool] = []
+    try:
+        with _sigint_handled_by(lambda *_: interrupted.append(True)):
+            yield
+    finally:
+        if interrupted:
+            raise KeyboardInterrupt
+
+
+def is_read_only(root: Path) -> bool:
+    """Return whether the project's filesystem is mounted read-only."""
+    return bool(os.statvfs(root).f_flag & os.ST_RDONLY)
+
+
+def _select_compression() -> str:
+    """Return the best Arrow codec this pyarrow build has: zstd, lz4, else none."""
+    import pyarrow as pa
+
+    for codec in ("zstd", "lz4"):
+        if pa.Codec.is_available(codec):
+            return codec
+    return "uncompressed"
+
+
+def _write_frame_file(
+    frame: Any, kind: FrameKind, path: Path, compression: str
+) -> None:
+    """Write the frame to `path` as an Arrow IPC file, with the library's own call.
+
+    Going through a Python file object makes `close()` raise upload errors, which
+    s3fs only surfaces there.
+    """
+    with path.open("wb") as sink:
+        if kind == "pandas":
+            frame.to_feather(sink, compression=compression)
+        else:
+            import polars
+
+            options = {}
+            if Version(polars.__version__) >= POLARS_BATCH_MIN_VERSION:
+                options["record_batch_size"] = POLARS_BATCH_ROWS
+            frame.write_ipc(sink, compression=compression, **options)
+
+
+def _delete_replaced_frames(
+    directory: Path, old: Optional[Manifest], new_file: str
+) -> None:
+    """Delete the frame the replaced manifest called `previous`, and stale strays.
+
+    A stray is a `<uuid>.arrow` that no manifest names, left by a writer that lost
+    a race or crashed. It must be a day old so that a write still in progress is
+    never mistaken for one.
+    """
+    named = {new_file}
+    if old is not None:
+        named.add(old.file)
+        if old.previous is not None and old.previous not in named:
+            (directory / old.previous).unlink(missing_ok=True)
+
+    now = time.time()
+    with os.scandir(directory) as entries:
+        for entry in entries:
+            if entry.name in named or not is_frame_name(entry.name):
+                continue
+            # another writer may prune the same stray first
+            with contextlib.suppress(FileNotFoundError):
+                if (
+                    entry.is_file(follow_symlinks=False)
+                    and now - entry.stat(follow_symlinks=False).st_mtime
+                    > STALE_FRAME_AGE_SECONDS
+                ):
+                    (directory / entry.name).unlink()
+
+
+def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str:
+    """Store the frame as the current one under `name` and return its reference ID.
+
+    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:
+    one after the manifest is renamed into place would delete the frame it now names.
+    """
+    directory = root / FRAMES_DIR / name
+    directory.mkdir(parents=True, exist_ok=True)
+
+    reference_id = str(uuid.uuid4())
+    new_file = f"{reference_id}.arrow"
+    committed = False
+    try:
+        _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))
+            committed = True
+    except BaseException:
+        if not committed:
+            with contextlib.suppress(OSError):
+                (directory / new_file).unlink(missing_ok=True)
+        raise
+    _delete_replaced_frames(directory, old, new_file)
+    return reference_id

deepnote_toolkit/artifacts/dataframe_storage_manifest.py

@@ -0,0 +1,99 @@
+"""Reads and writes the manifest that names a stored DataFrame's current file."""
+
+import contextlib
+import json
+import os
+import re
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, NamedTuple, Optional
+
+from deepnote_toolkit import __version__
+
+MANIFEST_NAME = "manifest.json"
+MANIFEST_VERSION = 1
+MANIFEST_MAX_BYTES = 4096
+
+_FRAME_NAME = re.compile(
+    r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\.arrow",
+    re.IGNORECASE,
+)
+
+
+class Manifest(NamedTuple):
+    """The frames a manifest names; `previous` is the one `file` replaced."""
+
+    file: str
+    previous: Optional[str]
+
+
+def is_frame_name(name: str) -> bool:
+    """Return whether `name` is `<uuid>.arrow`, the only form of frame file name."""
+    return _FRAME_NAME.fullmatch(name) is not None
+
+
+def read_manifest(directory: Path) -> Optional[Manifest]:
+    """Return the files the manifest names, or None if there is no usable one.
+
+    Any code in the project can write the manifest, so a manifest that is too big,
+    is not JSON or names anything but `<uuid>.arrow` counts as absent: the writer
+    then overwrites it and never deletes a path it supplied.
+    """
+    try:
+        with (directory / MANIFEST_NAME).open("rb") as handle:
+            raw = handle.read(MANIFEST_MAX_BYTES + 1)
+    except FileNotFoundError:
+        return None
+    if len(raw) > MANIFEST_MAX_BYTES:
+        return None
+    try:
+        data = json.loads(raw)
+    except (ValueError, RecursionError):
+        return None
+    if not isinstance(data, dict):
+        return None
+
+    file, previous = data.get("file"), data.get("previous")
+    if not isinstance(file, str) or not is_frame_name(file):
+        return None
+    if previous is not None and not (
+        isinstance(previous, str) and is_frame_name(previous)
+    ):
+        return None
+    return Manifest(file=file, previous=previous)
+
+
+def write_manifest(
+    directory: Path, new_file: str, previous: Optional[str], rows: int
+) -> None:
+    """Point the manifest at `new_file`.
+
+    The manifest is written under a temporary name and renamed over the old one, so
+    a write that fails part-way leaves the old manifest in place.
+    """
+    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__
+    # 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.
+    temporary = tempfile.NamedTemporaryFile(
+        "w",
+        encoding="utf-8",
+        dir=directory,
+        prefix=f".{MANIFEST_NAME}.",
+        suffix=".tmp",
+        delete=False,
+    )
+    try:
+        with temporary:
+            temporary.write(json.dumps(manifest))
+        os.replace(temporary.name, directory / MANIFEST_NAME)
+    except BaseException:
+        with contextlib.suppress(OSError):
+            os.unlink(temporary.name)
+        raise