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,212 @@
+"""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 import env
+from deepnote_toolkit.config import get_config
+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, mirroring set_notebook_path().
+
+    The working directory cannot be used: set_notebook_path() moves it into the
+    notebook's directory.
+    """
+    try:
+        cfg = get_config()
+        if cfg.paths.notebook_root:
+            return Path(cfg.paths.notebook_root)
+        if cfg.paths.home_dir:
+            return Path(cfg.paths.home_dir)
+        if cfg.runtime.dev_mode:
+            return Path("/work")
+        return Path(os.environ.get("HOME", str(Path.home()))) / "work"
+    except Exception:  # pylint: disable=broad-exception-caught
+        if env.get_env("DEEPNOTE_RUNNING_IN_DEV_MODE"):
+            return Path("/work")
+        return Path(os.environ.get("HOME", str(Path.home()))) / "work"
+
+
+@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."""
+    interrupted: list[bool] = []
+    with _sigint_handled_by(lambda *_: interrupted.append(True)):
+        yield
+    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"):
+        try:
+            if pa.Codec.is_available(codec):
+                return codec
+        except Exception:  # pylint: disable=broad-exception-caught
+            # pyarrow raises, rather than returning False, for names it does not know
+            continue
+    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 between truncating the manifest and writing it would upload an empty one,
+    and one after it 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,74 @@
+"""Reads and writes the manifest that names a stored DataFrame's current file."""
+
+import json
+import re
+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`."""
+    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__
+    (directory / MANIFEST_NAME).write_text(json.dumps(manifest), encoding="utf-8")

deepnote_toolkit/artifacts/dataframes.py

@@ -0,0 +1,81 @@
+"""Stores and deletes named DataFrames in the project's filesystem.
+
+A stored frame is an ordinary file under `deepnote_dataframes/<name>/`, so anyone
+who can edit the notebook could write it by hand. It holds the DataFrame last stored
+under its name, whichever notebook or execution stored it.
+"""
+
+import errno
+import shutil
+from typing import Any, Optional
+
+from . import dataframe_storage
+
+
+def _check_name(name: object) -> None:
+    if not dataframe_storage.is_safe_id(name):
+        raise ValueError(
+            f"A DataFrame name must match [A-Za-z0-9_-]{{1,128}}, got {name!r}"
+        )
+
+
+def store_dataframe(df: Any, name: str) -> Optional[str]:
+    """Store `df` as the current frame under `name`, replacing the one stored before.
+
+    Args:
+        df: A pandas or polars DataFrame, written whole with no row or size limit.
+        name: The frame's name, unique within the project. Must match
+            `[A-Za-z0-9_-]{1,128}`.
+
+    Returns:
+        The new frame's reference ID, or None if the project's filesystem is
+        read-only, in which case nothing is stored.
+
+    Raises:
+        ValueError: `name` is not valid.
+        TypeError: `df` is not a pandas or polars DataFrame. That includes PySpark
+            and pandas-on-Spark DataFrames, polars LazyFrames and SQL query previews.
+        OSError: The write failed.
+        Exception: `df` cannot be serialized to Arrow, as with `df.to_feather()`.
+
+    A failed call leaves the name's previous frame current.
+    """
+    _check_name(name)
+    kind = dataframe_storage.classify_frame(df)
+    if kind is None:
+        raise TypeError(
+            "store_dataframe needs a pandas or polars DataFrame, "
+            f"got {type(df).__qualname__}"
+        )
+
+    root = dataframe_storage.resolve_project_root()
+    if dataframe_storage.is_read_only(root):
+        return None
+    try:
+        return dataframe_storage.store_frame(df, kind, name, root)
+    except OSError as exc:
+        if exc.errno == errno.EROFS:
+            return None
+        raise
+
+
+def delete_dataframe(name: str) -> bool:
+    """Delete the frames stored under `name`.
+
+    Returns:
+        Whether there was anything to delete. False on a read-only filesystem, where
+        nothing is deleted.
+
+    Raises:
+        ValueError: `name` is not valid.
+    """
+    _check_name(name)
+    root = dataframe_storage.resolve_project_root()
+    if dataframe_storage.is_read_only(root):
+        return False
+
+    directory = root / dataframe_storage.FRAMES_DIR / name
+    if not directory.is_dir():
+        return False
+    shutil.rmtree(directory)
+    return True