diff --git a/src/dstrack/_cli.py b/src/dstrack/_cli.py index 70d7da3..1e77782 100644 --- a/src/dstrack/_cli.py +++ b/src/dstrack/_cli.py @@ -1,11 +1,12 @@ import logging -import shutil +import tempfile from pathlib import Path from typing import Annotated import typer from dstrack import __version__, console +from dstrack._store_template import STORE_TEMPLATE, materialize from dstrack.errors import StoreInitError from dstrack.utils import get_invocation_path @@ -20,59 +21,40 @@ ) -def _init_local_store_gitignore(path: Path) -> None: - """Create .gitignore file in provided path. - - creates and populates file only if it does not already exists. - - Args: - path: Path where .gitignore file should be created. - - Raises: - FileExistsError: When .gitignore file already exists in path. - - Returns: - None - """ - file_path = path / ".gitignore" - if file_path.is_file(): - raise FileExistsError(f"Store .gitignore file already exists: {file_path}") - - # Create and populate .gitignore file - console.info("Generating .gitignore file in local store.") - with file_path.open("w"): - file_path.write_text(".cache/") - - def init_local_store() -> Path: """Initializes local store structure. - creates a local store folder called `.dstrack`, a `datasets/` folder - within it, and a `.gitignore` file to ignore the cache directory. + Builds `STORE_TEMPLATE` (see `_store_template.py`) in a temporary + directory next to the destination, and only moves it into place once + every file and directory in the template has been created successfully. + The destination is therefore never touched by a failed or partial build. Returns: Path where the local store was created. Raises: FileExistsError: If the local store path already exists. - StoreInitError: If store creation starts but fails partway through, - e.g. because the `.gitignore` file cannot be written. Any - partially created state is rolled back before raising. + StoreInitError: If building the store from its template fails. The + temporary build directory is cleaned up before raising. """ - store_path = get_invocation_path() / ".dstrack" + invocation_path = get_invocation_path() + store_path = invocation_path / STORE_TEMPLATE.name if store_path.is_dir(): raise FileExistsError(f"Local store path already exists: {store_path}") - # Create local store path - store_path.mkdir() - try: - # Create dataset snapshots path - (store_path / "datasets").mkdir() - _init_local_store_gitignore(path=store_path) - except Exception as e: - shutil.rmtree(store_path) - raise StoreInitError(f"Failed to initialize local store at {store_path}") from e + with tempfile.TemporaryDirectory( + prefix=f"{STORE_TEMPLATE.name}-tmp-", dir=invocation_path + ) as tmp_dir: + try: + built_path = materialize(STORE_TEMPLATE, Path(tmp_dir)) + except Exception as e: + raise StoreInitError( + f"Failed to initialize local store at {store_path}" + ) from e + + console.info(f"Generating local store structure at {store_path}.") + built_path.rename(store_path) return store_path diff --git a/src/dstrack/_store_template.py b/src/dstrack/_store_template.py new file mode 100644 index 0000000..8ed6675 --- /dev/null +++ b/src/dstrack/_store_template.py @@ -0,0 +1,58 @@ +"""Declarative description of the local store's directory structure.""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Final + +from dstrack.paths import STORE_DIRNAME + + +@dataclass(frozen=True) +class TemplateFile: + """A file to create, with its literal content.""" + + name: str + content: str = "" + + +@dataclass(frozen=True) +class TemplateDir: + """A directory to create, optionally containing files and subdirectories.""" + + name: str + entries: "tuple[TemplateDir | TemplateFile, ...]" = field(default_factory=tuple) + + +STORE_TEMPLATE: Final[TemplateDir] = TemplateDir( + name=STORE_DIRNAME, + entries=( + TemplateDir(name="datasets"), + TemplateFile(name=".gitignore", content=".cache/"), + ), +) + + +def materialize(template: TemplateDir, parent: Path) -> Path: + """Recursively create `template`'s directory structure under `parent`. + + Args: + template: Directory template to materialize, including the + top-level directory itself. + parent: Existing directory that `template`'s directory is created in. + + Returns: + The path created for `template`. + + Raises: + FileExistsError: If any file or directory in the template already + exists under `parent`. + """ + path = parent / template.name + path.mkdir() + for entry in template.entries: + if isinstance(entry, TemplateDir): + materialize(entry, path) + else: + with open(path / entry.name, "x") as f: + f.write(entry.content) + return path diff --git a/src/dstrack/errors.py b/src/dstrack/errors.py index c774462..3c64dfd 100644 --- a/src/dstrack/errors.py +++ b/src/dstrack/errors.py @@ -1,2 +1,6 @@ class StoreInitError(Exception): """Raised when local store initialization fails partway through.""" + + +class StoreNotFoundError(Exception): + """Raised when no local store root can be resolved.""" diff --git a/src/dstrack/paths.py b/src/dstrack/paths.py new file mode 100644 index 0000000..25c2592 --- /dev/null +++ b/src/dstrack/paths.py @@ -0,0 +1,88 @@ +"""Resolution of the local store's directory, `STORE_DIRNAME/`. + +Mirrors how git resolves `.git/`: absent an explicit override, it is found by +walking up from the current directory until a `STORE_DIRNAME/` directory turns up. +""" + +import os +from pathlib import Path +from typing import Final + +from dstrack.errors import StoreNotFoundError +from dstrack.utils import get_invocation_path + +STORE_DIRNAME: Final = ".dstrack" +ROOT_PATH_ENV_VAR: Final = "DSTRACK_ROOT_PATH" + + +def _find_store_root(start: Path) -> Path: + """Walk upwards from `start` looking for a `STORE_DIRNAME/` directory. + + Stops, like git, at the filesystem root or as soon as it would cross + onto a different filesystem/mount than `start`, rather than wandering + indefinitely upwards into unrelated system directories. + + Args: + start: Directory to start the search from. + + Returns: + The path to the `STORE_DIRNAME/` directory itself, the first one found + at or above `start`. + + Raises: + StoreNotFoundError: If no directory at or above `start`, up to the + nearest filesystem boundary, contains a `STORE_DIRNAME/` directory. + """ + current = start.resolve() + boundary_dev = current.stat().st_dev + + while True: + candidate = current / STORE_DIRNAME + if candidate.is_dir(): + return candidate + + parent = current.parent + if parent == current or parent.stat().st_dev != boundary_dev: + break + current = parent + + raise StoreNotFoundError( + f"Could not find a `{STORE_DIRNAME}/` directory in `{start}` or any " + "of its parent directories up to the filesystem boundary. Run " + "`dstrack init` to create one, or point at one explicitly via " + f"`--root` or `{ROOT_PATH_ENV_VAR}`." + ) + + +def resolve_store_root(root: Path | str | None = None) -> Path: + """Resolve the path to the local store, `STORE_DIRNAME/`. + + Checked in order, the first one available wins: + 1. `root`, e.g. forwarded from a CLI `--root` option: the directory + `STORE_DIRNAME/` lives (or will be created) in. + 2. The `DSTRACK_ROOT_PATH` environment variable: same meaning as + `root`. + 3. Walking upwards from the current working directory until a + `STORE_DIRNAME/` directory is found, the same way git resolves `.git/`. + + Args: + root: Explicit parent directory of `STORE_DIRNAME/`, typically supplied + by a CLI `--root` option. Taken as-is, without checking that + `STORE_DIRNAME/` actually exists inside it. + + Returns: + The path to the `STORE_DIRNAME/` directory itself, as an absolute path. + + Raises: + StoreNotFoundError: `root` is `None`, `DSTRACK_ROOT_PATH` is + unset, and no `STORE_DIRNAME/` directory is found by walking up from + the current directory. + """ + if root is not None: + return Path(root).expanduser().resolve() / STORE_DIRNAME + + env_root = os.environ.get(ROOT_PATH_ENV_VAR) + if env_root: + return Path(env_root).expanduser().resolve() / STORE_DIRNAME + + return _find_store_root(get_invocation_path()) diff --git a/tests/test_cli.py b/tests/test_cli.py index b1ae65d..00c204e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -11,28 +11,6 @@ runner = CliRunner() -# --------------------------------------------------------------------------- -# _init_local_store_gitignore -# --------------------------------------------------------------------------- - - -def test_init_local_store_gitignore_creates_file(tmp_path: Path) -> None: - """Writes a .gitignore file that ignores the cache directory.""" - _cli._init_local_store_gitignore(tmp_path) - - gitignore = tmp_path / ".gitignore" - assert gitignore.is_file() - assert gitignore.read_text() == ".cache/" - - -def test_init_local_store_gitignore_raises_if_exists(tmp_path: Path) -> None: - """Refuses to overwrite an already-existing .gitignore file.""" - (tmp_path / ".gitignore").write_text("existing content") - - with pytest.raises(FileExistsError): - _cli._init_local_store_gitignore(tmp_path) - - # --------------------------------------------------------------------------- # init_local_store # --------------------------------------------------------------------------- @@ -63,25 +41,53 @@ def test_init_local_store_raises_if_store_exists( init_local_store() -def test_init_local_store_rolls_back_on_gitignore_failure( +def test_init_local_store_rolls_back_on_template_failure( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - """Removes the partially created store and raises StoreInitError. + """Raises StoreInitError and leaves no trace if building the template fails. - if writing the .gitignore file fails partway through. + partway through. Since the store is built in a temporary directory + first, the destination path should never even be created. """ monkeypatch.chdir(tmp_path) - def _boom(path: Path) -> None: + def _boom(template: object, parent: Path) -> Path: raise RuntimeError("disk on fire") - monkeypatch.setattr(_cli, "_init_local_store_gitignore", _boom) + monkeypatch.setattr(_cli, "materialize", _boom) with pytest.raises(StoreInitError) as exc_info: init_local_store() assert isinstance(exc_info.value.__cause__, RuntimeError) assert not (tmp_path / ".dstrack").exists() + assert list(tmp_path.iterdir()) == [] + + +def test_init_local_store_does_not_touch_destination_until_built( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Builds the store fully in a temp dir before moving it into place. + + at the moment the template is materialized, the final destination + must not exist yet. + """ + monkeypatch.chdir(tmp_path) + store_path = tmp_path / ".dstrack" + + from dstrack import _store_template + + original_materialize = _store_template.materialize + + def _check_and_materialize(template: object, parent: Path) -> Path: + assert not store_path.exists() + return original_materialize(template, parent) + + monkeypatch.setattr(_cli, "materialize", _check_and_materialize) + + init_local_store() + + assert store_path.is_dir() # --------------------------------------------------------------------------- diff --git a/tests/test_paths.py b/tests/test_paths.py new file mode 100644 index 0000000..c779a78 --- /dev/null +++ b/tests/test_paths.py @@ -0,0 +1,123 @@ +"""Tests for path resolution.""" + +import os +from pathlib import Path + +import pytest + +from dstrack import paths +from dstrack.errors import StoreNotFoundError +from dstrack.paths import ROOT_PATH_ENV_VAR, resolve_store_root + +# --------------------------------------------------------------------------- +# resolve_store_root +# --------------------------------------------------------------------------- + + +def test_resolve_store_root_uses_root_when_given( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An explicit `cli_root` wins, regardless of env var or cwd.""" + monkeypatch.setenv(ROOT_PATH_ENV_VAR, str(tmp_path / "env-root")) + monkeypatch.chdir(tmp_path) + + root = tmp_path / "cli-root" + + assert resolve_store_root(root=root) == root.resolve() / paths.STORE_DIRNAME + + +def test_resolve_store_root_uses_env_var_when_no_cli_root( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The env var is used when no `cli_root` is given.""" + env_root = tmp_path / "env-root" + monkeypatch.setenv(ROOT_PATH_ENV_VAR, str(env_root)) + monkeypatch.chdir(tmp_path) + + assert resolve_store_root() == env_root.resolve() / paths.STORE_DIRNAME + + +def test_resolve_store_root_finds_dstrack_in_cwd( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Falls back to walking up from cwd, finding `STORE_DIRNAME/` right there.""" + monkeypatch.delenv(ROOT_PATH_ENV_VAR, raising=False) + (tmp_path / paths.STORE_DIRNAME).mkdir() + monkeypatch.chdir(tmp_path) + + assert resolve_store_root() == tmp_path.resolve() / paths.STORE_DIRNAME + + +def test_resolve_store_root_finds_dstrack_in_ancestor( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Walks up through parent directories until `STORE_DIRNAME/` is found.""" + monkeypatch.delenv(ROOT_PATH_ENV_VAR, raising=False) + (tmp_path / paths.STORE_DIRNAME).mkdir() + nested = tmp_path / "a" / "b" / "c" + nested.mkdir(parents=True) + monkeypatch.chdir(nested) + + assert resolve_store_root() == tmp_path.resolve() / paths.STORE_DIRNAME + + +def test_resolve_store_root_raises_when_not_found( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Raises `StoreNotFoundError` if no `STORE_DIRNAME/` exists anywhere above cwd.""" + monkeypatch.delenv(ROOT_PATH_ENV_VAR, raising=False) + monkeypatch.chdir(tmp_path) + + with pytest.raises(StoreNotFoundError): + resolve_store_root() + + +# --------------------------------------------------------------------------- +# _find_store_root +# --------------------------------------------------------------------------- + + +def test_find_store_root_error_message_mentions_overrides( + tmp_path: Path, +) -> None: + """Error message points users at both override mechanisms.""" + with pytest.raises(StoreNotFoundError) as exc_info: + paths._find_store_root(tmp_path) + + assert "--root" in str(exc_info.value) + assert ROOT_PATH_ENV_VAR in str(exc_info.value) + + +def test_find_store_root_stops_at_filesystem_boundary( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Never searches past a directory that lives on a different device. + + Even if `STORE_DIRNAME/` exists further up, on the other side of the + boundary, like git it must not be found. + """ + (tmp_path / paths.STORE_DIRNAME).mkdir() + nested = tmp_path / "mounted" + nested.mkdir() + + real_stat = Path.stat + + class _FakeStat: + """Wraps a real `stat_result`, overriding only `st_dev`.""" + + def __init__(self, real_result: os.stat_result, st_dev: int) -> None: + self._real_result = real_result + self.st_dev = st_dev + + def __getattr__(self, name: str) -> object: + return getattr(self._real_result, name) + + def fake_stat(self: Path, *args: object, **kwargs: object) -> _FakeStat: + on_mount = self == nested or nested in self.parents + real_result = real_stat(self, *args, **kwargs) + return _FakeStat(real_result, st_dev=1 if on_mount else 2) + + monkeypatch.setattr(Path, "stat", fake_stat) + + with pytest.raises(StoreNotFoundError): + paths._find_store_root(nested) diff --git a/tests/test_store_template.py b/tests/test_store_template.py new file mode 100644 index 0000000..1d6197d --- /dev/null +++ b/tests/test_store_template.py @@ -0,0 +1,65 @@ +"""Tests for the local store template in src/dstrack/_store_template.py.""" + +from pathlib import Path + +import pytest + +from dstrack._store_template import ( + STORE_TEMPLATE, + TemplateDir, + TemplateFile, + materialize, +) +from dstrack.paths import STORE_DIRNAME + +# --------------------------------------------------------------------------- +# STORE_TEMPLATE +# --------------------------------------------------------------------------- + + +def test_store_template_name_matches_store_dirname() -> None: + """The template's root directory name matches the canonical store dirname.""" + assert STORE_TEMPLATE.name == STORE_DIRNAME + + +# --------------------------------------------------------------------------- +# materialize +# --------------------------------------------------------------------------- + + +def test_materialize_creates_expected_structure(tmp_path: Path) -> None: + """Builds `.dstrack/`, `.dstrack/datasets/`, and `.dstrack/.gitignore`.""" + store_path = materialize(STORE_TEMPLATE, tmp_path) + + assert store_path == tmp_path / STORE_DIRNAME + assert store_path.is_dir() + assert (store_path / "datasets").is_dir() + assert (store_path / ".gitignore").is_file() + assert (store_path / ".gitignore").read_text() == ".cache/" + + +def test_materialize_raises_if_target_already_exists(tmp_path: Path) -> None: + """Refuses to overwrite an already-existing directory.""" + (tmp_path / STORE_DIRNAME).mkdir() + + with pytest.raises(FileExistsError): + materialize(STORE_TEMPLATE, tmp_path) + + +def test_materialize_handles_nested_directories(tmp_path: Path) -> None: + """Recurses into nested `TemplateDir` entries, not just top-level ones.""" + template = TemplateDir( + name="root", + entries=( + TemplateDir( + name="child", + entries=(TemplateFile(name="leaf.txt", content="hello"),), + ), + ), + ) + + root_path = materialize(template, tmp_path) + + leaf_path = root_path / "child" / "leaf.txt" + assert leaf_path.is_file() + assert leaf_path.read_text() == "hello"