From 62a4d0ec7ba0aee167ed03669e5bd6793bd6d0fc Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 8 Oct 2026 12:33:58 +0000 Subject: [PATCH 01/41] 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 02/41] 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 03/41] 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} From 9e3fbb06f19ca5f9ded19d82ed9c003bacdbf063 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 09:17:42 +0000 Subject: [PATCH 04/41] fix(artifacts): store polars frames without string views From 1.0 polars writes strings and binaries as Utf8View/BinaryView by default, which pyarrow before 15 and apache-arrow before 21.2.0 can't read, while the same data written by pandas reads there. From polars 1.2.1, where `pl.CompatLevel` is first exported, the writer passes `compat_level=pl.CompatLevel.oldest()`, so they are stored as large_string and large_binary. polars 1.0 to 1.2.0 still write views. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage.py | 5 +++ tests/unit/artifacts/test_artifacts.py | 39 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 0d8cdd3f..2140f7b9 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -31,6 +31,9 @@ # batches, and polars only honours a batch size from 1.37. POLARS_BATCH_ROWS = 65_536 POLARS_BATCH_MIN_VERSION = Version("1.37") +# From 1.0 polars writes strings and binaries as views, which pyarrow before 15 and +# apache-arrow before 21.2 can't read; `pl.CompatLevel` is public from 1.2.1. +POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") FrameKind = Literal["pandas", "polars"] @@ -138,6 +141,8 @@ def _write_frame_file( options = {} if Version(polars.__version__) >= POLARS_BATCH_MIN_VERSION: options["record_batch_size"] = POLARS_BATCH_ROWS + if Version(polars.__version__) >= POLARS_COMPAT_LEVEL_MIN_VERSION: + options["compat_level"] = polars.CompatLevel.oldest() frame.write_ipc(sink, compression=compression, **options) diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index fcc4802c..d29fcbc8 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -603,6 +603,45 @@ def test_polars_batch_size_is_passed_to_polars_from_1_37( assert _stored_table(root).to_pydict() == {"a": [1, 2, 3]} +def test_polars_strings_are_stored_without_view_types(root: Path) -> None: + artifacts.store_dataframe( + pl.DataFrame( + { + "s": ["a", "a string longer than twelve bytes"], + "b": [b"x", b"y"], + "ls": [["a"], ["b", "c"]], + } + ), + NAME, + ) + + assert _stored_table(root).schema.types == [ + pa.large_string(), + pa.large_binary(), + pa.large_list(pa.large_string()), + ] + + +@pytest.mark.parametrize( + "polars_version, passed", + [("1.2.0", False), ("1.2.1", True)], + ids=["before-1.2.1", "from-1.2.1"], +) +def test_polars_compat_level_is_passed_to_polars_from_1_2_1( + root: Path, monkeypatch: pytest.MonkeyPatch, polars_version: str, passed: bool +) -> 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": ["x"]}), NAME) + + assert ("compat_level" in write_ipc.call_args.kwargs) is passed + assert _stored_table(root).to_pydict() == {"a": ["x"]} + + def _store_bad_frame_after_a_good_one( root: Path, make_frame: Callable[[], Any], expected_error: Optional[type] ) -> Dict[str, Any]: From aa40e0322a7970f3bf1304cc253e1d5ce701e4e1 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 09:19:54 +0000 Subject: [PATCH 05/41] fix(artifacts): keep the new frame once the manifest update starts - An s3fs rename copies the object and then deletes the source, so it can fail after the new manifest is already in place. store_frame then deleted the new frame, leaving the manifest naming a missing file. It now deletes the new frame only when the frame write or the manifest read fails; from write_manifest on, a failure keeps it, and a later write prunes it once it is a day old if no manifest names it. - With that, Stop at any point leaves a manifest that names a complete frame, so the SIGINT deferral around the manifest update is gone. Stop raises KeyboardInterrupt wherever it lands; the worst case is an unreferenced file. - Tests cover a rename that fails after the copy, a failed manifest read and an interrupted manifest write. The SIGINT test helpers and the tests of the held interrupt go with the deferral. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage.py | 59 ++------- .../artifacts/dataframe_storage_manifest.py | 3 +- deepnote_toolkit/artifacts/dataframes.py | 7 +- tests/unit/artifacts/test_artifacts.py | 118 ++++++++---------- tests/unit/helpers/dataframe_storage.py | 62 +-------- 5 files changed, 66 insertions(+), 183 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 2140f7b9..c31b0043 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -3,11 +3,10 @@ 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 typing import Any, Literal, Optional from packaging.version import Version @@ -75,40 +74,6 @@ def resolve_project_root() -> Path: return resolve_notebook_root() -@contextlib.contextmanager -def _sigint_handled_by(handler: Callable[..., Any]) -> Iterator[None]: - """Install `handler` for SIGINT, then restore the previous one. - - Off the main thread signal handlers cannot be set, and nothing needs to be. - """ - try: - previous = signal.signal(signal.SIGINT, handler) - except ValueError: - yield - return - try: - yield - finally: - if previous is not None: - signal.signal(signal.SIGINT, previous) - - -@contextlib.contextmanager -def _sigint_deferred() -> Iterator[None]: - """Hold SIGINT back until the block ends, then raise it as KeyboardInterrupt. - - It is raised even when the block fails, chained to that failure, so that an - error never swallows Stop. - """ - interrupted: list[bool] = [] - try: - with _sigint_handled_by(lambda *_: interrupted.append(True)): - yield - finally: - if interrupted: - raise KeyboardInterrupt - - def is_read_only(root: Path) -> bool: """Return whether the project's filesystem is mounted read-only.""" return bool(os.statvfs(root).f_flag & os.ST_RDONLY) @@ -177,29 +142,21 @@ def _delete_replaced_frames( def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> str: - """Store the frame as the current one under `name` and return its reference ID. - - The order matters because S3 replaces the manifest in one step: the frame is - complete before the manifest names it, and nothing is deleted before then. - An interrupt stops everything but the manifest update, which finishes first: - one after the manifest is renamed into place would delete the frame it now names. - """ + """Store the frame as the current one under `name` and return its reference ID.""" 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 + old = read_manifest(directory) except BaseException: - if not committed: - with contextlib.suppress(OSError): - (directory / new_file).unlink(missing_ok=True) + with contextlib.suppress(OSError): + (directory / new_file).unlink(missing_ok=True) raise + # A failure from here on keeps the new frame: an s3fs rename can put the new + # manifest in place and still fail. + write_manifest(directory, new_file, old.file if old else None, len(frame)) _delete_replaced_frames(directory, old, new_file) return reference_id diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index 036a328d..c1c91321 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -70,7 +70,8 @@ def write_manifest( """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. + a write that fails part-way leaves the old manifest in place, unless an s3fs + rename fails after copying the new one over it. """ manifest: dict[str, Any] = {"version": MANIFEST_VERSION, "file": new_file} if previous is not None: diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index 6290c992..e00dd813 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -38,9 +38,10 @@ 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, 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. + A failed or interrupted call leaves the name's previous frame current, unless + it stops after the manifest names the new frame: an interrupt can arrive then, + deleting the frames it replaced can fail, and so can an s3fs rename that has + already put the new manifest in place. """ _check_name(name) kind = dataframe_storage.classify_frame(df) diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index d29fcbc8..6499418f 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -7,7 +7,6 @@ import os import re import shutil -import signal import tempfile import threading import time @@ -30,7 +29,6 @@ from tests.unit.helpers.dataframe_storage import ( ALWAYS_FAILING, VERSION_DEPENDENT, - interrupt_at, probe_failure, ) @@ -743,7 +741,24 @@ def test_error_from_close_raises_and_stores_nothing( assert not (_frame_dir(root) / "manifest.json").exists() -def test_failure_before_the_manifest_step_deletes_nothing( +def test_failed_manifest_read_deletes_the_new_frame( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + frames = _arrow_files(root) + + def broken_read(*args: Any) -> Any: + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(dataframe_storage, "read_manifest", broken_read) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert _arrow_files(root) == frames + + +def test_failure_in_the_manifest_step_deletes_nothing( root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -764,7 +779,9 @@ def broken_manifest(*args: Any) -> Any: 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 + # the new frame stays, because a failed update may already have named it + entries = {path.name for path in directory.iterdir()} + assert set(listing) <= entries and len(entries - set(listing)) == 1 assert _stored_table(root).to_pydict() == {"a": [2]} @@ -784,7 +801,8 @@ def test_successful_write_leaves_only_the_manifest_beside_the_frames( class _FailingManifestHandle: - """A temporary manifest that fails in write(), or in close() like s3fs's upload.""" + """A temporary manifest that fails or is interrupted in write(), or fails in + close() like s3fs's upload.""" def __init__(self, handle: Any, failing: str) -> None: self._handle = handle @@ -793,6 +811,8 @@ def __init__(self, handle: Any, failing: str) -> None: def write(self, data: str) -> int: if self._failing == "write": raise OSError(errno.ENOSPC, "No space left on device") + if self._failing == "interrupt": + raise KeyboardInterrupt return self._handle.write(data) def __enter__(self) -> "_FailingManifestHandle": @@ -807,7 +827,7 @@ def __getattr__(self, name: str) -> Any: return getattr(self._handle, name) -@pytest.mark.parametrize("failing", ["write", "close"]) +@pytest.mark.parametrize("failing", ["write", "close", "interrupt"]) def test_failed_manifest_write_keeps_the_previous_frame( root: Path, monkeypatch: pytest.MonkeyPatch, failing: str ) -> None: @@ -824,7 +844,7 @@ def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: SimpleNamespace(NamedTemporaryFile=failing_temporary_file), ) - with pytest.raises(OSError): + with pytest.raises(KeyboardInterrupt if failing == "interrupt" else OSError): artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes @@ -832,6 +852,24 @@ def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: assert _entries_beside_the_frames(root) == ["manifest.json"] +def _copy_then_fail(source: str, destination: Path) -> None: + """s3fs's rename: the copy lands, then deleting the source fails.""" + shutil.copyfile(source, destination) + raise OSError(errno.EIO, "Input/output error") + + +def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + monkeypatch.setattr(dataframe_storage_manifest.os, "replace", _copy_then_fail) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + 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: @@ -901,76 +939,22 @@ def refuse(self: pd.DataFrame, sink: Any, **kwargs: Any) -> 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] + root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: + def write_some_then_interrupt(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + sink.write(b"ARROW1 partial") + raise KeyboardInterrupt + 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)) + monkeypatch.setattr(pd.DataFrame, "to_feather", write_some_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-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 -) -> 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 - - -@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: diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 1f062d33..dbc0f96f 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -1,10 +1,7 @@ -"""Frames and interrupts shared by the DataFrame storage tests.""" +"""Frames shared by the DataFrame storage tests.""" import io -import signal -import tempfile import uuid -from types import SimpleNamespace from typing import Any, Optional import numpy as np @@ -12,8 +9,6 @@ import pyarrow as pa import pytest -from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest - def probe_failure(frame: Any) -> Optional[type]: """The exception type pandas itself raises for this frame, if any.""" @@ -47,58 +42,3 @@ def dictionary_concat() -> pd.DataFrame: 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 opened, 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-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 - - 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 76d62cd5a20f0814457491f43dd3771a2510f7da Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 09:22:19 +0000 Subject: [PATCH 06/41] docs(artifacts): keep only comments the code can't carry Drop docstring paragraphs and comments that restate the code, describe behaviour the tests pin, or point at the RFC without a link. Keep the external facts: s3fs upload and rename semantics, polars and Python version quirks, the public store_dataframe and delete_dataframe contract. Fix the ones that were wrong: json raises RecursionError on deep nesting only before Python 3.12, polars has its own write_ipc errors, and the bad-frame helper returns the good call's manifest. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- deepnote_toolkit/artifacts/__init__.py | 5 +-- .../artifacts/dataframe_storage.py | 34 +++++-------------- .../artifacts/dataframe_storage_manifest.py | 10 +++--- deepnote_toolkit/artifacts/dataframes.py | 14 +++----- deepnote_toolkit/config.py | 5 +-- tests/unit/artifacts/conftest.py | 5 +-- tests/unit/artifacts/test_artifacts.py | 22 ++++++------ tests/unit/helpers/dataframe_storage.py | 3 -- 8 files changed, 31 insertions(+), 67 deletions(-) diff --git a/deepnote_toolkit/artifacts/__init__.py b/deepnote_toolkit/artifacts/__init__.py index d5498fc0..c425bf26 100644 --- a/deepnote_toolkit/artifacts/__init__.py +++ b/deepnote_toolkit/artifacts/__init__.py @@ -1,7 +1,4 @@ -"""Stores and deletes named DataFrames in the project's filesystem. - -Only the names in __all__ are supported public API. -""" +"""Stores and deletes named DataFrames in the project's filesystem.""" from .dataframes import delete_dataframe, store_dataframe diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index c31b0043..8803c03f 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -1,4 +1,4 @@ -"""Writes DataFrames into the project's filesystem as named Arrow files.""" +"""Writes named DataFrames into the project's filesystem as Arrow files.""" import contextlib import os @@ -26,8 +26,8 @@ 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. +# Each record batch is one page of the read endpoint; to_feather already writes +# 64k-row batches. POLARS_BATCH_ROWS = 65_536 POLARS_BATCH_MIN_VERSION = Version("1.37") # From 1.0 polars writes strings and binaries as views, which pyarrow before 15 and @@ -40,22 +40,12 @@ 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 whether `value` can be used as a single path segment.""" 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. + """Return "pandas" or "polars" for a frame worth storing, else None.""" if isinstance(value, DeepnoteQueryPreview): return None if is_pandas_dataframe(value): @@ -66,11 +56,7 @@ def classify_frame(value: object) -> Optional[FrameKind]: def resolve_project_root() -> Path: - """Return the project's root directory, the one set_notebook_path() uses. - - The working directory cannot be used: set_notebook_path() moves it into the - notebook's directory. - """ + """Return the project's root directory, the one set_notebook_path() uses.""" return resolve_notebook_root() @@ -80,7 +66,7 @@ def is_read_only(root: Path) -> bool: def _select_compression() -> str: - """Return the best Arrow codec this pyarrow build has: zstd, lz4, else none.""" + """Return the best codec this pyarrow build has: zstd, lz4, else uncompressed.""" import pyarrow as pa for codec in ("zstd", "lz4"): @@ -116,9 +102,8 @@ def _delete_replaced_frames( ) -> 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. + Another kernel may still be writing an unnamed frame here, so a stray must be a + day old before it is deleted. """ named = {new_file} if old is not None: @@ -131,7 +116,6 @@ def _delete_replaced_frames( 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) diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index c1c91321..a0cf1ffa 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -36,9 +36,8 @@ def is_frame_name(name: str) -> bool: 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. + Any code in the project can write the manifest, so a malformed one counts as + absent. """ try: with (directory / MANIFEST_NAME).open("rb") as handle: @@ -69,9 +68,8 @@ def write_manifest( ) -> None: """Point the manifest at `new_file`. - The manifest is written under a temporary name and renamed over the old one, so - a write that fails part-way leaves the old manifest in place, unless an s3fs - rename fails after copying the new one over it. + On s3fs a rename is a copy and then a delete, so a failed rename may already + have replaced the manifest. """ manifest: dict[str, Any] = {"version": MANIFEST_VERSION, "file": new_file} if previous is not None: diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index e00dd813..808e8620 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -1,9 +1,4 @@ -"""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. -""" +"""Stores and deletes named DataFrames in the project's filesystem.""" import errno import shutil @@ -24,8 +19,8 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: 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}`. + name: The frame's name, unique within the project, and its directory + under `deepnote_dataframes/`. Must match `[A-Za-z0-9_-]{1,128}`. Returns: The new frame's reference ID, or None if the project's filesystem is @@ -36,7 +31,8 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: 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()`. + Exception: `df` cannot be serialized to Arrow: whatever `df.to_feather()` + or polars' `df.write_ipc()` raises. A failed or interrupted call leaves the name's previous frame current, unless it stops after the manifest names the new frame: an interrupt can arrive then, diff --git a/deepnote_toolkit/config.py b/deepnote_toolkit/config.py index cfb7134f..ad4ba1a6 100644 --- a/deepnote_toolkit/config.py +++ b/deepnote_toolkit/config.py @@ -68,10 +68,7 @@ def get_config( 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. - """ + """Return the project's root directory, which notebook paths are relative to.""" try: cfg = get_config() if cfg.paths.notebook_root: diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index 973fa62f..b58da100 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -1,7 +1,4 @@ -"""Fixtures for the DataFrame storage tests. - -Kept out of tests/unit/conftest.py because the names are generic. -""" +"""Fixtures for the DataFrame storage tests.""" import time from pathlib import Path diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index 6499418f..487b0b56 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -288,8 +288,8 @@ def test_second_write_keeps_one_previous_frame(root: Path) -> None: 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) + replaced = directory / _manifest(root)["file"] + _backdate(replaced, 2 * DAY) def add(name: str, age_seconds: float) -> Path: path = directory / name @@ -320,8 +320,7 @@ def add(name: str, age_seconds: float) -> Path: 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() + assert replaced.exists() def _manifest_bytes(**fields: Any) -> bytes: @@ -394,7 +393,6 @@ def test_unusable_manifest_counts_as_no_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) @@ -417,7 +415,7 @@ def test_manifest_of_exactly_4096_bytes_is_honoured(root: Path) -> None: 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.""" + """Before 3.12, json raises RecursionError on the nesting 4 KiB allows.""" artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) def too_deep(raw: bytes) -> Any: @@ -471,7 +469,7 @@ def read_then_wait(directory: Path) -> Any: def write(value: int) -> None: try: artifacts.store_dataframe(pd.DataFrame({"a": [value]}), NAME) - except BaseException as exc: # surfaced by the assertion below + except BaseException as exc: errors.append(exc) with monkeypatch.context() as patch: @@ -643,7 +641,10 @@ def test_polars_compat_level_is_passed_to_polars_from_1_2_1( 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.""" + """Store a good frame, then a bad one. + + Return the good call's manifest bytes and the directory listing after the bad call. + """ artifacts.store_dataframe(pd.DataFrame({"ok": [1]}), NAME) directory = _frame_dir(root) manifest_bytes = (directory / "manifest.json").read_bytes() @@ -779,7 +780,6 @@ def broken_manifest(*args: Any) -> Any: artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) assert (directory / "manifest.json").read_bytes() == manifest_bytes - # the new frame stays, because a failed update may already have named it entries = {path.name for path in directory.iterdir()} assert set(listing) <= entries and len(entries - set(listing)) == 1 assert _stored_table(root).to_pydict() == {"a": [2]} @@ -801,8 +801,7 @@ def test_successful_write_leaves_only_the_manifest_beside_the_frames( class _FailingManifestHandle: - """A temporary manifest that fails or is interrupted in write(), or fails in - close() like s3fs's upload.""" + """A temporary manifest whose write() fails or is interrupted, or close() fails.""" def __init__(self, handle: Any, failing: str) -> None: self._handle = handle @@ -1109,7 +1108,6 @@ def test_delete_on_a_read_only_mount_deletes_nothing( 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: diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index dbc0f96f..17f5caf9 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -26,9 +26,6 @@ def dictionary_concat() -> pd.DataFrame: 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( From 15ac86898098c0ab8590e47aa796aade8c176228 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 14:16:06 +0000 Subject: [PATCH 07/41] refactor(artifacts): type the polars write options and explain them where used The options dict was untyped, so mypy inferred dict[str, int] and only the repo's disabled assignment check hid the CompatLevel value going into it. The polars version is now parsed once. The comments above the constants moved to the two options they explain, each cut to one line: batches are capped so a page read stays cheap, which polars before 1.37 can't do, and the oldest compat level avoids the string and binary views older Arrow readers can't read. How readers treat batch sizes needs no comment, because an Arrow IPC file lists its batches and their row counts. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- deepnote_toolkit/artifacts/dataframe_storage.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 8803c03f..14b60939 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -26,12 +26,8 @@ FRAMES_DIR = Path("deepnote_dataframes") STALE_FRAME_AGE_SECONDS = 24 * 60 * 60 -# Each record batch is one page of the read endpoint; to_feather already writes -# 64k-row batches. POLARS_BATCH_ROWS = 65_536 POLARS_BATCH_MIN_VERSION = Version("1.37") -# From 1.0 polars writes strings and binaries as views, which pyarrow before 15 and -# apache-arrow before 21.2 can't read; `pl.CompatLevel` is public from 1.2.1. POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") FrameKind = Literal["pandas", "polars"] @@ -89,10 +85,13 @@ def _write_frame_file( else: import polars - options = {} - if Version(polars.__version__) >= POLARS_BATCH_MIN_VERSION: + polars_version = Version(polars.__version__) + options: dict[str, Any] = {} + # Cap batches like write_feather so a page read stays cheap. + if polars_version >= POLARS_BATCH_MIN_VERSION: options["record_batch_size"] = POLARS_BATCH_ROWS - if Version(polars.__version__) >= POLARS_COMPAT_LEVEL_MIN_VERSION: + # pyarrow < 15, apache-arrow < 21.2 can't read 1.0+'s string/binary views. + if polars_version >= POLARS_COMPAT_LEVEL_MIN_VERSION: options["compat_level"] = polars.CompatLevel.oldest() frame.write_ipc(sink, compression=compression, **options) From f5a6dcbf24364b1659ab49ef2a328a367f173566 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 14:19:06 +0000 Subject: [PATCH 08/41] fix(artifacts): write pandas frames with pyarrow's write_feather DataFrame.to_feather before pandas 2.1 refuses any index but a default RangeIndex and any non-string column label, so filtered, sorted, grouped, concatenated and transposed frames raised ValueError on pandas 2.0 and 1.5, which the declared range still allows on Python 3.10 and 3.11. pandas 2.1+ calls pyarrow's write_feather itself, so the files on the locked versions are byte-for-byte the same. GeoDataFrames, which geopandas writes with its own to_feather, now raise ArrowTypeError on their geometry column. The tests that faked or spied on to_feather now target write_feather, and the test helper that predicts refusals asks pyarrow too. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage.py | 7 ++- deepnote_toolkit/artifacts/dataframes.py | 4 +- tests/unit/artifacts/test_artifacts.py | 55 +++++++++++++------ tests/unit/helpers/dataframe_storage.py | 5 +- 4 files changed, 48 insertions(+), 23 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 14b60939..61a634a9 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -74,14 +74,17 @@ def _select_compression() -> str: 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. + """Write the frame to `path` as an Arrow IPC file. 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) + from pyarrow import feather + + # pandas < 2.1's to_feather needs a default index and string column names. + feather.write_feather(frame, sink, compression=compression) else: import polars diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index 808e8620..c7567703 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -31,8 +31,8 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: 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: whatever `df.to_feather()` - or polars' `df.write_ipc()` raises. + Exception: `df` cannot be serialized to Arrow: whatever pyarrow's + `write_feather()` or polars' `df.write_ipc()` raises. A failed or interrupted call leaves the name's previous frame current, unless it stops after the manifest names the new frame: an interrupt can arrive then, diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index 487b0b56..7e2affd5 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -20,6 +20,7 @@ import polars as pl import pyarrow as pa import pytest +from pyarrow import feather import deepnote_toolkit from deepnote_toolkit import artifacts @@ -118,6 +119,24 @@ def test_writing_does_not_change_the_callers_frame(root: Path) -> None: pd.testing.assert_frame_equal(frame, before) +def test_filtered_frame_is_stored_with_its_row_labels( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def refuse(self: pd.DataFrame, *args: Any, **kwargs: Any) -> None: + raise ValueError("feather does not support serializing a non-default index") + + # pandas >= 2.1 accepts this frame, so make to_feather refuse it as < 2.1 does. + monkeypatch.setattr(pd.DataFrame, "to_feather", refuse) + frame = pd.DataFrame({"a": [5, 6, 7, 8]}) + + artifacts.store_dataframe(frame[frame.a != 7], NAME) + + assert _stored_table(root).to_pydict() == { + "a": [5, 6, 8], + "__index_level_0__": [0, 1, 3], + } + + 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() @@ -184,11 +203,11 @@ def test_valid_names_are_accepted(root: Path, name: str) -> 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 mock.patch.object(feather, "write_feather", autospec=True) as write: with pytest.raises(ValueError): artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) - to_feather.assert_not_called() + write.assert_not_called() assert list(root.iterdir()) == [] @@ -663,10 +682,10 @@ def _store_bad_frame_after_a_good_one( 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" + raised_by_pyarrow = probe_failure(make_frame()) + assert raised_by_pyarrow is not None, "pyarrow is expected to refuse this frame" - left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pandas) + left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) assert (_frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] assert left["listing"] == sorted(["manifest.json", _manifest(root)["file"]]) @@ -674,14 +693,14 @@ def test_unserializable_frame_raises_and_keeps_the_previous_frame( @pytest.mark.parametrize("make_frame, rows", VERSION_DEPENDENT) -def test_frames_pandas_may_refuse_follow_what_pandas_does( +def test_frames_pyarrow_may_refuse_follow_what_pyarrow_does( root: Path, make_frame: Callable[[], Any], rows: int ) -> None: - raised_by_pandas = probe_failure(make_frame()) + raised_by_pyarrow = probe_failure(make_frame()) - left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pandas) + left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) - if raised_by_pandas is None: + if raised_by_pyarrow is None: assert _stored_table(root).num_rows == rows return assert (_frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] @@ -704,11 +723,11 @@ def test_polars_object_column_raises_and_keeps_the_previous_frame(root: Path) -> 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: + def write_some_then_fail(frame: 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) + monkeypatch.setattr(feather, "write_feather", write_some_then_fail) with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -916,21 +935,21 @@ def test_read_only_mount_stores_nothing_and_does_not_convert( ) -> None: _read_only(monkeypatch) - with mock.patch.object(pd.DataFrame, "to_feather", autospec=True) as to_feather: + with mock.patch.object(feather, "write_feather", autospec=True) as write: result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) assert result is None - to_feather.assert_not_called() + write.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: + def refuse(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: raise OSError(errno.EROFS, "Read-only file system") - monkeypatch.setattr(pd.DataFrame, "to_feather", refuse) + monkeypatch.setattr(feather, "write_feather", refuse) result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -941,13 +960,15 @@ def refuse(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - def write_some_then_interrupt(self: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + def write_some_then_interrupt( + frame: pd.DataFrame, sink: Any, **kwargs: Any + ) -> None: sink.write(b"ARROW1 partial") raise KeyboardInterrupt artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() - monkeypatch.setattr(pd.DataFrame, "to_feather", write_some_then_interrupt) + monkeypatch.setattr(feather, "write_feather", write_some_then_interrupt) with pytest.raises(KeyboardInterrupt): artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 17f5caf9..a4ae35ae 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -8,12 +8,13 @@ import pandas as pd import pyarrow as pa import pytest +from pyarrow import feather def probe_failure(frame: Any) -> Optional[type]: - """The exception type pandas itself raises for this frame, if any.""" + """The exception type pyarrow itself raises for this frame, if any.""" try: - frame.to_feather(io.BytesIO()) + feather.write_feather(frame, io.BytesIO()) except Exception as exc: # pylint: disable=broad-exception-caught return type(exc) return None From b32661b1db90238e97bc354ebb23bcb5b287a3fd Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 14:20:22 +0000 Subject: [PATCH 09/41] test(artifacts): inject manifest read and cleanup failures at the filesystem Two tests replaced read_manifest and _delete_replaced_frames wholesale, so no test made the filesystem calls inside them fail. read_manifest treating any OSError as "no manifest", or the cleanup swallowing OSError, kept the whole suite green. The first lets a transient EIO on manifest.json drop the frame that was current until then from the manifest, and delete it in the same call if it is more than a day old. The tests now fail Path.open on manifest.json and Path.unlink on the replaced frame, and both of those regressions fail them. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_artifacts.py | 28 ++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index 7e2affd5..9f063136 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -766,11 +766,16 @@ def test_failed_manifest_read_deletes_the_new_frame( ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) frames = _arrow_files(root) + real_open = Path.open - def broken_read(*args: Any) -> Any: - raise OSError(errno.EIO, "Input/output error") + def open_failing_for_the_manifest( + self: Path, mode: str = "r", *args: Any, **kwargs: Any + ) -> Any: + if self.name == "manifest.json": + raise OSError(errno.EIO, "Input/output error") + return real_open(self, mode, *args, **kwargs) - monkeypatch.setattr(dataframe_storage, "read_manifest", broken_read) + monkeypatch.setattr(Path, "open", open_failing_for_the_manifest) with pytest.raises(OSError): artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) @@ -892,16 +897,23 @@ 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) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + replaced = _frame_dir(root) / _manifest(root)["previous"] + real_unlink = Path.unlink - def broken_cleanup(*args: Any) -> Any: - raise OSError(errno.EIO, "Input/output error") + def unlink_failing_for_the_replaced_frame( + self: Path, *args: Any, **kwargs: Any + ) -> None: + if self == replaced: + raise OSError(errno.EIO, "Input/output error") + real_unlink(self, *args, **kwargs) - monkeypatch.setattr(dataframe_storage, "_delete_replaced_frames", broken_cleanup) + monkeypatch.setattr(Path, "unlink", unlink_failing_for_the_replaced_frame) with pytest.raises(OSError): - artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) - assert _stored_table(root).to_pydict() == {"a": [2]} + assert _stored_table(root).to_pydict() == {"a": [3]} def test_stray_pruned_by_another_writer_meanwhile_is_not_a_failure( From 3a60709b8c1bd013d4c03420a4d260ff45837a1f Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 14:22:40 +0000 Subject: [PATCH 10/41] fix(artifacts): detect read-only mounts from EROFS alone The statvfs probe ran before anything was written, so when the project root did not exist, which is the default outside Deepnote Cloud, store_dataframe and delete_dataframe both raised FileNotFoundError instead of creating the root or returning False. The probe also never caught anything the writes don't: on a read-only mount the first write (creating the name's directory, opening the frame file, or rmtree's first unlink) fails with EROFS before any conversion. store_frame now maps EROFS to None itself, and only while nothing is stored yet: creating the directory moved into the step that deletes the new frame on failure. EROFS from the manifest step on raises like any other error, because by then the new frame may already be current and None would say nothing was stored. The read-only test helper now refuses writes with EROFS instead of faking statvfs. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage.py | 21 +++--- deepnote_toolkit/artifacts/dataframes.py | 12 +--- tests/unit/artifacts/test_artifacts.py | 70 +++++++++++++++++-- 3 files changed, 75 insertions(+), 28 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 61a634a9..4b662942 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -1,6 +1,7 @@ """Writes named DataFrames into the project's filesystem as Arrow files.""" import contextlib +import errno import os import re import time @@ -56,11 +57,6 @@ def resolve_project_root() -> Path: return resolve_notebook_root() -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 codec this pyarrow build has: zstd, lz4, else uncompressed.""" import pyarrow as pa @@ -127,19 +123,24 @@ def _delete_replaced_frames( (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.""" - directory = root / FRAMES_DIR / name - directory.mkdir(parents=True, exist_ok=True) +def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> Optional[str]: + """Store the frame as the current one under `name` and return its reference ID. + Returns None, having stored nothing, if the filesystem is read-only. EROFS from + the manifest step on raises like any other error: the new frame may be current. + """ + directory = root / FRAMES_DIR / name reference_id = str(uuid.uuid4()) new_file = f"{reference_id}.arrow" try: + directory.mkdir(parents=True, exist_ok=True) _write_frame_file(frame, kind, directory / new_file, _select_compression()) old = read_manifest(directory) - except BaseException: + except BaseException as exc: with contextlib.suppress(OSError): (directory / new_file).unlink(missing_ok=True) + if isinstance(exc, OSError) and exc.errno == errno.EROFS: + return None raise # A failure from here on keeps the new frame: an s3fs rename can put the new # manifest in place and still fail. diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index c7567703..d32d878b 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -48,14 +48,7 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: ) 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 + return dataframe_storage.store_frame(df, kind, name, root) def delete_dataframe(name: str) -> bool: @@ -72,9 +65,6 @@ def delete_dataframe(name: str) -> bool: """ _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 diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index 9f063136..7ee635c3 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -64,11 +64,29 @@ def _backdate(path: Path, age_seconds: float) -> None: def _read_only(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - dataframe_storage.os, - "statvfs", - lambda path: SimpleNamespace(f_flag=os.ST_RDONLY), - ) + """Fail Path.open writes, mkdir and deletes with EROFS, like a read-only mount.""" + + def refuse(*args: Any, **kwargs: Any) -> Any: + raise OSError(errno.EROFS, "Read-only file system") + + real_open = Path.open + + def open_for_reading(self: Path, mode: str = "r", *args: Any, **kwargs: Any) -> Any: + if mode.strip("rbt"): + refuse() + return real_open(self, mode, *args, **kwargs) + + # The kernel reports EEXIST before EROFS, so an existing directory passes. + def mkdir( + self: Path, mode: int = 0o777, parents: bool = False, exist_ok: bool = False + ) -> None: + if not (exist_ok and self.is_dir()): + refuse() + + monkeypatch.setattr(Path, "open", open_for_reading) + monkeypatch.setattr(Path, "mkdir", mkdir) + monkeypatch.setattr(os, "unlink", refuse) + monkeypatch.setattr(os, "rmdir", refuse) def test_stores_pandas_frame_and_manifest(root: Path) -> None: @@ -893,8 +911,9 @@ def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( assert _stored_table(root).to_pydict() == {"a": [2]} +@pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["EIO", "EROFS"]) def test_failure_after_the_manifest_step_keeps_the_new_frame( - root: Path, monkeypatch: pytest.MonkeyPatch + root: Path, monkeypatch: pytest.MonkeyPatch, error: int ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) @@ -905,7 +924,7 @@ def unlink_failing_for_the_replaced_frame( self: Path, *args: Any, **kwargs: Any ) -> None: if self == replaced: - raise OSError(errno.EIO, "Input/output error") + raise OSError(error, os.strerror(error)) real_unlink(self, *args, **kwargs) monkeypatch.setattr(Path, "unlink", unlink_failing_for_the_replaced_frame) @@ -955,6 +974,33 @@ def test_read_only_mount_stores_nothing_and_does_not_convert( assert list(root.iterdir()) == [] +def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() + _read_only(monkeypatch) + + with mock.patch.object(feather, "write_feather", autospec=True) as write: + result = artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert result is None + write.assert_not_called() + assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert _stored_table(root).to_pydict() == {"a": [1]} + + +def test_store_creates_a_missing_project_root( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "work" + monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) + + reference_id = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert _manifest(project)["file"] == f"{reference_id}.arrow" + + def test_erofs_from_the_write_stores_nothing( root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -1080,6 +1126,16 @@ def test_delete_returns_false_when_there_is_nothing_to_delete(root: Path) -> Non assert list(root.iterdir()) == [] +def test_delete_returns_false_when_the_project_root_is_missing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "work" + monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) + + assert artifacts.delete_dataframe(NAME) is False + assert not project.exists() + + def test_delete_twice_returns_false_the_second_time(root: Path) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) From 3c36f45740a6f42630eb9e414316b1ec12d57a8c Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 14:23:20 +0000 Subject: [PATCH 11/41] fix(artifacts): explain why a SQL query preview is refused A query preview is a pandas DataFrame subclass, and the frames derived from it keep the subclass, so "needs a pandas or polars DataFrame, got DeepnoteQueryPreview" was false on its face and said nothing about why or what to do instead. The message now says a preview holds at most 100 rows of the query's result and points to the DataFrame output mode or pd.DataFrame(df). classify_frame still refuses previews, which the post-run hook relies on to skip them silently; store_dataframe only picks the message. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- deepnote_toolkit/artifacts/dataframes.py | 14 ++++++++++++-- tests/unit/artifacts/test_artifacts.py | 15 ++++++++++++++- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index d32d878b..23bede8f 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -4,6 +4,8 @@ import shutil from typing import Any, Optional +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview + from . import dataframe_storage @@ -28,8 +30,9 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: 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. + TypeError: `df` is a SQL query preview or a frame derived from one, or is + not a pandas or polars DataFrame. That includes PySpark and + pandas-on-Spark DataFrames and polars LazyFrames. OSError: The write failed. Exception: `df` cannot be serialized to Arrow: whatever pyarrow's `write_feather()` or polars' `df.write_ipc()` raises. @@ -42,6 +45,13 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: _check_name(name) kind = dataframe_storage.classify_frame(df) if kind is None: + if isinstance(df, DeepnoteQueryPreview): + raise TypeError( + "store_dataframe refuses a SQL query preview or a frame derived from " + "one, which holds at most 100 rows of the query's result. Set the SQL " + "block's output mode to DataFrame, or pass pd.DataFrame(df) to store " + "these rows" + ) raise TypeError( "store_dataframe needs a pandas or polars DataFrame, " f"got {type(df).__qualname__}" diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py index 7ee635c3..ce881329 100644 --- a/tests/unit/artifacts/test_artifacts.py +++ b/tests/unit/artifacts/test_artifacts.py @@ -255,6 +255,9 @@ def _cleared_query_preview() -> DeepnoteQueryPreview: 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"), +] + +QUERY_PREVIEWS = [ pytest.param(_query_preview, id="query-preview"), pytest.param(_cleared_query_preview, id="cleared-query-preview"), pytest.param(lambda: _query_preview().copy(), id="copied-query-preview"), @@ -271,7 +274,17 @@ def test_unsupported_value_raises_type_error( assert list(root.iterdir()) == [] -@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES) +@pytest.mark.parametrize("make_preview", QUERY_PREVIEWS) +def test_query_preview_raises_type_error( + root: Path, make_preview: Callable[[], Any] +) -> None: + with pytest.raises(TypeError, match="SQL query preview"): + artifacts.store_dataframe(make_preview(), NAME) + + assert list(root.iterdir()) == [] + + +@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES + QUERY_PREVIEWS) def test_unsupported_value_raises_type_error_on_a_read_only_mount( root: Path, monkeypatch: pytest.MonkeyPatch, make_value: Callable[[], Any] ) -> None: From 1dc54cf459de9dbb6d1217b5e2254dd194c402bc Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 14 Oct 2026 19:57:11 +0000 Subject: [PATCH 12/41] test(artifacts): split test_artifacts.py by behaviour The 1,225-line test_artifacts.py is split into eight files by behaviour: storing, arguments, manifest and cleanup, file format, write failures, read-only mount, project root and delete_dataframe. The helpers and constants shared across them move to tests/unit/helpers/dataframe_storage.py under public names. No test changed: the same 150 tests are collected. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_arguments.py | 154 +++ tests/unit/artifacts/test_artifacts.py | 1225 ----------------- tests/unit/artifacts/test_delete_dataframe.py | 68 + tests/unit/artifacts/test_file_format.py | 161 +++ .../artifacts/test_manifest_and_cleanup.py | 293 ++++ tests/unit/artifacts/test_project_root.py | 102 ++ tests/unit/artifacts/test_read_only_mount.py | 87 ++ tests/unit/artifacts/test_storing.py | 136 ++ tests/unit/artifacts/test_write_failures.py | 296 ++++ tests/unit/helpers/dataframe_storage.py | 69 +- 10 files changed, 1364 insertions(+), 1227 deletions(-) create mode 100644 tests/unit/artifacts/test_arguments.py delete mode 100644 tests/unit/artifacts/test_artifacts.py create mode 100644 tests/unit/artifacts/test_delete_dataframe.py create mode 100644 tests/unit/artifacts/test_file_format.py create mode 100644 tests/unit/artifacts/test_manifest_and_cleanup.py create mode 100644 tests/unit/artifacts/test_project_root.py create mode 100644 tests/unit/artifacts/test_read_only_mount.py create mode 100644 tests/unit/artifacts/test_storing.py create mode 100644 tests/unit/artifacts/test_write_failures.py diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py new file mode 100644 index 00000000..c66005c6 --- /dev/null +++ b/tests/unit/artifacts/test_arguments.py @@ -0,0 +1,154 @@ +"""Tests for how the artifacts API checks the names and values it is given.""" + +from pathlib import Path +from typing import Any, Callable +from unittest import mock + +import pandas as pd +import polars as pl +import pytest +from pyarrow import feather + +from deepnote_toolkit import artifacts +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview +from tests.unit.helpers.dataframe_storage import ( + NAME, + arrow_files, + read_only, + stored_table, +) + + +@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(feather, "write_feather", autospec=True) as write: + with pytest.raises(ValueError): + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + + write.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"), +] + +QUERY_PREVIEWS = [ + pytest.param(_query_preview, id="query-preview"), + pytest.param(_cleared_query_preview, id="cleared-query-preview"), + pytest.param(lambda: _query_preview().copy(), id="copied-query-preview"), +] + + +@pytest.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_preview", QUERY_PREVIEWS) +def test_query_preview_raises_type_error( + root: Path, make_preview: Callable[[], Any] +) -> None: + with pytest.raises(TypeError, match="SQL query preview"): + artifacts.store_dataframe(make_preview(), NAME) + + assert list(root.iterdir()) == [] + + +@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES + QUERY_PREVIEWS) +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_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") + frame = spark_without_ansi_mode.range(3).pandas_api() + + with pytest.raises(TypeError): + artifacts.store_dataframe(frame, NAME) + + assert list(root.iterdir()) == [] + + +@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 diff --git a/tests/unit/artifacts/test_artifacts.py b/tests/unit/artifacts/test_artifacts.py deleted file mode 100644 index ce881329..00000000 --- a/tests/unit/artifacts/test_artifacts.py +++ /dev/null @@ -1,1225 +0,0 @@ -"""Tests for the artifacts API that stores and deletes named DataFrames.""" - -import contextlib -import copy -import errno -import json -import os -import re -import shutil -import tempfile -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 -from pyarrow import feather - -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 ( - ALWAYS_FAILING, - VERSION_DEPENDENT, - 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: - """Fail Path.open writes, mkdir and deletes with EROFS, like a read-only mount.""" - - def refuse(*args: Any, **kwargs: Any) -> Any: - raise OSError(errno.EROFS, "Read-only file system") - - real_open = Path.open - - def open_for_reading(self: Path, mode: str = "r", *args: Any, **kwargs: Any) -> Any: - if mode.strip("rbt"): - refuse() - return real_open(self, mode, *args, **kwargs) - - # The kernel reports EEXIST before EROFS, so an existing directory passes. - def mkdir( - self: Path, mode: int = 0o777, parents: bool = False, exist_ok: bool = False - ) -> None: - if not (exist_ok and self.is_dir()): - refuse() - - monkeypatch.setattr(Path, "open", open_for_reading) - monkeypatch.setattr(Path, "mkdir", mkdir) - monkeypatch.setattr(os, "unlink", refuse) - monkeypatch.setattr(os, "rmdir", refuse) - - -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_filtered_frame_is_stored_with_its_row_labels( - root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - def refuse(self: pd.DataFrame, *args: Any, **kwargs: Any) -> None: - raise ValueError("feather does not support serializing a non-default index") - - # pandas >= 2.1 accepts this frame, so make to_feather refuse it as < 2.1 does. - monkeypatch.setattr(pd.DataFrame, "to_feather", refuse) - frame = pd.DataFrame({"a": [5, 6, 7, 8]}) - - artifacts.store_dataframe(frame[frame.a != 7], NAME) - - assert _stored_table(root).to_pydict() == { - "a": [5, 6, 8], - "__index_level_0__": [0, 1, 3], - } - - -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(feather, "write_feather", autospec=True) as write: - with pytest.raises(ValueError): - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) - - write.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"), -] - -QUERY_PREVIEWS = [ - pytest.param(_query_preview, id="query-preview"), - pytest.param(_cleared_query_preview, id="cleared-query-preview"), - pytest.param(lambda: _query_preview().copy(), id="copied-query-preview"), -] - - -@pytest.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_preview", QUERY_PREVIEWS) -def test_query_preview_raises_type_error( - root: Path, make_preview: Callable[[], Any] -) -> None: - with pytest.raises(TypeError, match="SQL query preview"): - artifacts.store_dataframe(make_preview(), NAME) - - assert list(root.iterdir()) == [] - - -@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES + QUERY_PREVIEWS) -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_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") - frame = spark_without_ansi_mode.range(3).pandas_api() - - 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) - replaced = directory / _manifest(root)["file"] - _backdate(replaced, 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" - assert replaced.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" - 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: - """Before 3.12, json raises RecursionError on the nesting 4 KiB 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: - 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, bool]) -> None: - self._availability = availability - - def is_available(self, name: str) -> bool: - return self._availability[name] - - -@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": False, "lz4": False}, "uncompressed", id="none"), - ], -) -def test_compression_falls_back_from_zstd_to_lz4_to_none( - monkeypatch: pytest.MonkeyPatch, availability: Dict[str, bool], 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 test_polars_strings_are_stored_without_view_types(root: Path) -> None: - artifacts.store_dataframe( - pl.DataFrame( - { - "s": ["a", "a string longer than twelve bytes"], - "b": [b"x", b"y"], - "ls": [["a"], ["b", "c"]], - } - ), - NAME, - ) - - assert _stored_table(root).schema.types == [ - pa.large_string(), - pa.large_binary(), - pa.large_list(pa.large_string()), - ] - - -@pytest.mark.parametrize( - "polars_version, passed", - [("1.2.0", False), ("1.2.1", True)], - ids=["before-1.2.1", "from-1.2.1"], -) -def test_polars_compat_level_is_passed_to_polars_from_1_2_1( - root: Path, monkeypatch: pytest.MonkeyPatch, polars_version: str, passed: bool -) -> 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": ["x"]}), NAME) - - assert ("compat_level" in write_ipc.call_args.kwargs) is passed - assert _stored_table(root).to_pydict() == {"a": ["x"]} - - -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 the good call's manifest bytes and the directory listing after the bad call. - """ - 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_pyarrow = probe_failure(make_frame()) - assert raised_by_pyarrow is not None, "pyarrow is expected to refuse this frame" - - left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) - - 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_pyarrow_may_refuse_follow_what_pyarrow_does( - root: Path, make_frame: Callable[[], Any], rows: int -) -> None: - raised_by_pyarrow = probe_failure(make_frame()) - - left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) - - if raised_by_pyarrow 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(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: - sink.write(b"ARROW1 partial") - raise OSError(errno.EIO, "Input/output error") - - monkeypatch.setattr(feather, "write_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_failed_manifest_read_deletes_the_new_frame( - root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - frames = _arrow_files(root) - real_open = Path.open - - def open_failing_for_the_manifest( - self: Path, mode: str = "r", *args: Any, **kwargs: Any - ) -> Any: - if self.name == "manifest.json": - raise OSError(errno.EIO, "Input/output error") - return real_open(self, mode, *args, **kwargs) - - monkeypatch.setattr(Path, "open", open_failing_for_the_manifest) - - with pytest.raises(OSError): - artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - - assert _arrow_files(root) == frames - - -def test_failure_in_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 - entries = {path.name for path in directory.iterdir()} - assert set(listing) <= entries and len(entries - set(listing)) == 1 - 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 whose write() fails or is interrupted, or close() fails.""" - - 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") - if self._failing == "interrupt": - raise KeyboardInterrupt - 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", "interrupt"]) -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(KeyboardInterrupt if failing == "interrupt" else 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 _copy_then_fail(source: str, destination: Path) -> None: - """s3fs's rename: the copy lands, then deleting the source fails.""" - shutil.copyfile(source, destination) - raise OSError(errno.EIO, "Input/output error") - - -def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( - root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - monkeypatch.setattr(dataframe_storage_manifest.os, "replace", _copy_then_fail) - - with pytest.raises(OSError): - artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - - assert _stored_table(root).to_pydict() == {"a": [2]} - - -@pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["EIO", "EROFS"]) -def test_failure_after_the_manifest_step_keeps_the_new_frame( - root: Path, monkeypatch: pytest.MonkeyPatch, error: int -) -> None: - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - replaced = _frame_dir(root) / _manifest(root)["previous"] - real_unlink = Path.unlink - - def unlink_failing_for_the_replaced_frame( - self: Path, *args: Any, **kwargs: Any - ) -> None: - if self == replaced: - raise OSError(error, os.strerror(error)) - real_unlink(self, *args, **kwargs) - - monkeypatch.setattr(Path, "unlink", unlink_failing_for_the_replaced_frame) - - with pytest.raises(OSError): - artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) - - assert _stored_table(root).to_pydict() == {"a": [3]} - - -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(feather, "write_feather", autospec=True) as write: - result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - - assert result is None - write.assert_not_called() - assert list(root.iterdir()) == [] - - -def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( - root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() - _read_only(monkeypatch) - - with mock.patch.object(feather, "write_feather", autospec=True) as write: - result = artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - - assert result is None - write.assert_not_called() - assert (_frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes - assert _stored_table(root).to_pydict() == {"a": [1]} - - -def test_store_creates_a_missing_project_root( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - project = tmp_path / "work" - monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) - - reference_id = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - - assert _manifest(project)["file"] == f"{reference_id}.arrow" - - -def test_erofs_from_the_write_stores_nothing( - root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - def refuse(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: - raise OSError(errno.EROFS, "Read-only file system") - - monkeypatch.setattr(feather, "write_feather", refuse) - - result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - - assert result is None - assert _arrow_files(root) == [] - - -def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( - root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - def write_some_then_interrupt( - frame: pd.DataFrame, sink: Any, **kwargs: Any - ) -> None: - sink.write(b"ARROW1 partial") - raise KeyboardInterrupt - - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - manifest_bytes = (_frame_dir(root) / "manifest.json").read_bytes() - monkeypatch.setattr(feather, "write_feather", write_some_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 - - -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(toolkit_config, "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(toolkit_config, "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_returns_false_when_the_project_root_is_missing( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - project = tmp_path / "work" - monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) - - assert artifacts.delete_dataframe(NAME) is False - assert not project.exists() - - -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_delete_returns_false_when_the_filesystem_refuses_with_erofs( - root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - 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/artifacts/test_delete_dataframe.py b/tests/unit/artifacts/test_delete_dataframe.py new file mode 100644 index 00000000..1a46c475 --- /dev/null +++ b/tests/unit/artifacts/test_delete_dataframe.py @@ -0,0 +1,68 @@ +"""Tests for the artifacts API that deletes named DataFrames.""" + +from pathlib import Path + +import pandas as pd +import pytest + +from deepnote_toolkit import artifacts +from tests.unit.helpers.dataframe_storage import ( + NAME, + frame_dir, + stored_manifest, + stored_table, +) + + +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 stored_manifest(root) + assert stored_table(root).to_pydict() == {"a": [2]} + + +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" diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py new file mode 100644 index 00000000..a93b8450 --- /dev/null +++ b/tests/unit/artifacts/test_file_format.py @@ -0,0 +1,161 @@ +"""Tests for the compression, record batches and column types of stored DataFrames.""" + +from pathlib import Path +from typing import Any, Dict +from unittest import mock + +import pandas as pd +import polars as pl +import pyarrow as pa +import pytest + +from deepnote_toolkit import artifacts +from deepnote_toolkit.artifacts import dataframe_storage +from tests.unit.helpers.dataframe_storage import ( + NAME, + frame_dir, + stored_manifest, + stored_table, +) + + +@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) / stored_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, bool]) -> None: + self._availability = availability + + def is_available(self, name: str) -> bool: + return self._availability[name] + + +@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": False, "lz4": False}, "uncompressed", id="none"), + ], +) +def test_compression_falls_back_from_zstd_to_lz4_to_none( + monkeypatch: pytest.MonkeyPatch, availability: Dict[str, bool], 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) / stored_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) / stored_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) / stored_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 test_polars_strings_are_stored_without_view_types(root: Path) -> None: + artifacts.store_dataframe( + pl.DataFrame( + { + "s": ["a", "a string longer than twelve bytes"], + "b": [b"x", b"y"], + "ls": [["a"], ["b", "c"]], + } + ), + NAME, + ) + + assert stored_table(root).schema.types == [ + pa.large_string(), + pa.large_binary(), + pa.large_list(pa.large_string()), + ] + + +@pytest.mark.parametrize( + "polars_version, passed", + [("1.2.0", False), ("1.2.1", True)], + ids=["before-1.2.1", "from-1.2.1"], +) +def test_polars_compat_level_is_passed_to_polars_from_1_2_1( + root: Path, monkeypatch: pytest.MonkeyPatch, polars_version: str, passed: bool +) -> 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": ["x"]}), NAME) + + assert ("compat_level" in write_ipc.call_args.kwargs) is passed + assert stored_table(root).to_pydict() == {"a": ["x"]} diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py new file mode 100644 index 00000000..cf6acf33 --- /dev/null +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -0,0 +1,293 @@ +"""Tests for the manifest of a stored DataFrame and the cleanup of old frames.""" + +import contextlib +import json +import os +import threading +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Dict, Iterator, List + +import pandas as pd +import pyarrow as pa +import pytest + +from deepnote_toolkit import artifacts +from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest +from tests.unit.helpers.dataframe_storage import ( + DAY, + FRAME_A, + FRAME_B, + NAME, + arrow_files, + backdate, + entries_beside_the_frames, + frame_dir, + stored_manifest, + stored_table, +) + + +def test_second_write_keeps_one_previous_frame(root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + first = stored_manifest(root) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + second = stored_manifest(root) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + third = stored_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) + replaced = directory / stored_manifest(root)["file"] + backdate(replaced, 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" + assert replaced.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 = stored_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" + 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 stored_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: + """Before 3.12, json raises RecursionError on the nesting 4 KiB 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 stored_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: + 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 = stored_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]} + + +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"] + + +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]} diff --git a/tests/unit/artifacts/test_project_root.py b/tests/unit/artifacts/test_project_root.py new file mode 100644 index 00000000..d5cbc566 --- /dev/null +++ b/tests/unit/artifacts/test_project_root.py @@ -0,0 +1,102 @@ +"""Tests for the project root that stored DataFrames are written under.""" + +from pathlib import Path +from typing import Any + +import pandas as pd +import pytest + +from deepnote_toolkit import artifacts +from deepnote_toolkit import config as toolkit_config +from deepnote_toolkit.artifacts import dataframe_storage +from tests.unit.helpers.dataframe_storage import NAME, arrow_files, stored_manifest + + +def test_store_creates_a_missing_project_root( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "work" + monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) + + reference_id = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert stored_manifest(project)["file"] == f"{reference_id}.arrow" + + +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(toolkit_config, "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(toolkit_config, "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_returns_false_when_the_project_root_is_missing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + project = tmp_path / "work" + monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) + + assert artifacts.delete_dataframe(NAME) is False + assert not project.exists() diff --git a/tests/unit/artifacts/test_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py new file mode 100644 index 00000000..54c81995 --- /dev/null +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -0,0 +1,87 @@ +"""Tests for storing and deleting DataFrames on a read-only mount.""" + +import errno +import shutil +from pathlib import Path +from typing import Any +from unittest import mock + +import pandas as pd +import pytest +from pyarrow import feather + +from deepnote_toolkit import artifacts +from tests.unit.helpers.dataframe_storage import ( + NAME, + arrow_files, + frame_dir, + read_only, + stored_table, +) + + +def test_read_only_mount_stores_nothing_and_does_not_convert( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + read_only(monkeypatch) + + with mock.patch.object(feather, "write_feather", autospec=True) as write: + result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert result is None + write.assert_not_called() + assert list(root.iterdir()) == [] + + +def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + manifest_bytes = (frame_dir(root) / "manifest.json").read_bytes() + read_only(monkeypatch) + + with mock.patch.object(feather, "write_feather", autospec=True) as write: + result = artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert result is None + write.assert_not_called() + assert (frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes + assert stored_table(root).to_pydict() == {"a": [1]} + + +def test_erofs_from_the_write_stores_nothing( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def refuse(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + raise OSError(errno.EROFS, "Read-only file system") + + monkeypatch.setattr(feather, "write_feather", refuse) + + result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert result is None + assert arrow_files(root) == [] + + +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_delete_returns_false_when_the_filesystem_refuses_with_erofs( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + 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 diff --git a/tests/unit/artifacts/test_storing.py b/tests/unit/artifacts/test_storing.py new file mode 100644 index 00000000..cfaf37ee --- /dev/null +++ b/tests/unit/artifacts/test_storing.py @@ -0,0 +1,136 @@ +"""Tests for what the artifacts API writes when it stores a named DataFrame.""" + +import copy +import re +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import pandas as pd +import polars as pl +import pytest + +import deepnote_toolkit +from deepnote_toolkit import artifacts +from tests.unit.helpers.dataframe_storage import ( + NAME, + arrow_files, + stored_manifest, + stored_table, +) + +UUID_PATTERN = r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}" + + +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 = stored_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( + stored_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 stored_manifest(root)["file"] == f"{reference_id}.arrow" + assert stored_table(root).to_pydict() == {"a": [1, 2, 3], "b": ["x", "y", "z"]} + assert stored_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_filtered_frame_is_stored_with_its_row_labels( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def refuse(self: pd.DataFrame, *args: Any, **kwargs: Any) -> None: + raise ValueError("feather does not support serializing a non-default index") + + # pandas >= 2.1 accepts this frame, so make to_feather refuse it as < 2.1 does. + monkeypatch.setattr(pd.DataFrame, "to_feather", refuse) + frame = pd.DataFrame({"a": [5, 6, 7, 8]}) + + artifacts.store_dataframe(frame[frame.a != 7], NAME) + + assert stored_table(root).to_pydict() == { + "a": [5, 6, 8], + "__index_level_0__": [0, 1, 3], + } + + +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 stored_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]} + + +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/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py new file mode 100644 index 00000000..45f8c185 --- /dev/null +++ b/tests/unit/artifacts/test_write_failures.py @@ -0,0 +1,296 @@ +"""Tests for what a failed or interrupted DataFrame write leaves behind.""" + +import errno +import os +import shutil +import tempfile +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, Dict, Optional + +import pandas as pd +import polars as pl +import pytest +from pyarrow import feather + +from deepnote_toolkit import artifacts +from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest +from tests.unit.helpers.dataframe_storage import ( + ALWAYS_FAILING, + DAY, + FRAME_A, + NAME, + VERSION_DEPENDENT, + arrow_files, + backdate, + entries_beside_the_frames, + frame_dir, + probe_failure, + stored_manifest, + stored_table, +) + + +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 the good call's manifest bytes and the directory listing after the bad call. + """ + 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_pyarrow = probe_failure(make_frame()) + assert raised_by_pyarrow is not None, "pyarrow is expected to refuse this frame" + + left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) + + assert (frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] + assert left["listing"] == sorted(["manifest.json", stored_manifest(root)["file"]]) + assert stored_table(root).to_pydict() == {"ok": [1]} + + +@pytest.mark.parametrize("make_frame, rows", VERSION_DEPENDENT) +def test_frames_pyarrow_may_refuse_follow_what_pyarrow_does( + root: Path, make_frame: Callable[[], Any], rows: int +) -> None: + raised_by_pyarrow = probe_failure(make_frame()) + + left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) + + if raised_by_pyarrow 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", stored_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(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + sink.write(b"ARROW1 partial") + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(feather, "write_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_failed_manifest_read_deletes_the_new_frame( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + frames = arrow_files(root) + real_open = Path.open + + def open_failing_for_the_manifest( + self: Path, mode: str = "r", *args: Any, **kwargs: Any + ) -> Any: + if self.name == "manifest.json": + raise OSError(errno.EIO, "Input/output error") + return real_open(self, mode, *args, **kwargs) + + monkeypatch.setattr(Path, "open", open_failing_for_the_manifest) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert arrow_files(root) == frames + + +def test_failure_in_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 + entries = {path.name for path in directory.iterdir()} + assert set(listing) <= entries and len(entries - set(listing)) == 1 + assert stored_table(root).to_pydict() == {"a": [2]} + + +class _FailingManifestHandle: + """A temporary manifest whose write() fails or is interrupted, or close() fails.""" + + 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") + if self._failing == "interrupt": + raise KeyboardInterrupt + 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", "interrupt"]) +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(KeyboardInterrupt if failing == "interrupt" else 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 _copy_then_fail(source: str, destination: Path) -> None: + """s3fs's rename: the copy lands, then deleting the source fails.""" + shutil.copyfile(source, destination) + raise OSError(errno.EIO, "Input/output error") + + +def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + monkeypatch.setattr(dataframe_storage_manifest.os, "replace", _copy_then_fail) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert stored_table(root).to_pydict() == {"a": [2]} + + +@pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["EIO", "EROFS"]) +def test_failure_after_the_manifest_step_keeps_the_new_frame( + root: Path, monkeypatch: pytest.MonkeyPatch, error: int +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + replaced = frame_dir(root) / stored_manifest(root)["previous"] + real_unlink = Path.unlink + + def unlink_failing_for_the_replaced_frame( + self: Path, *args: Any, **kwargs: Any + ) -> None: + if self == replaced: + raise OSError(error, os.strerror(error)) + real_unlink(self, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", unlink_failing_for_the_replaced_frame) + + with pytest.raises(OSError): + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + + assert stored_table(root).to_pydict() == {"a": [3]} + + +def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( + root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def write_some_then_interrupt( + frame: pd.DataFrame, sink: Any, **kwargs: Any + ) -> None: + sink.write(b"ARROW1 partial") + raise KeyboardInterrupt + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + manifest_bytes = (frame_dir(root) / "manifest.json").read_bytes() + monkeypatch.setattr(feather, "write_feather", write_some_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 diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index a4ae35ae..28926bec 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -1,8 +1,13 @@ -"""Frames shared by the DataFrame storage tests.""" +"""Helpers and frames shared by the DataFrame storage tests.""" +import errno import io +import json +import os +import time import uuid -from typing import Any, Optional +from pathlib import Path +from typing import Any, Dict, List, Optional import numpy as np import pandas as pd @@ -10,6 +15,66 @@ import pytest from pyarrow import feather +NAME = "revenue" +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 stored_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) / stored_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: + """Fail Path.open writes, mkdir and deletes with EROFS, like a read-only mount.""" + + def refuse(*args: Any, **kwargs: Any) -> Any: + raise OSError(errno.EROFS, "Read-only file system") + + real_open = Path.open + + def open_for_reading(self: Path, mode: str = "r", *args: Any, **kwargs: Any) -> Any: + if mode.strip("rbt"): + refuse() + return real_open(self, mode, *args, **kwargs) + + # The kernel reports EEXIST before EROFS, so an existing directory passes. + def mkdir( + self: Path, mode: int = 0o777, parents: bool = False, exist_ok: bool = False + ) -> None: + if not (exist_ok and self.is_dir()): + refuse() + + monkeypatch.setattr(Path, "open", open_for_reading) + monkeypatch.setattr(Path, "mkdir", mkdir) + monkeypatch.setattr(os, "unlink", refuse) + monkeypatch.setattr(os, "rmdir", refuse) + + +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 probe_failure(frame: Any) -> Optional[type]: """The exception type pyarrow itself raises for this frame, if any.""" From 30daa3761557fa7ce99e255194f3c1f53d3b39b0 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 17:21:25 +0000 Subject: [PATCH 13/41] refactor(artifacts): rename the project root fixture and the regex constants The test fixture root becomes tmp_project_root, _SAFE_ID and _FRAME_NAME become _SAFE_ID_REGEX and _FRAME_NAME_REGEX, and _FailingClose becomes _FailingCloseS3Upload. The Spark fixtures are typed as SparkSession, and store_frame loses the comment before the manifest step. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage.py | 7 +- .../artifacts/dataframe_storage_manifest.py | 4 +- tests/unit/artifacts/conftest.py | 15 +-- tests/unit/artifacts/test_arguments.py | 44 ++++---- tests/unit/artifacts/test_delete_dataframe.py | 38 ++++--- tests/unit/artifacts/test_file_format.py | 35 +++--- .../artifacts/test_manifest_and_cleanup.py | 66 ++++++------ tests/unit/artifacts/test_read_only_mount.py | 24 +++-- tests/unit/artifacts/test_storing.py | 64 ++++++----- tests/unit/artifacts/test_write_failures.py | 102 +++++++++++------- 10 files changed, 228 insertions(+), 171 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index 4b662942..a8f70d90 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -33,12 +33,12 @@ FrameKind = Literal["pandas", "polars"] -_SAFE_ID = re.compile(r"[A-Za-z0-9_-]{1,128}") +_SAFE_ID_REGEX = 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.""" - return isinstance(value, str) and _SAFE_ID.fullmatch(value) is not None + return isinstance(value, str) and _SAFE_ID_REGEX.fullmatch(value) is not None def classify_frame(value: object) -> Optional[FrameKind]: @@ -142,8 +142,7 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> Optional[ if isinstance(exc, OSError) and exc.errno == errno.EROFS: return None raise - # A failure from here on keeps the new frame: an s3fs rename can put the new - # manifest in place and still fail. + write_manifest(directory, new_file, old.file if old else None, len(frame)) _delete_replaced_frames(directory, old, new_file) return reference_id diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index a0cf1ffa..b37e6ea5 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -15,7 +15,7 @@ MANIFEST_VERSION = 1 MANIFEST_MAX_BYTES = 4096 -_FRAME_NAME = re.compile( +_FRAME_NAME_REGEX = 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, ) @@ -30,7 +30,7 @@ class Manifest(NamedTuple): 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 + return _FRAME_NAME_REGEX.fullmatch(name) is not None def read_manifest(directory: Path) -> Optional[Manifest]: diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index b58da100..936e4ef4 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -5,12 +5,13 @@ from typing import Any, Dict, Iterator, List, Optional import pytest +from pyspark.sql import SparkSession from deepnote_toolkit.artifacts import dataframe_storage @pytest.fixture -def root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: +def tmp_project_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: """The project root, a fresh directory.""" project_root = tmp_path / "project" project_root.mkdir() @@ -32,11 +33,11 @@ def tokyo_time() -> Iterator[None]: @pytest.fixture(scope="module") -def spark() -> Any: +def spark() -> SparkSession: from pyspark.sql import SparkSession return ( - SparkSession.builder.master("local") # type: ignore + SparkSession.builder.master("local") .appName("Toolkit") .config("spark.driver.memory", "4g") .getOrCreate() @@ -44,7 +45,7 @@ def spark() -> Any: @pytest.fixture -def spark_without_ansi_mode(spark: Any) -> Iterator[Any]: +def spark_without_ansi_mode(spark: SparkSession) -> Iterator[SparkSession]: """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") @@ -68,7 +69,7 @@ def open_and_record(self: Path, mode: str = "r", *args: Any, **kwargs: Any) -> A return sinks -class _FailingClose: +class _FailingCloseS3Upload: """A write handle whose close() really closes, then fails like s3fs's upload.""" def __init__(self, handle: Any, error: OSError) -> None: @@ -79,7 +80,7 @@ def close(self) -> None: self._handle.close() raise self._error - def __enter__(self) -> "_FailingClose": + def __enter__(self) -> "_FailingCloseS3Upload": return self def __exit__(self, *exc_info: Any) -> None: @@ -102,7 +103,7 @@ def open_failing_on_close( error = state["error"] if error is None or self.suffix != ".arrow" or "w" not in mode: return handle - return _FailingClose(handle, error) + return _FailingCloseS3Upload(handle, error) monkeypatch.setattr(Path, "open", open_failing_on_close) return state diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index c66005c6..a6c748bd 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -20,10 +20,10 @@ @pytest.mark.parametrize("name", ["x", "x" * 128, "Rev-2026_Q3"]) -def test_valid_names_are_accepted(root: Path, name: str) -> None: +def test_valid_names_are_accepted(tmp_project_root: Path, name: str) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) - assert stored_table(root, name).to_pydict() == {"a": [1]} + assert stored_table(tmp_project_root, name).to_pydict() == {"a": [1]} INVALID_NAMES = [ @@ -43,18 +43,20 @@ def test_valid_names_are_accepted(root: Path, name: str) -> None: @pytest.mark.parametrize("name", INVALID_NAMES) -def test_invalid_name_raises_before_anything_is_written(root: Path, name: Any) -> None: +def test_invalid_name_raises_before_anything_is_written( + tmp_project_root: Path, name: Any +) -> None: with mock.patch.object(feather, "write_feather", autospec=True) as write: with pytest.raises(ValueError): artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) write.assert_not_called() - assert list(root.iterdir()) == [] + assert list(tmp_project_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 + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch, name: Any ) -> None: read_only(monkeypatch) @@ -89,27 +91,29 @@ def _cleared_query_preview() -> DeepnoteQueryPreview: @pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES) def test_unsupported_value_raises_type_error( - root: Path, make_value: Callable[[], Any] + tmp_project_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()) == [] + assert list(tmp_project_root.iterdir()) == [] @pytest.mark.parametrize("make_preview", QUERY_PREVIEWS) def test_query_preview_raises_type_error( - root: Path, make_preview: Callable[[], Any] + tmp_project_root: Path, make_preview: Callable[[], Any] ) -> None: with pytest.raises(TypeError, match="SQL query preview"): artifacts.store_dataframe(make_preview(), NAME) - assert list(root.iterdir()) == [] + assert list(tmp_project_root.iterdir()) == [] @pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES + QUERY_PREVIEWS) def test_unsupported_value_raises_type_error_on_a_read_only_mount( - root: Path, monkeypatch: pytest.MonkeyPatch, make_value: Callable[[], Any] + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + make_value: Callable[[], Any], ) -> None: read_only(monkeypatch) @@ -117,15 +121,17 @@ def test_unsupported_value_raises_type_error_on_a_read_only_mount( artifacts.store_dataframe(make_value(), NAME) -def test_pyspark_frame_raises_type_error(root: Path, spark: Any) -> None: +def test_pyspark_frame_raises_type_error(tmp_project_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()) == [] + assert list(tmp_project_root.iterdir()) == [] def test_pandas_on_spark_frame_raises_type_error( - root: Path, spark_without_ansi_mode: Any, monkeypatch: pytest.MonkeyPatch + tmp_project_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") @@ -136,19 +142,19 @@ def test_pandas_on_spark_frame_raises_type_error( with pytest.raises(TypeError): artifacts.store_dataframe(frame, NAME) - assert list(root.iterdir()) == [] + assert list(tmp_project_root.iterdir()) == [] @pytest.mark.parametrize("name", INVALID_NAMES) def test_delete_with_an_invalid_name_raises_and_deletes_nothing( - root: Path, name: Any + tmp_project_root: Path, name: Any ) -> None: - (root / "x").mkdir() - (root / "x" / "keep.txt").write_text("keep") + (tmp_project_root / "x").mkdir() + (tmp_project_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 + assert (tmp_project_root / "x" / "keep.txt").read_text() == "keep" + assert len(arrow_files(tmp_project_root)) == 1 diff --git a/tests/unit/artifacts/test_delete_dataframe.py b/tests/unit/artifacts/test_delete_dataframe.py index 1a46c475..d0891421 100644 --- a/tests/unit/artifacts/test_delete_dataframe.py +++ b/tests/unit/artifacts/test_delete_dataframe.py @@ -14,53 +14,59 @@ ) -def test_delete_removes_the_frames_of_the_name_only(root: Path) -> None: +def test_delete_removes_the_frames_of_the_name_only(tmp_project_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 = tmp_project_root / "deepnote_dataframes" / "revenue-2026" neighbour.mkdir() (neighbour / "keep.txt").write_text("keep") - (root / "notes.txt").write_text("keep") + (tmp_project_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 not frame_dir(tmp_project_root).exists() + assert stored_table(tmp_project_root, "other").to_pydict() == {"a": [3]} assert (neighbour / "keep.txt").read_text() == "keep" - assert (root / "notes.txt").read_text() == "keep" + assert (tmp_project_root / "notes.txt").read_text() == "keep" -def test_delete_returns_false_when_there_is_nothing_to_delete(root: Path) -> None: +def test_delete_returns_false_when_there_is_nothing_to_delete( + tmp_project_root: Path, +) -> None: assert artifacts.delete_dataframe(NAME) is False - assert list(root.iterdir()) == [] + assert list(tmp_project_root.iterdir()) == [] -def test_delete_twice_returns_false_the_second_time(root: Path) -> None: +def test_delete_twice_returns_false_the_second_time(tmp_project_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: +def test_a_name_can_be_stored_again_after_it_was_deleted( + tmp_project_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 stored_manifest(root) - assert stored_table(root).to_pydict() == {"a": [2]} + assert "previous" not in stored_manifest(tmp_project_root) + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} def test_delete_does_not_follow_a_link_out_of_the_frames_directory( - root: Path, + tmp_project_root: Path, ) -> None: - outside = root / "outside" + outside = tmp_project_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) + (tmp_project_root / "deepnote_dataframes").mkdir() + (tmp_project_root / "deepnote_dataframes" / NAME).symlink_to( + outside, target_is_directory=True + ) with pytest.raises(OSError): artifacts.delete_dataframe(NAME) diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index a93b8450..dea74805 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -22,10 +22,10 @@ @pytest.mark.parametrize( "make_frame", [pd.DataFrame, pl.DataFrame], ids=["pandas", "polars"] ) -def test_frames_are_compressed(root: Path, make_frame: Any) -> None: +def test_frames_are_compressed(tmp_project_root: Path, make_frame: Any) -> None: artifacts.store_dataframe(make_frame({"a": ["x" * 100] * 10_000}), NAME) - stored = frame_dir(root) / stored_manifest(root)["file"] + stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["file"] assert 0 < stored.stat().st_size < 100_000 @@ -66,29 +66,33 @@ def test_zstd_is_chosen_when_pyarrow_has_it() -> None: "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 + tmp_project_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) / stored_manifest(root)["file"] + stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_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: +def test_pandas_frame_is_written_in_64k_row_batches(tmp_project_root: Path) -> None: artifacts.store_dataframe(pd.DataFrame({"a": range(70_000)}), NAME) - reader = pa.ipc.open_file(frame_dir(root) / stored_manifest(root)["file"]) + reader = pa.ipc.open_file( + frame_dir(tmp_project_root) / stored_manifest(tmp_project_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: +def test_polars_frame_is_written_in_64k_row_batches(tmp_project_root: Path) -> None: artifacts.store_dataframe(pl.DataFrame({"a": range(200_000)}), NAME) - reader = pa.ipc.open_file(frame_dir(root) / stored_manifest(root)["file"]) + reader = pa.ipc.open_file( + frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["file"] + ) assert reader.num_record_batches == 4 assert reader.read_all().num_rows == 200_000 @@ -99,7 +103,7 @@ def test_polars_frame_is_written_in_64k_row_batches(root: Path) -> None: ids=["before-1.37", "from-1.37"], ) def test_polars_batch_size_is_passed_to_polars_from_1_37( - root: Path, + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch, polars_version: str, expected_options: Dict[str, int], @@ -119,10 +123,10 @@ def test_polars_batch_size_is_passed_to_polars_from_1_37( if name == "record_batch_size" } assert options == expected_options - assert stored_table(root).to_pydict() == {"a": [1, 2, 3]} + assert stored_table(tmp_project_root).to_pydict() == {"a": [1, 2, 3]} -def test_polars_strings_are_stored_without_view_types(root: Path) -> None: +def test_polars_strings_are_stored_without_view_types(tmp_project_root: Path) -> None: artifacts.store_dataframe( pl.DataFrame( { @@ -134,7 +138,7 @@ def test_polars_strings_are_stored_without_view_types(root: Path) -> None: NAME, ) - assert stored_table(root).schema.types == [ + assert stored_table(tmp_project_root).schema.types == [ pa.large_string(), pa.large_binary(), pa.large_list(pa.large_string()), @@ -147,7 +151,10 @@ def test_polars_strings_are_stored_without_view_types(root: Path) -> None: ids=["before-1.2.1", "from-1.2.1"], ) def test_polars_compat_level_is_passed_to_polars_from_1_2_1( - root: Path, monkeypatch: pytest.MonkeyPatch, polars_version: str, passed: bool + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + polars_version: str, + passed: bool, ) -> None: monkeypatch.setattr(pl, "__version__", polars_version) original = pl.DataFrame.write_ipc @@ -158,4 +165,4 @@ def test_polars_compat_level_is_passed_to_polars_from_1_2_1( artifacts.store_dataframe(pl.DataFrame({"a": ["x"]}), NAME) assert ("compat_level" in write_ipc.call_args.kwargs) is passed - assert stored_table(root).to_pydict() == {"a": ["x"]} + assert stored_table(tmp_project_root).to_pydict() == {"a": ["x"]} diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index cf6acf33..778ba61d 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -28,29 +28,31 @@ ) -def test_second_write_keeps_one_previous_frame(root: Path) -> None: +def test_second_write_keeps_one_previous_frame(tmp_project_root: Path) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - first = stored_manifest(root) + first = stored_manifest(tmp_project_root) artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - second = stored_manifest(root) + second = stored_manifest(tmp_project_root) artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) - third = stored_manifest(root) + third = stored_manifest(tmp_project_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( + assert [path.name for path in arrow_files(tmp_project_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 stored_table(tmp_project_root).to_pydict() == {"a": [3]} + previous = pa.ipc.open_file(frame_dir(tmp_project_root) / third["previous"]) assert previous.read_all().to_pydict() == {"a": [2]} -def test_unreferenced_frames_are_deleted_only_when_stale(root: Path) -> None: +def test_unreferenced_frames_are_deleted_only_when_stale( + tmp_project_root: Path, +) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - directory = frame_dir(root) - replaced = directory / stored_manifest(root)["file"] + directory = frame_dir(tmp_project_root) + replaced = directory / stored_manifest(tmp_project_root)["file"] backdate(replaced, 2 * DAY) def add(name: str, age_seconds: float) -> Path: @@ -68,7 +70,7 @@ def add(name: str, age_seconds: float) -> Path: add(f"backup-{FRAME_A}", 2 * DAY), add(f"{FRAME_A}.bak", 2 * DAY), ] - outside = root / "outside.arrow" + outside = tmp_project_root / "outside.arrow" outside.write_bytes(b"precious") link = directory / "77777777-6666-4555-8444-333333333333.arrow" link.symlink_to(outside) @@ -136,9 +138,9 @@ def _manifest_bytes(**fields: Any) -> bytes: ], ) def test_unusable_manifest_counts_as_no_manifest( - root: Path, manifest_bytes: bytes + tmp_project_root: Path, manifest_bytes: bytes ) -> None: - directory = frame_dir(root) + directory = frame_dir(tmp_project_root) directory.mkdir(parents=True) (directory / "manifest.json").write_bytes(manifest_bytes) named = [directory / FRAME_A, directory / FRAME_B] @@ -150,16 +152,16 @@ def test_unusable_manifest_counts_as_no_manifest( artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - manifest = stored_manifest(root) + manifest = stored_manifest(tmp_project_root) assert "previous" not in manifest assert manifest["file"] not in (FRAME_A, FRAME_B) - assert stored_table(root).to_pydict() == {"a": [1]} + assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} assert sentinel.read_text() == "keep" assert all(path.exists() for path in named) -def test_manifest_of_exactly_4096_bytes_is_honoured(root: Path) -> None: - directory = frame_dir(root) +def test_manifest_of_exactly_4096_bytes_is_honoured(tmp_project_root: Path) -> None: + directory = frame_dir(tmp_project_root) directory.mkdir(parents=True) raw = _manifest_bytes(file=FRAME_A, previous=FRAME_B).ljust(4096) assert len(raw) == 4096 @@ -169,13 +171,13 @@ def test_manifest_of_exactly_4096_bytes_is_honoured(root: Path) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - assert stored_manifest(root)["previous"] == FRAME_A + assert stored_manifest(tmp_project_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 + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """Before 3.12, json raises RecursionError on the nesting 4 KiB allows.""" artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -193,12 +195,12 @@ def too_deep(raw: bytes) -> Any: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - assert "previous" not in stored_manifest(root) - assert stored_table(root).to_pydict() == {"a": [2]} + assert "previous" not in stored_manifest(tmp_project_root) + assert stored_table(tmp_project_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 + tmp_project_root: Path, arrow_sinks: List[Any], monkeypatch: pytest.MonkeyPatch ) -> None: seen: List[Dict[str, Any]] = [] original = dataframe_storage.write_manifest @@ -216,7 +218,7 @@ def spy(directory: Path, new_file: str, *args: Any) -> None: def test_concurrent_writers_leave_a_complete_current_frame( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: barrier = threading.Barrier(2, timeout=10) original = dataframe_storage.read_manifest @@ -243,9 +245,9 @@ def write(value: int) -> None: thread.join(timeout=30) assert errors == [] - files = arrow_files(root) + files = arrow_files(tmp_project_root) assert len(files) == 2 - current = stored_manifest(root)["file"] + current = stored_manifest(tmp_project_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 @@ -254,24 +256,24 @@ def write(value: int) -> None: 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]} + assert (frame_dir(tmp_project_root) / current).exists() + assert stored_table(tmp_project_root).to_pydict() == {"a": [3]} def test_successful_write_leaves_only_the_manifest_beside_the_frames( - root: Path, + tmp_project_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"] + assert entries_beside_the_frames(tmp_project_root) == ["manifest.json"] def test_stray_pruned_by_another_writer_meanwhile_is_not_a_failure( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - directory = frame_dir(root) + directory = frame_dir(tmp_project_root) strays = [directory / FRAME_A, directory / FRAME_B] for stray in strays: stray.write_bytes(b"unreferenced") @@ -290,4 +292,4 @@ def scandir_then_lose_a_stray(path: Any) -> Iterator[Any]: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) assert not strays[1].exists() - assert stored_table(root).to_pydict() == {"a": [2]} + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} diff --git a/tests/unit/artifacts/test_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py index 54c81995..a0c45e74 100644 --- a/tests/unit/artifacts/test_read_only_mount.py +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -21,7 +21,7 @@ def test_read_only_mount_stores_nothing_and_does_not_convert( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: read_only(monkeypatch) @@ -30,14 +30,14 @@ def test_read_only_mount_stores_nothing_and_does_not_convert( assert result is None write.assert_not_called() - assert list(root.iterdir()) == [] + assert list(tmp_project_root.iterdir()) == [] def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - manifest_bytes = (frame_dir(root) / "manifest.json").read_bytes() + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() read_only(monkeypatch) with mock.patch.object(feather, "write_feather", autospec=True) as write: @@ -45,12 +45,14 @@ def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( assert result is None write.assert_not_called() - assert (frame_dir(root) / "manifest.json").read_bytes() == manifest_bytes - assert stored_table(root).to_pydict() == {"a": [1]} + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} def test_erofs_from_the_write_stores_nothing( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: def refuse(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: raise OSError(errno.EROFS, "Read-only file system") @@ -60,22 +62,22 @@ def refuse(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) assert result is None - assert arrow_files(root) == [] + assert arrow_files(tmp_project_root) == [] def test_delete_on_a_read_only_mount_deletes_nothing( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_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 + assert len(arrow_files(tmp_project_root)) == 1 def test_delete_returns_false_when_the_filesystem_refuses_with_erofs( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) diff --git a/tests/unit/artifacts/test_storing.py b/tests/unit/artifacts/test_storing.py index cfaf37ee..2486d612 100644 --- a/tests/unit/artifacts/test_storing.py +++ b/tests/unit/artifacts/test_storing.py @@ -22,16 +22,21 @@ UUID_PATTERN = r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}" -def test_stores_pandas_frame_and_manifest(root: Path) -> None: +def test_stores_pandas_frame_and_manifest(tmp_project_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 = stored_manifest(root) + assert [path.name for path in arrow_files(tmp_project_root)] == [ + f"{reference_id}.arrow" + ] + assert stored_table(tmp_project_root).to_pydict() == { + "a": [1, 2, 3], + "b": ["x", "y", "z"], + } + manifest = stored_manifest(tmp_project_root) assert set(manifest) == {"version", "file", "written_at", "rows", "toolkit_version"} assert manifest["version"] == 1 assert manifest["file"] == f"{reference_id}.arrow" @@ -41,39 +46,42 @@ def test_stores_pandas_frame_and_manifest(root: Path) -> None: @pytest.mark.usefixtures("tokyo_time") -def test_manifest_time_is_utc(root: Path) -> None: +def test_manifest_time_is_utc(tmp_project_root: Path) -> None: before = datetime.now(timezone.utc) artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) written = datetime.strptime( - stored_manifest(root)["written_at"], "%Y-%m-%dT%H:%M:%SZ" + stored_manifest(tmp_project_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: +def test_stores_polars_frame(tmp_project_root: Path) -> None: reference_id = artifacts.store_dataframe( pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}), NAME ) - assert stored_manifest(root)["file"] == f"{reference_id}.arrow" - assert stored_table(root).to_pydict() == {"a": [1, 2, 3], "b": ["x", "y", "z"]} - assert stored_manifest(root)["rows"] == 3 + assert stored_manifest(tmp_project_root)["file"] == f"{reference_id}.arrow" + assert stored_table(tmp_project_root).to_pydict() == { + "a": [1, 2, 3], + "b": ["x", "y", "z"], + } + assert stored_manifest(tmp_project_root)["rows"] == 3 -def test_writing_does_not_change_the_callers_frame(root: Path) -> None: +def test_writing_does_not_change_the_callers_frame(tmp_project_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 + assert len(arrow_files(tmp_project_root)) == 1 pd.testing.assert_frame_equal(frame, before) def test_filtered_frame_is_stored_with_its_row_labels( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: def refuse(self: pd.DataFrame, *args: Any, **kwargs: Any) -> None: raise ValueError("feather does not support serializing a non-default index") @@ -84,24 +92,28 @@ def refuse(self: pd.DataFrame, *args: Any, **kwargs: Any) -> None: artifacts.store_dataframe(frame[frame.a != 7], NAME) - assert stored_table(root).to_pydict() == { + assert stored_table(tmp_project_root).to_pydict() == { "a": [5, 6, 8], "__index_level_0__": [0, 1, 3], } -def test_groupby_keys_are_stored_after_the_data_columns(root: Path) -> None: +def test_groupby_keys_are_stored_after_the_data_columns(tmp_project_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"] + assert stored_table(tmp_project_root).schema.names == ["v", "k"] -def test_transposed_frame_is_stored_with_string_names(root: Path) -> None: +def test_transposed_frame_is_stored_with_string_names(tmp_project_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__"] + assert stored_table(tmp_project_root).schema.names == [ + "0", + "1", + "__index_level_0__", + ] @pytest.mark.parametrize( @@ -109,26 +121,26 @@ def test_transposed_frame_is_stored_with_string_names(root: Path) -> None: [pd.DataFrame({"a": []}), pl.DataFrame({"a": []})], ids=["pandas", "polars"], ) -def test_empty_frame_is_stored(root: Path, frame: Any) -> None: +def test_empty_frame_is_stored(tmp_project_root: Path, frame: Any) -> None: artifacts.store_dataframe(frame, NAME) - assert stored_table(root).num_rows == 0 - assert stored_manifest(root)["rows"] == 0 + assert stored_table(tmp_project_root).num_rows == 0 + assert stored_manifest(tmp_project_root)["rows"] == 0 -def test_names_are_stored_apart(root: Path) -> None: +def test_names_are_stored_apart(tmp_project_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]} + assert stored_table(tmp_project_root, "first").to_pydict() == {"a": [1]} + assert stored_table(tmp_project_root, "second").to_pydict() == {"a": [2]} -def test_the_last_call_with_a_name_wins(root: Path) -> None: +def test_the_last_call_with_a_name_wins(tmp_project_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]} + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} def test_functions_are_exported_from_the_package() -> None: diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 45f8c185..cecabef7 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -54,48 +54,64 @@ def _store_bad_frame_after_a_good_one( @pytest.mark.parametrize("make_frame", ALWAYS_FAILING) def test_unserializable_frame_raises_and_keeps_the_previous_frame( - root: Path, make_frame: Callable[[], Any] + tmp_project_root: Path, make_frame: Callable[[], Any] ) -> None: raised_by_pyarrow = probe_failure(make_frame()) assert raised_by_pyarrow is not None, "pyarrow is expected to refuse this frame" - left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) + left = _store_bad_frame_after_a_good_one( + tmp_project_root, make_frame, raised_by_pyarrow + ) - assert (frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] - assert left["listing"] == sorted(["manifest.json", stored_manifest(root)["file"]]) - assert stored_table(root).to_pydict() == {"ok": [1]} + assert (frame_dir(tmp_project_root) / "manifest.json").read_bytes() == left[ + "manifest_bytes" + ] + assert left["listing"] == sorted( + ["manifest.json", stored_manifest(tmp_project_root)["file"]] + ) + assert stored_table(tmp_project_root).to_pydict() == {"ok": [1]} @pytest.mark.parametrize("make_frame, rows", VERSION_DEPENDENT) def test_frames_pyarrow_may_refuse_follow_what_pyarrow_does( - root: Path, make_frame: Callable[[], Any], rows: int + tmp_project_root: Path, make_frame: Callable[[], Any], rows: int ) -> None: raised_by_pyarrow = probe_failure(make_frame()) - left = _store_bad_frame_after_a_good_one(root, make_frame, raised_by_pyarrow) + left = _store_bad_frame_after_a_good_one( + tmp_project_root, make_frame, raised_by_pyarrow + ) if raised_by_pyarrow is None: - assert stored_table(root).num_rows == rows + assert stored_table(tmp_project_root).num_rows == rows return - assert (frame_dir(root) / "manifest.json").read_bytes() == left["manifest_bytes"] - assert left["listing"] == sorted(["manifest.json", stored_manifest(root)["file"]]) + assert (frame_dir(tmp_project_root) / "manifest.json").read_bytes() == left[ + "manifest_bytes" + ] + assert left["listing"] == sorted( + ["manifest.json", stored_manifest(tmp_project_root)["file"]] + ) -def test_polars_object_column_raises_and_keeps_the_previous_frame(root: Path) -> None: +def test_polars_object_column_raises_and_keeps_the_previous_frame( + tmp_project_root: Path, +) -> None: artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) - manifest_bytes = (frame_dir(root) / "manifest.json").read_bytes() + manifest_bytes = (frame_dir(tmp_project_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 + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_root)) == 1 def test_partial_file_is_deleted_when_the_write_fails( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: def write_some_then_fail(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: sink.write(b"ARROW1 partial") @@ -107,8 +123,8 @@ def write_some_then_fail(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: 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() + assert arrow_files(tmp_project_root) == [] + assert not (frame_dir(tmp_project_root) / "manifest.json").exists() @pytest.mark.parametrize( @@ -122,7 +138,7 @@ def write_some_then_fail(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: ids=["EPERM", "EACCES", "EIO", "ENOSPC"], ) def test_error_from_close_raises_and_stores_nothing( - root: Path, close_error: Dict[str, Optional[OSError]], error: OSError + tmp_project_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 @@ -131,15 +147,15 @@ def test_error_from_close_raises_and_stores_nothing( 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() + assert arrow_files(tmp_project_root) == [] + assert not (frame_dir(tmp_project_root) / "manifest.json").exists() def test_failed_manifest_read_deletes_the_new_frame( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - frames = arrow_files(root) + frames = arrow_files(tmp_project_root) real_open = Path.open def open_failing_for_the_manifest( @@ -154,15 +170,15 @@ def open_failing_for_the_manifest( with pytest.raises(OSError): artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - assert arrow_files(root) == frames + assert arrow_files(tmp_project_root) == frames def test_failure_in_the_manifest_step_deletes_nothing( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_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) + directory = frame_dir(tmp_project_root) stale = directory / FRAME_A stale.write_bytes(b"unreferenced") backdate(stale, 2 * DAY) @@ -180,7 +196,7 @@ def broken_manifest(*args: Any) -> Any: assert (directory / "manifest.json").read_bytes() == manifest_bytes entries = {path.name for path in directory.iterdir()} assert set(listing) <= entries and len(entries - set(listing)) == 1 - assert stored_table(root).to_pydict() == {"a": [2]} + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} class _FailingManifestHandle: @@ -211,10 +227,10 @@ def __getattr__(self, name: str) -> Any: @pytest.mark.parametrize("failing", ["write", "close", "interrupt"]) def test_failed_manifest_write_keeps_the_previous_frame( - root: Path, monkeypatch: pytest.MonkeyPatch, failing: str + tmp_project_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() + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() real_temporary_file = tempfile.NamedTemporaryFile def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: @@ -229,9 +245,11 @@ def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: with pytest.raises(KeyboardInterrupt if failing == "interrupt" else 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"] + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} + assert entries_beside_the_frames(tmp_project_root) == ["manifest.json"] def _copy_then_fail(source: str, destination: Path) -> None: @@ -241,7 +259,7 @@ def _copy_then_fail(source: str, destination: Path) -> None: def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) monkeypatch.setattr(dataframe_storage_manifest.os, "replace", _copy_then_fail) @@ -249,16 +267,18 @@ def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( with pytest.raises(OSError): artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - assert stored_table(root).to_pydict() == {"a": [2]} + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} @pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["EIO", "EROFS"]) def test_failure_after_the_manifest_step_keeps_the_new_frame( - root: Path, monkeypatch: pytest.MonkeyPatch, error: int + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch, error: int ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - replaced = frame_dir(root) / stored_manifest(root)["previous"] + replaced = ( + frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["previous"] + ) real_unlink = Path.unlink def unlink_failing_for_the_replaced_frame( @@ -273,11 +293,11 @@ def unlink_failing_for_the_replaced_frame( with pytest.raises(OSError): artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) - assert stored_table(root).to_pydict() == {"a": [3]} + assert stored_table(tmp_project_root).to_pydict() == {"a": [3]} def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( - root: Path, monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: def write_some_then_interrupt( frame: pd.DataFrame, sink: Any, **kwargs: Any @@ -286,11 +306,13 @@ def write_some_then_interrupt( raise KeyboardInterrupt artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - manifest_bytes = (frame_dir(root) / "manifest.json").read_bytes() + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() monkeypatch.setattr(feather, "write_feather", write_some_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 ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_root)) == 1 From 9dc8d664c081709d93bdbf1e83159827efae2232 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 17:22:39 +0000 Subject: [PATCH 14/41] refactor(artifacts): validate the manifest with a pydantic model Manifest is now a pydantic model of the whole manifest, and read_manifest and write_manifest both go through it. Reading used to check only file and previous; now pydantic checks every field's type, all of them required but previous, with strict types so "1", 1.5 and true are not ints. read_manifest then refuses a version other than 1, negative rows, and a file or previous that is not .arrow. The name check stays on the read side because cleanup unlinks the replaced manifest's previous, and any code in the project can write the manifest. A refused manifest counts as absent, as before. Unknown keys are ignored. The written manifest is unchanged, previous still omitted when there is none. The pydantic_compat_helpers keep both pydantic 1 and 2 working; only known keys reach the model, since pydantic 1 raises TypeError on __pydantic_self__. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage_manifest.py | 68 +++++++++++----- .../artifacts/test_manifest_and_cleanup.py | 81 +++++++++++++------ 2 files changed, 104 insertions(+), 45 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index b37e6ea5..86d2f7b2 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -7,8 +7,15 @@ import tempfile from datetime import datetime, timezone from pathlib import Path -from typing import Any, NamedTuple, Optional +from typing import Optional +from pydantic import BaseModel, StrictInt, StrictStr, ValidationError + +from deepnote_core.pydantic_compat_helpers import ( + _model_dump_compat, + _model_fields_compat, + model_validate_compat, +) from deepnote_toolkit import __version__ MANIFEST_NAME = "manifest.json" @@ -21,20 +28,27 @@ ) -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_REGEX.fullmatch(name) is not None +class Manifest(BaseModel): + """The manifest file; `previous` is the frame `file` replaced. + + Unknown keys are ignored. + """ + + version: StrictInt + file: StrictStr + previous: Optional[StrictStr] = None + written_at: StrictStr + rows: StrictInt + toolkit_version: StrictStr + + def read_manifest(directory: Path) -> Optional[Manifest]: - """Return the files the manifest names, or None if there is no usable one. + """Return the directory's manifest, or None if there is no usable one. Any code in the project can write the manifest, so a malformed one counts as absent. @@ -50,17 +64,25 @@ def read_manifest(directory: Path) -> Optional[Manifest]: data = json.loads(raw) except (ValueError, RecursionError): return None + # pydantic v1 would turn a list of pairs into a dict. if not isinstance(data, dict): return None + # pydantic v1 passes every key to __init__, where __pydantic_self__ collides + # with its first parameter and raises TypeError. + fields = _model_fields_compat(Manifest) + known = {key: value for key, value in data.items() if key in fields} + try: + manifest = model_validate_compat(Manifest, known) + except ValidationError: + return None - file, previous = data.get("file"), data.get("previous") - if not isinstance(file, str) or not is_frame_name(file): + if manifest.version != MANIFEST_VERSION or manifest.rows < 0: + return None + if not is_frame_name(manifest.file): return None - if previous is not None and not ( - isinstance(previous, str) and is_frame_name(previous) - ): + if manifest.previous is not None and not is_frame_name(manifest.previous): return None - return Manifest(file=file, previous=previous) + return manifest def write_manifest( @@ -71,12 +93,14 @@ def write_manifest( On s3fs a rename is a copy and then a delete, so a failed rename may already have replaced the manifest. """ - 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__ + manifest = Manifest( + version=MANIFEST_VERSION, + file=new_file, + previous=previous, + written_at=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), + rows=rows, + toolkit_version=__version__, + ) # delete=False: s3fs uploads at close(), so the file must be closed before the # rename, and deleting it only when the context ends needs Python 3.12's # delete_on_close=False. @@ -90,7 +114,7 @@ def write_manifest( ) try: with temporary: - temporary.write(json.dumps(manifest)) + temporary.write(json.dumps(_model_dump_compat(manifest))) os.replace(temporary.name, directory / MANIFEST_NAME) except BaseException: with contextlib.suppress(OSError): diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 778ba61d..da01fe02 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -87,8 +87,20 @@ def add(name: str, age_seconds: float) -> Path: assert replaced.exists() -def _manifest_bytes(**fields: Any) -> bytes: - return json.dumps({"version": 1, **fields}).encode() +_VALID_MANIFEST: Dict[str, Any] = { + "version": 1, + "file": FRAME_A, + "previous": FRAME_B, + "written_at": "2026-10-15T09:30:00Z", + "rows": 1, + "toolkit_version": "1.0.0", +} + + +def _manifest_bytes(without: str = "", **fields: Any) -> bytes: + manifest = {**_VALID_MANIFEST, **fields} + manifest.pop(without, None) + return json.dumps(manifest).encode() @pytest.mark.parametrize( @@ -96,13 +108,18 @@ def _manifest_bytes(**fields: Any) -> 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( + json.dumps(list(_VALID_MANIFEST.items())).encode(), id="not-an-object" + ), + pytest.param(_manifest_bytes(without="version"), id="no-version"), + pytest.param(_manifest_bytes(version=2), id="unsupported-version"), + pytest.param(_manifest_bytes(version="1"), id="version-a-string"), + pytest.param(_manifest_bytes(version=True), id="version-a-bool"), + pytest.param(_manifest_bytes(without="file"), 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=7), id="file-not-a-string"), pytest.param( - _manifest_bytes(file=FRAME_A, previous="../../data/x.csv"), - id="previous-traversal", + _manifest_bytes(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"), @@ -111,29 +128,34 @@ def _manifest_bytes(**fields: Any) -> bytes: id="file-valid-name-then-traversal", ), pytest.param(_manifest_bytes(file=f"../{FRAME_A}"), id="file-in-parent"), + pytest.param(_manifest_bytes(previous=f"x{FRAME_B}"), id="previous-prefixed"), + pytest.param( + _manifest_bytes(previous=f"{FRAME_B}.bak"), id="previous-suffixed" + ), pytest.param( - _manifest_bytes(file=FRAME_A, previous=f"x{FRAME_B}"), - id="previous-prefixed", + _manifest_bytes(previous=f"{FRAME_B}/../../data/x.csv"), + id="previous-valid-name-then-traversal", ), pytest.param( - _manifest_bytes(file=FRAME_A, previous=f"{FRAME_B}.bak"), - id="previous-suffixed", + _manifest_bytes(previous=f"../{FRAME_B}"), id="previous-in-parent" ), + pytest.param(_manifest_bytes(without="written_at"), id="no-written-at"), + pytest.param(_manifest_bytes(written_at=7), id="written-at-not-a-string"), + pytest.param(_manifest_bytes(without="rows"), id="no-rows"), + pytest.param(_manifest_bytes(rows="1"), id="rows-a-string"), + pytest.param(_manifest_bytes(rows=1.5), id="rows-not-whole"), + pytest.param(_manifest_bytes(rows=-1), id="rows-negative"), pytest.param( - _manifest_bytes(file=FRAME_A, previous=f"{FRAME_B}/../../data/x.csv"), - id="previous-valid-name-then-traversal", + _manifest_bytes(without="toolkit_version"), id="no-toolkit-version" ), pytest.param( - _manifest_bytes(file=FRAME_A, previous=f"../{FRAME_B}"), - id="previous-in-parent", + _manifest_bytes(toolkit_version=7), id="toolkit-version-not-a-string" ), pytest.param( - _manifest_bytes(file=FRAME_A, previous=FRAME_B, pad="x" * 5000), - id="oversized-truncated-mid-string", + _manifest_bytes(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", + _manifest_bytes() + b" " * 5000, id="oversized-valid-json-then-spaces" ), ], ) @@ -160,12 +182,25 @@ def test_unusable_manifest_counts_as_no_manifest( assert all(path.exists() for path in named) -def test_manifest_of_exactly_4096_bytes_is_honoured(tmp_project_root: Path) -> None: +@pytest.mark.parametrize( + "manifest_bytes", + [ + pytest.param(_manifest_bytes().ljust(4096), id="exactly-4096-bytes"), + pytest.param( + _manifest_bytes(written_by_a_newer_toolkit=True), id="unknown-key" + ), + pytest.param( + _manifest_bytes(__pydantic_self__=1), id="unknown-key-named-like-self" + ), + ], +) +def test_usable_manifest_is_honoured( + tmp_project_root: Path, manifest_bytes: bytes +) -> None: directory = frame_dir(tmp_project_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) + assert len(manifest_bytes) <= 4096 + (directory / "manifest.json").write_bytes(manifest_bytes) for name in (FRAME_A, FRAME_B): (directory / name).write_bytes(b"named by the manifest") From aad5561b5b087828a14058023b47a663a4c256cc Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 17:23:10 +0000 Subject: [PATCH 15/41] test(artifacts): assert the exception each pytest.raises catches A pytest.raises that names only a type passes on any error of that type, such as an OSError from a different step than the one under test. Each one now checks what it caught: the errno of an OSError, the message of a ValueError or TypeError, the identity of an injected error or interrupt, and the exact type where a subclass would also match. The read-only TypeError test is split so a SQL query preview's message is checked apart from an unsupported value's. rmtree's symbolic link error is compared by its args: newer Pythons (3.12.x and up) set its filename, and str() then no longer shows the message. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_arguments.py | 33 +++++++++-- tests/unit/artifacts/test_delete_dataframe.py | 5 +- tests/unit/artifacts/test_write_failures.py | 57 +++++++++++++------ 3 files changed, 70 insertions(+), 25 deletions(-) diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index a6c748bd..22b9c903 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -47,9 +47,10 @@ def test_invalid_name_raises_before_anything_is_written( tmp_project_root: Path, name: Any ) -> None: with mock.patch.object(feather, "write_feather", autospec=True) as write: - with pytest.raises(ValueError): + with pytest.raises(ValueError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + assert str(raised.value).endswith(f"got {name!r}") write.assert_not_called() assert list(tmp_project_root.iterdir()) == [] @@ -60,9 +61,11 @@ def test_invalid_name_raises_on_a_read_only_mount( ) -> None: read_only(monkeypatch) - with pytest.raises(ValueError): + with pytest.raises(ValueError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + assert str(raised.value).endswith(f"got {name!r}") + def _query_preview() -> DeepnoteQueryPreview: return DeepnoteQueryPreview({"a": [1]}, deepnote_query="select 1") @@ -109,7 +112,7 @@ def test_query_preview_raises_type_error( assert list(tmp_project_root.iterdir()) == [] -@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES + QUERY_PREVIEWS) +@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES) def test_unsupported_value_raises_type_error_on_a_read_only_mount( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch, @@ -117,9 +120,25 @@ def test_unsupported_value_raises_type_error_on_a_read_only_mount( ) -> None: read_only(monkeypatch) - with pytest.raises(TypeError): + with pytest.raises(TypeError) as raised: artifacts.store_dataframe(make_value(), NAME) + raised.match("pandas or polars DataFrame") + + +@pytest.mark.parametrize("make_preview", QUERY_PREVIEWS) +def test_query_preview_raises_type_error_on_a_read_only_mount( + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + make_preview: Callable[[], Any], +) -> None: + read_only(monkeypatch) + + with pytest.raises(TypeError) as raised: + artifacts.store_dataframe(make_preview(), NAME) + + raised.match("SQL query preview") + def test_pyspark_frame_raises_type_error(tmp_project_root: Path, spark: Any) -> None: with pytest.raises(TypeError, match="pandas or polars DataFrame"): @@ -139,9 +158,10 @@ def test_pandas_on_spark_frame_raises_type_error( pytest.importorskip("pyspark.pandas") frame = spark_without_ansi_mode.range(3).pandas_api() - with pytest.raises(TypeError): + with pytest.raises(TypeError) as raised: artifacts.store_dataframe(frame, NAME) + raised.match("pandas or polars DataFrame") assert list(tmp_project_root.iterdir()) == [] @@ -153,8 +173,9 @@ def test_delete_with_an_invalid_name_raises_and_deletes_nothing( (tmp_project_root / "x" / "keep.txt").write_text("keep") artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - with pytest.raises(ValueError): + with pytest.raises(ValueError) as raised: artifacts.delete_dataframe(name) + assert str(raised.value).endswith(f"got {name!r}") assert (tmp_project_root / "x" / "keep.txt").read_text() == "keep" assert len(arrow_files(tmp_project_root)) == 1 diff --git a/tests/unit/artifacts/test_delete_dataframe.py b/tests/unit/artifacts/test_delete_dataframe.py index d0891421..1e09d941 100644 --- a/tests/unit/artifacts/test_delete_dataframe.py +++ b/tests/unit/artifacts/test_delete_dataframe.py @@ -68,7 +68,10 @@ def test_delete_does_not_follow_a_link_out_of_the_frames_directory( outside, target_is_directory=True ) - with pytest.raises(OSError): + with pytest.raises(OSError) as raised: artifacts.delete_dataframe(NAME) + # Not match=: newer Pythons (3.12.x and up) set this error's filename, and str() + # then shows "[Errno None] None: " instead of the message. + assert raised.value.args == ("Cannot call rmtree on a symbolic link",) assert (outside / "keep.txt").read_text() == "keep" diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index cecabef7..d2ffb423 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -44,8 +44,9 @@ def _store_bad_frame_after_a_good_one( if expected_error is None: artifacts.store_dataframe(make_frame(), NAME) else: - with pytest.raises(expected_error): + with pytest.raises(expected_error) as raised: artifacts.store_dataframe(make_frame(), NAME) + assert raised.type is expected_error return { "manifest_bytes": manifest_bytes, "listing": sorted(path.name for path in directory.iterdir()), @@ -99,11 +100,12 @@ def test_polars_object_column_raises_and_keeps_the_previous_frame( artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() - with pytest.raises(pl.exceptions.ComputeError): + with pytest.raises(pl.exceptions.ComputeError) as raised: artifacts.store_dataframe( pl.DataFrame({"a": pl.Series([object()], dtype=pl.Object)}), NAME ) + raised.match("'Object' datatype") assert ( frame_dir(tmp_project_root) / "manifest.json" ).read_bytes() == manifest_bytes @@ -167,9 +169,10 @@ def open_failing_for_the_manifest( monkeypatch.setattr(Path, "open", open_failing_for_the_manifest) - with pytest.raises(OSError): + with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + assert raised.value.errno == errno.EIO assert arrow_files(tmp_project_root) == frames @@ -190,9 +193,10 @@ def broken_manifest(*args: Any) -> Any: monkeypatch.setattr(dataframe_storage, "write_manifest", broken_manifest) - with pytest.raises(OSError): + with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + assert raised.value.errno == errno.EIO assert (directory / "manifest.json").read_bytes() == manifest_bytes entries = {path.name for path in directory.iterdir()} assert set(listing) <= entries and len(entries - set(listing)) == 1 @@ -200,17 +204,16 @@ def broken_manifest(*args: Any) -> Any: class _FailingManifestHandle: - """A temporary manifest whose write() fails or is interrupted, or close() fails.""" + """A temporary manifest whose write() or close() raises `error`.""" - def __init__(self, handle: Any, failing: str) -> None: + def __init__(self, handle: Any, failing: str, error: BaseException) -> None: self._handle = handle self._failing = failing + self._error = error def write(self, data: str) -> int: if self._failing == "write": - raise OSError(errno.ENOSPC, "No space left on device") - if self._failing == "interrupt": - raise KeyboardInterrupt + raise self._error return self._handle.write(data) def __enter__(self) -> "_FailingManifestHandle": @@ -219,22 +222,34 @@ def __enter__(self) -> "_FailingManifestHandle": def __exit__(self, *exc_info: Any) -> None: self._handle.close() if self._failing == "close": - raise OSError(errno.EIO, "Input/output error") + raise self._error def __getattr__(self, name: str) -> Any: return getattr(self._handle, name) -@pytest.mark.parametrize("failing", ["write", "close", "interrupt"]) +@pytest.mark.parametrize( + "failing, error", + [ + ("write", OSError(errno.ENOSPC, "No space left on device")), + ("close", OSError(errno.EIO, "Input/output error")), + ("write", KeyboardInterrupt()), + ], + ids=["write", "close", "interrupt"], +) def test_failed_manifest_write_keeps_the_previous_frame( - tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch, failing: str + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + failing: str, + error: BaseException, ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) manifest_bytes = (frame_dir(tmp_project_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) + handle = real_temporary_file(*args, **kwargs) + return _FailingManifestHandle(handle, failing, error) monkeypatch.setattr( dataframe_storage_manifest, @@ -242,9 +257,10 @@ def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: SimpleNamespace(NamedTemporaryFile=failing_temporary_file), ) - with pytest.raises(KeyboardInterrupt if failing == "interrupt" else OSError): + with pytest.raises(type(error)) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + assert raised.value is error assert ( frame_dir(tmp_project_root) / "manifest.json" ).read_bytes() == manifest_bytes @@ -264,9 +280,10 @@ def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) monkeypatch.setattr(dataframe_storage_manifest.os, "replace", _copy_then_fail) - with pytest.raises(OSError): + with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + assert raised.value.errno == errno.EIO assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} @@ -290,28 +307,32 @@ def unlink_failing_for_the_replaced_frame( monkeypatch.setattr(Path, "unlink", unlink_failing_for_the_replaced_frame) - with pytest.raises(OSError): + with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + assert raised.value.errno == error assert stored_table(tmp_project_root).to_pydict() == {"a": [3]} def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: + interrupt = KeyboardInterrupt() + def write_some_then_interrupt( frame: pd.DataFrame, sink: Any, **kwargs: Any ) -> None: sink.write(b"ARROW1 partial") - raise KeyboardInterrupt + raise interrupt artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() monkeypatch.setattr(feather, "write_feather", write_some_then_interrupt) - with pytest.raises(KeyboardInterrupt): + with pytest.raises(KeyboardInterrupt) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + assert raised.value is interrupt assert ( frame_dir(tmp_project_root) / "manifest.json" ).read_bytes() == manifest_bytes From d337a3c1bc8f91d98abe776bdbadee5ba491c6b1 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 18:23:18 +0000 Subject: [PATCH 16/41] fix(artifacts): correct parameter name in _store_bad_frame_after_a_good_one function --- tests/unit/artifacts/test_write_failures.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index d2ffb423..e76314ed 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -32,14 +32,14 @@ def _store_bad_frame_after_a_good_one( - root: Path, make_frame: Callable[[], Any], expected_error: Optional[type] + tmp_project_root: Path, make_frame: Callable[[], Any], expected_error: Optional[type] ) -> Dict[str, Any]: """Store a good frame, then a bad one. Return the good call's manifest bytes and the directory listing after the bad call. """ artifacts.store_dataframe(pd.DataFrame({"ok": [1]}), NAME) - directory = frame_dir(root) + directory = frame_dir(tmp_project_root) manifest_bytes = (directory / "manifest.json").read_bytes() if expected_error is None: artifacts.store_dataframe(make_frame(), NAME) From 329dc5113d1512d2276ca480aadc8bef1124736f Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 18:23:26 +0000 Subject: [PATCH 17/41] fix(artifacts): add comment to clarify mutation of query preview in _cleared_query_preview function --- tests/unit/artifacts/test_arguments.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index 22b9c903..e4478708 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -73,6 +73,7 @@ def _query_preview() -> DeepnoteQueryPreview: def _cleared_query_preview() -> DeepnoteQueryPreview: preview = _query_preview() + # Mutating the frame clears its attached SQL query metadata. preview["b"] = 2 return preview From 46b4884e55478425624385a17a460252c4811f82 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 19:14:21 +0000 Subject: [PATCH 18/41] test(artifacts): split the bad-frame helper so pytest.raises is unconditional _store_bad_frame_after_a_good_one took an optional expected_error and only wrapped the second store in pytest.raises when one was given. It now requires expected_error. _store_frame_after_a_good_one stores the second frame without expecting an error, and _store_a_good_frame holds the first store both share. test_frames_pyarrow_may_refuse_follow_what_pyarrow_does picks the helper by what pyarrow does. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_write_failures.py | 34 +++++++++++++-------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index e76314ed..b75fcdcf 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -31,25 +31,32 @@ ) +def _store_a_good_frame(tmp_project_root: Path) -> bytes: + artifacts.store_dataframe(pd.DataFrame({"ok": [1]}), NAME) + return (frame_dir(tmp_project_root) / "manifest.json").read_bytes() + + +def _store_frame_after_a_good_one( + tmp_project_root: Path, make_frame: Callable[[], Any] +) -> None: + _store_a_good_frame(tmp_project_root) + artifacts.store_dataframe(make_frame(), NAME) + + def _store_bad_frame_after_a_good_one( - tmp_project_root: Path, make_frame: Callable[[], Any], expected_error: Optional[type] + tmp_project_root: Path, make_frame: Callable[[], Any], expected_error: type ) -> Dict[str, Any]: """Store a good frame, then a bad one. Return the good call's manifest bytes and the directory listing after the bad call. """ - artifacts.store_dataframe(pd.DataFrame({"ok": [1]}), NAME) - directory = frame_dir(tmp_project_root) - manifest_bytes = (directory / "manifest.json").read_bytes() - if expected_error is None: + manifest_bytes = _store_a_good_frame(tmp_project_root) + with pytest.raises(expected_error) as raised: artifacts.store_dataframe(make_frame(), NAME) - else: - with pytest.raises(expected_error) as raised: - artifacts.store_dataframe(make_frame(), NAME) - assert raised.type is expected_error + assert raised.type is expected_error return { "manifest_bytes": manifest_bytes, - "listing": sorted(path.name for path in directory.iterdir()), + "listing": sorted(path.name for path in frame_dir(tmp_project_root).iterdir()), } @@ -79,13 +86,14 @@ def test_frames_pyarrow_may_refuse_follow_what_pyarrow_does( ) -> None: raised_by_pyarrow = probe_failure(make_frame()) + if raised_by_pyarrow is None: + _store_frame_after_a_good_one(tmp_project_root, make_frame) + assert stored_table(tmp_project_root).num_rows == rows + return left = _store_bad_frame_after_a_good_one( tmp_project_root, make_frame, raised_by_pyarrow ) - if raised_by_pyarrow is None: - assert stored_table(tmp_project_root).num_rows == rows - return assert (frame_dir(tmp_project_root) / "manifest.json").read_bytes() == left[ "manifest_bytes" ] From ca87765da25d7d6a80c242634881d9e57e13d164 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 19:14:46 +0000 Subject: [PATCH 19/41] test(artifacts): return the Manifest model from stored_manifest stored_manifest validated nothing and returned a Dict[str, Any]. It now validates the written JSON with the Manifest model and returns it, so a field of the wrong type fails there, naming the field, and callers read typed attributes. "previous" not in X became X.previous is None: on a pydantic model, in checks against (name, value) pairs and would always pass. That no longer tells an omitted previous from null, and the model ignores unknown keys, so test_stores_pandas_frame_and_manifest checks the written keys from the raw JSON. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_delete_dataframe.py | 2 +- tests/unit/artifacts/test_file_format.py | 8 +++---- .../artifacts/test_manifest_and_cleanup.py | 22 ++++++++--------- tests/unit/artifacts/test_project_root.py | 2 +- tests/unit/artifacts/test_storing.py | 24 +++++++++++-------- tests/unit/artifacts/test_write_failures.py | 10 ++++---- tests/unit/helpers/dataframe_storage.py | 12 ++++++---- 7 files changed, 44 insertions(+), 36 deletions(-) diff --git a/tests/unit/artifacts/test_delete_dataframe.py b/tests/unit/artifacts/test_delete_dataframe.py index 1e09d941..6e20634c 100644 --- a/tests/unit/artifacts/test_delete_dataframe.py +++ b/tests/unit/artifacts/test_delete_dataframe.py @@ -53,7 +53,7 @@ def test_a_name_can_be_stored_again_after_it_was_deleted( artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - assert "previous" not in stored_manifest(tmp_project_root) + assert stored_manifest(tmp_project_root).previous is None assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index dea74805..362db49d 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -25,7 +25,7 @@ def test_frames_are_compressed(tmp_project_root: Path, make_frame: Any) -> None: artifacts.store_dataframe(make_frame({"a": ["x" * 100] * 10_000}), NAME) - stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["file"] + stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file assert 0 < stored.stat().st_size < 100_000 @@ -72,7 +72,7 @@ def test_frame_is_stored_uncompressed_when_no_codec_is_available( artifacts.store_dataframe(make_frame({"a": ["x" * 100] * 10_000}), NAME) - stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["file"] + stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file assert pa.ipc.open_file(stored).read_all().num_rows == 10_000 assert stored.stat().st_size > 1_000_000 @@ -81,7 +81,7 @@ def test_pandas_frame_is_written_in_64k_row_batches(tmp_project_root: Path) -> N artifacts.store_dataframe(pd.DataFrame({"a": range(70_000)}), NAME) reader = pa.ipc.open_file( - frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["file"] + frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file ) assert reader.num_record_batches == 2 assert reader.read_all().num_rows == 70_000 @@ -91,7 +91,7 @@ def test_polars_frame_is_written_in_64k_row_batches(tmp_project_root: Path) -> N artifacts.store_dataframe(pl.DataFrame({"a": range(200_000)}), NAME) reader = pa.ipc.open_file( - frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["file"] + frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file ) assert reader.num_record_batches == 4 assert reader.read_all().num_rows == 200_000 diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index da01fe02..8541c728 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -36,14 +36,14 @@ def test_second_write_keeps_one_previous_frame(tmp_project_root: Path) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) third = stored_manifest(tmp_project_root) - assert "previous" not in first - assert second["previous"] == first["file"] - assert third["previous"] == second["file"] + assert first.previous is None + assert second.previous == first.file + assert third.previous == second.file assert [path.name for path in arrow_files(tmp_project_root)] == sorted( - [second["file"], third["file"]] + [second.file, third.file] ) assert stored_table(tmp_project_root).to_pydict() == {"a": [3]} - previous = pa.ipc.open_file(frame_dir(tmp_project_root) / third["previous"]) + previous = pa.ipc.open_file(frame_dir(tmp_project_root) / third.previous) assert previous.read_all().to_pydict() == {"a": [2]} @@ -52,7 +52,7 @@ def test_unreferenced_frames_are_deleted_only_when_stale( ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) directory = frame_dir(tmp_project_root) - replaced = directory / stored_manifest(tmp_project_root)["file"] + replaced = directory / stored_manifest(tmp_project_root).file backdate(replaced, 2 * DAY) def add(name: str, age_seconds: float) -> Path: @@ -175,8 +175,8 @@ def test_unusable_manifest_counts_as_no_manifest( artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) manifest = stored_manifest(tmp_project_root) - assert "previous" not in manifest - assert manifest["file"] not in (FRAME_A, FRAME_B) + assert manifest.previous is None + assert manifest.file not in (FRAME_A, FRAME_B) assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} assert sentinel.read_text() == "keep" assert all(path.exists() for path in named) @@ -206,7 +206,7 @@ def test_usable_manifest_is_honoured( artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - assert stored_manifest(tmp_project_root)["previous"] == FRAME_A + assert stored_manifest(tmp_project_root).previous == FRAME_A assert (directory / FRAME_A).exists() assert not (directory / FRAME_B).exists() @@ -230,7 +230,7 @@ def too_deep(raw: bytes) -> Any: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - assert "previous" not in stored_manifest(tmp_project_root) + assert stored_manifest(tmp_project_root).previous is None assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} @@ -282,7 +282,7 @@ def write(value: int) -> None: assert errors == [] files = arrow_files(tmp_project_root) assert len(files) == 2 - current = stored_manifest(tmp_project_root)["file"] + current = stored_manifest(tmp_project_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 diff --git a/tests/unit/artifacts/test_project_root.py b/tests/unit/artifacts/test_project_root.py index d5cbc566..d425ac38 100644 --- a/tests/unit/artifacts/test_project_root.py +++ b/tests/unit/artifacts/test_project_root.py @@ -20,7 +20,7 @@ def test_store_creates_a_missing_project_root( reference_id = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - assert stored_manifest(project)["file"] == f"{reference_id}.arrow" + assert stored_manifest(project).file == f"{reference_id}.arrow" def _broken_config() -> Any: diff --git a/tests/unit/artifacts/test_storing.py b/tests/unit/artifacts/test_storing.py index 2486d612..32bed08d 100644 --- a/tests/unit/artifacts/test_storing.py +++ b/tests/unit/artifacts/test_storing.py @@ -1,6 +1,7 @@ """Tests for what the artifacts API writes when it stores a named DataFrame.""" import copy +import json import re from datetime import datetime, timezone from pathlib import Path @@ -15,6 +16,7 @@ from tests.unit.helpers.dataframe_storage import ( NAME, arrow_files, + frame_dir, stored_manifest, stored_table, ) @@ -36,13 +38,15 @@ def test_stores_pandas_frame_and_manifest(tmp_project_root: Path) -> None: "a": [1, 2, 3], "b": ["x", "y", "z"], } + # Manifest ignores unknown keys and defaults previous, so the keys are checked raw. + written = json.loads((frame_dir(tmp_project_root) / "manifest.json").read_text()) + assert set(written) == {"version", "file", "written_at", "rows", "toolkit_version"} manifest = stored_manifest(tmp_project_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"]) + 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") @@ -52,7 +56,7 @@ def test_manifest_time_is_utc(tmp_project_root: Path) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) written = datetime.strptime( - stored_manifest(tmp_project_root)["written_at"], "%Y-%m-%dT%H:%M:%SZ" + stored_manifest(tmp_project_root).written_at, "%Y-%m-%dT%H:%M:%SZ" ) assert abs((written.replace(tzinfo=timezone.utc) - before).total_seconds()) < 60 @@ -62,12 +66,12 @@ def test_stores_polars_frame(tmp_project_root: Path) -> None: pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}), NAME ) - assert stored_manifest(tmp_project_root)["file"] == f"{reference_id}.arrow" + assert stored_manifest(tmp_project_root).file == f"{reference_id}.arrow" assert stored_table(tmp_project_root).to_pydict() == { "a": [1, 2, 3], "b": ["x", "y", "z"], } - assert stored_manifest(tmp_project_root)["rows"] == 3 + assert stored_manifest(tmp_project_root).rows == 3 def test_writing_does_not_change_the_callers_frame(tmp_project_root: Path) -> None: @@ -125,7 +129,7 @@ def test_empty_frame_is_stored(tmp_project_root: Path, frame: Any) -> None: artifacts.store_dataframe(frame, NAME) assert stored_table(tmp_project_root).num_rows == 0 - assert stored_manifest(tmp_project_root)["rows"] == 0 + assert stored_manifest(tmp_project_root).rows == 0 def test_names_are_stored_apart(tmp_project_root: Path) -> None: diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index b75fcdcf..93838947 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -75,7 +75,7 @@ def test_unserializable_frame_raises_and_keeps_the_previous_frame( "manifest_bytes" ] assert left["listing"] == sorted( - ["manifest.json", stored_manifest(tmp_project_root)["file"]] + ["manifest.json", stored_manifest(tmp_project_root).file] ) assert stored_table(tmp_project_root).to_pydict() == {"ok": [1]} @@ -98,7 +98,7 @@ def test_frames_pyarrow_may_refuse_follow_what_pyarrow_does( "manifest_bytes" ] assert left["listing"] == sorted( - ["manifest.json", stored_manifest(tmp_project_root)["file"]] + ["manifest.json", stored_manifest(tmp_project_root).file] ) @@ -301,9 +301,9 @@ def test_failure_after_the_manifest_step_keeps_the_new_frame( ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - replaced = ( - frame_dir(tmp_project_root) / stored_manifest(tmp_project_root)["previous"] - ) + previous = stored_manifest(tmp_project_root).previous + assert previous is not None + replaced = frame_dir(tmp_project_root) / previous real_unlink = Path.unlink def unlink_failing_for_the_replaced_frame( diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 28926bec..ff582369 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -7,7 +7,7 @@ import time import uuid from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, List, Optional import numpy as np import pandas as pd @@ -15,6 +15,9 @@ import pytest from pyarrow import feather +from deepnote_core.pydantic_compat_helpers import model_validate_compat +from deepnote_toolkit.artifacts.dataframe_storage_manifest import Manifest + NAME = "revenue" DAY = 24 * 60 * 60 FRAME_A = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee.arrow" @@ -29,13 +32,14 @@ def arrow_files(root: Path, name: str = NAME) -> List[Path]: return sorted(frame_dir(root, name).glob("*.arrow")) -def stored_manifest(root: Path, name: str = NAME) -> Dict[str, Any]: - return json.loads((frame_dir(root, name) / "manifest.json").read_text()) +def stored_manifest(root: Path, name: str = NAME) -> Manifest: + raw = json.loads((frame_dir(root, name) / "manifest.json").read_text()) + return model_validate_compat(Manifest, raw) def stored_table(root: Path, name: str = NAME) -> pa.Table: """The frame the name's manifest currently names.""" - path = frame_dir(root, name) / stored_manifest(root, name)["file"] + path = frame_dir(root, name) / stored_manifest(root, name).file return pa.ipc.open_file(path).read_all() From 1c89789f5530d6ffb110ff9f299953f0cc8b0975 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 15 Oct 2026 19:15:12 +0000 Subject: [PATCH 20/41] test(artifacts): check the accepted frame replaced the good one When pyarrow accepts a version-dependent frame, the test only compared row counts. The uuid case has one row, as does the good frame stored before it, so a store that silently kept the previous frame passed. The test now also checks that the stored table is not the good frame. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_write_failures.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 93838947..9f4fa1fc 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -88,7 +88,9 @@ def test_frames_pyarrow_may_refuse_follow_what_pyarrow_does( if raised_by_pyarrow is None: _store_frame_after_a_good_one(tmp_project_root, make_frame) - assert stored_table(tmp_project_root).num_rows == rows + stored = stored_table(tmp_project_root) + assert stored.column_names != ["ok"] + assert stored.num_rows == rows return left = _store_bad_frame_after_a_good_one( tmp_project_root, make_frame, raised_by_pyarrow From 6cd3bb9c55384e5aef780475cdf1dca801351b54 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 06:50:05 +0000 Subject: [PATCH 21/41] test(artifacts): check that EROFS from the frame write leaves no manifest The test claimed the store left nothing but checked only the frame files, so a store that wrote the manifest before the frame still passed it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_read_only_mount.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/artifacts/test_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py index a0c45e74..35808012 100644 --- a/tests/unit/artifacts/test_read_only_mount.py +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -63,6 +63,7 @@ def refuse(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: assert result is None assert arrow_files(tmp_project_root) == [] + assert not (frame_dir(tmp_project_root) / "manifest.json").exists() def test_delete_on_a_read_only_mount_deletes_nothing( From 398fe374d4025045cb17a00b1e4fdbf6b1bc6c9e Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 06:50:05 +0000 Subject: [PATCH 22/41] test(artifacts): check that a failed manifest rename leaves no temporary file This is the only test where the rename fails, so nothing caught a write_manifest whose cleanup covered only the write and not the rename. That version leaves .manifest.json.*.tmp behind for good, because the cleanup prunes only .arrow names. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- tests/unit/artifacts/test_write_failures.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 9f4fa1fc..dfd9869f 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -295,6 +295,7 @@ def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( assert raised.value.errno == errno.EIO assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} + assert entries_beside_the_frames(tmp_project_root) == ["manifest.json"] @pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["EIO", "EROFS"]) From c042c29f89e7e4990c77eef5ef13c2c9f1c5d8da Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 06:55:11 +0000 Subject: [PATCH 23/41] refactor(artifacts): replace entries_beside_the_frames with non_arrow_files in tests --- tests/unit/artifacts/test_manifest_and_cleanup.py | 4 ++-- tests/unit/artifacts/test_write_failures.py | 6 +++--- tests/unit/helpers/dataframe_storage.py | 10 ++++------ 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 8541c728..c9d7ddd6 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -21,8 +21,8 @@ NAME, arrow_files, backdate, - entries_beside_the_frames, frame_dir, + non_arrow_files, stored_manifest, stored_table, ) @@ -301,7 +301,7 @@ def test_successful_write_leaves_only_the_manifest_beside_the_frames( artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) - assert entries_beside_the_frames(tmp_project_root) == ["manifest.json"] + assert non_arrow_files(tmp_project_root) == ["manifest.json"] def test_stray_pruned_by_another_writer_meanwhile_is_not_a_failure( diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index dfd9869f..087a1455 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -23,8 +23,8 @@ VERSION_DEPENDENT, arrow_files, backdate, - entries_beside_the_frames, frame_dir, + non_arrow_files, probe_failure, stored_manifest, stored_table, @@ -275,7 +275,7 @@ def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: frame_dir(tmp_project_root) / "manifest.json" ).read_bytes() == manifest_bytes assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} - assert entries_beside_the_frames(tmp_project_root) == ["manifest.json"] + assert non_arrow_files(tmp_project_root) == ["manifest.json"] def _copy_then_fail(source: str, destination: Path) -> None: @@ -295,7 +295,7 @@ def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( assert raised.value.errno == errno.EIO assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} - assert entries_beside_the_frames(tmp_project_root) == ["manifest.json"] + assert non_arrow_files(tmp_project_root) == ["manifest.json"] @pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["EIO", "EROFS"]) diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index ff582369..d8f0f8a2 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -31,6 +31,10 @@ def frame_dir(root: Path, name: str = NAME) -> Path: def arrow_files(root: Path, name: str = NAME) -> List[Path]: return sorted(frame_dir(root, name).glob("*.arrow")) +def non_arrow_files(root: Path) -> List[str]: + return sorted( + path.name for path in frame_dir(root).iterdir() if path.suffix != ".arrow" + ) def stored_manifest(root: Path, name: str = NAME) -> Manifest: raw = json.loads((frame_dir(root, name) / "manifest.json").read_text()) @@ -74,12 +78,6 @@ def mkdir( monkeypatch.setattr(os, "rmdir", refuse) -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 probe_failure(frame: Any) -> Optional[type]: """The exception type pyarrow itself raises for this frame, if any.""" try: From 95760d4d2aebb201e155427f1e5ec1bacb115bd9 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 07:05:07 +0000 Subject: [PATCH 24/41] Reformat code --- tests/unit/helpers/dataframe_storage.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index d8f0f8a2..52355c96 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -31,11 +31,13 @@ def frame_dir(root: Path, name: str = NAME) -> Path: def arrow_files(root: Path, name: str = NAME) -> List[Path]: return sorted(frame_dir(root, name).glob("*.arrow")) + def non_arrow_files(root: Path) -> List[str]: return sorted( path.name for path in frame_dir(root).iterdir() if path.suffix != ".arrow" ) + def stored_manifest(root: Path, name: str = NAME) -> Manifest: raw = json.loads((frame_dir(root, name) / "manifest.json").read_text()) return model_validate_compat(Manifest, raw) From 5b149c15b47c6717d5437c1b2f5e1f764bd46cec Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 12:21:58 +0000 Subject: [PATCH 25/41] refactor(artifacts): pass the typed frame to store_frame instead of a kind classify_frame mapped a frame to a "pandas" or "polars" string, and store_frame and _write_frame_file trusted that string without checking it against the frame. store_dataframe now checks the frame with the ocelots TypeGuards and passes on only the narrowed frame, typed as StorableFrame. _write_frame_file dispatches on it with isinstance and ends in assert_never. classify_frame also refused query previews while is_pandas_dataframe listed DeepnoteQueryPreview as pandas, and store_dataframe checked for a preview a second time to pick its error message. The refusal now lives only in store_dataframe. DeepnoteQueryPreview subclasses pd.DataFrame, so dropping it from is_pandas_dataframe changes no result. Stored files are byte-identical before and after, and every refused input raises the same TypeError. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage.py | 43 ++++++++----------- deepnote_toolkit/artifacts/dataframes.py | 23 +++++----- deepnote_toolkit/ocelots/utils.py | 3 +- 3 files changed, 31 insertions(+), 38 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index a8f70d90..b9e2078d 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -7,16 +7,14 @@ import time import uuid from pathlib import Path -from typing import Any, Literal, Optional +from typing import Any, Optional, Union +import pandas as pd from packaging.version import Version +from typing_extensions import TypeAlias, assert_never from deepnote_toolkit.config import resolve_notebook_root -from deepnote_toolkit.ocelots.utils import ( - is_pandas_dataframe, - is_polars_eager_dataframe, -) -from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview +from deepnote_toolkit.ocelots.types import PandasDF, PolarsEagerDF from .dataframe_storage_manifest import ( Manifest, @@ -31,7 +29,7 @@ POLARS_BATCH_MIN_VERSION = Version("1.37") POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") -FrameKind = Literal["pandas", "polars"] +StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF] _SAFE_ID_REGEX = re.compile(r"[A-Za-z0-9_-]{1,128}") @@ -41,17 +39,6 @@ def is_safe_id(value: object) -> bool: return isinstance(value, str) and _SAFE_ID_REGEX.fullmatch(value) is not None -def classify_frame(value: object) -> Optional[FrameKind]: - """Return "pandas" or "polars" for a frame worth storing, else None.""" - if isinstance(value, DeepnoteQueryPreview): - return None - if is_pandas_dataframe(value): - return "pandas" - if is_polars_eager_dataframe(value): - return "polars" - return None - - def resolve_project_root() -> Path: """Return the project's root directory, the one set_notebook_path() uses.""" return resolve_notebook_root() @@ -67,23 +54,24 @@ def _select_compression() -> str: return "uncompressed" -def _write_frame_file( - frame: Any, kind: FrameKind, path: Path, compression: str -) -> None: +def _write_frame_file(frame: StorableFrame, path: Path, compression: str) -> None: """Write the frame to `path` as an Arrow IPC file. 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": + if isinstance(frame, pd.DataFrame): from pyarrow import feather # pandas < 2.1's to_feather needs a default index and string column names. feather.write_feather(frame, sink, compression=compression) - else: - import polars + return + + # polars is an optional dependency, so import it only for a non-pandas frame. + import polars + if isinstance(frame, polars.DataFrame): polars_version = Version(polars.__version__) options: dict[str, Any] = {} # Cap batches like write_feather so a page read stays cheap. @@ -93,6 +81,9 @@ def _write_frame_file( if polars_version >= POLARS_COMPAT_LEVEL_MIN_VERSION: options["compat_level"] = polars.CompatLevel.oldest() frame.write_ipc(sink, compression=compression, **options) + return + + assert_never(frame) def _delete_replaced_frames( @@ -123,7 +114,7 @@ def _delete_replaced_frames( (directory / entry.name).unlink() -def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> Optional[str]: +def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]: """Store the frame as the current one under `name` and return its reference ID. Returns None, having stored nothing, if the filesystem is read-only. EROFS from @@ -134,7 +125,7 @@ def store_frame(frame: Any, kind: FrameKind, name: str, root: Path) -> Optional[ new_file = f"{reference_id}.arrow" try: directory.mkdir(parents=True, exist_ok=True) - _write_frame_file(frame, kind, directory / new_file, _select_compression()) + _write_frame_file(frame, directory / new_file, _select_compression()) old = read_manifest(directory) except BaseException as exc: with contextlib.suppress(OSError): diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index 23bede8f..46e4ccfe 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -4,6 +4,10 @@ import shutil from typing import Any, Optional +from deepnote_toolkit.ocelots.utils import ( + is_pandas_dataframe, + is_polars_eager_dataframe, +) from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview from . import dataframe_storage @@ -43,22 +47,21 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: already put the new manifest in place. """ _check_name(name) - kind = dataframe_storage.classify_frame(df) - if kind is None: - if isinstance(df, DeepnoteQueryPreview): - raise TypeError( - "store_dataframe refuses a SQL query preview or a frame derived from " - "one, which holds at most 100 rows of the query's result. Set the SQL " - "block's output mode to DataFrame, or pass pd.DataFrame(df) to store " - "these rows" - ) + if isinstance(df, DeepnoteQueryPreview): + raise TypeError( + "store_dataframe refuses a SQL query preview or a frame derived from " + "one, which holds at most 100 rows of the query's result. Set the SQL " + "block's output mode to DataFrame, or pass pd.DataFrame(df) to store " + "these rows" + ) + if not (is_pandas_dataframe(df) or is_polars_eager_dataframe(df)): raise TypeError( "store_dataframe needs a pandas or polars DataFrame, " f"got {type(df).__qualname__}" ) root = dataframe_storage.resolve_project_root() - return dataframe_storage.store_frame(df, kind, name, root) + return dataframe_storage.store_frame(df, name, root) def delete_dataframe(name: str) -> bool: diff --git a/deepnote_toolkit/ocelots/utils.py b/deepnote_toolkit/ocelots/utils.py index 685f1ff1..4e25bbc5 100644 --- a/deepnote_toolkit/ocelots/utils.py +++ b/deepnote_toolkit/ocelots/utils.py @@ -9,11 +9,10 @@ PolarsEagerDF, PysparkDF, ) -from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview def is_pandas_dataframe(df) -> TypeGuard[PandasDF]: - return isinstance(df, (pd.DataFrame, DeepnoteQueryPreview)) + return isinstance(df, pd.DataFrame) def is_pyspark_dataframe(df) -> TypeGuard[PysparkDF]: From a68c0c625e2e82d8f4ded6798f62491077d87582 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 12:21:58 +0000 Subject: [PATCH 26/41] chore(mypy): enable arg-type for the artifacts package The global config disables arg-type, the code mypy reports a non-exhaustive assert_never under, so a missing branch in _write_frame_file's isinstance chain went unreported. pandas has no stubs here, so PandasDF is Any and only the polars branch is checked. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 447f2e5e..f22244cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -339,6 +339,11 @@ exclude = [ "(^|[\\/])deepnote_core[\\/]resources[\\/]" ] +# An isinstance chain that misses a typed union member fails its assert_never with arg-type. +[[tool.mypy.overrides]] +module = ["deepnote_toolkit.artifacts.*"] +enable_error_code = ["arg-type"] + [tool.coverage.run] branch = true source = ["deepnote_toolkit", "installer", "deepnote_core"] From 1b3e28161ecb9a61327974a80f30da04188999ef Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 12:21:58 +0000 Subject: [PATCH 27/41] test(artifacts): rename the failing-close fixture after the s3fs upload handle --- tests/unit/artifacts/conftest.py | 12 ++++++------ tests/unit/artifacts/test_write_failures.py | 6 ++++-- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index 936e4ef4..0f2268dd 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -2,7 +2,7 @@ import time from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional +from typing import IO, Any, Dict, Iterator, List, Optional import pytest from pyspark.sql import SparkSession @@ -69,10 +69,10 @@ def open_and_record(self: Path, mode: str = "r", *args: Any, **kwargs: Any) -> A return sinks -class _FailingCloseS3Upload: +class _FailingCloseS3UploadFileHandle: """A write handle whose close() really closes, then fails like s3fs's upload.""" - def __init__(self, handle: Any, error: OSError) -> None: + def __init__(self, handle: IO[Any], error: OSError) -> None: self._handle = handle self._error = error @@ -80,7 +80,7 @@ def close(self) -> None: self._handle.close() raise self._error - def __enter__(self) -> "_FailingCloseS3Upload": + def __enter__(self) -> "_FailingCloseS3UploadFileHandle": return self def __exit__(self, *exc_info: Any) -> None: @@ -91,7 +91,7 @@ def __getattr__(self, name: str) -> Any: @pytest.fixture -def close_error(monkeypatch: pytest.MonkeyPatch) -> Dict[str, Optional[OSError]]: +def close_upload_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 @@ -103,7 +103,7 @@ def open_failing_on_close( error = state["error"] if error is None or self.suffix != ".arrow" or "w" not in mode: return handle - return _FailingCloseS3Upload(handle, error) + return _FailingCloseS3UploadFileHandle(handle, error) monkeypatch.setattr(Path, "open", open_failing_on_close) return state diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 087a1455..0c611137 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -150,10 +150,12 @@ def write_some_then_fail(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: ids=["EPERM", "EACCES", "EIO", "ENOSPC"], ) def test_error_from_close_raises_and_stores_nothing( - tmp_project_root: Path, close_error: Dict[str, Optional[OSError]], error: OSError + tmp_project_root: Path, + close_upload_error: Dict[str, Optional[OSError]], + error: OSError, ) -> None: """s3fs surfaces a failed upload only when the file is closed.""" - close_error["error"] = error + close_upload_error["error"] = error with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) From bd3b7424fce08d904171cab1e9d79a6bb7edabbe Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 12:21:58 +0000 Subject: [PATCH 28/41] test(artifacts): drop the match= note from the symlink delete test --- tests/unit/artifacts/test_delete_dataframe.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/artifacts/test_delete_dataframe.py b/tests/unit/artifacts/test_delete_dataframe.py index 6e20634c..291d1ee9 100644 --- a/tests/unit/artifacts/test_delete_dataframe.py +++ b/tests/unit/artifacts/test_delete_dataframe.py @@ -71,7 +71,5 @@ def test_delete_does_not_follow_a_link_out_of_the_frames_directory( with pytest.raises(OSError) as raised: artifacts.delete_dataframe(NAME) - # Not match=: newer Pythons (3.12.x and up) set this error's filename, and str() - # then shows "[Errno None] None: " instead of the message. assert raised.value.args == ("Cannot call rmtree on a symbolic link",) assert (outside / "keep.txt").read_text() == "keep" From 6613a6f2fdd67c012f7e67c0a88cd2ed15119358 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 12:21:58 +0000 Subject: [PATCH 29/41] chore(vscode): recommend the Python extensions and use poetry environments --- .vscode/extensions.json | 10 +++++++++- .vscode/settings.json | 4 +++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/.vscode/extensions.json b/.vscode/extensions.json index 2d96d383..b809de4d 100644 --- a/.vscode/extensions.json +++ b/.vscode/extensions.json @@ -1,3 +1,11 @@ { - "recommendations": ["ms-python.isort"] + "recommendations": [ + "ms-python.python", + "ms-python.vscode-python-envs", + "ms-python.debugpy", + "ms-python.vscode-pylance", + "ms-python.isort", + "ms-python.black-formatter", + "redhat.vscode-yaml" + ] } diff --git a/.vscode/settings.json b/.vscode/settings.json index 580d0411..ecf3c95f 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -22,5 +22,7 @@ ], "python.analysis.extraPaths": [ "./installer" - ] + ], + "python-envs.defaultEnvManager": "ms-python.python:poetry", + "python-envs.defaultPackageManager": "ms-python.python:poetry" } From e29aabb3857a8b2c113f194a40bffd8e04d5d2c9 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 13:18:03 +0000 Subject: [PATCH 30/41] refactor(artifacts): pass the parsed manifest straight to pydantic pydantic ignores unknown keys by default on both v1 and v2. The filter to the model's fields only worked around v1's parse_obj, which calls cls(**obj): a key named __pydantic_self__, or a non-string key from a list of pairs, makes it raise TypeError rather than ValidationError. read_manifest now catches TypeError instead. The isinstance(dict) guard is gone too. On v2 every result is unchanged. On v1 a list of pairs is now honoured, still through every field and file-name check, and a __pydantic_self__ key counts as absent. The unknown-key-named-like-self case only covered the removed workaround, and on v2 it duplicates unknown-key. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/dataframe_storage_manifest.py | 17 +++-------------- .../unit/artifacts/test_manifest_and_cleanup.py | 3 --- 2 files changed, 3 insertions(+), 17 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index 86d2f7b2..3478ac25 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -13,7 +13,6 @@ from deepnote_core.pydantic_compat_helpers import ( _model_dump_compat, - _model_fields_compat, model_validate_compat, ) from deepnote_toolkit import __version__ @@ -34,10 +33,7 @@ def is_frame_name(name: str) -> bool: class Manifest(BaseModel): - """The manifest file; `previous` is the frame `file` replaced. - - Unknown keys are ignored. - """ + """The manifest file; `previous` is the frame `file` replaced.""" version: StrictInt file: StrictStr @@ -64,16 +60,9 @@ def read_manifest(directory: Path) -> Optional[Manifest]: data = json.loads(raw) except (ValueError, RecursionError): return None - # pydantic v1 would turn a list of pairs into a dict. - if not isinstance(data, dict): - return None - # pydantic v1 passes every key to __init__, where __pydantic_self__ collides - # with its first parameter and raises TypeError. - fields = _model_fields_compat(Manifest) - known = {key: value for key, value in data.items() if key in fields} try: - manifest = model_validate_compat(Manifest, known) - except ValidationError: + manifest = model_validate_compat(Manifest, data) + except (ValidationError, TypeError): return None if manifest.version != MANIFEST_VERSION or manifest.rows < 0: diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index c9d7ddd6..51ac7589 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -189,9 +189,6 @@ def test_unusable_manifest_counts_as_no_manifest( pytest.param( _manifest_bytes(written_by_a_newer_toolkit=True), id="unknown-key" ), - pytest.param( - _manifest_bytes(__pydantic_self__=1), id="unknown-key-named-like-self" - ), ], ) def test_usable_manifest_is_honoured( From d31c16055f69db2e740c4d1de566ae4a902162fc Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 13:18:03 +0000 Subject: [PATCH 31/41] refactor(core): infer model_validate_compat's return type from the model class Callers got Any back, so mypy checked nothing they did with the result. The TypeVar is bound to BaseModel, so passing a class that isn't a model is a type-var error, which the repo's mypy config reports. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- deepnote_core/pydantic_compat_helpers.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/deepnote_core/pydantic_compat_helpers.py b/deepnote_core/pydantic_compat_helpers.py index 8d42b06d..94f87fe1 100644 --- a/deepnote_core/pydantic_compat_helpers.py +++ b/deepnote_core/pydantic_compat_helpers.py @@ -1,4 +1,8 @@ -from typing import Any, Mapping +from typing import Any, Mapping, TypeVar + +from pydantic import BaseModel + +Model = TypeVar("Model", bound=BaseModel) def _model_dump_compat(model_cls) -> dict: @@ -45,7 +49,7 @@ def _get_description_compat(field: Any) -> str: return getattr(field_info, "description", "") if field_info else "" -def model_validate_compat(model_cls, data: Mapping[str, Any]) -> Any: +def model_validate_compat(model_cls: type[Model], data: Mapping[str, Any]) -> Model: """Validate using Pydantic v2 if available, else fall back to v1. Pydantic v2 exposes `model_validate`, while v1 uses `parse_obj` on the class. From 6c1fe413fe177b3e984ce061d84b1a197a903ea3 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 14:19:04 +0000 Subject: [PATCH 32/41] refactor(artifacts): simplify manifest validation in read_manifest function --- deepnote_toolkit/artifacts/dataframe_storage.py | 6 +----- deepnote_toolkit/artifacts/dataframe_storage_manifest.py | 6 ------ 2 files changed, 1 insertion(+), 11 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index b9e2078d..b78830ad 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -55,11 +55,7 @@ def _select_compression() -> str: def _write_frame_file(frame: StorableFrame, path: Path, compression: str) -> None: - """Write the frame to `path` as an Arrow IPC file. - - Going through a Python file object makes `close()` raise upload errors, which - s3fs only surfaces there. - """ + """Write the frame to `path` as an Arrow IPC file.""" with path.open("wb") as sink: if isinstance(frame, pd.DataFrame): from pyarrow import feather diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index 3478ac25..bdf94c3e 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -65,12 +65,6 @@ def read_manifest(directory: Path) -> Optional[Manifest]: except (ValidationError, TypeError): return None - if manifest.version != MANIFEST_VERSION or manifest.rows < 0: - return None - if not is_frame_name(manifest.file): - return None - if manifest.previous is not None and not is_frame_name(manifest.previous): - return None return manifest From e31d9a0c518fb478b558a2571f89c0199f029565 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 15:29:43 +0000 Subject: [PATCH 33/41] test(artifacts): drop the manifest cases for the removed version, rows and name checks 6c1fe41 removed the version, rows and file-name checks from read_manifest, so a manifest with version 2, negative rows, or a file or previous that isn't .arrow is now honoured. The 12 cases that expected it to count as absent go, and so does the ../../data/x.csv sentinel only the traversal cases used. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .../artifacts/test_manifest_and_cleanup.py | 30 +------------------ 1 file changed, 1 insertion(+), 29 deletions(-) diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 51ac7589..17ed9b43 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -112,39 +112,15 @@ def _manifest_bytes(without: str = "", **fields: Any) -> bytes: json.dumps(list(_VALID_MANIFEST.items())).encode(), id="not-an-object" ), pytest.param(_manifest_bytes(without="version"), id="no-version"), - pytest.param(_manifest_bytes(version=2), id="unsupported-version"), pytest.param(_manifest_bytes(version="1"), id="version-a-string"), pytest.param(_manifest_bytes(version=True), id="version-a-bool"), pytest.param(_manifest_bytes(without="file"), id="no-file"), - pytest.param(_manifest_bytes(file="../../data/x.csv"), id="file-traversal"), pytest.param(_manifest_bytes(file=7), id="file-not-a-string"), - pytest.param( - _manifest_bytes(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(previous=f"x{FRAME_B}"), id="previous-prefixed"), - pytest.param( - _manifest_bytes(previous=f"{FRAME_B}.bak"), id="previous-suffixed" - ), - pytest.param( - _manifest_bytes(previous=f"{FRAME_B}/../../data/x.csv"), - id="previous-valid-name-then-traversal", - ), - pytest.param( - _manifest_bytes(previous=f"../{FRAME_B}"), id="previous-in-parent" - ), pytest.param(_manifest_bytes(without="written_at"), id="no-written-at"), pytest.param(_manifest_bytes(written_at=7), id="written-at-not-a-string"), pytest.param(_manifest_bytes(without="rows"), id="no-rows"), pytest.param(_manifest_bytes(rows="1"), id="rows-a-string"), pytest.param(_manifest_bytes(rows=1.5), id="rows-not-whole"), - pytest.param(_manifest_bytes(rows=-1), id="rows-negative"), pytest.param( _manifest_bytes(without="toolkit_version"), id="no-toolkit-version" ), @@ -167,10 +143,7 @@ def test_unusable_manifest_counts_as_no_manifest( (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") + path.write_bytes(b"named by the unusable manifest") artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -178,7 +151,6 @@ def test_unusable_manifest_counts_as_no_manifest( assert manifest.previous is None assert manifest.file not in (FRAME_A, FRAME_B) assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} - assert sentinel.read_text() == "keep" assert all(path.exists() for path in named) From 77f8d94f270b16d8bca252fb91796f21b4370829 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 16 Oct 2026 19:49:23 +0000 Subject: [PATCH 34/41] chore(vscode): keep the editor settings out of this branch Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01FhzfrytCUojy8Vw1gbX3Bx --- .vscode/extensions.json | 10 +--------- .vscode/settings.json | 4 +--- 2 files changed, 2 insertions(+), 12 deletions(-) diff --git a/.vscode/extensions.json b/.vscode/extensions.json index b809de4d..2d96d383 100644 --- a/.vscode/extensions.json +++ b/.vscode/extensions.json @@ -1,11 +1,3 @@ { - "recommendations": [ - "ms-python.python", - "ms-python.vscode-python-envs", - "ms-python.debugpy", - "ms-python.vscode-pylance", - "ms-python.isort", - "ms-python.black-formatter", - "redhat.vscode-yaml" - ] + "recommendations": ["ms-python.isort"] } diff --git a/.vscode/settings.json b/.vscode/settings.json index ecf3c95f..580d0411 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -22,7 +22,5 @@ ], "python.analysis.extraPaths": [ "./installer" - ], - "python-envs.defaultEnvManager": "ms-python.python:poetry", - "python-envs.defaultPackageManager": "ms-python.python:poetry" + ] } From 5a1ea1d044dae200fea1bc711ad4fab906274240 Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 18 Oct 2026 20:14:28 +0000 Subject: [PATCH 35/41] fix(artifacts): count a manifest that names anything but a frame file as absent Cleanup unlinks the replaced manifest's previous. Since 6c1fe41 removed the name checks, a manifest with "previous": "../../../outside.txt" made the next store_dataframe delete a file outside the project. Any code in the project can write the manifest. read_manifest again requires file and previous to be .arrow. Checking file too stops the toolkit from writing a previous that its own next read would reject. The file-traversal and previous-traversal cases come back, with a sentinel outside the project that must survive. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VD15J21xfHn4E8p1nZAPFB --- deepnote_toolkit/artifacts/dataframe_storage_manifest.py | 6 ++++++ tests/unit/artifacts/test_manifest_and_cleanup.py | 8 ++++++++ 2 files changed, 14 insertions(+) diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index bdf94c3e..6ef37f86 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -65,6 +65,12 @@ def read_manifest(directory: Path) -> Optional[Manifest]: except (ValidationError, TypeError): return None + # Cleanup unlinks `previous`, and `file` becomes the next `previous`, so any + # other name could delete a file outside this directory. + if not is_frame_name(manifest.file): + return None + if manifest.previous is not None and not is_frame_name(manifest.previous): + return None return manifest diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 17ed9b43..91c06686 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -116,6 +116,10 @@ def _manifest_bytes(without: str = "", **fields: Any) -> bytes: pytest.param(_manifest_bytes(version=True), id="version-a-bool"), pytest.param(_manifest_bytes(without="file"), id="no-file"), pytest.param(_manifest_bytes(file=7), id="file-not-a-string"), + pytest.param(_manifest_bytes(file="../../data/x.csv"), id="file-traversal"), + pytest.param( + _manifest_bytes(previous="../../data/x.csv"), id="previous-traversal" + ), pytest.param(_manifest_bytes(without="written_at"), id="no-written-at"), pytest.param(_manifest_bytes(written_at=7), id="written-at-not-a-string"), pytest.param(_manifest_bytes(without="rows"), id="no-rows"), @@ -144,6 +148,9 @@ def test_unusable_manifest_counts_as_no_manifest( named = [directory / FRAME_A, directory / FRAME_B] for path in named: path.write_bytes(b"named by the unusable manifest") + sentinel = (directory / "../../data/x.csv").resolve() + sentinel.parent.mkdir(parents=True) + sentinel.write_text("keep") artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -151,6 +158,7 @@ def test_unusable_manifest_counts_as_no_manifest( assert manifest.previous is None assert manifest.file not in (FRAME_A, FRAME_B) assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} + assert sentinel.read_text() == "keep" assert all(path.exists() for path in named) From 4d64274e17c1eb886e37e56cfe6d2c93d1a073cc Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 18 Oct 2026 20:14:58 +0000 Subject: [PATCH 36/41] test(artifacts): drop cases that another test already catches Each case goes only if a remaining test fails under the same source mutation, so the suite still catches all 18 mutations tried. - test_arguments: the read-only variants of the name, type and query-preview checks, which run before any filesystem call. Of the invalid delete names, "", ".", ".." and "../x" stay: each would make delete_dataframe remove a directory that isn't the name's own. - test_manifest_and_cleanup: missing-key and type cases that test pydantic rather than read_manifest. version-a-bool and rows-a-string still catch a StrictInt loosened to int. - test_write_failures: the case that asserts whatever pyarrow does, four of the five unserializable frames, and the EPERM and ENOSPC close errors, which take the same paths as EACCES and EIO. The interrupt and polars-object tests still catch a narrowed except. - test_file_format: the real-zstd test, which fails only under the mutation the fake codec's [zstd] case already catches. - test_read_only_mount: the EROFS delete test through an rmtree mock. The test through the real rmtree covers the branch and checks the files too. - test_storing: deepnote_toolkit.artifacts is artifacts, which the import makes true. The __all__ check stays. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VD15J21xfHn4E8p1nZAPFB --- tests/unit/artifacts/test_arguments.py | 59 ++++--------------- tests/unit/artifacts/test_file_format.py | 7 --- .../artifacts/test_manifest_and_cleanup.py | 13 ---- tests/unit/artifacts/test_read_only_mount.py | 14 ----- tests/unit/artifacts/test_storing.py | 3 +- tests/unit/artifacts/test_write_failures.py | 48 +++------------ tests/unit/helpers/dataframe_storage.py | 27 +-------- 7 files changed, 22 insertions(+), 149 deletions(-) diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index e4478708..97e36268 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -11,12 +11,7 @@ from deepnote_toolkit import artifacts from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview -from tests.unit.helpers.dataframe_storage import ( - NAME, - arrow_files, - read_only, - stored_table, -) +from tests.unit.helpers.dataframe_storage import NAME, arrow_files, stored_table @pytest.mark.parametrize("name", ["x", "x" * 128, "Rev-2026_Q3"]) @@ -55,18 +50,6 @@ def test_invalid_name_raises_before_anything_is_written( assert list(tmp_project_root.iterdir()) == [] -@pytest.mark.parametrize("name", INVALID_NAMES) -def test_invalid_name_raises_on_a_read_only_mount( - tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch, name: Any -) -> None: - read_only(monkeypatch) - - with pytest.raises(ValueError) as raised: - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) - - assert str(raised.value).endswith(f"got {name!r}") - - def _query_preview() -> DeepnoteQueryPreview: return DeepnoteQueryPreview({"a": [1]}, deepnote_query="select 1") @@ -113,34 +96,6 @@ def test_query_preview_raises_type_error( assert list(tmp_project_root.iterdir()) == [] -@pytest.mark.parametrize("make_value", UNSUPPORTED_VALUES) -def test_unsupported_value_raises_type_error_on_a_read_only_mount( - tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - make_value: Callable[[], Any], -) -> None: - read_only(monkeypatch) - - with pytest.raises(TypeError) as raised: - artifacts.store_dataframe(make_value(), NAME) - - raised.match("pandas or polars DataFrame") - - -@pytest.mark.parametrize("make_preview", QUERY_PREVIEWS) -def test_query_preview_raises_type_error_on_a_read_only_mount( - tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - make_preview: Callable[[], Any], -) -> None: - read_only(monkeypatch) - - with pytest.raises(TypeError) as raised: - artifacts.store_dataframe(make_preview(), NAME) - - raised.match("SQL query preview") - - def test_pyspark_frame_raises_type_error(tmp_project_root: Path, spark: Any) -> None: with pytest.raises(TypeError, match="pandas or polars DataFrame"): artifacts.store_dataframe(spark.range(3), NAME) @@ -166,7 +121,17 @@ def test_pandas_on_spark_frame_raises_type_error( assert list(tmp_project_root.iterdir()) == [] -@pytest.mark.parametrize("name", INVALID_NAMES) +# Let through, each of these would make delete_dataframe remove a directory that is +# not the name's own. +@pytest.mark.parametrize( + "name", + [ + pytest.param("", id="empty"), + pytest.param(".", id="dot"), + pytest.param("..", id="dotdot"), + pytest.param("../x", id="traversal"), + ], +) def test_delete_with_an_invalid_name_raises_and_deletes_nothing( tmp_project_root: Path, name: Any ) -> None: diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index 362db49d..7b817313 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -55,13 +55,6 @@ def test_compression_falls_back_from_zstd_to_lz4_to_none( 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"] ) diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 91c06686..66337b25 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -111,26 +111,13 @@ def _manifest_bytes(without: str = "", **fields: Any) -> bytes: pytest.param( json.dumps(list(_VALID_MANIFEST.items())).encode(), id="not-an-object" ), - pytest.param(_manifest_bytes(without="version"), id="no-version"), - pytest.param(_manifest_bytes(version="1"), id="version-a-string"), pytest.param(_manifest_bytes(version=True), id="version-a-bool"), pytest.param(_manifest_bytes(without="file"), id="no-file"), - pytest.param(_manifest_bytes(file=7), id="file-not-a-string"), pytest.param(_manifest_bytes(file="../../data/x.csv"), id="file-traversal"), pytest.param( _manifest_bytes(previous="../../data/x.csv"), id="previous-traversal" ), - pytest.param(_manifest_bytes(without="written_at"), id="no-written-at"), - pytest.param(_manifest_bytes(written_at=7), id="written-at-not-a-string"), - pytest.param(_manifest_bytes(without="rows"), id="no-rows"), pytest.param(_manifest_bytes(rows="1"), id="rows-a-string"), - pytest.param(_manifest_bytes(rows=1.5), id="rows-not-whole"), - pytest.param( - _manifest_bytes(without="toolkit_version"), id="no-toolkit-version" - ), - pytest.param( - _manifest_bytes(toolkit_version=7), id="toolkit-version-not-a-string" - ), pytest.param( _manifest_bytes(pad="x" * 5000), id="oversized-truncated-mid-string" ), diff --git a/tests/unit/artifacts/test_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py index 35808012..608823b8 100644 --- a/tests/unit/artifacts/test_read_only_mount.py +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -1,7 +1,6 @@ """Tests for storing and deleting DataFrames on a read-only mount.""" import errno -import shutil from pathlib import Path from typing import Any from unittest import mock @@ -75,16 +74,3 @@ def test_delete_on_a_read_only_mount_deletes_nothing( assert artifacts.delete_dataframe(NAME) is False assert len(arrow_files(tmp_project_root)) == 1 - - -def test_delete_returns_false_when_the_filesystem_refuses_with_erofs( - tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - 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 diff --git a/tests/unit/artifacts/test_storing.py b/tests/unit/artifacts/test_storing.py index 32bed08d..6c517c76 100644 --- a/tests/unit/artifacts/test_storing.py +++ b/tests/unit/artifacts/test_storing.py @@ -147,6 +147,5 @@ def test_the_last_call_with_a_name_wins(tmp_project_root: Path) -> None: assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} -def test_functions_are_exported_from_the_package() -> None: - assert deepnote_toolkit.artifacts is artifacts +def test_artifacts_is_in_the_package_all() -> None: assert "artifacts" in deepnote_toolkit.__all__ diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 0c611137..a42a0f55 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -16,11 +16,9 @@ from deepnote_toolkit import artifacts from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest from tests.unit.helpers.dataframe_storage import ( - ALWAYS_FAILING, DAY, FRAME_A, NAME, - VERSION_DEPENDENT, arrow_files, backdate, frame_dir, @@ -36,13 +34,6 @@ def _store_a_good_frame(tmp_project_root: Path) -> bytes: return (frame_dir(tmp_project_root) / "manifest.json").read_bytes() -def _store_frame_after_a_good_one( - tmp_project_root: Path, make_frame: Callable[[], Any] -) -> None: - _store_a_good_frame(tmp_project_root) - artifacts.store_dataframe(make_frame(), NAME) - - def _store_bad_frame_after_a_good_one( tmp_project_root: Path, make_frame: Callable[[], Any], expected_error: type ) -> Dict[str, Any]: @@ -60,15 +51,18 @@ def _store_bad_frame_after_a_good_one( } -@pytest.mark.parametrize("make_frame", ALWAYS_FAILING) +def _mixed_types() -> pd.DataFrame: + return pd.DataFrame({"a": [1, "two"]}) + + def test_unserializable_frame_raises_and_keeps_the_previous_frame( - tmp_project_root: Path, make_frame: Callable[[], Any] + tmp_project_root: Path, ) -> None: - raised_by_pyarrow = probe_failure(make_frame()) + raised_by_pyarrow = probe_failure(_mixed_types()) assert raised_by_pyarrow is not None, "pyarrow is expected to refuse this frame" left = _store_bad_frame_after_a_good_one( - tmp_project_root, make_frame, raised_by_pyarrow + tmp_project_root, _mixed_types, raised_by_pyarrow ) assert (frame_dir(tmp_project_root) / "manifest.json").read_bytes() == left[ @@ -80,30 +74,6 @@ def test_unserializable_frame_raises_and_keeps_the_previous_frame( assert stored_table(tmp_project_root).to_pydict() == {"ok": [1]} -@pytest.mark.parametrize("make_frame, rows", VERSION_DEPENDENT) -def test_frames_pyarrow_may_refuse_follow_what_pyarrow_does( - tmp_project_root: Path, make_frame: Callable[[], Any], rows: int -) -> None: - raised_by_pyarrow = probe_failure(make_frame()) - - if raised_by_pyarrow is None: - _store_frame_after_a_good_one(tmp_project_root, make_frame) - stored = stored_table(tmp_project_root) - assert stored.column_names != ["ok"] - assert stored.num_rows == rows - return - left = _store_bad_frame_after_a_good_one( - tmp_project_root, make_frame, raised_by_pyarrow - ) - - assert (frame_dir(tmp_project_root) / "manifest.json").read_bytes() == left[ - "manifest_bytes" - ] - assert left["listing"] == sorted( - ["manifest.json", stored_manifest(tmp_project_root).file] - ) - - def test_polars_object_column_raises_and_keeps_the_previous_frame( tmp_project_root: Path, ) -> None: @@ -142,12 +112,10 @@ def write_some_then_fail(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: @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"], + ids=["EACCES", "EIO"], ) def test_error_from_close_raises_and_stores_nothing( tmp_project_root: Path, diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 52355c96..8ab3619c 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -1,16 +1,13 @@ -"""Helpers and frames shared by the DataFrame storage tests.""" +"""Helpers shared by the DataFrame storage tests.""" import errno import io import json import os import time -import uuid from pathlib import Path from typing import Any, List, Optional -import numpy as np -import pandas as pd import pyarrow as pa import pytest from pyarrow import feather @@ -87,25 +84,3 @@ def probe_failure(frame: Any) -> Optional[type]: 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]) - - -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"), -] From f731429cca65d2d675b40feff1caeb9d76aa73f4 Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 18 Oct 2026 20:15:06 +0000 Subject: [PATCH 37/41] test(artifacts): fail the concurrent-writers test if a writer hangs join(timeout=30) returns quietly while a thread is still running, and the test then asserted against a write that hadn't finished. Also read the stored manifest once in test_stores_polars_frame, and drop the spark fixture's duplicate SparkSession import. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VD15J21xfHn4E8p1nZAPFB --- tests/unit/artifacts/conftest.py | 2 -- tests/unit/artifacts/test_manifest_and_cleanup.py | 1 + tests/unit/artifacts/test_storing.py | 5 +++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index 0f2268dd..05844d10 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -34,8 +34,6 @@ def tokyo_time() -> Iterator[None]: @pytest.fixture(scope="module") def spark() -> SparkSession: - from pyspark.sql import SparkSession - return ( SparkSession.builder.master("local") .appName("Toolkit") diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 66337b25..7396ae2c 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -242,6 +242,7 @@ def write(value: int) -> None: thread.start() for thread in threads: thread.join(timeout=30) + assert not thread.is_alive() assert errors == [] files = arrow_files(tmp_project_root) diff --git a/tests/unit/artifacts/test_storing.py b/tests/unit/artifacts/test_storing.py index 6c517c76..c6eebc9f 100644 --- a/tests/unit/artifacts/test_storing.py +++ b/tests/unit/artifacts/test_storing.py @@ -66,12 +66,13 @@ def test_stores_polars_frame(tmp_project_root: Path) -> None: pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}), NAME ) - assert stored_manifest(tmp_project_root).file == f"{reference_id}.arrow" + manifest = stored_manifest(tmp_project_root) + assert manifest.file == f"{reference_id}.arrow" + assert manifest.rows == 3 assert stored_table(tmp_project_root).to_pydict() == { "a": [1, 2, 3], "b": ["x", "y", "z"], } - assert stored_manifest(tmp_project_root).rows == 3 def test_writing_does_not_change_the_callers_frame(tmp_project_root: Path) -> None: From 229b7975fb03979d5d6113578f47e776794e6001 Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 18 Oct 2026 20:15:42 +0000 Subject: [PATCH 38/41] fix(artifacts): count a manifest with a version other than 1 as absent version was written but, since 6c1fe41, never read back, so a toolkit would act on the file and previous of a future manifest format. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VD15J21xfHn4E8p1nZAPFB --- deepnote_toolkit/artifacts/dataframe_storage_manifest.py | 2 ++ tests/unit/artifacts/test_manifest_and_cleanup.py | 1 + 2 files changed, 3 insertions(+) diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index 6ef37f86..4e841f64 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -65,6 +65,8 @@ def read_manifest(directory: Path) -> Optional[Manifest]: except (ValidationError, TypeError): return None + if manifest.version != MANIFEST_VERSION: + return None # Cleanup unlinks `previous`, and `file` becomes the next `previous`, so any # other name could delete a file outside this directory. if not is_frame_name(manifest.file): diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 7396ae2c..3fca8f16 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -111,6 +111,7 @@ def _manifest_bytes(without: str = "", **fields: Any) -> bytes: pytest.param( json.dumps(list(_VALID_MANIFEST.items())).encode(), id="not-an-object" ), + pytest.param(_manifest_bytes(version=2), id="unsupported-version"), pytest.param(_manifest_bytes(version=True), id="version-a-bool"), pytest.param(_manifest_bytes(without="file"), id="no-file"), pytest.param(_manifest_bytes(file="../../data/x.csv"), id="file-traversal"), From e83a979e76c5d034111f027569334d005ec47859 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 19 Oct 2026 06:16:35 +0000 Subject: [PATCH 39/41] test(artifacts): reject a delete name that traverses to a frame file The new case points delete_dataframe at a real .arrow file in the project root. Its target is the only one that isn't a directory, so it alone fails if the name check moves below the is_dir() early return. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VD15J21xfHn4E8p1nZAPFB --- tests/unit/artifacts/test_arguments.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index 97e36268..9d1c28e8 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -11,7 +11,12 @@ from deepnote_toolkit import artifacts from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview -from tests.unit.helpers.dataframe_storage import NAME, arrow_files, stored_table +from tests.unit.helpers.dataframe_storage import ( + FRAME_A, + NAME, + arrow_files, + stored_table, +) @pytest.mark.parametrize("name", ["x", "x" * 128, "Rev-2026_Q3"]) @@ -121,8 +126,9 @@ def test_pandas_on_spark_frame_raises_type_error( assert list(tmp_project_root.iterdir()) == [] -# Let through, each of these would make delete_dataframe remove a directory that is -# not the name's own. +# Let through, each of these would point delete_dataframe outside the name's own +# directory. Only the frame file is not a directory, so only it catches a name check +# moved below the is_dir() return. @pytest.mark.parametrize( "name", [ @@ -130,6 +136,7 @@ def test_pandas_on_spark_frame_raises_type_error( pytest.param(".", id="dot"), pytest.param("..", id="dotdot"), pytest.param("../x", id="traversal"), + pytest.param(f"../{FRAME_A}", id="traversal-to-a-frame-file"), ], ) def test_delete_with_an_invalid_name_raises_and_deletes_nothing( @@ -137,6 +144,7 @@ def test_delete_with_an_invalid_name_raises_and_deletes_nothing( ) -> None: (tmp_project_root / "x").mkdir() (tmp_project_root / "x" / "keep.txt").write_text("keep") + (tmp_project_root / FRAME_A).write_text("keep") artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) with pytest.raises(ValueError) as raised: @@ -144,4 +152,5 @@ def test_delete_with_an_invalid_name_raises_and_deletes_nothing( assert str(raised.value).endswith(f"got {name!r}") assert (tmp_project_root / "x" / "keep.txt").read_text() == "keep" + assert (tmp_project_root / FRAME_A).read_text() == "keep" assert len(arrow_files(tmp_project_root)) == 1 From 91ae178ed90b5566f686411f05f6f8a921464b64 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 19 Oct 2026 08:43:03 +0000 Subject: [PATCH 40/41] refactor(artifacts): check previous where cleanup unlinks it, not in read_manifest read_manifest no longer checks version or the file names, so such a manifest is honoured. _delete_replaced_frames unlinks previous only if it is .arrow. That is the only place a name from the manifest reaches the filesystem: file only becomes the next store's previous. The traversal cases move out of "counts as no manifest" into a test that a frame-named file in the parent directory survives a store, whether previous names it by a relative or an absolute path. The unsupported-version case goes with the version check. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VD15J21xfHn4E8p1nZAPFB --- .../artifacts/dataframe_storage.py | 6 ++++- .../artifacts/dataframe_storage_manifest.py | 8 ------ .../artifacts/test_manifest_and_cleanup.py | 26 ++++++++++++------- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py index b78830ad..b571b7f7 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -93,7 +93,11 @@ def _delete_replaced_frames( named = {new_file} if old is not None: named.add(old.file) - if old.previous is not None and old.previous not in named: + if ( + old.previous is not None + and is_frame_name(old.previous) + and old.previous not in named + ): (directory / old.previous).unlink(missing_ok=True) now = time.time() diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py index 4e841f64..bdf94c3e 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -65,14 +65,6 @@ def read_manifest(directory: Path) -> Optional[Manifest]: except (ValidationError, TypeError): return None - if manifest.version != MANIFEST_VERSION: - return None - # Cleanup unlinks `previous`, and `file` becomes the next `previous`, so any - # other name could delete a file outside this directory. - if not is_frame_name(manifest.file): - return None - if manifest.previous is not None and not is_frame_name(manifest.previous): - return None return manifest diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 3fca8f16..94ebe170 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -111,13 +111,8 @@ def _manifest_bytes(without: str = "", **fields: Any) -> bytes: pytest.param( json.dumps(list(_VALID_MANIFEST.items())).encode(), id="not-an-object" ), - pytest.param(_manifest_bytes(version=2), id="unsupported-version"), pytest.param(_manifest_bytes(version=True), id="version-a-bool"), pytest.param(_manifest_bytes(without="file"), id="no-file"), - pytest.param(_manifest_bytes(file="../../data/x.csv"), id="file-traversal"), - pytest.param( - _manifest_bytes(previous="../../data/x.csv"), id="previous-traversal" - ), pytest.param(_manifest_bytes(rows="1"), id="rows-a-string"), pytest.param( _manifest_bytes(pad="x" * 5000), id="oversized-truncated-mid-string" @@ -136,9 +131,6 @@ def test_unusable_manifest_counts_as_no_manifest( named = [directory / FRAME_A, directory / FRAME_B] for path in named: path.write_bytes(b"named by the unusable manifest") - sentinel = (directory / "../../data/x.csv").resolve() - sentinel.parent.mkdir(parents=True) - sentinel.write_text("keep") artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -146,7 +138,6 @@ def test_unusable_manifest_counts_as_no_manifest( assert manifest.previous is None assert manifest.file not in (FRAME_A, FRAME_B) assert stored_table(tmp_project_root).to_pydict() == {"a": [1]} - assert sentinel.read_text() == "keep" assert all(path.exists() for path in named) @@ -176,6 +167,23 @@ def test_usable_manifest_is_honoured( assert not (directory / FRAME_B).exists() +@pytest.mark.parametrize("absolute", [False, True], ids=["relative", "absolute"]) +def test_previous_in_a_parent_directory_is_not_deleted( + tmp_project_root: Path, absolute: bool +) -> None: + directory = frame_dir(tmp_project_root) + directory.mkdir(parents=True) + # Named like a frame, so checking only the last path component would delete it. + in_the_parent = directory.parent / FRAME_B + in_the_parent.write_text("keep") + previous = str(in_the_parent) if absolute else f"../{FRAME_B}" + (directory / "manifest.json").write_bytes(_manifest_bytes(previous=previous)) + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert in_the_parent.read_text() == "keep" + + def test_manifest_nested_too_deeply_to_parse_counts_as_no_manifest( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: From 880dc2201eb2de71df7402f50f81fe33fadd00f6 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 19 Oct 2026 12:40:40 +0000 Subject: [PATCH 41/41] refactor(artifacts): mark the storage and manifest modules private deepnote_toolkit imports artifacts into every kernel, and importing a submodule binds it on the package, so tab completion on deepnote_toolkit.artifacts listed dataframe_storage and dataframe_storage_manifest next to store_dataframe and delete_dataframe. __all__ does not hide submodules. The modules stay flat, as notebooks/_schemas.py is. A _dataframe/ subpackage, or splitting out types and constants, would move each definition away from the one module that uses it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VD15J21xfHn4E8p1nZAPFB --- ...aframe_storage.py => _dataframe_storage.py} | 2 +- ...ifest.py => _dataframe_storage_manifest.py} | 0 deepnote_toolkit/artifacts/dataframes.py | 12 ++++++------ tests/unit/artifacts/conftest.py | 6 ++++-- tests/unit/artifacts/test_file_format.py | 4 ++-- .../artifacts/test_manifest_and_cleanup.py | 14 +++++++------- tests/unit/artifacts/test_project_root.py | 18 +++++++++--------- tests/unit/artifacts/test_write_failures.py | 8 ++++---- tests/unit/helpers/dataframe_storage.py | 2 +- 9 files changed, 34 insertions(+), 32 deletions(-) rename deepnote_toolkit/artifacts/{dataframe_storage.py => _dataframe_storage.py} (99%) rename deepnote_toolkit/artifacts/{dataframe_storage_manifest.py => _dataframe_storage_manifest.py} (100%) diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py similarity index 99% rename from deepnote_toolkit/artifacts/dataframe_storage.py rename to deepnote_toolkit/artifacts/_dataframe_storage.py index b571b7f7..95efac17 100644 --- a/deepnote_toolkit/artifacts/dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -16,7 +16,7 @@ from deepnote_toolkit.config import resolve_notebook_root from deepnote_toolkit.ocelots.types import PandasDF, PolarsEagerDF -from .dataframe_storage_manifest import ( +from ._dataframe_storage_manifest import ( Manifest, is_frame_name, read_manifest, diff --git a/deepnote_toolkit/artifacts/dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py similarity index 100% rename from deepnote_toolkit/artifacts/dataframe_storage_manifest.py rename to deepnote_toolkit/artifacts/_dataframe_storage_manifest.py diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index 46e4ccfe..aa11fd82 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -10,11 +10,11 @@ ) from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview -from . import dataframe_storage +from . import _dataframe_storage def _check_name(name: object) -> None: - if not dataframe_storage.is_safe_id(name): + 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}" ) @@ -60,8 +60,8 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: f"got {type(df).__qualname__}" ) - root = dataframe_storage.resolve_project_root() - return dataframe_storage.store_frame(df, name, root) + root = _dataframe_storage.resolve_project_root() + return _dataframe_storage.store_frame(df, name, root) def delete_dataframe(name: str) -> bool: @@ -77,8 +77,8 @@ def delete_dataframe(name: str) -> bool: files. A symbolic link in place of the name's directory raises too. """ _check_name(name) - root = dataframe_storage.resolve_project_root() - directory = root / dataframe_storage.FRAMES_DIR / name + root = _dataframe_storage.resolve_project_root() + directory = root / _dataframe_storage.FRAMES_DIR / name if not directory.is_dir(): return False try: diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index 05844d10..acab15c5 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -7,7 +7,7 @@ import pytest from pyspark.sql import SparkSession -from deepnote_toolkit.artifacts import dataframe_storage +from deepnote_toolkit.artifacts import _dataframe_storage @pytest.fixture @@ -15,7 +15,9 @@ def tmp_project_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) + monkeypatch.setattr( + _dataframe_storage, "resolve_project_root", lambda: project_root + ) return project_root diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index 7b817313..96f1eb12 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -10,7 +10,7 @@ import pytest from deepnote_toolkit import artifacts -from deepnote_toolkit.artifacts import dataframe_storage +from deepnote_toolkit.artifacts import _dataframe_storage from tests.unit.helpers.dataframe_storage import ( NAME, frame_dir, @@ -52,7 +52,7 @@ def test_compression_falls_back_from_zstd_to_lz4_to_none( ) -> None: monkeypatch.setattr(pa, "Codec", _FakeCodec(availability)) - assert dataframe_storage._select_compression() == expected + assert _dataframe_storage._select_compression() == expected @pytest.mark.parametrize( diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 94ebe170..9a35a337 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -13,7 +13,7 @@ import pytest from deepnote_toolkit import artifacts -from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest +from deepnote_toolkit.artifacts import _dataframe_storage, _dataframe_storage_manifest from tests.unit.helpers.dataframe_storage import ( DAY, FRAME_A, @@ -196,7 +196,7 @@ def too_deep(raw: bytes) -> Any: ) monkeypatch.setattr( - dataframe_storage_manifest, + _dataframe_storage_manifest, "json", SimpleNamespace(loads=too_deep, dumps=json.dumps), ) @@ -211,14 +211,14 @@ def test_manifest_is_written_after_the_frame_file_is_complete( tmp_project_root: Path, arrow_sinks: List[Any], monkeypatch: pytest.MonkeyPatch ) -> None: seen: List[Dict[str, Any]] = [] - original = dataframe_storage.write_manifest + 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) + monkeypatch.setattr(_dataframe_storage, "write_manifest", spy) artifacts.store_dataframe(pd.DataFrame({"a": [1, 2, 3]}), NAME) @@ -229,7 +229,7 @@ def test_concurrent_writers_leave_a_complete_current_frame( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: barrier = threading.Barrier(2, timeout=10) - original = dataframe_storage.read_manifest + original = _dataframe_storage.read_manifest def read_then_wait(directory: Path) -> Any: manifest = original(directory) @@ -245,7 +245,7 @@ def write(value: int) -> None: errors.append(exc) with monkeypatch.context() as patch: - patch.setattr(dataframe_storage, "read_manifest", read_then_wait) + 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() @@ -296,7 +296,7 @@ def scandir_then_lose_a_stray(path: Any) -> Iterator[Any]: strays[0].unlink() yield iter(listing) - monkeypatch.setattr(dataframe_storage.os, "scandir", scandir_then_lose_a_stray) + monkeypatch.setattr(_dataframe_storage.os, "scandir", scandir_then_lose_a_stray) artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) diff --git a/tests/unit/artifacts/test_project_root.py b/tests/unit/artifacts/test_project_root.py index d425ac38..be5c53e5 100644 --- a/tests/unit/artifacts/test_project_root.py +++ b/tests/unit/artifacts/test_project_root.py @@ -8,7 +8,7 @@ from deepnote_toolkit import artifacts from deepnote_toolkit import config as toolkit_config -from deepnote_toolkit.artifacts import dataframe_storage +from deepnote_toolkit.artifacts import _dataframe_storage from tests.unit.helpers.dataframe_storage import NAME, arrow_files, stored_manifest @@ -16,7 +16,7 @@ def test_store_creates_a_missing_project_root( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: project = tmp_path / "work" - monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) + monkeypatch.setattr(_dataframe_storage, "resolve_project_root", lambda: project) reference_id = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -37,26 +37,26 @@ def test_notebook_root_wins( 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" + 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" + 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") + 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" + assert _dataframe_storage.resolve_project_root() == tmp_path / "work" def test_unreadable_config_falls_back_to_dev_mode_env( self, monkeypatch: pytest.MonkeyPatch @@ -64,7 +64,7 @@ def test_unreadable_config_falls_back_to_dev_mode_env( monkeypatch.setattr(toolkit_config, "get_config", _broken_config) monkeypatch.setenv("DEEPNOTE_RUNNING_IN_DEV_MODE", "true") - assert dataframe_storage.resolve_project_root() == Path("/work") + 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 @@ -72,7 +72,7 @@ def test_unreadable_config_falls_back_to_work_under_home( monkeypatch.setattr(toolkit_config, "get_config", _broken_config) monkeypatch.setenv("HOME", str(tmp_path)) - assert dataframe_storage.resolve_project_root() == tmp_path / "work" + assert _dataframe_storage.resolve_project_root() == tmp_path / "work" @pytest.mark.usefixtures("no_root_settings") @@ -96,7 +96,7 @@ def test_delete_returns_false_when_the_project_root_is_missing( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: project = tmp_path / "work" - monkeypatch.setattr(dataframe_storage, "resolve_project_root", lambda: project) + monkeypatch.setattr(_dataframe_storage, "resolve_project_root", lambda: project) assert artifacts.delete_dataframe(NAME) is False assert not project.exists() diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index a42a0f55..9137c5f5 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -14,7 +14,7 @@ from pyarrow import feather from deepnote_toolkit import artifacts -from deepnote_toolkit.artifacts import dataframe_storage, dataframe_storage_manifest +from deepnote_toolkit.artifacts import _dataframe_storage, _dataframe_storage_manifest from tests.unit.helpers.dataframe_storage import ( DAY, FRAME_A, @@ -171,7 +171,7 @@ def test_failure_in_the_manifest_step_deletes_nothing( def broken_manifest(*args: Any) -> Any: raise OSError(errno.EIO, "Input/output error") - monkeypatch.setattr(dataframe_storage, "write_manifest", broken_manifest) + monkeypatch.setattr(_dataframe_storage, "write_manifest", broken_manifest) with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) @@ -232,7 +232,7 @@ def failing_temporary_file(*args: Any, **kwargs: Any) -> Any: return _FailingManifestHandle(handle, failing, error) monkeypatch.setattr( - dataframe_storage_manifest, + _dataframe_storage_manifest, "tempfile", SimpleNamespace(NamedTemporaryFile=failing_temporary_file), ) @@ -258,7 +258,7 @@ def test_rename_failing_after_the_copy_keeps_the_frame_it_put_in_place( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) - monkeypatch.setattr(dataframe_storage_manifest.os, "replace", _copy_then_fail) + monkeypatch.setattr(_dataframe_storage_manifest.os, "replace", _copy_then_fail) with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 8ab3619c..9a7ecdcf 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -13,7 +13,7 @@ from pyarrow import feather from deepnote_core.pydantic_compat_helpers import model_validate_compat -from deepnote_toolkit.artifacts.dataframe_storage_manifest import Manifest +from deepnote_toolkit.artifacts._dataframe_storage_manifest import Manifest NAME = "revenue" DAY = 24 * 60 * 60