databricks-sdk 0.29.0__py3-none-any.whl → 0.30.0__py3-none-any.whl

databricks/sdk/data_plane.py

@@ -0,0 +1,65 @@
+import threading
+from dataclasses import dataclass
+from typing import Callable, List
+
+from databricks.sdk.oauth import Token
+from databricks.sdk.service.oauth2 import DataPlaneInfo
+
+
+@dataclass
+class DataPlaneDetails:
+    """
+    Contains details required to query a DataPlane endpoint.
+    """
+    endpoint_url: str
+    """URL used to query the endpoint through the DataPlane."""
+    token: Token
+    """Token to query the DataPlane endpoint."""
+
+
+class DataPlaneService:
+    """Helper class to fetch and manage DataPlane details."""
+
+    def __init__(self):
+        self._data_plane_info = {}
+        self._tokens = {}
+        self._lock = threading.Lock()
+
+    def get_data_plane_details(self, method: str, params: List[str], info_getter: Callable[[], DataPlaneInfo],
+                               refresh: Callable[[str], Token]):
+        """Get and cache information required to query a Data Plane endpoint using the provided methods.
+
+        Returns a cached DataPlaneDetails if the details have already been fetched previously and are still valid.
+        If not, it uses the provided functions to fetch the details.
+
+        :param method: method name. Used to construct a unique key for the cache.
+        :param params: path params used in the "get" operation which uniquely determine the object. Used to construct a unique key for the cache.
+        :param info_getter: function which returns the DataPlaneInfo. It will only be called if the information is not already present in the cache.
+        :param refresh: function to refresh the token. It will only be called if the token is missing or expired.
+        """
+        all_elements = params.copy()
+        all_elements.insert(0, method)
+        map_key = "/".join(all_elements)
+        info = self._data_plane_info.get(map_key)
+        if not info:
+            self._lock.acquire()
+            try:
+                info = self._data_plane_info.get(map_key)
+                if not info:
+                    info = info_getter()
+                    self._data_plane_info[map_key] = info
+            finally:
+                self._lock.release()
+
+        token = self._tokens.get(map_key)
+        if not token or not token.valid:
+            self._lock.acquire()
+            token = self._tokens.get(map_key)
+            try:
+                if not token or not token.valid:
+                    token = refresh(info.authorization_details)
+                    self._tokens[map_key] = token
+            finally:
+                self._lock.release()
+
+        return DataPlaneDetails(endpoint_url=info.endpoint_url, token=token)

databricks/sdk/mixins/files.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import base64
 import os
 import pathlib
+import platform
 import shutil
 import sys
 from abc import ABC, abstractmethod
@@ -266,8 +267,9 @@ class _VolumesIO(BinaryIO):
 
 class _Path(ABC):
 
-    def __init__(self, path: str):
-        self._path = pathlib.Path(str(path).replace('dbfs:', '').replace('file:', ''))
+    @abstractmethod
+    def __init__(self):
+        ...
 
     @property
     def is_local(self) -> bool:
@@ -327,6 +329,12 @@ class _Path(ABC):
 
 class _LocalPath(_Path):
 
+    def __init__(self, path: str):
+        if platform.system() == "Windows":
+            self._path = pathlib.Path(str(path).replace('file:///', '').replace('file:', ''))
+        else:
+            self._path = pathlib.Path(str(path).replace('file:', ''))
+
     def _is_local(self) -> bool:
         return True
 
@@ -393,7 +401,7 @@ class _LocalPath(_Path):
 class _VolumesPath(_Path):
 
     def __init__(self, api: files.FilesAPI, src: Union[str, pathlib.Path]):
-        super().__init__(src)
+        self._path = pathlib.PurePosixPath(str(src).replace('dbfs:', '').replace('file:', ''))
         self._api = api
 
     def _is_local(self) -> bool:
@@ -462,7 +470,7 @@ class _VolumesPath(_Path):
 class _DbfsPath(_Path):
 
     def __init__(self, api: files.DbfsAPI, src: str):
-        super().__init__(src)
+        self._path = pathlib.PurePosixPath(str(src).replace('dbfs:', '').replace('file:', ''))
         self._api = api
 
     def _is_local(self) -> bool: