ygg 0.1.30__py3-none-any.whl → 0.1.32__py3-none-any.whl

yggdrasil/databricks/workspaces/workspace.py

@@ -1,3 +1,5 @@
+"""Workspace configuration and Databricks SDK helpers."""
+
 import dataclasses
 import logging
 import os
@@ -16,14 +18,14 @@ from typing import (
 if TYPE_CHECKING:
     from ..compute.cluster import Cluster
 
-from .databricks_path import DatabricksPath, DatabricksPathKind
+from .path import DatabricksPath, DatabricksPathKind
+from ...version import __version__ as YGGDRASIL_VERSION
 from ...libs.databrickslib import require_databricks_sdk, databricks_sdk
 
 if databricks_sdk is not None:
     from databricks.sdk import WorkspaceClient
-    from databricks.sdk.errors import ResourceDoesNotExist, NotFound
+    from databricks.sdk.errors import ResourceDoesNotExist
     from databricks.sdk.service.workspace import ExportFormat, ObjectInfo
-    from databricks.sdk.service import catalog as catalog_svc
     from databricks.sdk.dbutils import FileInfo
     from databricks.sdk.service.files import DirectoryEntry
 
@@ -45,7 +47,7 @@ def _get_env_product():
     v = os.getenv("DATABRICKS_PRODUCT")
 
     if not v:
-        return None
+        return "yggdrasil"
     return v.strip().lower()
 
 
@@ -53,7 +55,7 @@ def _get_env_product_version():
     v = os.getenv("DATABRICKS_PRODUCT_VERSION")
 
     if not v:
-        return None
+        return YGGDRASIL_VERSION
     return v.strip().lower()
 
 
@@ -67,6 +69,7 @@ def _get_env_product_tag():
 
 @dataclass
 class Workspace:
+    """Configuration wrapper for connecting to a Databricks workspace."""
     # Databricks / generic
     host: Optional[str] = None
     account_id: Optional[str] = None
@@ -113,6 +116,11 @@ class Workspace:
     # Pickle support
     # -------------------------
     def __getstate__(self):
+        """Serialize the workspace state for pickling.
+
+        Returns:
+            A pickle-ready state dictionary.
+        """
         state = self.__dict__.copy()
         state.pop("_sdk", None)
 
@@ -122,6 +130,11 @@ class Workspace:
         return state
 
     def __setstate__(self, state):
+        """Restore workspace state after unpickling.
+
+        Args:
+            state: Serialized state dictionary.
+        """
         self.__dict__.update(state)
         self._sdk = None
 
@@ -132,21 +145,46 @@ class Workspace:
             self.connect(reset=True)
 
     def __enter__(self) -> "Workspace":
+        """Enter a context manager and connect to the workspace.
+
+        Returns:
+            The connected Workspace instance.
+        """
         self._was_connected = self._sdk is not None
-        self.connect()
-        return self
+        return self.connect()
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context manager and close if newly connected.
+
+        Args:
+            exc_type: Exception type, if raised.
+            exc_val: Exception value, if raised.
+            exc_tb: Exception traceback, if raised.
+
+        Returns:
+            None.
+        """
         if not self._was_connected:
             self.close()
 
+    def __del__(self):
+        self.close()
+
     # -------------------------
     # Clone
     # -------------------------
-    def clone(
+    def clone_instance(
         self,
         **kwargs
     ) -> "Workspace":
+        """Clone the workspace config with overrides.
+
+        Args:
+            **kwargs: Field overrides for the clone.
+
+        Returns:
+            A new Workspace instance with updated fields.
+        """
         state = self.__getstate__()
         state.update(kwargs)
         return Workspace().__setstate__(state)
@@ -154,86 +192,108 @@ class Workspace:
     # -------------------------
     # SDK connection
     # -------------------------
-    def connect(self, reset: bool = False) -> "Workspace":
+    @property
+    def connected(self):
+        """Return True when a WorkspaceClient is cached.
+
+        Returns:
+            True if connected, otherwise False.
+        """
+        return self._sdk is not None
+
+    def connect(self, reset: bool = False, clone: bool = False) -> "Workspace":
+        """Connect to the workspace and cache the SDK client.
+
+        Args:
+            reset: Whether to reset the cached client before connecting.
+            clone: Whether to connect a cloned instance.
+
+        Returns:
+            The connected Workspace instance.
+        """
         if reset:
             self._sdk = None
 
-        if self._sdk is None:
-            require_databricks_sdk()
-            logger.debug("Connecting %s", self)
-
-            # Build Config from config_dict if available, else from fields.
-            kwargs = {
-                "host": self.host,
-                "account_id": self.account_id,
-                "token": self.token,
-                "client_id": self.client_id,
-                "client_secret": self.client_secret,
-                "token_audience": self.token_audience,
-                "azure_workspace_resource_id": self.azure_workspace_resource_id,
-                "azure_use_msi": self.azure_use_msi,
-                "azure_client_secret": self.azure_client_secret,
-                "azure_client_id": self.azure_client_id,
-                "azure_tenant_id": self.azure_tenant_id,
-                "azure_environment": self.azure_environment,
-                "google_credentials": self.google_credentials,
-                "google_service_account": self.google_service_account,
-                "profile": self.profile,
-                "config_file": self.config_file,
-                "auth_type": self.auth_type,
-                "http_timeout_seconds": self.http_timeout_seconds,
-                "retry_timeout_seconds": self.retry_timeout_seconds,
-                "debug_truncate_bytes": self.debug_truncate_bytes,
-                "debug_headers": self.debug_headers,
-                "rate_limit": self.rate_limit,
-                "product": self.product,
-                "product_version": self.product_version,
-            }
-
-            build_kwargs = {k: v for k, v in kwargs.items() if v is not None}
+        if self._sdk is not None:
+            return self
+
+        instance = self.clone_instance() if clone else self
+
+        require_databricks_sdk()
+        logger.debug("Connecting %s", self)
+
+        # Build Config from config_dict if available, else from fields.
+        kwargs = {
+            "host": instance.host,
+            "account_id": instance.account_id,
+            "token": instance.token,
+            "client_id": instance.client_id,
+            "client_secret": instance.client_secret,
+            "token_audience": instance.token_audience,
+            "azure_workspace_resource_id": instance.azure_workspace_resource_id,
+            "azure_use_msi": instance.azure_use_msi,
+            "azure_client_secret": instance.azure_client_secret,
+            "azure_client_id": instance.azure_client_id,
+            "azure_tenant_id": instance.azure_tenant_id,
+            "azure_environment": instance.azure_environment,
+            "google_credentials": instance.google_credentials,
+            "google_service_account": instance.google_service_account,
+            "profile": instance.profile,
+            "config_file": instance.config_file,
+            "auth_type": instance.auth_type,
+            "http_timeout_seconds": instance.http_timeout_seconds,
+            "retry_timeout_seconds": instance.retry_timeout_seconds,
+            "debug_truncate_bytes": instance.debug_truncate_bytes,
+            "debug_headers": instance.debug_headers,
+            "rate_limit": instance.rate_limit,
+            "product": instance.product,
+            "product_version": instance.product_version,
+        }
 
-            try:
-                self._sdk = WorkspaceClient(**build_kwargs)
-            except ValueError as e:
-                if "cannot configure default credentials" in str(e) and self.auth_type is None:
-                    last_error = e
+        build_kwargs = {k: v for k, v in kwargs.items() if v is not None}
+
+        try:
+            instance._sdk = WorkspaceClient(**build_kwargs)
+        except ValueError as e:
+            if "cannot configure default credentials" in str(e) and instance.auth_type is None:
+                last_error = e
 
-                    auth_types = ["runtime"] if self.is_in_databricks_environment() else ["external-browser"]
+                auth_types = ["runtime"] if instance.is_in_databricks_environment() else ["external-browser"]
 
-                    for auth_type in auth_types:
-                        build_kwargs["auth_type"] = auth_type
+                for auth_type in auth_types:
+                    build_kwargs["auth_type"] = auth_type
+
+                    try:
+                        instance._sdk = WorkspaceClient(**build_kwargs)
+                        break
+                    except Exception as se:
+                        last_error = se
+                        build_kwargs.pop("auth_type")
+
+                if instance._sdk is None:
+                    if instance.is_in_databricks_environment() and instance._cached_token:
+                        build_kwargs["token"] = instance._cached_token
 
                         try:
-                            self._sdk = WorkspaceClient(**build_kwargs)
-                            break
+                            instance._sdk = WorkspaceClient(**build_kwargs)
                         except Exception as se:
                             last_error = se
-                            build_kwargs.pop("auth_type")
-
-                    if self._sdk is None:
-                        if self.is_in_databricks_environment() and self._cached_token:
-                            build_kwargs["token"] = self._cached_token
 
-                            try:
-                                self._sdk = WorkspaceClient(**build_kwargs)
-                            except Exception as se:
-                                last_error = se
+                if instance._sdk is None:
+                    raise last_error
+            else:
+                raise e
 
-                    if self._sdk is None:
-                        raise last_error
-                else:
-                    raise e
+        # backfill resolved config values
+        for key in list(kwargs.keys()):
+            if getattr(instance, key, None) is None:
+                v = getattr(instance._sdk.config, key, None)
+                if v is not None:
+                    setattr(instance, key, v)
 
-            # backfill resolved config values
-            for key in list(kwargs.keys()):
-                if getattr(self, key, None) is None:
-                    v = getattr(self._sdk.config, key, None)
-                    if v is not None:
-                        setattr(self, key, v)
+        logger.info("Connected %s", instance)
 
-            logger.info("Connected %s", self)
-
-        return self
+        return instance
 
     # ------------------------------------------------------------------ #
     # Context manager + lifecycle
@@ -260,6 +320,11 @@ class Workspace:
         return str(files[0]) if files else None
 
     def reset_local_cache(self):
+        """Remove cached browser OAuth tokens.
+
+        Returns:
+            None.
+        """
         local_cache = self._local_cache_token_path()
 
         if local_cache:
@@ -267,6 +332,11 @@ class Workspace:
 
     @property
     def current_user(self):
+        """Return the current Databricks user.
+
+        Returns:
+            The current user object from the SDK.
+        """
         try:
             return self.sdk().current_user.me()
         except:
@@ -275,6 +345,11 @@ class Workspace:
             raise
 
     def current_token(self) -> str:
+        """Return the active API token for this workspace.
+
+        Returns:
+            The bearer token string.
+        """
         if self.token:
             return self.token
 
@@ -287,24 +362,56 @@ class Workspace:
     # ------------------------------------------------------------------ #
     # Path helpers
     # ------------------------------------------------------------------ #
+    def filesytem(
+        self,
+        workspace: Optional["Workspace"] = None,
+    ):
+        """Return a PyArrow filesystem for Databricks paths.
+
+        Args:
+            workspace: Optional workspace override.
+
+        Returns:
+            A DatabricksFileSystem instance.
+        """
+        from .filesytem import DatabricksFileSystem, DatabricksFileSystemHandler
+
+        handler = DatabricksFileSystemHandler(
+            workspace=self if workspace is None else workspace
+        )
+
+        return DatabricksFileSystem(
+            handler=handler
+        )
+
     def dbfs_path(
         self,
         parts: Union[List[str], str],
         kind: Optional[DatabricksPathKind] = None,
         workspace: Optional["Workspace"] = None
     ):
+        """Create a DatabricksPath in this workspace.
+
+        Args:
+            parts: Path parts or string to parse.
+            kind: Optional path kind override.
+            workspace: Optional workspace override.
+
+        Returns:
+            A DatabricksPath instance.
+        """
         workspace = self if workspace is None else workspace
 
         if kind is None or isinstance(parts, str):
             return DatabricksPath.parse(
-                parts=parts,
+                obj=parts,
                 workspace=workspace
             )
 
         return DatabricksPath(
             kind=kind,
             parts=parts,
-            workspace=workspace
+            _workspace=workspace
         )
 
     def shared_cache_path(
@@ -313,6 +420,12 @@ class Workspace:
     ) -> DatabricksPath:
         """
         Shared cache base under Volumes for the current user.
+
+        Args:
+            suffix: Optional path suffix to append.
+
+        Returns:
+            A DatabricksPath pointing at the shared cache location.
         """
         base = "/Workspace/Shared/.ygg/cache"
 
@@ -327,61 +440,12 @@ class Workspace:
     # ------------------------------------------------------------------ #
 
     def sdk(self) -> "WorkspaceClient":
-        return self.connect()._sdk
+        """Return the connected WorkspaceClient.
 
-    # ------------------------------------------------------------------ #
-    # UC volume + directory management
-    # ------------------------------------------------------------------ #
-
-    def ensure_uc_volume_and_dir(
-        self,
-        target_path: str,
-    ) -> None:
-        """
-        Ensure catalog, schema, volume exist for a UC volume path
-        like /Volumes/<catalog>/<schema>/<volume>/...,
-        then create the directory.
+        Returns:
+            The WorkspaceClient instance.
         """
-        sdk = self.sdk()
-        parts = target_path.split("/")
-
-        # basic sanity check
-        if len(parts) < 5 or parts[1] != "Volumes":
-            raise ValueError(
-                f"Unexpected UC volume path: {target_path!r}. "
-                "Expected /Volumes/<catalog>/<schema>/<volume>/..."
-            )
-
-        # /Volumes/<catalog>/<schema>/<volume>/...
-        _, _, catalog_name, schema_name, volume_name, *subpath = parts
-
-        # 1) ensure catalog
-        try:
-            sdk.catalogs.get(name=catalog_name)
-        except NotFound:
-            sdk.catalogs.create(name=catalog_name)
-
-        # 2) ensure schema
-        schema_full_name = f"{catalog_name}.{schema_name}"
-        try:
-            sdk.schemas.get(full_name=schema_full_name)
-        except NotFound:
-            sdk.schemas.create(name=schema_name, catalog_name=catalog_name)
-
-        # 3) ensure volume (managed volume is simplest)
-        volume_full_name = f"{catalog_name}.{schema_name}.{volume_name}"
-        try:
-            sdk.volumes.read(name=volume_full_name)
-        except NotFound:
-            sdk.volumes.create(
-                catalog_name=catalog_name,
-                schema_name=schema_name,
-                name=volume_name,
-                volume_type=catalog_svc.VolumeType.MANAGED,
-            )
-
-        # 4) finally create the directory path itself
-        sdk.files.create_directory(target_path)
+        return self.connect()._sdk
 
     # ------------------------------------------------------------------ #
     # List / open / delete / SQL
@@ -400,6 +464,13 @@ class Workspace:
           - other paths      -> Workspace paths (sdk.workspace.list)
 
         If recursive=True, yield all nested files/directories.
+
+        Args:
+            path: Path string to list.
+            recursive: Whether to list recursively.
+
+        Returns:
+            An iterator of workspace/DBFS/volume entries.
         """
         sdk = self.sdk()
 
@@ -452,6 +523,13 @@ class Workspace:
           via workspace.download(...).
 
         Returned object is a BinaryIO context manager.
+
+        Args:
+            path: Path to open.
+            workspace_format: Optional export format for workspace paths.
+
+        Returns:
+            A BinaryIO stream for reading.
         """
         sdk = self.sdk()
 
@@ -462,37 +540,24 @@ class Workspace:
 
         # Workspace path
         fmt = workspace_format or ExportFormat.AUTO
-        return sdk.workspace.download(path=path, format=fmt)
-
-    def delete_path(
-        self,
-        target_path: str,
-        recursive: bool = True,
-        ignore_missing: bool = True,
-    ) -> None:
-        """
-        Delete a path in Databricks Workspace (file or directory).
-
-        - If recursive=True and target_path is a directory, deletes entire tree.
-        - If ignore_missing=True, missing paths won't raise.
-        """
-        sdk = self.sdk()
 
-        try:
-            sdk.workspace.delete(
-                path=target_path,
-                recursive=recursive,
-            )
-        except ResourceDoesNotExist:
-            if ignore_missing:
-                return
-            raise
+        return sdk.workspace.download(path=path, format=fmt)
 
     @staticmethod
     def is_in_databricks_environment():
+        """Return True when running on a Databricks runtime.
+
+        Returns:
+            True if running on Databricks, otherwise False.
+        """
         return os.getenv("DATABRICKS_RUNTIME_VERSION") is not None
 
     def default_tags(self):
+        """Return default resource tags for Databricks assets.
+
+        Returns:
+            A dict of default tags.
+        """
         return {
             k: v
             for k, v in (
@@ -504,6 +569,14 @@ class Workspace:
         }
 
     def merge_tags(self, existing: dict | None = None):
+        """Merge default tags with an existing set.
+
+        Args:
+            existing: Optional existing tags.
+
+        Returns:
+            A dict of merged tags.
+        """
         if existing:
             return self.default_tags()
 
@@ -514,6 +587,17 @@ class Workspace:
         schema_name: Optional[str] = None,
         **kwargs
     ):
+        """Return a SQLEngine configured for this workspace.
+
+        Args:
+            workspace: Optional workspace override.
+            catalog_name: Optional catalog name.
+            schema_name: Optional schema name.
+            **kwargs: Additional SQLEngine parameters.
+
+        Returns:
+            A SQLEngine instance.
+        """
         from ..sql import SQLEngine
 
         return SQLEngine(
@@ -529,11 +613,20 @@ class Workspace:
         cluster_name: Optional[str] = None,
         **kwargs
     ) -> "Cluster":
+        """Return a Cluster helper bound to this workspace.
+
+        Args:
+            cluster_id: Optional cluster id.
+            cluster_name: Optional cluster name.
+            **kwargs: Additional Cluster parameters.
+
+        Returns:
+            A Cluster instance.
+        """
         from ..compute.cluster import Cluster
 
         return Cluster(workspace=self, cluster_id=cluster_id, cluster_name=cluster_name, **kwargs)
 
-
 # ---------------------------------------------------------------------------
 # Workspace-bound base class
 # ---------------------------------------------------------------------------
@@ -543,23 +636,54 @@ DBXWorkspace = Workspace
 
 @dataclass
 class WorkspaceService(ABC):
+    """Base class for helpers that depend on a Workspace."""
     workspace: Workspace = dataclasses.field(default_factory=Workspace)
 
     def __post_init__(self):
+        """Ensure a Workspace instance is available.
+
+        Returns:
+            None.
+        """
         if self.workspace is None:
             self.workspace = Workspace()
 
     def __enter__(self):
+        """Enter a context manager and connect the workspace.
+
+        Returns:
+            The current WorkspaceService instance.
+        """
         self.workspace.__enter__()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the context manager and close the workspace.
+
+        Args:
+            exc_type: Exception type, if raised.
+            exc_val: Exception value, if raised.
+            exc_tb: Exception traceback, if raised.
+
+        Returns:
+            None.
+        """
         self.workspace.__exit__(exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb)
 
     def is_in_databricks_environment(self):
+        """Return True when running on a Databricks runtime.
+
+        Returns:
+            True if running on Databricks, otherwise False.
+        """
         return self.workspace.is_in_databricks_environment()
 
     def connect(self):
+        """Connect the underlying workspace.
+
+        Returns:
+            The current WorkspaceService instance.
+        """
         self.workspace = self.workspace.connect()
         return self
 
@@ -569,6 +693,16 @@ class WorkspaceService(ABC):
         kind: Optional[DatabricksPathKind] = None,
         workspace: Optional["Workspace"] = None
     ):
+        """Create a DatabricksPath in the underlying workspace.
+
+        Args:
+            parts: Path parts or string to parse.
+            kind: Optional path kind override.
+            workspace: Optional workspace override.
+
+        Returns:
+            A DatabricksPath instance.
+        """
         return self.workspace.dbfs_path(
             kind=kind,
             parts=parts,
@@ -576,8 +710,18 @@ class WorkspaceService(ABC):
         )
 
     def sdk(self):
+        """Return the WorkspaceClient for the underlying workspace.
+
+        Returns:
+            The WorkspaceClient instance.
+        """
         return self.workspace.sdk()
 
     @property
     def current_user(self):
+        """Return the current Databricks user.
+
+        Returns:
+            The current user object from the SDK.
+        """
         return self.workspace.current_user

yggdrasil/dataclasses/__init__.py

@@ -1,3 +1,5 @@
+"""Enhanced dataclass helpers with Arrow awareness."""
+
 from .dataclass import yggdataclass
 
 __all__ = ["yggdataclass"]