star

venv/lib/python3.9/site-packages/databricks/sdk/mixins/workspace.py

@@ -0,0 +1,117 @@
+from typing import Any, BinaryIO, Iterator, Optional, Union
+
+from ..core import DatabricksError
+from ..service.workspace import (ExportFormat, ImportFormat, Language,
+                                 ObjectInfo, ObjectType, WorkspaceAPI)
+
+
+def _fqcn(x: Any) -> str:
+    return f"{x.__module__}.{x.__name__}"
+
+
+class WorkspaceExt(WorkspaceAPI):
+    __doc__ = WorkspaceAPI.__doc__
+
+    def list(
+        self,
+        path: str,
+        *,
+        notebooks_modified_after: Optional[int] = None,
+        recursive: Optional[bool] = False,
+        **kwargs,
+    ) -> Iterator[ObjectInfo]:
+        """List workspace objects
+
+        :param recursive: bool
+            Optionally invoke recursive traversal
+
+        :returns: Iterator of workspaceObjectInfo
+        """
+        parent_list = super().list
+        queue = [path]
+        while queue:
+            path, queue = queue[0], queue[1:]
+            for object_info in parent_list(path, notebooks_modified_after=notebooks_modified_after):
+                if recursive and object_info.object_type == ObjectType.DIRECTORY:
+                    queue.append(object_info.path)
+                    continue
+                yield object_info
+
+    def upload(
+        self,
+        path: str,
+        content: Union[bytes, BinaryIO],
+        *,
+        format: Optional[ImportFormat] = None,
+        language: Optional[Language] = None,
+        overwrite: Optional[bool] = False,
+    ) -> None:
+        """
+        Uploads a workspace object (for example, a notebook or file) or the contents of an entire
+        directory (`DBC` format).
+
+        Errors:
+         * `RESOURCE_ALREADY_EXISTS`: if `path` already exists no `overwrite=True`.
+         * `INVALID_PARAMETER_VALUE`: if `format` and `content` values are not compatible.
+
+        :param path:     target location of the file on workspace.
+        :param content:  the contents as either raw binary data `bytes` or a file-like the file-like `io.BinaryIO` of the `path` contents.
+        :param format:   By default, `ImportFormat.SOURCE`. If using `ImportFormat.AUTO` the `path`
+                         is imported or exported as either a workspace file or a notebook, depending
+                         on an analysis of the `item`’s extension and the header content provided in
+                         the request. In addition, if the `path` is imported as a notebook, then
+                         the `item`’s extension is automatically removed.
+        :param language: Only required if using `ExportFormat.SOURCE`.
+        """
+        if format is not None and not isinstance(format, ImportFormat):
+            raise ValueError(f"format is expected to be {_fqcn(ImportFormat)}, but got {_fqcn(format.__class__)}")
+        if (not format or format == ImportFormat.SOURCE) and not language:
+            suffixes = {
+                ".py": Language.PYTHON,
+                ".sql": Language.SQL,
+                ".scala": Language.SCALA,
+                ".R": Language.R,
+            }
+            for sfx, lang in suffixes.items():
+                if path.endswith(sfx):
+                    language = lang
+                    break
+        if language is not None and not isinstance(language, Language):
+            raise ValueError(f"language is expected to be {_fqcn(Language)}, but got {_fqcn(language.__class__)}")
+        data = {"path": path}
+        if format:
+            data["format"] = format.value
+        if language:
+            data["language"] = language.value
+        if overwrite:
+            data["overwrite"] = "true"
+        try:
+            return self._api.do(
+                "POST",
+                "/api/2.0/workspace/import",
+                files={"content": content},
+                data=data,
+            )
+        except DatabricksError as e:
+            if e.error_code == "INVALID_PARAMETER_VALUE":
+                msg = f"Perhaps you forgot to specify the `format=ImportFormat.AUTO`. {e}"
+                raise DatabricksError(message=msg, error_code=e.error_code)
+            else:
+                raise e
+
+    def download(self, path: str, *, format: Optional[ExportFormat] = None) -> BinaryIO:
+        """
+        Downloads notebook or file from the workspace
+
+        :param path:     location of the file or notebook on workspace.
+        :param format:   By default, `ExportFormat.SOURCE`. If using `ExportFormat.AUTO` the `path`
+                         is imported or exported as either a workspace file or a notebook, depending
+                         on an analysis of the `item`’s extension and the header content provided in
+                         the request.
+        :return:         file-like `io.BinaryIO` of the `path` contents.
+        """
+        query = {"path": path, "direct_download": "true"}
+        if format:
+            query["format"] = format.value
+        response = self._api.do("GET", "/api/2.0/workspace/export", query=query, raw=True)
+        return response["contents"]