review and refactor export, fix tests again

Author: shcheklein

src/datachain/client/local.py

@@ -22,53 +22,6 @@ class FileClient(Client):
     PREFIX = "file://"
     protocol = "file"
 
-    @classmethod
-    def is_path_in(
-        cls,
-        output: str | os.PathLike[str],
-        dst: str,
-    ) -> bool:
-        """Return whether `dst` is safe to write under local `output`.
-
-        Accepts both plain OS paths and `file://` urlpaths. For other schemes,
-        returns False.
-        """
-
-        from fsspec.utils import stringify_path
-
-        output_str = stringify_path(output)
-
-        # Only handle local paths + file:// urlpaths.
-        if "://" in dst and not dst.startswith(cls.PREFIX):
-            return False
-        if "://" in output_str and not output_str.startswith(cls.PREFIX):
-            return False
-
-        output_os = (
-            LocalFileSystem._strip_protocol(output_str)
-            if output_str.startswith(cls.PREFIX)
-            else output_str
-        )
-        dst_os = (
-            LocalFileSystem._strip_protocol(dst) if dst.startswith(cls.PREFIX) else dst
-        )
-
-        # Use abspath (makes relative paths absolute based on CWD, collapses
-        # .. and .) followed by normcase (lowercases on Windows) for deterministic,
-        # filesystem-independent comparison.  resolve(strict=False) is avoided
-        # because its symlink / junction behaviour varies across Python versions
-        # and can differ depending on which path components already exist.
-        output_normed = os.path.normcase(os.path.abspath(output_os))
-        dst_normed = os.path.normcase(os.path.abspath(dst_os))
-
-        # Destination must be a file path under output, not the output dir itself.
-        if dst_normed == output_normed:
-            return False
-
-        # Ensure dst starts with the output prefix followed by a separator,
-        # so that output="/foo/bar" does not match dst="/foo/bar2".
-        return dst_normed.startswith(output_normed + os.sep)
-
     def __init__(
         self,
         name: str,
@@ -171,7 +124,10 @@ async def get_current_etag(self, file: "File") -> str:
     @classmethod
     def rel_path_for_file(cls, file: "File") -> str:
         path = file.path
-        cls.validate_local_relpath(file.source, path)
+        try:
+            cls.validate_file_relpath(path)
+        except ValueError as exc:
+            raise FileError(str(exc), file.source, path) from None
 
         normpath = os.path.normpath(path)
         normpath = Path(normpath).as_posix()
@@ -184,43 +140,57 @@ def rel_path_for_file(cls, file: "File") -> str:
 
         return normpath
 
-    @classmethod
-    def validate_local_relpath(cls, source: str, path: str) -> None:
-        """Validate a local relative path string.
+    @staticmethod
+    def validate_file_relpath(path: str) -> None:
+        """Ensure a file's relative path is safe to use on the local filesystem.
 
-        Used both for local `File.path` values and for local export destination
-        suffixes. Reject inputs that would require normalization (e.g. empty
-        segments) or that could be unsafe on the local filesystem.
+        Rejects absolute, traversal, and malformed paths that could escape the
+        intended directory or require implicit normalization. On Windows,
+        backslashes and drive-letter prefixes are handled as separators/absolute.
         """
 
         if not path:
-            raise FileError("path must not be empty", source, path)
+            raise ValueError(f"unsafe file path {path!r}: must not be empty")
 
-        if path.endswith("/"):
-            raise FileError("path must not be a directory", source, path)
+        # On Windows, backslash is a path separator — normalize so the
+        # checks below work uniformly.  On Linux/macOS, backslash is a legal
+        # filename character and must not be reinterpreted as a separator.
+        if os.name == "nt":
+            canonical = path.replace("\\", "/")
+        else:
+            canonical = path
 
-        raw_posix = path.replace("\\", "/")
+        if canonical.endswith("/"):
+            raise ValueError(f"unsafe file path {path!r}: must not be a directory")
 
         # Disallow absolute paths; local file paths are interpreted relative to
         # the source/output prefix.
-        if raw_posix.startswith("/"):
-            raise FileError("path must not be absolute", source, path)
+        if canonical.startswith("/"):
+            raise ValueError(f"unsafe file path {path!r}: must not be absolute")
 
         # On Windows, a drive-letter prefix like "C:/" is absolute even
-        # without a leading "/".
-        if len(raw_posix) >= 2 and raw_posix[0].isalpha() and raw_posix[1] == ":":
-            raise FileError("path must not be absolute", source, path)
+        # without a leading "/".  On Unix, colons are legal in filenames,
+        # so only enforce this on Windows.
+        if (
+            os.name == "nt"
+            and len(canonical) >= 2
+            and canonical[0].isalpha()
+            and canonical[1] == ":"
+        ):
+            raise ValueError(f"unsafe file path {path!r}: must not be absolute")
 
         # Disallow empty segments (e.g. 'dir//file.txt') to avoid implicit
         # normalization.
-        if "//" in raw_posix:
-            raise FileError("path must not contain empty segments", source, path)
+        if "//" in canonical:
+            raise ValueError(
+                f"unsafe file path {path!r}: must not contain empty segments"
+            )
 
         # Disallow dot segments like '.' or '..' (even if they could be
         # normalized away) because they can make I/O and exports unsafe.
-        raw_parts = raw_posix.split("/")
+        raw_parts = canonical.split("/")
         if any(part in (".", "..") for part in raw_parts):
-            raise FileError("path must not contain '.' or '..'", source, path)
+            raise ValueError(f"unsafe file path {path!r}: must not contain '.' or '..'")
 
     @staticmethod
     def _has_drive_letter(source: str) -> bool:

src/datachain/fs/utils.py

@@ -1,3 +1,4 @@
+import os
 from typing import TYPE_CHECKING
 
 from fsspec.implementations.local import LocalFileSystem
@@ -6,6 +7,23 @@
     from fsspec import AbstractFileSystem
 
 
+def is_subpath(parent: str, child: str) -> bool:
+    """True iff *child* is strictly inside *parent* (path-traversal guard).
+
+    Both paths must be absolute OS paths. Comparison is case-insensitive on
+    Windows.
+    """
+    assert os.path.isabs(parent), f"parent must be absolute: {parent!r}"
+    assert os.path.isabs(child), f"child must be absolute: {child!r}"
+
+    parent_normed = os.path.normcase(parent)
+    child_normed = os.path.normcase(child)
+
+    if child_normed == parent_normed:
+        return False
+    return child_normed.startswith(parent_normed + os.sep)
+
+
 def _isdir(fs: "AbstractFileSystem", path: str) -> bool:
     info = fs.info(path)
     return info["type"] == "directory" or (

src/datachain/lib/file.py

@@ -633,45 +633,84 @@ def export(
         link_type: Literal["copy", "symlink"] = "copy",
         client_config: dict | None = None,
     ) -> None:
-        """Export file to new location."""
+        """Copy or link this file into an output directory.
+
+        Args:
+            output: Destination directory (local path or cloud prefix).
+            placement: How to build the path under *output*:
+
+                - ``"fullpath"`` (default) — ``output/bucket/dir/file.txt``
+                - ``"filepath"`` — ``output/dir/file.txt``
+                - ``"filename"`` — ``output/file.txt``
+                - ``"etag"`` — ``output/<etag>.txt``
+            use_cache: If True, download to local cache first.  Also
+                required for symlinking remote files.
+            link_type: ``"copy"`` (default) or ``"symlink"``.
+                Symlink falls back to copy for virtual files and for
+                remote files when *use_cache* is False.
+            client_config: Extra kwargs forwarded to the storage client.
+
+        Example:
+            ```py
+            # flat export by filename
+            f.export("./export", placement="filename")
+
+            # export to a cloud prefix
+            f.export("s3://output-bucket/results", placement="filepath")
+
+            # pass storage credentials via client_config
+            f.export("s3://bucket/out", client_config={"aws_access_key_id": "…"})
+
+            # symlink from local cache (avoids re-downloading)
+            f.export("./local_out", use_cache=True, link_type="symlink")
+            ```
+        """
         if self._catalog is None:
             raise RuntimeError("Cannot export file: catalog is not set")
 
         self._caching_enabled = use_cache
 
         suffix = self._get_destination_suffix(placement)
-        # Normalize backslashes to forward slashes so posixpath.join produces
-        # consistent separator paths on Windows.
-        output_fspath = os.fspath(output).replace("\\", "/")
-        dst = posixpath.join(output_fspath, suffix)
-        dst_dir = posixpath.dirname(dst)
-        client: Client = self._catalog.get_client(dst_dir, **(client_config or {}))
-
-        # If we're exporting to a local directory, ensure the resolved destination
-        # stays within the chosen output directory. This protects against traversal
-        # and absolute-path suffixes even when the source key is cloud/opaque.
+        output_fspath = os.fspath(output)
+        # On Windows, normalize backslash separators to forward slashes so
+        # posixpath.join produces consistent paths. On Linux/macOS backslash
+        # is a legal filename character and must not be replaced.
+        if os.name == "nt":
+            output_fspath = output_fspath.replace("\\", "/")
+
+        # Resolve output to absolute immediately so all derived paths are
+        # absolute and don't need further normalization.
+        output_abs = os.path.abspath(output_fspath)
+        client: Client = self._catalog.get_client(output_abs, **(client_config or {}))
+        dst_abs = posixpath.join(output_abs, suffix)
+
+        # Traversal safety: for local exports, validate the suffix
         if client.PREFIX == "file://":
             from datachain.client.local import FileClient
+            from datachain.fs.utils import is_subpath
 
-            FileClient.validate_local_relpath(os.fspath(output), suffix)
+            try:
+                FileClient.validate_file_relpath(suffix)
+            except ValueError as exc:
+                raise FileError(str(exc), stringify_path(output), suffix) from None
 
-            if not FileClient.is_path_in(output, dst):
+            if not is_subpath(output_abs, dst_abs):
                 raise FileError(
                     "destination is not within output directory",
                     stringify_path(output),
-                    dst,
+                    dst_abs,
                 )
 
-        client.fs.makedirs(dst_dir, exist_ok=True)
+        client.fs.makedirs(posixpath.dirname(dst_abs), exist_ok=True)
 
         if link_type == "symlink":
             try:
-                return self._symlink_to(dst)
+                return self._symlink_to(dst_abs)
             except OSError as exc:
                 if exc.errno not in (errno.ENOTSUP, errno.EXDEV, errno.ENOSYS):
                     raise
 
-        self.save(dst, client_config=client_config)
+        self.save(dst_abs, client_config=client_config)
 
     def _set_stream(
         self,

tests/unit/lib/test_file.py

@@ -152,9 +152,10 @@ def test_export_placements_build_expected_destination_for_each_source(
     file.export(output, placement="filepath", use_cache=False)
     file.export(output, placement="fullpath", use_cache=False)
 
-    # export() normalises backslashes to forward slashes via
-    # os.fspath(output).replace("\\", "/"), so the expected prefix must match.
-    expected_prefix = os.fspath(output).replace("\\", "/")
+    # export() resolves output to absolute (os.path.abspath) and normalizes
+    # backslashes to forward slashes on Windows, so build the expected prefix
+    # the same way.
+    expected_prefix = os.path.abspath(os.fspath(output)).replace("\\", "/")
 
     expected_fullpath_suffix = (
         f"{expected_fullpath_prefix}/dir1/dir2/test.txt"
@@ -220,19 +221,19 @@ def test_export_local_output_allows_literal_percent_encoded_traversal(
 @pytest.mark.parametrize(
     "source,path,is_error,expected",
     [
-        ("", "", True, r"path must not be empty"),
-        ("", ".", True, r"path must not contain"),
-        ("", "..", True, r"path must not contain"),
-        ("", "/abs/file.txt", True, r"path must not be absolute"),
+        ("", "", True, r"must not be empty"),
+        ("", ".", True, r"must not contain"),
+        ("", "..", True, r"must not contain"),
+        ("", "/abs/file.txt", True, r"must not be absolute"),
         ("", "file#hash.txt", False, "file#hash.txt"),
         ("", "./dir/../file#hash.txt", True, r"must not contain"),
         ("", "../escape.txt", True, r"must not contain"),
         ("", "dir//file.txt", True, r"must not contain empty segments"),
         ("", "dir/", True, r"must not be a directory"),
-        ("file:///bucket", "", True, r"path must not be empty"),
-        ("file:///bucket", ".", True, r"path must not contain"),
-        ("file:///bucket", "..", True, r"path must not contain"),
-        ("file:///bucket", "/abs/file.txt", True, r"path must not be absolute"),
+        ("file:///bucket", "", True, r"must not be empty"),
+        ("file:///bucket", ".", True, r"must not contain"),
+        ("file:///bucket", "..", True, r"must not contain"),
+        ("file:///bucket", "/abs/file.txt", True, r"must not be absolute"),
         ("file:///bucket", "file#hash.txt", False, "file:///bucket/file#hash.txt"),
         (
             "file:///bucket",
@@ -306,17 +307,17 @@ def test_get_uri_contract(
 @pytest.mark.parametrize(
     "source,path,is_error,expected",
     [
-        ("", "", True, r"path must not be empty"),
-        ("", ".", True, r"path must not contain"),
-        ("", "..", True, r"path must not contain"),
-        ("", "/abs/file.txt", True, r"path must not be absolute"),
+        ("", "", True, r"must not be empty"),
+        ("", ".", True, r"must not contain"),
+        ("", "..", True, r"must not contain"),
+        ("", "/abs/file.txt", True, r"must not be absolute"),
         ("", "file.txt", False, "file.txt"),
         ("", "../escape.txt", True, r"must not contain"),
         ("", "dir/", True, r"must not be a directory"),
-        ("file:///bucket", "", True, r"path must not be empty"),
-        ("file:///bucket", ".", True, r"path must not contain"),
-        ("file:///bucket", "..", True, r"path must not contain"),
-        ("file:///bucket", "/abs/file.txt", True, r"path must not be absolute"),
+        ("file:///bucket", "", True, r"must not be empty"),
+        ("file:///bucket", ".", True, r"must not contain"),
+        ("file:///bucket", "..", True, r"must not contain"),
+        ("file:///bucket", "/abs/file.txt", True, r"must not be absolute"),
         pytest.param(
             "file:///bucket",
             "file#hash.txt",