From 62a4d0ec7ba0aee167ed03669e5bd6793bd6d0fc Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 8 Oct 2026 12:33:58 +0000 Subject: [PATCH 1/3] feat(artifacts): add store_dataframe and delete_dataframe `artifacts.store_dataframe(df, name)` writes a whole pandas or polars DataFrame to `deepnote_dataframes//.arrow` in the project's filesystem, with `to_feather` / `write_ipc`, zstd falling back to lz4 and then uncompressed. A `manifest.json` names the current file and the one it replaced, and every write rewrites it. `artifacts.delete_dataframe(name)` deletes a name's directory. The function raises like the user's own `df.to_feather()` would, and returns None on a read-only mount. An interrupt stops the write and removes the partial file, except during the manifest update, which finishes first. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01B441SpUqRn7QTzcnWxjNa5 --- deepnote_toolkit/__init__.py | 1 + deepnote_toolkit/artifacts.py | 81 ++ deepnote_toolkit/dataframe_storage.py | 211 ++++ .../dataframe_storage_manifest.py | 74 ++ tests/unit/dataframe_storage/__init__.py | 0 tests/unit/dataframe_storage/conftest.py | 116 ++ .../unit/dataframe_storage/test_artifacts.py | 1019 +++++++++++++++++ tests/unit/helpers/dataframe_storage.py | 104 ++ 8 files changed, 1606 insertions(+) create mode 100644 deepnote_toolkit/artifacts.py create mode 100644 deepnote_toolkit/dataframe_storage.py create mode 100644 deepnote_toolkit/dataframe_storage_manifest.py create mode 100644 tests/unit/dataframe_storage/__init__.py create mode 100644 tests/unit/dataframe_storage/conftest.py create mode 100644 tests/unit/dataframe_storage/test_artifacts.py create mode 100644 tests/unit/helpers/dataframe_storage.py diff --git a/deepnote_toolkit/__init__.py b/deepnote_toolkit/__init__.py index 39688789..ec178ad7 100644 --- a/deepnote_toolkit/__init__.py +++ b/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 diff --git a/deepnote_toolkit/artifacts.py b/deepnote_toolkit/artifacts.py new file mode 100644 index 00000000..6c359a94 --- /dev/null +++ b/deepnote_toolkit/artifacts.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//`, 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 deepnote_toolkit 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 diff --git a/deepnote_toolkit/dataframe_storage.py b/deepnote_toolkit/dataframe_storage.py new file mode 100644 index 00000000..df2e4329 --- /dev/null +++ b/deepnote_toolkit/dataframe_storage.py @@ -0,0 +1,211 @@ +"""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.dataframe_storage_manifest import ( + Manifest, + is_frame_name, + read_manifest, + write_manifest, +) +from deepnote_toolkit.ocelots.utils import ( + is_pandas_dataframe, + is_polars_eager_dataframe, +) +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview + +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 `.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 diff --git a/deepnote_toolkit/dataframe_storage_manifest.py b/deepnote_toolkit/dataframe_storage_manifest.py new file mode 100644 index 00000000..0bd535b7 --- /dev/null +++ b/deepnote_toolkit/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 `.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 `.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") diff --git a/tests/unit/dataframe_storage/__init__.py b/tests/unit/dataframe_storage/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/dataframe_storage/conftest.py b/tests/unit/dataframe_storage/conftest.py new file mode 100644 index 00000000..83879d80 --- /dev/null +++ b/tests/unit/dataframe_storage/conftest.py @@ -0,0 +1,116 @@ +"""Fixtures for the DataFrame storage tests. + +Kept out of tests/unit/conftest.py because the names are generic. +""" + +import time +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional + +import pytest + +from deepnote_toolkit import dataframe_storage + + +@pytest.fixture +def root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """The project root, a fresh directory.""" + project_root = tmp_path / "project" + project_root.mkdir() + monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project_root) + return project_root + + +@pytest.fixture +def tokyo_time() -> Iterator[None]: + """Run in a local time zone nine hours ahead of UTC.""" + if not hasattr(time, "tzset"): + pytest.skip("time zones cannot be switched on this platform") + with pytest.MonkeyPatch.context() as patch: + patch.setenv("TZ", "Asia/Tokyo") + time.tzset() + yield + # the C library caches the zone, so restoring the variable is not enough + time.tzset() + + +@pytest.fixture(scope="module") +def spark() -> Any: + from pyspark.sql import SparkSession + + return ( + SparkSession.builder.master("local") # type: ignore + .appName("Toolkit") + .config("spark.driver.memory", "4g") + .getOrCreate() + ) + + +@pytest.fixture +def arrow_sinks(monkeypatch: pytest.MonkeyPatch) -> List[Any]: + """Every file object opened for writing a frame, in order.""" + sinks: List[Any] = [] + real_open = Path.open + + def open_and_record(self: Path, mode: str = "r", *args: Any, **kwargs: Any) -> Any: + handle = real_open(self, mode, *args, **kwargs) + if self.suffix == ".arrow" and "w" in mode: + sinks.append(handle) + return handle + + monkeypatch.setattr(Path, "open", open_and_record) + return sinks + + +class _FailingClose: + """A write handle whose close() really closes, then fails like s3fs's upload.""" + + def __init__(self, handle: Any, error: OSError) -> None: + self._handle = handle + self._error = error + + def close(self) -> None: + self._handle.close() + raise self._error + + def __enter__(self) -> "_FailingClose": + return self + + def __exit__(self, *exc_info: Any) -> None: + self.close() + + def __getattr__(self, name: str) -> Any: + return getattr(self._handle, name) + + +@pytest.fixture +def close_error(monkeypatch: pytest.MonkeyPatch) -> Dict[str, Optional[OSError]]: + """Set `["error"]` to make closing a frame file fail with that error.""" + state: Dict[str, Optional[OSError]] = {"error": None} + real_open = Path.open + + def open_failing_on_close( + self: Path, mode: str = "r", *args: Any, **kwargs: Any + ) -> Any: + handle = real_open(self, mode, *args, **kwargs) + error = state["error"] + if error is None or self.suffix != ".arrow" or "w" not in mode: + return handle + return _FailingClose(handle, error) + + monkeypatch.setattr(Path, "open", open_failing_on_close) + return state + + +@pytest.fixture +def no_root_settings(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "DEEPNOTE_PATHS__NOTEBOOK_ROOT", + "DEEPNOTE_PATHS__HOME_DIR", + "DEEPNOTE_HOME_DIR", + "HOME_DIR", + "DEEPNOTE_RUNTIME__DEV_MODE", + "DEEPNOTE_RUNNING_IN_DEV_MODE", + "RUNNING_IN_DEV_MODE", + ): + monkeypatch.delenv(name, raising=False) diff --git a/tests/unit/dataframe_storage/test_artifacts.py b/tests/unit/dataframe_storage/test_artifacts.py new file mode 100644 index 00000000..7675f886 --- /dev/null +++ b/tests/unit/dataframe_storage/test_artifacts.py @@ -0,0 +1,1019 @@ +"""Tests for the artifacts API that stores and deletes named DataFrames.""" + +import contextlib +import copy +import errno +import json +import os +import re +import signal +import threading +import time +from datetime import datetime, timezone +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, Iterator, List, Optional +from unittest import mock + +import pandas as pd +import polars as pl +import pyarrow as pa +import pytest + +import deepnote_toolkit +from deepnote_toolkit import artifacts, dataframe_storage, dataframe_storage_manifest +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview +from tests.unit.helpers.dataframe_storage import ( + ALWAYS_FAILING, + VERSION_DEPENDENT, + interrupt_at, + probe_failure, +) + +NAME = "revenue" +UUID_PATTERN = r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}" +DAY = 24 * 60 * 60 +FRAME_A = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee.arrow" +FRAME_B = "11111111-2222-4333-8444-555555555555.arrow" + + +def _frame_dir(root: Path, name: str = NAME) -> Path: + return root / "deepnote_dataframes" / name + + +def _arrow_files(root: Path, name: str = NAME) -> List[Path]: + return sorted(_frame_dir(root, name).glob("*.arrow")) + + +def _manifest(root: Path, name: str = NAME) -> Dict[str, Any]: + return json.loads((_frame_dir(root, name) / "manifest.json").read_text()) + + +def _stored_table(root: Path, name: str = NAME) -> pa.Table: + """The frame the name's manifest currently names.""" + path = _frame_dir(root, name) / _manifest(root, name)["file"] + return pa.ipc.open_file(path).read_all() + + +def _backdate(path: Path, age_seconds: float) -> None: + mtime = time.time() - age_seconds + os.utime(path, (mtime, mtime), follow_symlinks=False) + + +def _read_only(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + dataframe_storage.os, + "statvfs", + lambda path: SimpleNamespace(f_flag=os.ST_RDONLY), + ) + + +def test_stores_pandas_frame_and_manifest(root: Path) -> None: + reference_id = artifacts.store_dataframe( + pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}), NAME + ) + + assert reference_id is not None + assert re.fullmatch(UUID_PATTERN, reference_id) + assert [path.name for path in _arrow_files(root)] == [f"{reference_id}.arrow"] + assert _stored_table(root).to_pydict() == {"a": [1, 2, 3], "b": ["x", "y", "z"]} + manifest = _manifest(root) + assert set(manifest) == {"version", "file", "written_at", "rows", "toolkit_version"} + assert manifest["version"] == 1 + assert manifest["file"] == f"{reference_id}.arrow" + assert manifest["rows"] == 3 + assert manifest["toolkit_version"] == deepnote_toolkit.__version__ + assert re.fullmatch(r"\d{4}-\d\d-\d\dT\d\d:\d\d:\d\dZ", manifest["written_at"]) + + +@pytest.mark.usefixtures("tokyo_time") +def test_manifest_time_is_utc(root: Path) -> None: + before = datetime.now(timezone.utc) + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + written = datetime.strptime(_manifest(root)["written_at"], "%Y-%m-%dT%H:%M:%SZ") + assert abs((written.replace(tzinfo=timezone.utc) - before).total_seconds()) < 60 + + +def test_stores_polars_frame(root: Path) -> None: + reference_id = artifacts.store_dataframe( + pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}), NAME + ) + + assert _manifest(root)["file"] == f"{reference_id}.arrow" + assert _stored_table(root).to_pydict() == {"a": [1, 2, 3], "b": ["x", "y", "z"]} + assert _manifest(root)["rows"] == 3 + + +def test_writing_does_not_change_the_callers_frame(root: Path) -> None: + frame = pd.DataFrame({1: [1, 2], 2: [3, 4]}, index=pd.Index(["a", "b"], name="idx")) + before = copy.deepcopy(frame) + + artifacts.store_dataframe(frame, NAME) + + assert len(_arrow_files(root)) == 1 + pd.testing.assert_frame_equal(frame, before) + + +def test_groupby_keys_are_stored_after_the_data_columns(root: Path) -> None: + frame = pd.DataFrame({"k": ["a", "a", "b"], "v": [1, 2, 3]}).groupby("k").sum() + + artifacts.store_dataframe(frame, NAME) + + assert _stored_table(root).schema.names == ["v", "k"] + + +def test_transposed_frame_is_stored_with_string_names(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"x": [1, 2], "y": [3, 4]}).T, NAME) + + assert _stored_table(root).schema.names == ["0", "1", "__index_level_0__"] + + +@pytest.mark.parametrize( + "frame", + [pd.DataFrame({"a": []}), pl.DataFrame({"a": []})], + ids=["pandas", "polars"], +) +def test_empty_frame_is_stored(root: Path, frame: Any) -> None: + artifacts.store_dataframe(frame, NAME) + + assert _stored_table(root).num_rows == 0 + assert _manifest(root)["rows"] == 0 + + +def test_names_are_stored_apart(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), "first") + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), "second") + + assert _stored_table(root, "first").to_pydict() == {"a": [1]} + assert _stored_table(root, "second").to_pydict() == {"a": [2]} + + +def test_the_last_call_with_a_name_wins(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.store_dataframe(pl.DataFrame({"a": [2]}), NAME) + + assert _stored_table(root).to_pydict() == {"a": [2]} + + +@pytest.mark.parametrize("name", ["x", "x" * 128, "Rev-2026_Q3"]) +def test_valid_names_are_accepted(root: Path, name: str) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + + assert _stored_table(root, name).to_pydict() == {"a": [1]} + + +INVALID_NAMES = [ + pytest.param("", id="empty"), + pytest.param(".", id="dot"), + pytest.param("..", id="dotdot"), + pytest.param("../x", id="traversal"), + pytest.param("a/b", id="slash"), + pytest.param("a\\b", id="backslash"), + pytest.param("a b", id="space"), + pytest.param("revenue\n", id="trailing-newline"), + pytest.param("é", id="non-ascii"), + pytest.param("x" * 129, id="too-long"), + pytest.param(5, id="int"), + pytest.param(None, id="none"), +] + + +@pytest.mark.parametrize("name", INVALID_NAMES) +def test_invalid_name_raises_before_anything_is_written(root: Path, name: Any) -> None: + with mock.patch.object(pd.DataFrame, "to_feather", autospec=True) as to_feather: + with pytest.raises(ValueError): + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + + to_feather.assert_not_called() + assert list(root.iterdir()) == [] + + +@pytest.mark.parametrize("name", INVALID_NAMES) +def test_invalid_name_raises_on_a_read_only_mount( + root: Path, monkeypatch: pytest.MonkeyPatch, name: Any +) -> None: + _read_only(monkeypatch) + + with pytest.raises(ValueError): + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + + +def _query_preview() -> DeepnoteQueryPreview: + return DeepnoteQueryPreview({"a": [1]}, deepnote_query="select 1") + + +def _cleared_query_preview() -> DeepnoteQueryPreview: + preview = _query_preview() + preview["b"] = 2 + return preview + + +UNSUPPORTED_VALUES = [ + pytest.param(lambda: 1, id="int"), + pytest.param(lambda: None, id="none"), + pytest.param(lambda: pd.Series([1, 2]), id="series"), + pytest.param(lambda: {"a": [1]}, id="dict"), + pytest.param(lambda: pl.DataFrame({"a": [1]}).lazy(), id="lazyframe"), + 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.mark.parametrize("make_value", UNSUPPORTED_VALUES) +def test_unsupported_value_raises_type_error( + root: Path, make_value: Callable[[], Any] +) -> None: + with pytest.raises(TypeError, match="pandas or polars DataFrame"): + artifacts.store_dataframe(make_value(), NAME) + + assert list(root.iterdir()) == [] + + +@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES) +def test_unsupported_value_raises_type_error_on_a_read_only_mount( + root: Path, monkeypatch: pytest.MonkeyPatch, make_value: Callable[[], Any] +) -> None: + _read_only(monkeypatch) + + with pytest.raises(TypeError): + artifacts.store_dataframe(make_value(), NAME) + + +def test_pyspark_frame_raises_type_error(root: Path, spark: Any) -> None: + with pytest.raises(TypeError, match="pandas or polars DataFrame"): + artifacts.store_dataframe(spark.range(3), NAME) + + assert list(root.iterdir()) == [] + + +def test_pandas_on_spark_frame_raises_type_error( + root: Path, spark: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # pyspark.pandas warns on import unless told that pyarrow may ignore time zones + monkeypatch.setenv("PYARROW_IGNORE_TIMEZONE", "1") + # pyspark 3.5, which Python < 3.13 gets, imports distutils, gone since Python 3.12 + pytest.importorskip("pyspark.pandas") + # pandas-on-Spark refuses to run under ANSI mode, which Spark 4 enables by default + ansi_mode = spark.conf.get("spark.sql.ansi.enabled") + spark.conf.set("spark.sql.ansi.enabled", "false") + try: + frame = spark.range(3).pandas_api() + finally: + spark.conf.set("spark.sql.ansi.enabled", ansi_mode) + + with pytest.raises(TypeError): + artifacts.store_dataframe(frame, NAME) + + assert list(root.iterdir()) == [] + + +def test_second_write_keeps_one_previous_frame(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + first = _manifest(root) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + second = _manifest(root) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + third = _manifest(root) + + assert "previous" not in first + assert second["previous"] == first["file"] + assert third["previous"] == second["file"] + assert [path.name for path in _arrow_files(root)] == sorted( + [second["file"], third["file"]] + ) + assert _stored_table(root).to_pydict() == {"a": [3]} + previous = pa.ipc.open_file(_frame_dir(root) / third["previous"]) + assert previous.read_all().to_pydict() == {"a": [2]} + + +def test_unreferenced_frames_are_deleted_only_when_stale(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + directory = _frame_dir(root) + current = directory / _manifest(root)["file"] + _backdate(current, 2 * DAY) + + def add(name: str, age_seconds: float) -> Path: + path = directory / name + path.write_bytes(b"unreferenced") + _backdate(path, age_seconds) + return path + + stale = add(FRAME_A, 2 * DAY) + almost_stale = add(FRAME_B, DAY - 3600) + fresh = add("99999999-8888-4777-8666-555555555555.arrow", 0) + not_a_frame = [ + add("notes.arrow", 2 * DAY), + add("x.txt", 2 * DAY), + add(f"backup-{FRAME_A}", 2 * DAY), + add(f"{FRAME_A}.bak", 2 * DAY), + ] + outside = root / "outside.arrow" + outside.write_bytes(b"precious") + link = directory / "77777777-6666-4555-8444-333333333333.arrow" + link.symlink_to(outside) + _backdate(link, 2 * DAY) + + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert not stale.exists() + assert almost_stale.exists() + assert fresh.exists() + assert all(path.exists() for path in not_a_frame) + assert link.is_symlink() + assert outside.read_bytes() == b"precious" + # the frame the replaced manifest named stays as `previous`, however old it is + assert current.exists() + + +def _manifest_bytes(**fields: Any) -> bytes: + return json.dumps({"version": 1, **fields}).encode() + + +@pytest.mark.parametrize( + "manifest_bytes", + [ + pytest.param(b"{not json", id="invalid-json"), + pytest.param(b"\x80\x81\x82", id="not-utf8"), + pytest.param(b"[]", id="not-an-object"), + pytest.param(_manifest_bytes(), id="no-file"), + pytest.param(_manifest_bytes(file="../../data/x.csv"), id="file-traversal"), + pytest.param(_manifest_bytes(file=7, previous=FRAME_B), id="file-not-a-string"), + pytest.param( + _manifest_bytes(file=FRAME_A, previous="../../data/x.csv"), + id="previous-traversal", + ), + pytest.param(_manifest_bytes(file=f"x{FRAME_A}"), id="file-prefixed"), + pytest.param(_manifest_bytes(file=f"{FRAME_A}.bak"), id="file-suffixed"), + pytest.param( + _manifest_bytes(file=f"{FRAME_A}/../../data/x.csv"), + id="file-valid-name-then-traversal", + ), + pytest.param(_manifest_bytes(file=f"../{FRAME_A}"), id="file-in-parent"), + pytest.param( + _manifest_bytes(file=FRAME_A, previous=f"x{FRAME_B}"), + id="previous-prefixed", + ), + pytest.param( + _manifest_bytes(file=FRAME_A, previous=f"{FRAME_B}.bak"), + id="previous-suffixed", + ), + pytest.param( + _manifest_bytes(file=FRAME_A, previous=f"{FRAME_B}/../../data/x.csv"), + id="previous-valid-name-then-traversal", + ), + pytest.param( + _manifest_bytes(file=FRAME_A, previous=f"../{FRAME_B}"), + id="previous-in-parent", + ), + pytest.param( + _manifest_bytes(file=FRAME_A, previous=FRAME_B, pad="x" * 5000), + id="oversized-truncated-mid-string", + ), + pytest.param( + _manifest_bytes(file=FRAME_A, previous=FRAME_B) + b" " * 5000, + id="oversized-valid-json-then-spaces", + ), + ], +) +def test_unusable_manifest_counts_as_no_manifest( + root: Path, manifest_bytes: bytes +) -> None: + directory = _frame_dir(root) + directory.mkdir(parents=True) + (directory / "manifest.json").write_bytes(manifest_bytes) + named = [directory / FRAME_A, directory / FRAME_B] + for path in named: + path.write_bytes(b"named by the hostile manifest") + sentinel = (directory / "../../data/x.csv").resolve() + sentinel.parent.mkdir(parents=True) + sentinel.write_text("keep") + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + manifest = _manifest(root) + assert "previous" not in manifest + assert manifest["file"] not in (FRAME_A, FRAME_B) + assert _stored_table(root).to_pydict() == {"a": [1]} + assert sentinel.read_text() == "keep" + # a manifest that counts as absent names nothing, so nothing it names is deleted + assert all(path.exists() for path in named) + + +def test_manifest_of_exactly_4096_bytes_is_honoured(root: Path) -> None: + directory = _frame_dir(root) + directory.mkdir(parents=True) + raw = _manifest_bytes(file=FRAME_A, previous=FRAME_B).ljust(4096) + assert len(raw) == 4096 + (directory / "manifest.json").write_bytes(raw) + for name in (FRAME_A, FRAME_B): + (directory / name).write_bytes(b"named by the manifest") + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert _manifest(root)["previous"] == FRAME_A + assert (directory / FRAME_A).exists() + assert not (directory / FRAME_B).exists() + + +def test_manifest_nested_too_deeply_to_parse_counts_as_no_manifest( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Python before 3.13 raises RecursionError for the nesting a 4 KiB manifest allows.""" + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + def too_deep(raw: bytes) -> Any: + raise RecursionError( + "maximum recursion depth exceeded while decoding a JSON array" + ) + + monkeypatch.setattr( + dataframe_storage_manifest, + "json", + SimpleNamespace(loads=too_deep, dumps=json.dumps), + ) + + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert "previous" not in _manifest(root) + assert _stored_table(root).to_pydict() == {"a": [2]} + + +def test_manifest_is_written_after_the_frame_file_is_complete( + root: Path, arrow_sinks: List[Any], monkeypatch: pytest.MonkeyPatch +) -> None: + seen: List[Dict[str, Any]] = [] + original = dataframe_storage.write_manifest + + def spy(directory: Path, new_file: str, *args: Any) -> None: + table = pa.ipc.open_file(directory / new_file).read_all() + seen.append({"closed": arrow_sinks[0].closed, "rows": table.num_rows}) + original(directory, new_file, *args) + + monkeypatch.setattr(dataframe_storage, "write_manifest", spy) + + artifacts.store_dataframe(pd.DataFrame({"a": [1, 2, 3]}), NAME) + + assert seen == [{"closed": True, "rows": 3}] + + +def test_concurrent_writers_leave_a_complete_current_frame( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + barrier = threading.Barrier(2, timeout=10) + original = dataframe_storage.read_manifest + + def read_then_wait(directory: Path) -> Any: + manifest = original(directory) + barrier.wait() + return manifest + + errors: List[BaseException] = [] + + def write(value: int) -> None: + try: + artifacts.store_dataframe(pd.DataFrame({"a": [value]}), NAME) + except BaseException as exc: # surfaced by the assertion below + errors.append(exc) + + with monkeypatch.context() as patch: + patch.setattr(dataframe_storage, "read_manifest", read_then_wait) + threads = [threading.Thread(target=write, args=(value,)) for value in (1, 2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + + assert errors == [] + files = _arrow_files(root) + assert len(files) == 2 + current = _manifest(root)["file"] + stray = next(path for path in files if path.name != current) + for path in files: + assert pa.ipc.open_file(path).read_all().num_rows == 1 + + _backdate(stray, 2 * DAY) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + + assert not stray.exists() + assert (_frame_dir(root) / current).exists() + assert _stored_table(root).to_pydict() == {"a": [3]} + + +@pytest.mark.parametrize( + "make_frame", [pd.DataFrame, pl.DataFrame], ids=["pandas", "polars"] +) +def test_frames_are_compressed(root: Path, make_frame: Any) -> None: + artifacts.store_dataframe(make_frame({"a": ["x" * 100] * 10_000}), NAME) + + stored = _frame_dir(root) / _manifest(root)["file"] + assert 0 < stored.stat().st_size < 100_000 + + +class _FakeCodec: + """Stands in for pyarrow.Codec, whose extension type cannot be patched.""" + + def __init__(self, availability: Dict[str, Any]) -> None: + self._availability = availability + + def is_available(self, name: str) -> bool: + outcome = self._availability[name] + if isinstance(outcome, Exception): + raise outcome + return outcome + + +@pytest.mark.parametrize( + "availability, expected", + [ + pytest.param({"zstd": True, "lz4": True}, "zstd", id="zstd"), + pytest.param({"zstd": False, "lz4": True}, "lz4", id="zstd-missing"), + pytest.param( + {"zstd": ValueError("Unsupported codec"), "lz4": True}, + "lz4", + id="zstd-unrecognised", + ), + pytest.param({"zstd": False, "lz4": False}, "uncompressed", id="none"), + pytest.param( + {"zstd": False, "lz4": ValueError("Unsupported codec")}, + "uncompressed", + id="none-unrecognised", + ), + ], +) +def test_compression_falls_back_from_zstd_to_lz4_to_none( + monkeypatch: pytest.MonkeyPatch, availability: Dict[str, Any], expected: str +) -> None: + monkeypatch.setattr(pa, "Codec", _FakeCodec(availability)) + + assert dataframe_storage._select_compression() == expected + + +@pytest.mark.skipif( + not pa.Codec.is_available("zstd"), reason="this pyarrow build has no zstd" +) +def test_zstd_is_chosen_when_pyarrow_has_it() -> None: + assert dataframe_storage._select_compression() == "zstd" + + +@pytest.mark.parametrize( + "make_frame", [pd.DataFrame, pl.DataFrame], ids=["pandas", "polars"] +) +def test_frame_is_stored_uncompressed_when_no_codec_is_available( + root: Path, monkeypatch: pytest.MonkeyPatch, make_frame: Any +) -> None: + monkeypatch.setattr(pa, "Codec", _FakeCodec({"zstd": False, "lz4": False})) + + artifacts.store_dataframe(make_frame({"a": ["x" * 100] * 10_000}), NAME) + + stored = _frame_dir(root) / _manifest(root)["file"] + assert pa.ipc.open_file(stored).read_all().num_rows == 10_000 + assert stored.stat().st_size > 1_000_000 + + +def test_pandas_frame_is_written_in_64k_row_batches(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": range(70_000)}), NAME) + + reader = pa.ipc.open_file(_frame_dir(root) / _manifest(root)["file"]) + assert reader.num_record_batches == 2 + assert reader.read_all().num_rows == 70_000 + + +def test_polars_frame_is_written_in_64k_row_batches(root: Path) -> None: + artifacts.store_dataframe(pl.DataFrame({"a": range(200_000)}), NAME) + + reader = pa.ipc.open_file(_frame_dir(root) / _manifest(root)["file"]) + assert reader.num_record_batches == 4 + assert reader.read_all().num_rows == 200_000 + + +@pytest.mark.parametrize( + "polars_version, expected_options", + [("1.36.9", {}), ("1.37.0", {"record_batch_size": 65_536})], + ids=["before-1.37", "from-1.37"], +) +def test_polars_batch_size_is_passed_to_polars_from_1_37( + root: Path, + monkeypatch: pytest.MonkeyPatch, + polars_version: str, + expected_options: Dict[str, int], +) -> None: + monkeypatch.setattr(pl, "__version__", polars_version) + original = pl.DataFrame.write_ipc + + with mock.patch.object( + pl.DataFrame, "write_ipc", autospec=True, side_effect=original + ) as write_ipc: + artifacts.store_dataframe(pl.DataFrame({"a": [1, 2, 3]}), NAME) + + write_ipc.assert_called_once() + options = { + name: value + for name, value in write_ipc.call_args.kwargs.items() + if name == "record_batch_size" + } + assert options == expected_options + assert _stored_table(root).to_pydict() == {"a": [1, 2, 3]} + + +def _store_bad_frame_after_a_good_one( + root: Path, make_frame: Callable[[], Any], expected_error: Optional[type] +) -> Dict[str, Any]: + """Store a good frame, then a bad one; return what the bad call left behind.""" + artifacts.store_dataframe(pd.DataFrame({"ok": [1]}), NAME) + directory = _frame_dir(root) + manifest_bytes = (directory / "manifest.json").read_bytes() + if expected_error is None: + artifacts.store_dataframe(make_frame(), NAME) + else: + with pytest.raises(expected_error): + artifacts.store_dataframe(make_frame(), NAME) + return { + "manifest_bytes": manifest_bytes, + "listing": sorted(path.name for path in directory.iterdir()), + } + + +@pytest.mark.parametrize("make_frame", ALWAYS_FAILING) +def test_unserializable_frame_raises_and_keeps_the_previous_frame( + root: Path, make_frame: Callable[[], Any] +) -> None: + raised_by_pandas = probe_failure(make_frame()) + assert raised_by_pandas is not None, "pandas is expected to refuse this frame" + + left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pandas) + + assert (_frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] + assert left["listing"] == sorted(["manifest.json", _manifest(root)["file"]]) + assert _stored_table(root).to_pydict() == {"ok": [1]} + + +@pytest.mark.parametrize("make_frame, rows", VERSION_DEPENDENT) +def test_frames_pandas_may_refuse_follow_what_pandas_does( + root: Path, make_frame: Callable[[], Any], rows: int +) -> None: + raised_by_pandas = probe_failure(make_frame()) + + left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pandas) + + if raised_by_pandas is None: + assert _stored_table(root).num_rows == rows + return + assert (_frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] + assert left["listing"] == sorted(["manifest.json", _manifest(root)["file"]]) + + +def test_polars_object_column_raises_and_keeps_the_previous_frame(root: Path) -> None: + artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) + manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() + + with pytest.raises(pl.exceptions.ComputeError): + artifacts.store_dataframe( + pl.DataFrame({"a": pl.Series([object()], dtype=pl.Object)}), NAME + ) + + assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert len(_arrow_files(root)) == 1 + + +def test_partial_file_is_deleted_when_the_write_fails( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def write_some_then_fail(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + sink.write(b"ARROW1 partial") + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(pd.DataFrame, "to_feather", write_some_then_fail) + + with pytest.raises(OSError) as raised: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert raised.value.errno == errno.EIO + assert _arrow_files(root) == [] + assert not (_frame_dir(root) / "manifest.json").exists() + + +@pytest.mark.parametrize( + "error", + [ + PermissionError(errno.EPERM, "Operation not permitted"), + PermissionError(errno.EACCES, "Permission denied"), + OSError(errno.EIO, "Input/output error"), + OSError(errno.ENOSPC, "No space left on device"), + ], + ids=["EPERM", "EACCES", "EIO", "ENOSPC"], +) +def test_error_from_close_raises_and_stores_nothing( + root: Path, close_error: Dict[str, Optional[OSError]], error: OSError +) -> None: + """s3fs surfaces a failed upload only when the file is closed.""" + close_error["error"] = error + + with pytest.raises(OSError) as raised: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert raised.value is error + assert _arrow_files(root) == [] + assert not (_frame_dir(root) / "manifest.json").exists() + + +def test_failure_before_the_manifest_step_deletes_nothing( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + directory = _frame_dir(root) + stale = directory / FRAME_A + stale.write_bytes(b"unreferenced") + _backdate(stale, 2 * DAY) + manifest_bytes = (directory / "manifest.json").read_bytes() + listing = sorted(path.name for path in directory.iterdir()) + + def broken_manifest(*args: Any) -> Any: + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(dataframe_storage, "write_manifest", broken_manifest) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + + assert (directory / "manifest.json").read_bytes() == manifest_bytes + assert sorted(path.name for path in directory.iterdir()) == listing + assert _stored_table(root).to_pydict() == {"a": [2]} + + +def test_failure_after_the_manifest_step_keeps_the_new_frame( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + def broken_cleanup(*args: Any) -> Any: + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(dataframe_storage, "_delete_replaced_frames", broken_cleanup) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert _stored_table(root).to_pydict() == {"a": [2]} + + +def test_stray_pruned_by_another_writer_meanwhile_is_not_a_failure( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + directory = _frame_dir(root) + strays = [directory / FRAME_A, directory / FRAME_B] + for stray in strays: + stray.write_bytes(b"unreferenced") + _backdate(stray, 2 * DAY) + real_scandir = os.scandir + + @contextlib.contextmanager + def scandir_then_lose_a_stray(path: Any) -> Iterator[Any]: + with real_scandir(path) as entries: + listing = list(entries) + strays[0].unlink() + yield iter(listing) + + monkeypatch.setattr(dataframe_storage.os, "scandir", scandir_then_lose_a_stray) + + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert not strays[1].exists() + assert _stored_table(root).to_pydict() == {"a": [2]} + + +def test_read_only_mount_stores_nothing_and_does_not_convert( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _read_only(monkeypatch) + + with mock.patch.object(pd.DataFrame, "to_feather", autospec=True) as to_feather: + result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert result is None + to_feather.assert_not_called() + assert list(root.iterdir()) == [] + + +def test_erofs_from_the_write_stores_nothing( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def refuse(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + raise OSError(errno.EROFS, "Read-only file system") + + monkeypatch.setattr(pd.DataFrame, "to_feather", refuse) + + result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert result is None + assert _arrow_files(root) == [] + + +def _write_partial_file_then(interrupt: Callable[[], None]) -> Callable[..., None]: + def to_feather(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + sink.write(b"ARROW1 partial") + interrupt() + + return to_feather + + +def _raise_keyboard_interrupt() -> None: + raise KeyboardInterrupt + + +@pytest.mark.parametrize( + "interrupt", + [lambda: signal.raise_signal(signal.SIGINT), _raise_keyboard_interrupt], + ids=["sigint", "keyboard-interrupt"], +) +def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( + root: Path, monkeypatch: pytest.MonkeyPatch, interrupt: Callable[[], None] +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() + handler_before = signal.getsignal(signal.SIGINT) + monkeypatch.setattr(pd.DataFrame, "to_feather", _write_partial_file_then(interrupt)) + + with pytest.raises(KeyboardInterrupt): + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert len(_arrow_files(root)) == 1 + assert signal.getsignal(signal.SIGINT) is handler_before + + +@pytest.mark.parametrize( + "point", ["manifest-truncated", "manifest-written", "old-frames-deleted"] +) +def test_interrupt_after_the_frame_is_written_leaves_a_usable_manifest( + root: Path, monkeypatch: pytest.MonkeyPatch, point: str +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + handler_before = signal.getsignal(signal.SIGINT) + interrupt_at(point, monkeypatch) + + with pytest.raises(KeyboardInterrupt): + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + + # the manifest parses and names the new frame, which is there and complete + assert _stored_table(root).to_pydict() == {"a": [3]} + assert signal.getsignal(signal.SIGINT) is handler_before + + +def _broken_config() -> Any: + raise ValueError("invalid config") + + +@pytest.mark.usefixtures("no_root_settings") +class TestProjectRoot: + """resolve_project_root() follows the precedence of set_notebook_path().""" + + def test_notebook_root_wins( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("DEEPNOTE_PATHS__NOTEBOOK_ROOT", str(tmp_path / "notebooks")) + monkeypatch.setenv("DEEPNOTE_HOME_DIR", str(tmp_path / "home")) + + assert dataframe_storage.resolve_project_root() == tmp_path / "notebooks" + + def test_home_dir_is_used_as_is( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("DEEPNOTE_HOME_DIR", str(tmp_path / "home")) + + assert dataframe_storage.resolve_project_root() == tmp_path / "home" + + def test_dev_mode_uses_slash_work(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DEEPNOTE_RUNNING_IN_DEV_MODE", "true") + + assert dataframe_storage.resolve_project_root() == Path("/work") + + def test_default_is_work_under_home( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("HOME", str(tmp_path)) + + assert dataframe_storage.resolve_project_root() == tmp_path / "work" + + def test_unreadable_config_falls_back_to_dev_mode_env( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(dataframe_storage, "get_config", _broken_config) + monkeypatch.setenv("DEEPNOTE_RUNNING_IN_DEV_MODE", "true") + + assert dataframe_storage.resolve_project_root() == Path("/work") + + def test_unreadable_config_falls_back_to_work_under_home( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(dataframe_storage, "get_config", _broken_config) + monkeypatch.setenv("HOME", str(tmp_path)) + + assert dataframe_storage.resolve_project_root() == tmp_path / "work" + + +@pytest.mark.usefixtures("no_root_settings") +def test_frame_lands_under_the_project_root_not_the_working_directory( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "project" + notebook_directory = tmp_path / "notebooks" + project.mkdir() + notebook_directory.mkdir() + monkeypatch.setenv("DEEPNOTE_PATHS__NOTEBOOK_ROOT", str(project)) + monkeypatch.chdir(notebook_directory) + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert len(_arrow_files(project)) == 1 + assert list(notebook_directory.iterdir()) == [] + + +def test_delete_removes_the_frames_of_the_name_only(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), "other") + neighbour = root / "deepnote_dataframes" / "revenue-2026" + neighbour.mkdir() + (neighbour / "keep.txt").write_text("keep") + (root / "notes.txt").write_text("keep") + + assert artifacts.delete_dataframe(NAME) is True + + assert not _frame_dir(root).exists() + assert _stored_table(root, "other").to_pydict() == {"a": [3]} + assert (neighbour / "keep.txt").read_text() == "keep" + assert (root / "notes.txt").read_text() == "keep" + + +def test_delete_returns_false_when_there_is_nothing_to_delete(root: Path) -> None: + assert artifacts.delete_dataframe(NAME) is False + assert list(root.iterdir()) == [] + + +def test_delete_twice_returns_false_the_second_time(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert artifacts.delete_dataframe(NAME) is True + assert artifacts.delete_dataframe(NAME) is False + + +def test_a_name_can_be_stored_again_after_it_was_deleted(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.delete_dataframe(NAME) + + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert "previous" not in _manifest(root) + assert _stored_table(root).to_pydict() == {"a": [2]} + + +@pytest.mark.parametrize("name", INVALID_NAMES) +def test_delete_with_an_invalid_name_raises_and_deletes_nothing( + root: Path, name: Any +) -> None: + (root / "x").mkdir() + (root / "x" / "keep.txt").write_text("keep") + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + with pytest.raises(ValueError): + artifacts.delete_dataframe(name) + + assert (root / "x" / "keep.txt").read_text() == "keep" + assert len(_arrow_files(root)) == 1 + + +def test_delete_does_not_follow_a_link_out_of_the_frames_directory( + root: Path, +) -> None: + outside = root / "outside" + outside.mkdir() + (outside / "keep.txt").write_text("keep") + (root / "deepnote_dataframes").mkdir() + (root / "deepnote_dataframes" / NAME).symlink_to(outside, target_is_directory=True) + + with pytest.raises(OSError): + artifacts.delete_dataframe(NAME) + + assert (outside / "keep.txt").read_text() == "keep" + + +def test_delete_on_a_read_only_mount_deletes_nothing( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + _read_only(monkeypatch) + + assert artifacts.delete_dataframe(NAME) is False + + assert len(_arrow_files(root)) == 1 + + +def test_functions_are_exported_from_the_package() -> None: + assert deepnote_toolkit.artifacts is artifacts + assert "artifacts" in deepnote_toolkit.__all__ diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py new file mode 100644 index 00000000..f7fd1fbb --- /dev/null +++ b/tests/unit/helpers/dataframe_storage.py @@ -0,0 +1,104 @@ +"""Frames and interrupts shared by the DataFrame storage tests.""" + +import io +import signal +import uuid +from pathlib import Path +from typing import Any, Optional + +import numpy as np +import pandas as pd +import pyarrow as pa +import pytest + +from deepnote_toolkit import dataframe_storage + + +def probe_failure(frame: Any) -> Optional[type]: + """The exception type pandas itself raises for this frame, if any.""" + try: + frame.to_feather(io.BytesIO()) + except Exception as exc: # pylint: disable=broad-exception-caught + return type(exc) + return None + + +def dictionary_concat() -> pd.DataFrame: + dictionary_type = pd.ArrowDtype(pa.dictionary(pa.int32(), pa.string())) + first = pd.DataFrame({"c": pd.Series(["x", "y"], dtype=dictionary_type)}) + second = pd.DataFrame({"c": pd.Series(["z", "w"], dtype=dictionary_type)}) + return pd.concat([first, second]) + + +# Frames the RFC lists as failing on every supported pyarrow. The UUID column and the +# dictionary concatenation only fail on some versions, so for them whether the frame +# is stored is decided by asking pandas directly. +ALWAYS_FAILING = [ + pytest.param(lambda: pd.DataFrame({"a": [1, "two"]}), id="mixed-types"), + pytest.param( + lambda: pd.DataFrame([[1, 2]], columns=["a", "a"]), id="duplicate-labels" + ), + pytest.param(lambda: pd.DataFrame({"a": [2**70]}), id="beyond-int64"), + pytest.param(lambda: pd.DataFrame({"a": np.array([1 + 2j])}), id="complex"), + pytest.param(lambda: pd.DataFrame({"a": ["\ud800"]}), id="lone-surrogate"), +] +VERSION_DEPENDENT = [ + pytest.param(lambda: pd.DataFrame({"a": [uuid.UUID(int=1)]}), 1, id="uuid"), + pytest.param(dictionary_concat, 4, id="dictionary-concat"), +] + + +class SigintBeforeWrite: + """A handle that gets SIGINT after the file was truncated, before it is written.""" + + def __init__(self, handle: Any) -> None: + self._handle = handle + + def write(self, data: str) -> int: + signal.raise_signal(signal.SIGINT) + return self._handle.write(data) + + def __enter__(self) -> "SigintBeforeWrite": + return self + + def __exit__(self, *exc_info: Any) -> None: + self._handle.close() + + def __getattr__(self, name: str) -> Any: + return getattr(self._handle, name) + + +def interrupt_at(point: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Make the next write receive SIGINT at the named point after its frame file.""" + if point == "manifest-truncated": + real_open = Path.open + + def open_manifest( + self: Path, mode: str = "r", *args: Any, **kwargs: Any + ) -> Any: + handle = real_open(self, mode, *args, **kwargs) + if self.name == "manifest.json" and "w" in mode: + return SigintBeforeWrite(handle) + return handle + + monkeypatch.setattr(Path, "open", open_manifest) + elif point == "manifest-written": + write_manifest = dataframe_storage.write_manifest + + def write_manifest_then_interrupt(*args: Any, **kwargs: Any) -> None: + write_manifest(*args, **kwargs) + signal.raise_signal(signal.SIGINT) + + monkeypatch.setattr( + dataframe_storage, "write_manifest", write_manifest_then_interrupt + ) + else: + delete_replaced_frames = dataframe_storage._delete_replaced_frames + + def interrupt_then_delete(*args: Any, **kwargs: Any) -> None: + signal.raise_signal(signal.SIGINT) + delete_replaced_frames(*args, **kwargs) + + monkeypatch.setattr( + dataframe_storage, "_delete_replaced_frames", interrupt_then_delete + ) From f5820e163e4e96e0ee5558f64a681a1f973eca32 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 12 Oct 2026 09:38:15 +0000 Subject: [PATCH 2/3] refactor(artifacts): move DataFrame storage into an artifacts package The public API (store_dataframe, delete_dataframe) moves to deepnote_toolkit/artifacts/dataframes.py and is re-exported from the package, so `deepnote_toolkit.artifacts.store_dataframe` is unchanged. The writer and its manifest move alongside it, and their tests move to tests/unit/artifacts/. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- deepnote_toolkit/artifacts/__init__.py | 8 ++++++++ .../{ => artifacts}/dataframe_storage.py | 13 +++++++------ .../{ => artifacts}/dataframe_storage_manifest.py | 0 .../{artifacts.py => artifacts/dataframes.py} | 2 +- .../{dataframe_storage => artifacts}/__init__.py | 0 .../{dataframe_storage => artifacts}/conftest.py | 2 +- .../test_artifacts.py | 3 ++- tests/unit/helpers/dataframe_storage.py | 2 +- 8 files changed, 20 insertions(+), 10 deletions(-) create mode 100644 deepnote_toolkit/artifacts/__init__.py rename deepnote_toolkit/{ => artifacts}/dataframe_storage.py (99%) rename deepnote_toolkit/{ => artifacts}/dataframe_storage_manifest.py (100%) rename deepnote_toolkit/{artifacts.py => artifacts/dataframes.py} (98%) rename tests/unit/{dataframe_storage => artifacts}/__init__.py (100%) rename tests/unit/{dataframe_storage => artifacts}/conftest.py (98%) rename tests/unit/{dataframe_storage => artifacts}/test_artifacts.py (99%) diff --git a/deepnote_toolkit/artifacts/__init__.py b/deepnote_toolkit/artifacts/__init__.py new file mode 100644 index 00000000..d5498fc0 --- /dev/null +++ b/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"] diff --git a/deepnote_toolkit/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py similarity index 99% rename from deepnote_toolkit/dataframe_storage.py rename to deepnote_toolkit/artifacts/dataframe_storage.py index df2e4329..61f872fc 100644 --- a/deepnote_toolkit/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -13,18 +13,19 @@ from deepnote_toolkit import env from deepnote_toolkit.config import get_config -from deepnote_toolkit.dataframe_storage_manifest import ( - Manifest, - is_frame_name, - read_manifest, - write_manifest, -) 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 diff --git a/deepnote_toolkit/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py similarity index 100% rename from deepnote_toolkit/dataframe_storage_manifest.py rename to deepnote_toolkit/artifacts/dataframe_storage_manifest.py diff --git a/deepnote_toolkit/artifacts.py b/deepnote_toolkit/artifacts/dataframes.py similarity index 98% rename from deepnote_toolkit/artifacts.py rename to deepnote_toolkit/artifacts/dataframes.py index 6c359a94..623138a8 100644 --- a/deepnote_toolkit/artifacts.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -9,7 +9,7 @@ import shutil from typing import Any, Optional -from deepnote_toolkit import dataframe_storage +from . import dataframe_storage def _check_name(name: object) -> None: diff --git a/tests/unit/dataframe_storage/__init__.py b/tests/unit/artifacts/__init__.py similarity index 100% rename from tests/unit/dataframe_storage/__init__.py rename to tests/unit/artifacts/__init__.py diff --git a/tests/unit/dataframe_storage/conftest.py b/tests/unit/artifacts/conftest.py similarity index 98% rename from tests/unit/dataframe_storage/conftest.py rename to tests/unit/artifacts/conftest.py index 83879d80..eeb461c3 100644 --- a/tests/unit/dataframe_storage/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -9,7 +9,7 @@ import pytest -from deepnote_toolkit import dataframe_storage +from deepnote_toolkit.artifacts import dataframe_storage @pytest.fixture diff --git a/tests/unit/dataframe_storage/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py similarity index 99% rename from tests/unit/dataframe_storage/test_artifacts.py rename to tests/unit/artifacts/test_artifacts.py index 7675f886..7cb84a05 100644 --- a/tests/unit/dataframe_storage/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -21,7 +21,8 @@ import pytest import deepnote_toolkit -from deepnote_toolkit import artifacts, dataframe_storage, dataframe_storage_manifest +from deepnote_toolkit import artifacts +from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview from tests.unit.helpers.dataframe_storage import ( ALWAYS_FAILING, diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index f7fd1fbb..5f19c3b9 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -11,7 +11,7 @@ import pyarrow as pa import pytest -from deepnote_toolkit import dataframe_storage +from deepnote_toolkit.artifacts import dataframe_storage def probe_failure(frame: Any) -> Optional[type]: From 00a8f6f0f3e753de3df638bea7904047106426c0 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 12 Oct 2026 19:17:53 +0000 Subject: [PATCH 3/3] fix(artifacts): rename the manifest into place and keep Stop on failures - Write manifest.json to a temporary file and rename it over the old one. s3fs uploads an empty object when a file is opened with O_TRUNC, so the in-place write left the name without a current frame during every store, and a failed write or upload left the manifest empty or naming a frame the call then deleted. - Raise a SIGINT held during the manifest update even when the update fails, so an I/O error can't swallow Stop. - delete_dataframe returns False on EROFS, as store_dataframe returns None, and documents the OSError it raises. - store_dataframe's docstring covers failures after the manifest names the new frame. - Drop the codec guard for an exception pyarrow never raises for zstd or lz4, with the two tests of that impossible state. - set_notebook_path() and the writer share config.resolve_notebook_root(), which keeps ipykernel out of the artifacts package's imports. - Move the Spark ANSI-mode toggle into a fixture. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage.py | 46 +++--- .../artifacts/dataframe_storage_manifest.py | 29 +++- deepnote_toolkit/artifacts/dataframes.py | 13 +- deepnote_toolkit/config.py | 22 +++ deepnote_toolkit/set_notebook_path.py | 24 +--- tests/unit/artifacts/conftest.py | 9 ++ tests/unit/artifacts/test_artifacts.py | 136 ++++++++++++++---- tests/unit/helpers/dataframe_storage.py | 30 ++-- tests/unit/test_set_notebook_path.py | 10 +- 9 files changed, 219 insertions(+), 100 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 61f872fc..0d8cdd3f 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -11,8 +11,7 @@ from packaging.version import Version -from deepnote_toolkit import env -from deepnote_toolkit.config import get_config +from deepnote_toolkit.config import resolve_notebook_root from deepnote_toolkit.ocelots.utils import ( is_pandas_dataframe, is_polars_eager_dataframe, @@ -65,24 +64,12 @@ def classify_frame(value: object) -> Optional[FrameKind]: def resolve_project_root() -> Path: - """Return the project's root directory, mirroring set_notebook_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. """ - 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" + return resolve_notebook_root() @contextlib.contextmanager @@ -105,12 +92,18 @@ def _sigint_handled_by(handler: Callable[..., Any]) -> Iterator[None]: @contextlib.contextmanager def _sigint_deferred() -> Iterator[None]: - """Hold SIGINT back until the block ends, then raise it as KeyboardInterrupt.""" + """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] = [] - with _sigint_handled_by(lambda *_: interrupted.append(True)): - yield - if interrupted: - raise KeyboardInterrupt + try: + with _sigint_handled_by(lambda *_: interrupted.append(True)): + yield + finally: + if interrupted: + raise KeyboardInterrupt def is_read_only(root: Path) -> bool: @@ -123,12 +116,8 @@ def _select_compression() -> str: 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 + if pa.Codec.is_available(codec): + return codec return "uncompressed" @@ -188,8 +177,7 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str: 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. + 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) diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index 0bd535b7..036a328d 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -1,7 +1,10 @@ """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 @@ -64,11 +67,33 @@ def read_manifest(directory: Path) -> Optional[Manifest]: def write_manifest( directory: Path, new_file: str, previous: Optional[str], rows: int ) -> None: - """Point the manifest at `new_file`.""" + """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__ - (directory / MANIFEST_NAME).write_text(json.dumps(manifest), encoding="utf-8") + # 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 diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index 623138a8..6290c992 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -38,7 +38,9 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: 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. + A failed call leaves the name's previous frame current, unless it fails after + the manifest names the new frame: an interrupt held during the manifest update + is raised then, and so is a failure while deleting the frames it replaced. """ _check_name(name) kind = dataframe_storage.classify_frame(df) @@ -68,6 +70,8 @@ def delete_dataframe(name: str) -> bool: Raises: ValueError: `name` is not valid. + OSError: The delete failed, possibly part-way, leaving some of the name's + files. A symbolic link in place of the name's directory raises too. """ _check_name(name) root = dataframe_storage.resolve_project_root() @@ -77,5 +81,10 @@ def delete_dataframe(name: str) -> bool: directory = root / dataframe_storage.FRAMES_DIR / name if not directory.is_dir(): return False - shutil.rmtree(directory) + try: + shutil.rmtree(directory) + except OSError as exc: + if exc.errno == errno.EROFS: + return False + raise return True diff --git a/deepnote_toolkit/config.py b/deepnote_toolkit/config.py index b7936a2d..cfb7134f 100644 --- a/deepnote_toolkit/config.py +++ b/deepnote_toolkit/config.py @@ -15,6 +15,8 @@ from deepnote_core.config.loader import ConfigurationLoader from deepnote_core.config.models import DeepnoteConfig +from . import env + @lru_cache(maxsize=128) def _load_config_cached(config_path: Optional[Path]) -> DeepnoteConfig: @@ -63,3 +65,23 @@ def get_config( # Call the cached function return _load_config_cached(cfg_path) + + +def resolve_notebook_root() -> Path: + """Return the project's root directory, which notebook paths are relative to. + + Falls back to environment variables when the configuration cannot be loaded. + """ + 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" diff --git a/deepnote_toolkit/set_notebook_path.py b/deepnote_toolkit/set_notebook_path.py index d3e7c4e1..d3f9ab89 100644 --- a/deepnote_toolkit/set_notebook_path.py +++ b/deepnote_toolkit/set_notebook_path.py @@ -2,13 +2,12 @@ import os import sys import traceback -from pathlib import Path import requests from ipykernel import connect from . import env -from .config import get_config +from .config import get_config, resolve_notebook_root from .logging import get_logger @@ -22,26 +21,7 @@ def set_notebook_path() -> None: - Changes the working directory to the notebook's directory - Sets DEEPNOTE_PROJECT_ID in detached mode """ - # Determine notebooks common root using config first, with env fallback - try: - cfg = get_config() - # Prefer explicit notebook_root when provided - if cfg.paths.notebook_root: - root_path = Path(cfg.paths.notebook_root) - elif cfg.paths.home_dir: - root_path = Path(cfg.paths.home_dir) - elif cfg.runtime.dev_mode: - root_path = Path("/work") - else: - # Respect HOME if provided for compatibility with tests and environments - root_path = Path(os.environ.get("HOME", str(Path.home()))) / "work" - except Exception: - # Fallback: respect HOME when available, but honor dev-mode env toggle - root_path = ( - Path("/work") - if env.get_env("DEEPNOTE_RUNNING_IN_DEV_MODE") - else Path(os.environ.get("HOME", str(Path.home()))) / "work" - ) + root_path = resolve_notebook_root() try: connection_file = os.path.basename(connect.get_connection_file()) diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index eeb461c3..973fa62f 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -46,6 +46,15 @@ def spark() -> Any: ) +@pytest.fixture +def spark_without_ansi_mode(spark: Any) -> Iterator[Any]: + """pandas-on-Spark refuses to run under ANSI mode, which Spark 4 enables.""" + ansi_mode = spark.conf.get("spark.sql.ansi.enabled") + spark.conf.set("spark.sql.ansi.enabled", "false") + yield spark + spark.conf.set("spark.sql.ansi.enabled", ansi_mode) + + @pytest.fixture def arrow_sinks(monkeypatch: pytest.MonkeyPatch) -> List[Any]: """Every file object opened for writing a frame, in order.""" diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index 7cb84a05..fcc4802c 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -6,7 +6,9 @@ import json import os import re +import shutil import signal +import tempfile import threading import time from datetime import datetime, timezone @@ -22,6 +24,7 @@ import deepnote_toolkit from deepnote_toolkit import artifacts +from deepnote_toolkit import config as toolkit_config from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview from tests.unit.helpers.dataframe_storage import ( @@ -251,19 +254,13 @@ def test_pyspark_frame_raises_type_error(root: Path, spark: Any) -> None: def test_pandas_on_spark_frame_raises_type_error( - root: Path, spark: Any, monkeypatch: pytest.MonkeyPatch + root: Path, spark_without_ansi_mode: Any, monkeypatch: pytest.MonkeyPatch ) -> None: # pyspark.pandas warns on import unless told that pyarrow may ignore time zones monkeypatch.setenv("PYARROW_IGNORE_TIMEZONE", "1") # pyspark 3.5, which Python < 3.13 gets, imports distutils, gone since Python 3.12 pytest.importorskip("pyspark.pandas") - # pandas-on-Spark refuses to run under ANSI mode, which Spark 4 enables by default - ansi_mode = spark.conf.get("spark.sql.ansi.enabled") - spark.conf.set("spark.sql.ansi.enabled", "false") - try: - frame = spark.range(3).pandas_api() - finally: - spark.conf.set("spark.sql.ansi.enabled", ansi_mode) + frame = spark_without_ansi_mode.range(3).pandas_api() with pytest.raises(TypeError): artifacts.store_dataframe(frame, NAME) @@ -516,14 +513,11 @@ def test_frames_are_compressed(root: Path, make_frame: Any) -> None: class _FakeCodec: """Stands in for pyarrow.Codec, whose extension type cannot be patched.""" - def __init__(self, availability: Dict[str, Any]) -> None: + def __init__(self, availability: Dict[str, bool]) -> None: self._availability = availability def is_available(self, name: str) -> bool: - outcome = self._availability[name] - if isinstance(outcome, Exception): - raise outcome - return outcome + return self._availability[name] @pytest.mark.parametrize( @@ -531,21 +525,11 @@ def is_available(self, name: str) -> bool: [ pytest.param({"zstd": True, "lz4": True}, "zstd", id="zstd"), pytest.param({"zstd": False, "lz4": True}, "lz4", id="zstd-missing"), - pytest.param( - {"zstd": ValueError("Unsupported codec"), "lz4": True}, - "lz4", - id="zstd-unrecognised", - ), pytest.param({"zstd": False, "lz4": False}, "uncompressed", id="none"), - pytest.param( - {"zstd": False, "lz4": ValueError("Unsupported codec")}, - "uncompressed", - id="none-unrecognised", - ), ], ) def test_compression_falls_back_from_zstd_to_lz4_to_none( - monkeypatch: pytest.MonkeyPatch, availability: Dict[str, Any], expected: str + monkeypatch: pytest.MonkeyPatch, availability: Dict[str, bool], expected: str ) -> None: monkeypatch.setattr(pa, "Codec", _FakeCodec(availability)) @@ -745,6 +729,70 @@ def broken_manifest(*args: Any) -> Any: assert _stored_table(root).to_pydict() == {"a": [2]} +def _entries_beside_the_frames(root: Path) -> List[str]: + return sorted( + path.name for path in _frame_dir(root).iterdir() if path.suffix != ".arrow" + ) + + +def test_successful_write_leaves_only_the_manifest_beside_the_frames( + root: Path, +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert _entries_beside_the_frames(root) == ["manifest.json"] + + +class _FailingManifestHandle: + """A temporary manifest that fails in write(), or in close() like s3fs's upload.""" + + def __init__(self, handle: Any, failing: str) -> None: + self._handle = handle + self._failing = failing + + def write(self, data: str) -> int: + if self._failing == "write": + raise OSError(errno.ENOSPC, "No space left on device") + return self._handle.write(data) + + def __enter__(self) -> "_FailingManifestHandle": + return self + + def __exit__(self, *exc_info: Any) -> None: + self._handle.close() + if self._failing == "close": + raise OSError(errno.EIO, "Input/output error") + + def __getattr__(self, name: str) -> Any: + return getattr(self._handle, name) + + +@pytest.mark.parametrize("failing", ["write", "close"]) +def test_failed_manifest_write_keeps_the_previous_frame( + root: Path, monkeypatch: pytest.MonkeyPatch, failing: str +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() + real_temporary_file = tempfile.NamedTemporaryFile + + def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: + return _FailingManifestHandle(real_temporary_file(*args, **kwargs), failing) + + monkeypatch.setattr( + dataframe_storage_manifest, + "tempfile", + SimpleNamespace(NamedTemporaryFile=failing_temporary_file), + ) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert _stored_table(root).to_pydict() == {"a": [1]} + assert _entries_beside_the_frames(root) == ["manifest.json"] + + def test_failure_after_the_manifest_step_keeps_the_new_frame( root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -848,7 +896,7 @@ def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( @pytest.mark.parametrize( - "point", ["manifest-truncated", "manifest-written", "old-frames-deleted"] + "point", ["manifest-being-written", "manifest-written", "old-frames-deleted"] ) def test_interrupt_after_the_frame_is_written_leaves_a_usable_manifest( root: Path, monkeypatch: pytest.MonkeyPatch, point: str @@ -866,6 +914,26 @@ def test_interrupt_after_the_frame_is_written_leaves_a_usable_manifest( assert signal.getsignal(signal.SIGINT) is handler_before +@pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["eio", "erofs"]) +def test_interrupt_during_a_failing_manifest_write_is_raised( + root: Path, monkeypatch: pytest.MonkeyPatch, error: int +) -> None: + handler_before = signal.getsignal(signal.SIGINT) + + def interrupt_then_fail(*args: Any) -> None: + signal.raise_signal(signal.SIGINT) + raise OSError(error, os.strerror(error)) + + monkeypatch.setattr(dataframe_storage, "write_manifest", interrupt_then_fail) + + with pytest.raises(KeyboardInterrupt) as raised: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert isinstance(raised.value.__context__, OSError) + assert _arrow_files(root) == [] + assert signal.getsignal(signal.SIGINT) is handler_before + + def _broken_config() -> Any: raise ValueError("invalid config") @@ -904,7 +972,7 @@ def test_default_is_work_under_home( def test_unreadable_config_falls_back_to_dev_mode_env( self, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(dataframe_storage, "get_config", _broken_config) + monkeypatch.setattr(toolkit_config, "get_config", _broken_config) monkeypatch.setenv("DEEPNOTE_RUNNING_IN_DEV_MODE", "true") assert dataframe_storage.resolve_project_root() == Path("/work") @@ -912,7 +980,7 @@ def test_unreadable_config_falls_back_to_dev_mode_env( def test_unreadable_config_falls_back_to_work_under_home( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(dataframe_storage, "get_config", _broken_config) + monkeypatch.setattr(toolkit_config, "get_config", _broken_config) monkeypatch.setenv("HOME", str(tmp_path)) assert dataframe_storage.resolve_project_root() == tmp_path / "work" @@ -1015,6 +1083,20 @@ def test_delete_on_a_read_only_mount_deletes_nothing( assert len(_arrow_files(root)) == 1 +def test_delete_returns_false_when_the_filesystem_refuses_with_erofs( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """store_dataframe() returns None for the same error.""" + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + def refuse(path: Any) -> None: + raise OSError(errno.EROFS, "Read-only file system") + + monkeypatch.setattr(shutil, "rmtree", refuse) + + assert artifacts.delete_dataframe(NAME) is False + + def test_functions_are_exported_from_the_package() -> None: assert deepnote_toolkit.artifacts is artifacts assert "artifacts" in deepnote_toolkit.__all__ diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 5f19c3b9..1f062d33 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -2,8 +2,9 @@ import io import signal +import tempfile import uuid -from pathlib import Path +from types import SimpleNamespace from typing import Any, Optional import numpy as np @@ -11,7 +12,7 @@ import pyarrow as pa import pytest -from deepnote_toolkit.artifacts import dataframe_storage +from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest def probe_failure(frame: Any) -> Optional[type]: @@ -49,7 +50,7 @@ def dictionary_concat() -> pd.DataFrame: class SigintBeforeWrite: - """A handle that gets SIGINT after the file was truncated, before it is written.""" + """A handle that gets SIGINT after the file was opened, before it is written.""" def __init__(self, handle: Any) -> None: self._handle = handle @@ -70,18 +71,17 @@ def __getattr__(self, name: str) -> Any: def interrupt_at(point: str, monkeypatch: pytest.MonkeyPatch) -> None: """Make the next write receive SIGINT at the named point after its frame file.""" - if point == "manifest-truncated": - real_open = Path.open - - def open_manifest( - self: Path, mode: str = "r", *args: Any, **kwargs: Any - ) -> Any: - handle = real_open(self, mode, *args, **kwargs) - if self.name == "manifest.json" and "w" in mode: - return SigintBeforeWrite(handle) - return handle - - monkeypatch.setattr(Path, "open", open_manifest) + if point == "manifest-being-written": + real_temporary_file = tempfile.NamedTemporaryFile + + def temporary_file_getting_sigint(*args: Any, **kwargs: Any) -> Any: + return SigintBeforeWrite(real_temporary_file(*args, **kwargs)) + + monkeypatch.setattr( + dataframe_storage_manifest, + "tempfile", + SimpleNamespace(NamedTemporaryFile=temporary_file_getting_sigint), + ) elif point == "manifest-written": write_manifest = dataframe_storage.write_manifest diff --git a/tests/unit/test_set_notebook_path.py b/tests/unit/test_set_notebook_path.py index 79e9c1e9..26390f5a 100644 --- a/tests/unit/test_set_notebook_path.py +++ b/tests/unit/test_set_notebook_path.py @@ -6,6 +6,7 @@ from pathlib import Path from unittest import mock +from deepnote_toolkit import config as toolkit_config from deepnote_toolkit import env as dnenv from deepnote_toolkit.set_notebook_path import set_notebook_path @@ -58,9 +59,12 @@ def recorder(url, headers): monkeypatch.setattr(mod.requests, "get", recorder) notebook_root = tmp_path / "root" notebook_root.mkdir(parents=True, exist_ok=True) - monkeypatch.setattr( - mod, "get_config", lambda: _cfg(notebook_root=str(notebook_root)) - ) + + def fake_config(): + return _cfg(notebook_root=str(notebook_root)) + + monkeypatch.setattr(mod, "get_config", fake_config) + monkeypatch.setattr(toolkit_config, "get_config", fake_config) # Capture chdir target changedir = {"path": None}