Infra: allow for multiple install_hooks() running at the same time

dagshub/common/api/repo.py

@@ -3,6 +3,8 @@
 from pathlib import Path, PurePosixPath
 
 import rich.progress
+from httpx import Response
+from tenacity import stop_after_attempt, wait_exponential, before_sleep_log, retry, retry_if_exception
 
 from dagshub.common.api.responses import (
     RepoAPIResponse,
@@ -54,6 +56,22 @@ class PathNotFoundError(Exception):
     pass
 
 
+class DagsHubHTTPError(Exception):
+    def __init__(self, msg: str, response: Response):
+        super().__init__()
+        self.msg = msg
+        self.response = response
+
+    def __str__(self):
+        return self.msg
+
+
+def _is_server_error_exception(exception: BaseException) -> bool:
+    if not isinstance(exception, DagsHubHTTPError):
+        return False
+    return exception.response.status_code >= 500
+
+
 class RepoAPI:
     def __init__(self, repo: str, host: Optional[str] = None, auth: Optional[Any] = None):
         """
@@ -89,7 +107,7 @@ def get_repo_info(self) -> RepoAPIResponse:
             error_msg = f"Got status code {res.status_code} when getting repository info."
             logger.error(error_msg)
             logger.debug(res.content)
-            raise RuntimeError(error_msg)
+            raise DagsHubHTTPError(error_msg, res)
         return dacite.from_dict(RepoAPIResponse, res.json())
 
     def get_branch_info(self, branch: str) -> BranchAPIResponse:
@@ -107,7 +125,7 @@ def get_branch_info(self, branch: str) -> BranchAPIResponse:
             error_msg = f"Got status code {res.status_code} when getting branch."
             logger.error(error_msg)
             logger.debug(res.content)
-            raise RuntimeError(error_msg)
+            raise DagsHubHTTPError(error_msg, res)
 
         return dacite.from_dict(BranchAPIResponse, res.json())
 
@@ -126,7 +144,7 @@ def get_commit_info(self, sha: str) -> CommitAPIResponse:
             error_msg = f"Got status code {res.status_code} when getting commit."
             logger.error(error_msg)
             logger.debug(res.content)
-            raise RuntimeError(error_msg)
+            raise DagsHubHTTPError(error_msg, res)
 
         return dacite.from_dict(CommitAPIResponse, res.json()["commit"])
 
@@ -142,7 +160,7 @@ def get_connected_storages(self) -> List[StorageAPIEntry]:
             error_msg = f"Got status code {res.status_code} when getting repository info."
             logger.error(error_msg)
             logger.debug(res.content)
-            raise RuntimeError(error_msg)
+            raise DagsHubHTTPError(error_msg, res)
 
         return [dacite.from_dict(StorageAPIEntry, storage_entry) for storage_entry in res.json()]
 
@@ -164,7 +182,7 @@ def list_path(self, path: str, revision: Optional[str] = None, include_size: boo
             error_msg = f"Got status code {res.status_code} when listing path {path}"
             logger.error(error_msg)
             logger.debug(res.content)
-            raise RuntimeError(error_msg)
+            raise DagsHubHTTPError(error_msg, res)
 
         content = res.json()
         if type(content) is dict:
@@ -194,7 +212,7 @@ def _get():
                 error_msg = f"Got status code {res.status_code} when listing path {path}"
                 logger.error(error_msg)
                 logger.debug(res.content)
-                raise RuntimeError(error_msg)
+                raise DagsHubHTTPError(error_msg, res)
 
             content = res.json()
             if "entries" not in content:
@@ -218,6 +236,12 @@ def _get():
 
         return entries
 
+    @retry(
+        retry=retry_if_exception(_is_server_error_exception),
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
     def get_file(self, path: str, revision: Optional[str] = None) -> bytes:
         """
         Download file from repo.
@@ -229,16 +253,22 @@ def get_file(self, path: str, revision: Optional[str] = None) -> bytes:
         Returns:
             bytes: The content of the file.
         """
-        res = self._http_request("GET", self.raw_api_url(path, revision))
+        res = self._http_request("GET", self.raw_api_url(path, revision), timeout=None)
         if res.status_code == 404:
             raise PathNotFoundError(f"Path {path} not found")
         elif res.status_code >= 400:
             error_msg = f"Got status code {res.status_code} when getting file {path}"
             logger.error(error_msg)
             logger.debug(res.content)
-            raise RuntimeError(error_msg)
+            raise DagsHubHTTPError(error_msg, res)
         return res.content
 
+    @retry(
+        retry=retry_if_exception(_is_server_error_exception),
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
     def get_storage_file(self, path: str) -> bytes:
         """
         Download file from a connected storage bucket.
@@ -258,7 +288,7 @@ def get_storage_file(self, path: str) -> bytes:
             error_msg = f"Got status code {res.status_code} when getting file {path}"
             logger.error(error_msg)
             logger.debug(res.content)
-            raise RuntimeError(error_msg)
+            raise DagsHubHTTPError(error_msg, res)
         return res.content
 
     def _get_files_in_path(

dagshub/common/helpers.py

@@ -97,7 +97,7 @@ def prompt_user(prompt, default=False) -> bool:
     return prompt_response == "y"
 
 
-def log_message(msg, logger=None):
+def log_message(msg: str, logger=None):
     """
     Logs message to the info of the logger + prints, unless the printing was suppressed
     """

dagshub/streaming/__init__.py

@@ -1,4 +1,4 @@
-from .filesystem import DagsHubFilesystem, install_hooks, uninstall_hooks
+from .filesystem import DagsHubFilesystem, install_hooks, uninstall_hooks, get_mounted_filesystems
 
 try:
     from .mount import mount
@@ -15,4 +15,10 @@ def mount(*args, **kwargs):
         print(error)
 
 
-__all__ = [DagsHubFilesystem.__name__, install_hooks.__name__, mount.__name__, uninstall_hooks.__name__]
+__all__ = [
+    DagsHubFilesystem.__name__,
+    install_hooks.__name__,
+    mount.__name__,
+    uninstall_hooks.__name__,
+    get_mounted_filesystems.__name__,
+]

dagshub/streaming/dataclasses.py

@@ -1,14 +1,16 @@
+import os
 from dataclasses import dataclass
+from os import PathLike
 from pathlib import Path
-from typing import Optional, TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING, Union, Tuple
 
 try:
     from functools import cached_property
 except ImportError:
     from cached_property import cached_property
 
 if TYPE_CHECKING:
-    from dagshub.streaming import DagsHubFilesystem
+    from filesystem import DagsHubFilesystem
 
 storage_schemas = ["s3", "gs", "azure"]
 
@@ -24,15 +26,35 @@ class DagshubPath:
         relative_path (Optional[Path]): Path relative to the root of the encapsulating FileSystem.
                                         If None, path is outside the FS
         original_path (Path): Original path as it was accessed by the user
+        is_binary_path_requested (bool): For functions like scandir and listdir that have
+            different behaviour whether user requested a string or a binary path
     """
 
-    # TODO: this couples this class hard to the fs, need to decouple later
-    fs: "DagsHubFilesystem"  # Actual type is DagsHubFilesystem, but imports are wonky
-    absolute_path: Optional[Path]
-    relative_path: Optional[Path]
-    original_path: Optional[Path]
+    def __init__(self, fs: "DagsHubFilesystem", file_path: Union[str, bytes, PathLike, "DagshubPath"]):
+        self.fs = fs
+        self.is_binary_path_requested = isinstance(file_path, bytes)
+        self.absolute_path, self.relative_path, self.original_path = self.parse_path(file_path)
 
-    def __post_init__(self):
+    def parse_path(self, file_path: Union[str, bytes, PathLike, "DagshubPath"]) -> Tuple[Path, Optional[Path], Path]:
+        if isinstance(file_path, DagshubPath):
+            self.is_binary_path_requested = file_path.is_binary_path_requested
+            if file_path.fs != self.fs:
+                relativized = DagshubPath(self.fs, file_path.absolute_path)
+                return relativized.absolute_path, relativized.relative_path, relativized.original_path
+            return file_path.absolute_path, file_path.relative_path, file_path.original_path
+        if isinstance(file_path, bytes):
+            file_path = os.fsdecode(file_path)
+        orig_path = Path(file_path)
+        abspath = Path(os.path.abspath(file_path))
+        try:
+            relpath = abspath.relative_to(os.path.abspath(self.fs.project_root))
+            if str(relpath).startswith("<"):
+                return abspath, None, orig_path
+            return abspath, relpath, orig_path
+        except ValueError:
+            return abspath, None, orig_path
+
+    def handle_storages(self):
         # Handle storage paths - translate s3:/bla-bla to .dagshub/storage/s3/bla-bla
         if self.relative_path is not None:
             str_path = self.relative_path.as_posix()
@@ -41,9 +63,11 @@ def __post_init__(self):
                     str_path = str_path[len(storage_schema) + 2 :]
                     self.relative_path = Path(".dagshub/storage") / storage_schema / str_path
                     self.absolute_path = self.fs.project_root / self.relative_path
+                    break
 
     @cached_property
     def name(self):
+        assert self.absolute_path is not None
         return self.absolute_path.name
 
     @cached_property
@@ -56,10 +80,11 @@ def is_storage_path(self):
         Is path a storage path (stored in a bucket)
         Those paths are accessible via a path like `.dagshub/storage/s3/bucket/...`
         """
+        if self.relative_path is None:
+            return False
         return self.relative_path.as_posix().startswith(".dagshub/storage")
 
-    @cached_property
-    def is_passthrough_path(self):
+    def is_passthrough_path(self, fs: "DagsHubFilesystem"):
         """
         Is path a "passthrough" path
         A passthrough path is a path that the FS ignores when trying to look up if the file exists on DagsHub
@@ -68,17 +93,144 @@ def is_passthrough_path(self):
                 If you need to read with streaming from a .dvc folder (to read config for example), please pull the repo
             - Any /site-packages/ folder - if you have a venv in your repo, python will try to find packages there.
         """
+        if self.relative_path is None:
+            return True
         str_path = self.relative_path.as_posix()
         if "/site-packages/" in str_path or str_path.endswith("/site-packages"):
             return True
         if str_path.startswith((".git/", ".dvc/")) or str_path in (".git", ".dvc"):
             return True
-        return any((self.relative_path.match(glob) for glob in self.fs.exclude_globs))
+        return any((self.relative_path.match(glob) for glob in fs.exclude_globs))
 
     def __truediv__(self, other):
-        return DagshubPath(
-            absolute_path=self.absolute_path / other,
-            relative_path=self.relative_path / other,
-            original_path=self.original_path / other,
-            fs=self.fs,
+        new = DagshubPath(
+            self.fs,
+            self.original_path / other,
         )
+        new.is_binary_path_requested = self.is_binary_path_requested
+        return new
+
+
+class DagshubScandirIterator:
+    def __init__(self, iterator):
+        self._iterator = iterator
+
+    def __iter__(self):
+        return self._iterator
+
+    def __next__(self):
+        return self._iterator.__next__()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        return self
+
+
+class DagshubStatResult:
+    def __init__(
+        self, fs: "DagsHubFilesystem", path: DagshubPath, is_directory: bool, custom_size: Optional[int] = None
+    ):
+        self._fs = fs
+        self._path = path
+        self._is_directory = is_directory
+        self._custom_size = custom_size
+        self._true_stat: Optional[os.stat_result] = None
+        assert not self._is_directory  # TODO make folder stats lazy?
+
+    def __getattr__(self, name: str):
+        if not name.startswith("st_"):
+            raise AttributeError
+        if self._true_stat is not None:
+            return os.stat_result.__getattribute__(self._true_stat, name)
+        if name == "st_uid":
+            return os.getuid()
+        elif name == "st_gid":
+            return os.getgid()
+        elif name == "st_atime" or name == "st_mtime" or name == "st_ctime":
+            return 0
+        elif name == "st_mode":
+            return 0o100644
+        elif name == "st_size":
+            if self._custom_size is not None:
+                return self._custom_size
+            return 1100  # hardcoded size because size requests take a disproportionate amount of time
+        self._fs.open(self._path)
+        self._true_stat = self._fs.original_stat(self._path.absolute_path)
+        return os.stat_result.__getattribute__(self._true_stat, name)
+
+    def __repr__(self):
+        inner = repr(self._true_stat) if self._true_stat is not None else "pending..."
+        return f"dagshub_stat_result({inner}, path={self._path})"
+
+
+class DagshubDirEntry:
+    def __init__(self, fs: "DagsHubFilesystem", path: DagshubPath, is_directory: bool = False, is_binary: bool = False):
+        self._fs = fs
+        self._path = path
+        self._is_directory = is_directory
+        self._is_binary = is_binary
+        self._true_direntry: Optional[os.DirEntry] = None
+
+    @property
+    def name(self):
+        if self._true_direntry is not None:
+            name = self._true_direntry.name
+        else:
+            name = self._path.name
+        return os.fsencode(name) if self._is_binary else name
+
+    @property
+    def path(self):
+        if self._true_direntry is not None:
+            path = self._true_direntry.path
+        else:
+            path = str(self._path.original_path)
+        return os.fsencode(path) if self._is_binary else path
+
+    def is_dir(self):
+        if self._true_direntry is not None:
+            return self._true_direntry.is_dir()
+        else:
+            return self._is_directory
+
+    def is_file(self):
+        if self._true_direntry is not None:
+            return self._true_direntry.is_file()
+        else:
+            # TODO: Symlinks should return false
+            return not self._is_directory
+
+    def stat(self):
+        if self._true_direntry is not None:
+            return self._true_direntry.stat()
+        else:
+            return self._fs.stat(self._path.original_path)
+
+    def __getattr__(self, name: str):
+        if name == "_true_direntry":
+            raise AttributeError
+        if self._true_direntry is not None:
+            return os.DirEntry.__getattribute__(self._true_direntry, name)
+
+        # Either create a dir, or download the file
+        if self._is_directory:
+            self._fs.mkdirs(self._path.absolute_path)
+        else:
+            self._fs.open(self._path.absolute_path)
+
+        for direntry in self._fs.original_stat(self._path.original_path):
+            if direntry.name == self._path.name:
+                self._true_direntry = direntry
+                return os.DirEntry.__getattribute__(self._true_direntry, name)
+        else:
+            raise FileNotFoundError
+
+    def __repr__(self):
+        cached = " (cached)" if self._true_direntry is not None else ""
+        return f"<dagshub_DirEntry '{self.name}'{cached}>"
+
+
+PathType = Union[str, int, bytes, PathLike]
+PathTypeWithDagshubPath = Union[PathType, DagshubPath]