Adds Databricks backend (#325)

* Initial DB framework, ls works (no details)

* Switches to using folder_url instead of path

* pin_exists() works now

* pin_versions() works

* pin_open() works

* Makes _list_folders into _list_items

* Adds mkdir & put, gets pin_write() working

* Formatting improvements

* Removes conditional to use a custom board object

* Improvements to file reading, and adds initial rm

* pin_delete() works, new 'exists' aproach

* Adds some test pieces

* Partially runs tests

* Adds support for `detail` in ls

* Figures out how to properly cache the fs

* Fixes teardown for tests

* extends `rm` one more level (fix later)

* No errors in tests, moving on to addressing failures

* Adds recursive file/folder mapper

* fully recursive put (clean later)

* Improvements to _map_put

* Starts moving to discrete functions

* Finishes moving everything to discrete functions

* Fixes typo

* Creates custom cache class, fixes issue with reading wrong pin

* Removes _open from PinsDBCache after confirming that it's not needed

* Removes data.csv file

* Adds get() and _databricks_get()

* Removes test data

* Fixes issue with not handling board versioning for test boards

* Passes constructor test

* Removes PinsDBCache and manual registration of dbc

* Emulates the structure of the other test boards

* Adds notes to the function, cleans up arguments

* Restores BoardRsConnect to helpers

* Adds full BoardRSconnect call

* Gets fs via calling the function for the tests

* Removes Databricks from CI tests

* Adds databricks-sdk to requirements dev file

* Updates rest of dev reqs versions

* Attempts to avoid double forward slash in Linux

* Checks if path_to_version is a str before removing trailing slash

* Fixes typo on isinstance call

* Removes protocol assignment for DatabricksFs

* Passes pre-commit

* Adds databricks-sdk to minimum reqs

* Addresses additional issue from precommit and attempts to solve pyright error

* Update pins/databricks/fs.py

Co-authored-by: Isabel Zimmerman <[email protected]>

* Update requirements/minimum.txt

Co-authored-by: Isabel Zimmerman <[email protected]>

* Update pins/databricks/fs.py

Co-authored-by: Isabel Zimmerman <[email protected]>

* Update pins/databricks/fs.py

Co-authored-by: Isabel Zimmerman <[email protected]>

* Switches from os to pathlib

* Converts functions to staticmethods

* add docs for board_databricks

* load in quartodoc

* update all tests

* clean up class, make true optional import

* run dbc tests

* add databricks into pyright deps

* load databricks creds

* error earlier

* pass for windows cache

* resolve dependencies

* resolve dependencies again

* cannot clean up constructors, skip

---------

Co-authored-by: Isabel Zimmerman <[email protected]>
Co-authored-by: isabel zimmerman <[email protected]>

Author: edgararuiz

.env.dev

@@ -30,3 +30,8 @@ RSC_LICENSE=
 # (Note that the local file backend always uses a temporary directory.)
 #
 # PINS_TEST_S3__PATH="ci-pins"
+
+# Databricks backend ----
+DATABRICKS_HOST=
+DATABRICKS_TOKEN=
+DATABRICKS_VOLUME=

.github/workflows/ci.yml

@@ -71,6 +71,8 @@ jobs:
           AWS_REGION: "us-east-1"
           AZURE_STORAGE_ACCOUNT_NAME: ${{ secrets.AZURE_STORAGE_ACCOUNT_NAME }}
           AZURE_STORAGE_ACCOUNT_KEY: ${{ secrets.AZURE_STORAGE_ACCOUNT_KEY }}
+          DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }}
+          DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN }}
           PYTEST_OPTS: ${{ matrix.pytest_opts }}
           REQUIREMENTS: ${{ matrix.requirements }}
           ACTION_OS: ${{ matrix.os }}

docs/_quarto.yml

@@ -49,6 +49,8 @@ website:
             href: reference/board_gcs.qmd
           - text: "`board_azure`"
             href: reference/board_azure.qmd
+          - text: "`board_databricks`"
+            href: reference/board_databricks.qmd
           - text: "`board_connect`"
             href: reference/board_connect.qmd
           - text: "`board_url`"
@@ -99,6 +101,7 @@ quartodoc:
         - board_s3
         - board_gcs
         - board_azure
+        - board_databricks
         - board_connect
         - board_url
         - board

pins/__init__.py

@@ -22,6 +22,7 @@
     board_azure,
     board_s3,
     board_gcs,
+    board_databricks,
     board,
 )
 from .boards import board_deparse

pins/boards.py

@@ -870,6 +870,8 @@ def board_deparse(board: BaseBoard):
         return f"board_gcs({repr(board.board)}{allow_pickle})"
     elif prot == "http":
         return f"board_url({repr(board.board)}, {board.pin_paths}{allow_pickle})"
+    elif prot == "dbc":
+        return f"board_databricks({repr(board.board)}{allow_pickle})"
     else:
         raise NotImplementedError(
             f"board deparsing currently not supported for protocol: {prot}"

pins/constructors.py

@@ -10,6 +10,7 @@
 from .boards import BaseBoard, BoardManual, BoardRsConnect, board_deparse
 from .cache import PinsAccessTimeCache, PinsCache, PinsRscCacheMapper, prefix_cache
 from .config import get_cache_dir, get_data_dir
+from .errors import PinsError
 
 # Kept here for backward-compatibility reasons
 # Note that this is not a constructor, but a function to represent them.
@@ -87,6 +88,11 @@ def board(
 
         fs = RsConnectFs(**storage_options)
 
+    elif protocol == "dbc":
+        from pins.databricks.fs import DatabricksFs
+
+        fs = DatabricksFs(**storage_options)
+
     else:
         fs = fsspec.filesystem(protocol, **storage_options)
 
@@ -569,3 +575,61 @@ def board_azure(path, versioned=True, cache=DEFAULT, allow_pickle_read=None):
 
     opts = {"use_listings_cache": False}
     return board("abfs", path, versioned, cache, allow_pickle_read, storage_options=opts)
+
+
+def board_databricks(path, versioned=True, cache=DEFAULT, allow_pickle_read=None):
+    """Create a board to read and write pins from an Databricks Volume folder.
+
+    Parameters
+    ----------
+    path:
+        The path to the target folder inside Unity Catalog. The path must include the
+        catalog, schema, and volume names, preceded by 'Volumes/', for example:
+        "/Volumes/my-catalog/my-schema/my-volume".
+    versioned:
+        Whether or not pins should be versioned.
+    cache:
+        Whether to use a cache. By default, pins attempts to select the right cache
+        directory, given your filesystem. If `None` is passed, then no cache will be
+        used. You can set the cache using the `PINS_CACHE_DIR` environment variable.
+    allow_pickle_read: optional, bool
+        Whether to allow reading pins that use the pickle protocol. Pickles are unsafe,
+        and can execute arbitrary code. Only allow reading pickles if you trust the
+        board to execute Python code on your computer.
+
+        You can enable reading pickles by setting this to `True`, or by setting the
+        environment variable `PINS_ALLOW_PICKLE_READ`. If both are set, this argument
+        takes precedence.
+
+    Notes
+    -----
+    The Databricks board uses the `databricks-sdk` library to authenticate and interact
+    with the Databricks Volume.
+
+    See <https://docs.databricks.com/aws/en/dev-tools/sdk-python>
+
+
+    Examples
+    --------
+
+    >>> import pytest; pytest.skip()
+
+    >>> import pins
+    >>> from dotenv import load_dotenv
+    >>> load_dotenv() # eg, for a .env file with DATABRICKS_HOST and DATABRICKS_TOKEN set
+    >>> board = pins.board_databricks("/Volumes/examples/my-board/test-volume")
+    >>> board.pin_list()
+    ['df_csv']
+
+    >>> board.pin_read("df_csv")
+       x  y  z
+    0  1  a  3
+    1  2  b  4
+    """
+    try:
+        import databricks.sdk  # noqa: F401
+    except ModuleNotFoundError:
+        raise PinsError(
+            "Install the `databricks-sdk` package for Databricks board support."
+        )
+    return board("dbc", path, versioned, cache, allow_pickle_read)

pins/databricks/__init__.py

@@ -0,0 +1,196 @@
+import shutil
+from io import BytesIO
+from pathlib import Path, PurePath
+
+from fsspec import AbstractFileSystem
+
+from pins.errors import PinsError
+
+
+class DatabricksFs(AbstractFileSystem):
+    protocol = "dbc"
+
+    def ls(self, path, detail=False, **kwargs):
+        return self._databricks_ls(path, detail)
+
+    def exists(self, path: str, **kwargs):
+        return self._databricks_exists(path)
+
+    def open(self, path: str, mode: str = "rb", *args, **kwargs):
+        if mode != "rb":
+            raise NotImplementedError
+        return self._databricks_open(path)
+
+    def get(self, rpath, lpath, recursive=False, **kwargs):
+        self._databricks_get(rpath, lpath, recursive, **kwargs)
+
+    def mkdir(self, path, create_parents=True, **kwargs):
+        if not create_parents:
+            raise NotImplementedError
+        self._databricks_mkdir(path)
+
+    def put(
+        self,
+        lpath,
+        rpath,
+        recursive=True,
+        maxdepth=None,
+        **kwargs,
+    ):
+        if not recursive:
+            raise NotImplementedError
+        if maxdepth is not None:
+            raise NotImplementedError
+        self._databricks_put(lpath, rpath)
+
+    def rm(self, path, recursive=True, maxdepth=None) -> None:
+        if not recursive:
+            raise NotImplementedError
+        if maxdepth is not None:
+            raise NotImplementedError
+        if self._databricks_exists(path):
+            self._databricks_rm_dir(path)
+
+    @staticmethod
+    def _databricks_put(lpath, rpath):
+        from databricks.sdk import WorkspaceClient
+
+        w = WorkspaceClient()
+        path = Path(lpath).absolute()
+        orig_path = path
+
+        def _upload_files(path):
+            contents = Path(path)
+            for item in contents.iterdir():
+                abs_path = PurePath(path).joinpath(item)
+                is_file = Path(abs_path).is_file()
+                if is_file:
+                    rel_path = abs_path.relative_to(orig_path)
+                    db_path = PurePath(rpath).joinpath(rel_path)
+                    file = open(abs_path, "rb")
+                    w.files.upload(str(db_path), BytesIO(file.read()), overwrite=True)
+                else:
+                    _upload_files(abs_path)
+
+        _upload_files(path)
+
+    def _databricks_get(self, board, rpath, lpath, recursive=False, **kwargs):
+        from databricks.sdk import WorkspaceClient
+
+        w = WorkspaceClient()
+        file_type = self._databricks_is_type(rpath)
+        if file_type == "file":
+            board.fs.get(rpath, lpath, **kwargs)
+            return
+
+        def _get_files(path, recursive, **kwargs):
+            raw_contents = w.files.list_directory_contents(path)
+            contents = list(raw_contents)
+            details = list(map(self._databricks_content_details, contents))
+            for item in details:
+                item_path = item.get("path")
+                if item.get("is_directory"):
+                    if recursive:
+                        _get_files(item_path, recursive=recursive, **kwargs)
+                else:
+                    rel_path = PurePath(item_path).relative_to(rpath)
+                    target_path = PurePath(lpath).joinpath(rel_path)
+                    board.fs.get(item_path, str(target_path))
+
+        _get_files(rpath, recursive, **kwargs)
+
+    def _databricks_open(self, path):
+        from databricks.sdk import WorkspaceClient
+
+        if not self._databricks_exists(path):
+            raise PinsError(f"File or directory does not exist at path: {path}")
+        w = WorkspaceClient()
+        resp = w.files.download(path)
+        f = BytesIO()
+        shutil.copyfileobj(resp.contents, f)
+        f.seek(0)
+        return f
+
+    def _databricks_exists(self, path: str):
+        if self._databricks_is_type(path) == "nothing":
+            return False
+        else:
+            return True
+
+    @staticmethod
+    def _databricks_is_type(path: str):
+        from databricks.sdk import WorkspaceClient
+        from databricks.sdk.errors import NotFound
+
+        w = WorkspaceClient()
+        try:
+            w.files.get_metadata(path)
+        except NotFound:
+            try:
+                w.files.get_directory_metadata(path)
+            except NotFound:
+                return "nothing"
+            else:
+                return "directory"
+        else:
+            return "file"
+
+    def _databricks_ls(self, path, detail):
+        from databricks.sdk import WorkspaceClient
+
+        if not self._databricks_exists(path):
+            raise PinsError(f"File or directory does not exist at path: {path}")
+        w = WorkspaceClient()
+        if self._databricks_is_type(path) == "file":
+            if detail:
+                return [dict(name=path, size=None, type="file")]
+            else:
+                return path
+
+        contents_raw = w.files.list_directory_contents(path)
+        contents = list(contents_raw)
+        items = []
+        for item in contents:
+            item = self._databricks_content_details(item)
+            item_path = item.get("path")
+            item_path = item_path.rstrip("/")
+            if detail:
+                if item.get("is_directory"):
+                    item_type = "directory"
+                else:
+                    item_type = "file"
+                items.append(dict(name=item_path, size=None, type=item_type))
+            else:
+                items.append(item_path)
+        return items
+
+    def _databricks_rm_dir(self, path):
+        from databricks.sdk import WorkspaceClient
+
+        w = WorkspaceClient()
+        raw_contents = w.files.list_directory_contents(path)
+        contents = list(raw_contents)
+        details = list(map(self._databricks_content_details, contents))
+        for item in details:
+            item_path = item.get("path")
+            if item.get("is_directory"):
+                self._databricks_rm_dir(item_path)
+            else:
+                w.files.delete(item_path)
+        w.files.delete_directory(path)
+
+    @staticmethod
+    def _databricks_mkdir(path):
+        from databricks.sdk import WorkspaceClient
+
+        w = WorkspaceClient()
+        w.files.create_directory(path)
+
+    @staticmethod
+    def _databricks_content_details(item):
+        details = {
+            "path": item.path,
+            "name": item.name,
+            "is_directory": item.is_directory,
+        }
+        return details

pins/databricks/fs.py

@@ -24,6 +24,8 @@ def load_path(filename: str, path_to_version, pin_type=None):
         filename = "data.csv"
 
     if path_to_version is not None:
+        if isinstance(path_to_version, str):
+            path_to_version = path_to_version.rstrip("/")
         path_to_file = f"{path_to_version}/{filename}"
     else:
         # BoardUrl doesn't have versions, and the file is the full url

pins/drivers.py

@@ -6,10 +6,17 @@
 from importlib_resources import files
 from pytest import mark as m
 
-from pins.tests.helpers import BoardBuilder, RscBoardBuilder, Snapshot, rm_env
+from pins.tests.helpers import (
+    BoardBuilder,
+    DbcBoardBuilder,
+    RscBoardBuilder,
+    Snapshot,
+    rm_env,
+)
 
 EXAMPLE_REL_PATH = "pins/tests/pins-compat"
 PATH_TO_EXAMPLE_BOARD = files("pins") / "tests/pins-compat"
+PATH_TO_EXAMPLE_BOARD_DBC = "/Volumes/workshops/my-board/my-volume/test"
 PATH_TO_EXAMPLE_VERSION = PATH_TO_EXAMPLE_BOARD / "df_csv/20220214T163720Z-9bfad/"
 EXAMPLE_PIN_NAME = "df_csv"
 
@@ -21,6 +28,7 @@
     pytest.param(lambda: BoardBuilder("s3"), id="s3", marks=m.fs_s3),
     pytest.param(lambda: BoardBuilder("gcs"), id="gcs", marks=m.fs_gcs),
     pytest.param(lambda: BoardBuilder("abfs"), id="abfs", marks=m.fs_abfs),
+    pytest.param(lambda: DbcBoardBuilder("dbc"), id="dbc", marks=m.fs_dbc),
 ]
 
 # rsc should only be used once, because users are created at docker setup time