Aluode/PerceptionLabPortable

1from __future__ import annotations2 3import base644import urllib5 6import requests7from requests.adapters import HTTPAdapter, Retry8from typing_extensions import override9 10from fsspec import AbstractFileSystem11from fsspec.spec import AbstractBufferedFile12 13 14class DatabricksException(Exception):15    """16    Helper class for exceptions raised in this module.17    """18 19    def __init__(self, error_code, message, details=None):20        """Create a new DatabricksException"""21        super().__init__(message)22 23        self.error_code = error_code24        self.message = message25        self.details = details26 27 28class DatabricksFileSystem(AbstractFileSystem):29    """30    Get access to the Databricks filesystem implementation over HTTP.31    Can be used inside and outside of a databricks cluster.32    """33 34    def __init__(self, instance, token, **kwargs):35        """36        Create a new DatabricksFileSystem.37 38        Parameters39        ----------40        instance: str41            The instance URL of the databricks cluster.42            For example for an Azure databricks cluster, this43            has the form adb-<some-number>.<two digits>.azuredatabricks.net.44        token: str45            Your personal token. Find out more46            here: https://docs.databricks.com/dev-tools/api/latest/authentication.html47        """48        self.instance = instance49        self.token = token50        self.session = requests.Session()51        self.retries = Retry(52            total=10,53            backoff_factor=0.05,54            status_forcelist=[408, 429, 500, 502, 503, 504],55        )56 57        self.session.mount("https://", HTTPAdapter(max_retries=self.retries))58        self.session.headers.update({"Authorization": f"Bearer {self.token}"})59 60        super().__init__(**kwargs)61 62    @override63    def _ls_from_cache(self, path) -> list[dict[str, str | int]] | None:64        """Check cache for listing65 66        Returns listing, if found (may be empty list for a directory that67        exists but contains nothing), None if not in cache.68        """69        self.dircache.pop(path.rstrip("/"), None)70 71        parent = self._parent(path)72        if parent in self.dircache:73            for entry in self.dircache[parent]:74                if entry["name"] == path.rstrip("/"):75                    if entry["type"] != "directory":76                        return [entry]77                    return []78            raise FileNotFoundError(path)79 80    def ls(self, path, detail=True, **kwargs):81        """82        List the contents of the given path.83 84        Parameters85        ----------86        path: str87            Absolute path88        detail: bool89            Return not only the list of filenames,90            but also additional information on file sizes91            and types.92        """93        try:94            out = self._ls_from_cache(path)95        except FileNotFoundError:96            # This happens if the `path`'s parent was cached, but `path` is not97            # there. This suggests that `path` is new since the parent was98            # cached. Attempt to invalidate parent's cache before continuing.99            self.dircache.pop(self._parent(path), None)100            out = None101 102        if not out:103            try:104                r = self._send_to_api(105                    method="get", endpoint="list", json={"path": path}106                )107            except DatabricksException as e:108                if e.error_code == "RESOURCE_DOES_NOT_EXIST":109                    raise FileNotFoundError(e.message) from e110 111                raise112            files = r.get("files", [])113            out = [114                {115                    "name": o["path"],116                    "type": "directory" if o["is_dir"] else "file",117                    "size": o["file_size"],118                }119                for o in files120            ]121            self.dircache[path] = out122 123        if detail:124            return out125        return [o["name"] for o in out]126 127    def makedirs(self, path, exist_ok=True):128        """129        Create a given absolute path and all of its parents.130 131        Parameters132        ----------133        path: str134            Absolute path to create135        exist_ok: bool136            If false, checks if the folder137            exists before creating it (and raises an138            Exception if this is the case)139        """140        if not exist_ok:141            try:142                # If the following succeeds, the path is already present143                self._send_to_api(144                    method="get", endpoint="get-status", json={"path": path}145                )146                raise FileExistsError(f"Path {path} already exists")147            except DatabricksException as e:148                if e.error_code == "RESOURCE_DOES_NOT_EXIST":149                    pass150 151        try:152            self._send_to_api(method="post", endpoint="mkdirs", json={"path": path})153        except DatabricksException as e:154            if e.error_code == "RESOURCE_ALREADY_EXISTS":155                raise FileExistsError(e.message) from e156 157            raise158        self.invalidate_cache(self._parent(path))159 160    def mkdir(self, path, create_parents=True, **kwargs):161        """162        Create a given absolute path and all of its parents.163 164        Parameters165        ----------166        path: str167            Absolute path to create168        create_parents: bool169            Whether to create all parents or not.170            "False" is not implemented so far.171        """172        if not create_parents:173            raise NotImplementedError174 175        self.mkdirs(path, **kwargs)176 177    def rm(self, path, recursive=False, **kwargs):178        """179        Remove the file or folder at the given absolute path.180 181        Parameters182        ----------183        path: str184            Absolute path what to remove185        recursive: bool186            Recursively delete all files in a folder.187        """188        try:189            self._send_to_api(190                method="post",191                endpoint="delete",192                json={"path": path, "recursive": recursive},193            )194        except DatabricksException as e:195            # This is not really an exception, it just means196            # not everything was deleted so far197            if e.error_code == "PARTIAL_DELETE":198                self.rm(path=path, recursive=recursive)199            elif e.error_code == "IO_ERROR":200                # Using the same exception as the os module would use here201                raise OSError(e.message) from e202 203            raise204        self.invalidate_cache(self._parent(path))205 206    def mv(207        self, source_path, destination_path, recursive=False, maxdepth=None, **kwargs208    ):209        """210        Move a source to a destination path.211 212        A note from the original [databricks API manual]213        (https://docs.databricks.com/dev-tools/api/latest/dbfs.html#move).214 215        When moving a large number of files the API call will time out after216        approximately 60s, potentially resulting in partially moved data.217        Therefore, for operations that move more than 10k files, we strongly218        discourage using the DBFS REST API.219 220        Parameters221        ----------222        source_path: str223            From where to move (absolute path)224        destination_path: str225            To where to move (absolute path)226        recursive: bool227            Not implemented to far.228        maxdepth:229            Not implemented to far.230        """231        if recursive:232            raise NotImplementedError233        if maxdepth:234            raise NotImplementedError235 236        try:237            self._send_to_api(238                method="post",239                endpoint="move",240                json={"source_path": source_path, "destination_path": destination_path},241            )242        except DatabricksException as e:243            if e.error_code == "RESOURCE_DOES_NOT_EXIST":244                raise FileNotFoundError(e.message) from e245            elif e.error_code == "RESOURCE_ALREADY_EXISTS":246                raise FileExistsError(e.message) from e247 248            raise249        self.invalidate_cache(self._parent(source_path))250        self.invalidate_cache(self._parent(destination_path))251 252    def _open(self, path, mode="rb", block_size="default", **kwargs):253        """254        Overwrite the base class method to make sure to create a DBFile.255        All arguments are copied from the base method.256 257        Only the default blocksize is allowed.258        """259        return DatabricksFile(self, path, mode=mode, block_size=block_size, **kwargs)260 261    def _send_to_api(self, method, endpoint, json):262        """263        Send the given json to the DBFS API264        using a get or post request (specified by the argument `method`).265 266        Parameters267        ----------268        method: str269            Which http method to use for communication; "get" or "post".270        endpoint: str271            Where to send the request to (last part of the API URL)272        json: dict273            Dictionary of information to send274        """275        if method == "post":276            session_call = self.session.post277        elif method == "get":278            session_call = self.session.get279        else:280            raise ValueError(f"Do not understand method {method}")281 282        url = urllib.parse.urljoin(f"https://{self.instance}/api/2.0/dbfs/", endpoint)283 284        r = session_call(url, json=json)285 286        # The DBFS API will return a json, also in case of an exception.287        # We want to preserve this information as good as possible.288        try:289            r.raise_for_status()290        except requests.HTTPError as e:291            # try to extract json error message292            # if that fails, fall back to the original exception293            try:294                exception_json = e.response.json()295            except Exception:296                raise e from None297 298            raise DatabricksException(**exception_json) from e299 300        return r.json()301 302    def _create_handle(self, path, overwrite=True):303        """304        Internal function to create a handle, which can be used to305        write blocks of a file to DBFS.306        A handle has a unique identifier which needs to be passed307        whenever written during this transaction.308        The handle is active for 10 minutes - after that a new309        write transaction needs to be created.310        Make sure to close the handle after you are finished.311 312        Parameters313        ----------314        path: str315            Absolute path for this file.316        overwrite: bool317            If a file already exist at this location, either overwrite318            it or raise an exception.319        """320        try:321            r = self._send_to_api(322                method="post",323                endpoint="create",324                json={"path": path, "overwrite": overwrite},325            )326            return r["handle"]327        except DatabricksException as e:328            if e.error_code == "RESOURCE_ALREADY_EXISTS":329                raise FileExistsError(e.message) from e330 331            raise332 333    def _close_handle(self, handle):334        """335        Close a handle, which was opened by :func:`_create_handle`.336 337        Parameters338        ----------339        handle: str340            Which handle to close.341        """342        try:343            self._send_to_api(method="post", endpoint="close", json={"handle": handle})344        except DatabricksException as e:345            if e.error_code == "RESOURCE_DOES_NOT_EXIST":346                raise FileNotFoundError(e.message) from e347 348            raise349 350    def _add_data(self, handle, data):351        """352        Upload data to an already opened file handle353        (opened by :func:`_create_handle`).354        The maximal allowed data size is 1MB after355        conversion to base64.356        Remember to close the handle when you are finished.357 358        Parameters359        ----------360        handle: str361            Which handle to upload data to.362        data: bytes363            Block of data to add to the handle.364        """365        data = base64.b64encode(data).decode()366        try:367            self._send_to_api(368                method="post",369                endpoint="add-block",370                json={"handle": handle, "data": data},371            )372        except DatabricksException as e:373            if e.error_code == "RESOURCE_DOES_NOT_EXIST":374                raise FileNotFoundError(e.message) from e375            elif e.error_code == "MAX_BLOCK_SIZE_EXCEEDED":376                raise ValueError(e.message) from e377 378            raise379 380    def _get_data(self, path, start, end):381        """382        Download data in bytes from a given absolute path in a block383        from [start, start+length].384        The maximum number of allowed bytes to read is 1MB.385 386        Parameters387        ----------388        path: str389            Absolute path to download data from390        start: int391            Start position of the block392        end: int393            End position of the block394        """395        try:396            r = self._send_to_api(397                method="get",398                endpoint="read",399                json={"path": path, "offset": start, "length": end - start},400            )401            return base64.b64decode(r["data"])402        except DatabricksException as e:403            if e.error_code == "RESOURCE_DOES_NOT_EXIST":404                raise FileNotFoundError(e.message) from e405            elif e.error_code in ["INVALID_PARAMETER_VALUE", "MAX_READ_SIZE_EXCEEDED"]:406                raise ValueError(e.message) from e407 408            raise409 410    def invalidate_cache(self, path=None):411        if path is None:412            self.dircache.clear()413        else:414            self.dircache.pop(path, None)415        super().invalidate_cache(path)416 417 418class DatabricksFile(AbstractBufferedFile):419    """420    Helper class for files referenced in the DatabricksFileSystem.421    """422 423    DEFAULT_BLOCK_SIZE = 1 * 2**20  # only allowed block size424 425    def __init__(426        self,427        fs,428        path,429        mode="rb",430        block_size="default",431        autocommit=True,432        cache_type="readahead",433        cache_options=None,434        **kwargs,435    ):436        """437        Create a new instance of the DatabricksFile.438 439        The blocksize needs to be the default one.440        """441        if block_size is None or block_size == "default":442            block_size = self.DEFAULT_BLOCK_SIZE443 444        assert block_size == self.DEFAULT_BLOCK_SIZE, (445            f"Only the default block size is allowed, not {block_size}"446        )447 448        super().__init__(449            fs,450            path,451            mode=mode,452            block_size=block_size,453            autocommit=autocommit,454            cache_type=cache_type,455            cache_options=cache_options or {},456            **kwargs,457        )458 459    def _initiate_upload(self):460        """Internal function to start a file upload"""461        self.handle = self.fs._create_handle(self.path)462 463    def _upload_chunk(self, final=False):464        """Internal function to add a chunk of data to a started upload"""465        self.buffer.seek(0)466        data = self.buffer.getvalue()467 468        data_chunks = [469            data[start:end] for start, end in self._to_sized_blocks(len(data))470        ]471 472        for data_chunk in data_chunks:473            self.fs._add_data(handle=self.handle, data=data_chunk)474 475        if final:476            self.fs._close_handle(handle=self.handle)477            return True478 479    def _fetch_range(self, start, end):480        """Internal function to download a block of data"""481        return_buffer = b""482        length = end - start483        for chunk_start, chunk_end in self._to_sized_blocks(length, start):484            return_buffer += self.fs._get_data(485                path=self.path, start=chunk_start, end=chunk_end486            )487 488        return return_buffer489 490    def _to_sized_blocks(self, length, start=0):491        """Helper function to split a range from 0 to total_length into blocksizes"""492        end = start + length493        for data_chunk in range(start, end, self.blocksize):494            data_start = data_chunk495            data_end = min(end, data_chunk + self.blocksize)496            yield data_start, data_end497