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/__init__.py b/deepnote_toolkit/artifacts/__init__.py new file mode 100644 index 00000000..c425bf26 --- /dev/null +++ b/deepnote_toolkit/artifacts/__init__.py @@ -0,0 +1,5 @@ +"""Stores and deletes named DataFrames in the project's filesystem.""" + +from .dataframes import delete_dataframe, store_dataframe + +__all__ = ["delete_dataframe", "store_dataframe"] diff --git a/deepnote_toolkit/artifacts/dataframe_storage.py b/deepnote_toolkit/artifacts/dataframe_storage.py new file mode 100644 index 00000000..a8f70d90 --- /dev/null +++ b/deepnote_toolkit/artifacts/dataframe_storage.py @@ -0,0 +1,148 @@ +"""Writes named DataFrames into the project's filesystem as Arrow files.""" + +import contextlib +import errno +import os +import re +import time +import uuid +from pathlib import Path +from typing import Any, Literal, Optional + +from packaging.version import Version + +from deepnote_toolkit.config import resolve_notebook_root +from deepnote_toolkit.ocelots.utils import ( + is_pandas_dataframe, + is_polars_eager_dataframe, +) +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview + +from .dataframe_storage_manifest import ( + Manifest, + is_frame_name, + read_manifest, + write_manifest, +) + +FRAMES_DIR = Path("deepnote_dataframes") +STALE_FRAME_AGE_SECONDS = 24 * 60 * 60 +POLARS_BATCH_ROWS = 65_536 +POLARS_BATCH_MIN_VERSION = Version("1.37") +POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") + +FrameKind = Literal["pandas", "polars"] + +_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_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() + + +def _select_compression() -> str: + """Return the best codec this pyarrow build has: zstd, lz4, else uncompressed.""" + import pyarrow as pa + + for codec in ("zstd", "lz4"): + if pa.Codec.is_available(codec): + return codec + return "uncompressed" + + +def _write_frame_file( + frame: Any, kind: FrameKind, path: Path, compression: str +) -> None: + """Write the frame to `path` as an Arrow IPC file. + + 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": + 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 + + 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 + # 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) + + +def _delete_replaced_frames( + directory: Path, old: Optional[Manifest], new_file: str +) -> None: + """Delete the frame the replaced manifest called `previous`, and stale strays. + + 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: + 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 + 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) -> 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 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 + + 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 new file mode 100644 index 00000000..86d2f7b2 --- /dev/null +++ b/deepnote_toolkit/artifacts/dataframe_storage_manifest.py @@ -0,0 +1,122 @@ +"""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 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" +MANIFEST_VERSION = 1 +MANIFEST_MAX_BYTES = 4096 + +_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, +) + + +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 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. + """ + 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 + # 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 + + 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 + + +def write_manifest( + directory: Path, new_file: str, previous: Optional[str], rows: int +) -> None: + """Point the manifest at `new_file`. + + On s3fs a rename is a copy and then a delete, so a failed rename may already + have replaced the manifest. + """ + 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. + temporary = tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=directory, + prefix=f".{MANIFEST_NAME}.", + suffix=".tmp", + delete=False, + ) + try: + with temporary: + temporary.write(json.dumps(_model_dump_compat(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 new file mode 100644 index 00000000..23bede8f --- /dev/null +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -0,0 +1,87 @@ +"""Stores and deletes named DataFrames in the project's filesystem.""" + +import errno +import shutil +from typing import Any, Optional + +from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview + +from . import dataframe_storage + + +def _check_name(name: object) -> None: + if not dataframe_storage.is_safe_id(name): + raise ValueError( + f"A DataFrame name must match [A-Za-z0-9_-]{{1,128}}, got {name!r}" + ) + + +def store_dataframe(df: Any, name: str) -> Optional[str]: + """Store `df` as the current frame under `name`, replacing the one stored before. + + Args: + df: A pandas or polars DataFrame, written whole with no row or size limit. + name: The frame's name, unique within the project, 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 + read-only, in which case nothing is stored. + + Raises: + ValueError: `name` is not valid. + 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. + + 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) + 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__}" + ) + + root = dataframe_storage.resolve_project_root() + return dataframe_storage.store_frame(df, kind, name, root) + + +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. + 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() + directory = root / dataframe_storage.FRAMES_DIR / name + if not directory.is_dir(): + return False + 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..ad4ba1a6 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,20 @@ 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.""" + 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/__init__.py b/tests/unit/artifacts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py new file mode 100644 index 00000000..936e4ef4 --- /dev/null +++ b/tests/unit/artifacts/conftest.py @@ -0,0 +1,123 @@ +"""Fixtures for the DataFrame storage tests.""" + +import time +from pathlib import Path +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 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) + 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() -> SparkSession: + from pyspark.sql import SparkSession + + return ( + SparkSession.builder.master("local") + .appName("Toolkit") + .config("spark.driver.memory", "4g") + .getOrCreate() + ) + + +@pytest.fixture +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") + 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.""" + 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 _FailingCloseS3Upload: + """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) -> "_FailingCloseS3Upload": + 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 _FailingCloseS3Upload(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/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py new file mode 100644 index 00000000..e4478708 --- /dev/null +++ b/tests/unit/artifacts/test_arguments.py @@ -0,0 +1,182 @@ +"""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(tmp_project_root: Path, name: str) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + + assert stored_table(tmp_project_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( + tmp_project_root: Path, name: Any +) -> None: + with mock.patch.object(feather, "write_feather", autospec=True) as write: + 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()) == [] + + +@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") + + +def _cleared_query_preview() -> DeepnoteQueryPreview: + preview = _query_preview() + # Mutating the frame clears its attached SQL query metadata. + 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( + 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(tmp_project_root.iterdir()) == [] + + +@pytest.mark.parametrize("make_preview", QUERY_PREVIEWS) +def test_query_preview_raises_type_error( + 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(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) + + assert list(tmp_project_root.iterdir()) == [] + + +def test_pandas_on_spark_frame_raises_type_error( + 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") + # 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) as raised: + artifacts.store_dataframe(frame, NAME) + + raised.match("pandas or polars DataFrame") + assert list(tmp_project_root.iterdir()) == [] + + +@pytest.mark.parametrize("name", INVALID_NAMES) +def test_delete_with_an_invalid_name_raises_and_deletes_nothing( + tmp_project_root: Path, name: Any +) -> None: + (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) 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 new file mode 100644 index 00000000..6e20634c --- /dev/null +++ b/tests/unit/artifacts/test_delete_dataframe.py @@ -0,0 +1,77 @@ +"""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(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 = tmp_project_root / "deepnote_dataframes" / "revenue-2026" + neighbour.mkdir() + (neighbour / "keep.txt").write_text("keep") + (tmp_project_root / "notes.txt").write_text("keep") + + assert artifacts.delete_dataframe(NAME) is True + + 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 (tmp_project_root / "notes.txt").read_text() == "keep" + + +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(tmp_project_root.iterdir()) == [] + + +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( + 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 stored_manifest(tmp_project_root).previous is None + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} + + +def test_delete_does_not_follow_a_link_out_of_the_frames_directory( + tmp_project_root: Path, +) -> None: + outside = tmp_project_root / "outside" + outside.mkdir() + (outside / "keep.txt").write_text("keep") + (tmp_project_root / "deepnote_dataframes").mkdir() + (tmp_project_root / "deepnote_dataframes" / NAME).symlink_to( + outside, target_is_directory=True + ) + + 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_file_format.py b/tests/unit/artifacts/test_file_format.py new file mode 100644 index 00000000..362db49d --- /dev/null +++ b/tests/unit/artifacts/test_file_format.py @@ -0,0 +1,168 @@ +"""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(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 + 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( + 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(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(tmp_project_root: Path) -> None: + 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 + ) + assert reader.num_record_batches == 2 + assert reader.read_all().num_rows == 70_000 + + +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(tmp_project_root) / stored_manifest(tmp_project_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( + tmp_project_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(tmp_project_root).to_pydict() == {"a": [1, 2, 3]} + + +def test_polars_strings_are_stored_without_view_types(tmp_project_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(tmp_project_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( + tmp_project_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(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 new file mode 100644 index 00000000..8541c728 --- /dev/null +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -0,0 +1,330 @@ +"""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(tmp_project_root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + first = stored_manifest(tmp_project_root) + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + second = stored_manifest(tmp_project_root) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + third = stored_manifest(tmp_project_root) + + 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] + ) + 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( + tmp_project_root: Path, +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + 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: + 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 = tmp_project_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() + + +_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( + "manifest_bytes", + [ + pytest.param(b"{not json", id="invalid-json"), + pytest.param(b"\x80\x81\x82", id="not-utf8"), + 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), 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" + ), + 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" + ), + pytest.param( + _manifest_bytes() + b" " * 5000, id="oversized-valid-json-then-spaces" + ), + ], +) +def test_unusable_manifest_counts_as_no_manifest( + tmp_project_root: Path, manifest_bytes: bytes +) -> None: + directory = frame_dir(tmp_project_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(tmp_project_root) + 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) + + +@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) + 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") + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + 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( + 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) + + 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 stored_manifest(tmp_project_root).previous is None + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} + + +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 + + 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( + tmp_project_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(tmp_project_root) + assert len(files) == 2 + 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 + + backdate(stray, 2 * DAY) + artifacts.store_dataframe(pd.DataFrame({"a": [3]}), NAME) + + assert not stray.exists() + 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( + 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(tmp_project_root) == ["manifest.json"] + + +def test_stray_pruned_by_another_writer_meanwhile_is_not_a_failure( + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + directory = frame_dir(tmp_project_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(tmp_project_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..d425ac38 --- /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..a0c45e74 --- /dev/null +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -0,0 +1,89 @@ +"""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( + tmp_project_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(tmp_project_root.iterdir()) == [] + + +def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + 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: + result = artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert result is None + write.assert_not_called() + 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( + 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") + + monkeypatch.setattr(feather, "write_feather", refuse) + + result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert result is None + assert arrow_files(tmp_project_root) == [] + + +def test_delete_on_a_read_only_mount_deletes_nothing( + 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(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 new file mode 100644 index 00000000..32bed08d --- /dev/null +++ b/tests/unit/artifacts/test_storing.py @@ -0,0 +1,152 @@ +"""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 +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, + frame_dir, + 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(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(tmp_project_root)] == [ + f"{reference_id}.arrow" + ] + assert stored_table(tmp_project_root).to_pydict() == { + "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 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(tmp_project_root: Path) -> None: + before = datetime.now(timezone.utc) + + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + written = datetime.strptime( + 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(tmp_project_root: Path) -> None: + reference_id = artifacts.store_dataframe( + pl.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]}), NAME + ) + + 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(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(tmp_project_root)) == 1 + pd.testing.assert_frame_equal(frame, before) + + +def test_filtered_frame_is_stored_with_its_row_labels( + 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") + + # 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(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(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(tmp_project_root).schema.names == ["v", "k"] + + +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(tmp_project_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(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 + + +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(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(tmp_project_root: Path) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + artifacts.store_dataframe(pl.DataFrame({"a": [2]}), NAME) + + 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 + 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..9f4fa1fc --- /dev/null +++ b/tests/unit/artifacts/test_write_failures.py @@ -0,0 +1,349 @@ +"""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_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: 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. + """ + manifest_bytes = _store_a_good_frame(tmp_project_root) + 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 frame_dir(tmp_project_root).iterdir()), + } + + +@pytest.mark.parametrize("make_frame", ALWAYS_FAILING) +def test_unserializable_frame_raises_and_keeps_the_previous_frame( + 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( + 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] + ) + 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: + 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) 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 + assert len(arrow_files(tmp_project_root)) == 1 + + +def test_partial_file_is_deleted_when_the_write_fails( + 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") + 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(tmp_project_root) == [] + assert not (frame_dir(tmp_project_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( + 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 + + with pytest.raises(OSError) as raised: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + + assert raised.value is error + 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( + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) + frames = arrow_files(tmp_project_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) as raised: + artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) + + assert raised.value.errno == errno.EIO + assert arrow_files(tmp_project_root) == frames + + +def test_failure_in_the_manifest_step_deletes_nothing( + 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(tmp_project_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) 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 + assert stored_table(tmp_project_root).to_pydict() == {"a": [2]} + + +class _FailingManifestHandle: + """A temporary manifest whose write() or close() raises `error`.""" + + 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 self._error + 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 self._error + + def __getattr__(self, name: str) -> Any: + return getattr(self._handle, name) + + +@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, + 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: + handle = real_temporary_file(*args, **kwargs) + return _FailingManifestHandle(handle, failing, error) + + monkeypatch.setattr( + dataframe_storage_manifest, + "tempfile", + SimpleNamespace(NamedTemporaryFile=failing_temporary_file), + ) + + 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 + 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: + """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( + 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) + + 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]} + + +@pytest.mark.parametrize("error", [errno.EIO, errno.EROFS], ids=["EIO", "EROFS"]) +def test_failure_after_the_manifest_step_keeps_the_new_frame( + 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) + 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( + 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) 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 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) 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 + assert len(arrow_files(tmp_project_root)) == 1 diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py new file mode 100644 index 00000000..ff582369 --- /dev/null +++ b/tests/unit/helpers/dataframe_storage.py @@ -0,0 +1,111 @@ +"""Helpers and frames 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 + +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" +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) -> 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 + 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.""" + try: + feather.write_feather(frame, 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]) + + +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"), +] 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}