ygg 0.1.57__py3-none-any.whl → 0.1.64__py3-none-any.whl

yggdrasil/databricks/workspaces/workspace.py

@@ -4,29 +4,40 @@ import dataclasses
 import logging
 import os
 import posixpath
+import time
 from abc import ABC
 from dataclasses import dataclass
 from pathlib import Path
+from threading import Thread
 from typing import (
     BinaryIO,
     Iterator,
     Optional,
-    Union, TYPE_CHECKING, List
+    Union, TYPE_CHECKING, List, Set
 )
 
-if TYPE_CHECKING:
-    from ..compute.cluster import Cluster
-
 from .path import DatabricksPath, DatabricksPathKind
+from ...libs.databrickslib import databricks_sdk, WorkspaceClient, DatabricksDummyClass
+from ...pyutils.expiring_dict import ExpiringDict
+from ...pyutils.waiting_config import WaitingConfig, WaitingConfigArg
 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
     from databricks.sdk.service.workspace import ExportFormat, ObjectInfo
     from databricks.sdk.dbutils import FileInfo
     from databricks.sdk.service.files import DirectoryEntry
+else:
+    ResourceDoesNotExist = DatabricksDummyClass
+    ExportFormat = DatabricksDummyClass
+    ObjectInfo = DatabricksDummyClass
+    FileInfo = DatabricksDummyClass
+    DirectoryEntry = DatabricksDummyClass
+
+
+if TYPE_CHECKING:
+    from ..sql.warehouse import SQLWarehouse
+    from ..compute.cluster import Cluster
 
 
 __all__ = [
@@ -36,7 +47,26 @@ __all__ = [
 ]
 
 
-logger = logging.getLogger(__name__)
+LOGGER = logging.getLogger(__name__)
+CHECKED_TMP_WORKSPACES: ExpiringDict[str, Set[str]] = ExpiringDict()
+
+def is_checked_tmp_path(
+    host: str,
+    base_path: str
+):
+    existing = CHECKED_TMP_WORKSPACES.get(host)
+
+    if existing is None:
+        CHECKED_TMP_WORKSPACES[host] = set(base_path)
+
+        return False
+
+    if base_path in existing:
+        return True
+
+    existing.add(base_path)
+
+    return False
 
 
 # ---------------------------------------------------------------------------
@@ -47,7 +77,7 @@ def _get_env_product():
 
     if not v:
         return "yggdrasil"
-    return v.strip().lower()
+    return v
 
 
 def _get_env_product_version():
@@ -57,7 +87,7 @@ def _get_env_product_version():
         if _get_env_product() == "yggdrasil":
             return YGGDRASIL_VERSION
         return None
-    return v.strip().lower()
+    return v
 
 
 def _get_env_product_tag():
@@ -65,7 +95,7 @@ def _get_env_product_tag():
 
     if not v:
         return None
-    return v.strip().lower()
+    return v
 
 
 @dataclass
@@ -115,8 +145,12 @@ class Workspace:
     _cached_token: Optional[str] = dataclasses.field(default=None, repr=False, compare=False, hash=False)
 
     # -------------------------
-    # Pickle support
+    # Python methods
     # -------------------------
+    def __post_init__(self):
+        self.product = self.product.strip().lower() if self.product else "yggdrasil"
+        self.product_tag = self.product_tag.strip().lower() if self.product_tag else "main"
+
     def __getstate__(self):
         """Serialize the workspace state for pickling.
 
@@ -245,8 +279,6 @@ class Workspace:
 
         instance = self.clone_instance() if clone else self
 
-        require_databricks_sdk()
-
         # Build Config from config_dict if available, else from fields.
         kwargs = {
             "host": instance.host,
@@ -357,7 +389,7 @@ class Workspace:
     @property
     def safe_host(self):
         if not self.host:
-            return self.connect().host
+            self.host = self.sdk().config.host
         return self.host
 
     @property
@@ -392,7 +424,7 @@ class Workspace:
     # ------------------------------------------------------------------ #
     # Path helpers
     # ------------------------------------------------------------------ #
-    def filesytem(
+    def filesystem(
         self,
         workspace: Optional["Workspace"] = None,
     ):
@@ -418,7 +450,8 @@ class Workspace:
         self,
         parts: Union[List[str], str],
         kind: Optional[DatabricksPathKind] = None,
-        workspace: Optional["Workspace"] = None
+        workspace: Optional["Workspace"] = None,
+        temporary: bool = False
     ):
         """Create a DatabricksPath in this workspace.
 
@@ -426,6 +459,7 @@ class Workspace:
             parts: Path parts or string to parse.
             kind: Optional path kind override.
             workspace: Optional workspace override.
+            temporary: Temporary path
 
         Returns:
             A DatabricksPath instance.
@@ -435,15 +469,137 @@ class Workspace:
         if kind is None or isinstance(parts, str):
             return DatabricksPath.parse(
                 obj=parts,
-                workspace=workspace
+                workspace=workspace,
+                temporary=temporary
             )
 
         return DatabricksPath(
             kind=kind,
             parts=parts,
+            temporary=temporary,
             _workspace=workspace
         )
 
+    @staticmethod
+    def _base_tmp_path(
+        catalog_name: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        volume_name: Optional[str] = None,
+    ) -> str:
+        if catalog_name and schema_name:
+            base_path = "/Volumes/%s/%s/%s" % (
+                catalog_name, schema_name, volume_name or "tmp"
+            )
+        else:
+            base_path = "/Workspace/Shared/.ygg/tmp"
+
+        return base_path
+
+    def tmp_path(
+        self,
+        suffix: Optional[str] = None,
+        extension: Optional[str] = None,
+        max_lifetime: Optional[float] = None,
+        catalog_name: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        volume_name: Optional[str] = None,
+        base_path: Optional[str] = None,
+    ) -> DatabricksPath:
+        """
+        Shared cache base under Volumes for the current user.
+
+        Args:
+            suffix: Optional suffix
+            extension: Optional extension suffix to append.
+            max_lifetime: Max lifetime of temporary path
+            catalog_name: Unity catalog name for volume path
+            schema_name: Unity schema name for volume path
+            volume_name: Unity volume name for volume path
+            base_path: Base temporary path
+
+        Returns:
+            A DatabricksPath pointing at the shared cache location.
+        """
+        start = int(time.time())
+        max_lifetime = int(max_lifetime or 48 * 3600)
+        end = max(0, int(start + max_lifetime))
+
+        base_path = base_path or self._base_tmp_path(
+            catalog_name=catalog_name,
+            schema_name=schema_name,
+            volume_name=volume_name
+        )
+
+        rnd = os.urandom(4).hex()
+        temp_path = f"tmp-{start}-{end}-{rnd}"
+
+        if suffix:
+            temp_path += suffix
+
+        if extension:
+            temp_path += ".%s" % extension
+
+        self.clean_tmp_folder(
+            raise_error=False,
+            wait=False,
+            base_path=base_path
+        )
+
+        return self.dbfs_path(f"{base_path}/{temp_path}")
+
+    def clean_tmp_folder(
+        self,
+        raise_error: bool = True,
+        wait: Optional[WaitingConfigArg] = True,
+        catalog_name: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        volume_name: Optional[str] = None,
+        base_path: Optional[str] = None,
+    ):
+        wait = WaitingConfig.check_arg(wait)
+
+        base_path = base_path or self._base_tmp_path(
+            catalog_name=catalog_name,
+            schema_name=schema_name,
+            volume_name=volume_name
+        )
+
+        if is_checked_tmp_path(host=self.safe_host, base_path=base_path):
+            return self
+
+        if wait.timeout:
+            base_path = self.dbfs_path(base_path)
+
+            LOGGER.debug(
+                "Cleaning temp path %s",
+                base_path
+            )
+
+            for path in base_path.ls(recursive=False, allow_not_found=True):
+                if path.name.startswith("tmp"):
+                    parts = path.name.split("-")
+
+                    if len(parts) > 2 and parts[0] == "tmp" and parts[1].isdigit() and parts[2].isdigit():
+                        end = int(parts[2])
+
+                        if end and time.time() > end:
+                            path.remove(recursive=True)
+
+            LOGGER.info(
+                "Cleaned temp path %s",
+                base_path
+            )
+        else:
+            Thread(
+                target=self.clean_tmp_folder,
+                kwargs={
+                    "raise_error": raise_error,
+                    "base_path": base_path
+                }
+            ).start()
+
+        return self
+
     def shared_cache_path(
         self,
         suffix: Optional[str] = None
@@ -469,13 +625,13 @@ class Workspace:
     # SDK access / connection
     # ------------------------------------------------------------------ #
 
-    def sdk(self) -> "WorkspaceClient":
+    def sdk(self) -> WorkspaceClient:
         """Return the connected WorkspaceClient.
 
         Returns:
             The WorkspaceClient instance.
         """
-        return self.connect()._sdk
+        return self.connect(clone=False)._sdk
 
     # ------------------------------------------------------------------ #
     # List / open / delete / SQL
@@ -529,14 +685,15 @@ class Workspace:
                     yield from self.list_path(child_path, recursive=True)
             return
 
-        # Workspace files / notebooks
-        try:
-            entries = list(sdk.workspace.list(path, recursive=recursive))
-        except ResourceDoesNotExist:
-            return
+        else:
+            # Workspace files / notebooks
+            try:
+                entries = list(sdk.workspace.list(path, recursive=recursive))
+            except ResourceDoesNotExist:
+                return
 
-        for obj in entries:
-            yield obj
+            for obj in entries:
+                yield obj
 
     def open_path(
         self,
@@ -582,20 +739,29 @@ class Workspace:
         """
         return os.getenv("DATABRICKS_RUNTIME_VERSION") is not None
 
-    def default_tags(self):
+    def default_tags(self, update: bool = True):
         """Return default resource tags for Databricks assets.
 
         Returns:
             A dict of default tags.
         """
-        base = {
-            k: v
-            for k, v in (
-                ("Product", self.product),
-                ("ProductTag", self.product_tag),
-            )
-            if v
-        }
+        if update:
+            base = {
+                k: v
+                for k, v in (
+                    ("Product", self.product),
+                )
+                if v
+            }
+        else:
+            base = {
+                k: v
+                for k, v in (
+                    ("Product", self.product),
+                    ("ProductTag", self.product_tag),
+                )
+                if v
+            }
 
         if self.custom_tags:
             base.update(self.custom_tags)
@@ -605,7 +771,7 @@ class Workspace:
     def sql(
         self,
         workspace: Optional["Workspace"] = None,
-        warehouse_id: Optional[str] = None,
+        warehouse: Optional["SQLWarehouse"] = None,
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
     ):
@@ -613,7 +779,7 @@ class Workspace:
 
         Args:
             workspace: Optional workspace override.
-            warehouse_id: Optional SQL warehouse id.
+            warehouse: Optional SQL warehouse.
             catalog_name: Optional catalog name.
             schema_name: Optional schema name.
 
@@ -622,11 +788,13 @@ class Workspace:
         """
         from ..sql import SQLEngine
 
+        workspace = self if workspace is None else workspace
+
         return SQLEngine(
-            workspace=self if workspace is None else workspace,
-            warehouse_id=warehouse_id,
+            workspace=workspace,
             catalog_name=catalog_name,
             schema_name=schema_name,
+            _warehouse=warehouse,
         )
 
     def warehouses(
@@ -652,9 +820,9 @@ class Workspace:
         """Return a Cluster helper bound to this workspace.
 
         Args:
+            workspace: Optional workspace override.
             cluster_id: Optional cluster id.
             cluster_name: Optional cluster name.
-            **kwargs: Additional Cluster parameters.
 
         Returns:
             A Cluster instance.
@@ -662,21 +830,11 @@ class Workspace:
         from ..compute.cluster import Cluster
 
         return Cluster(
-            workspace=self,
+            workspace=self if workspace is None else workspace,
             cluster_id=cluster_id,
             cluster_name=cluster_name,
         )
 
-    def loki(
-        self,
-        workspace: Optional["Workspace"] = None,
-    ):
-        from ..ai.loki import Loki
-
-        return Loki(
-            workspace=self,
-        )
-
 # ---------------------------------------------------------------------------
 # Workspace-bound base class
 # ---------------------------------------------------------------------------
@@ -689,15 +847,6 @@ 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.
 
@@ -742,7 +891,7 @@ class WorkspaceService(ABC):
         parts: Union[List[str], str],
         kind: Optional[DatabricksPathKind] = None,
         workspace: Optional["Workspace"] = None
-    ):
+    ) -> "DatabricksPath":
         """Create a DatabricksPath in the underlying workspace.
 
         Args:

yggdrasil/exceptions.py

@@ -0,0 +1,7 @@
+__all__ = [
+    "YGGException"
+]
+
+
+class YGGException(Exception):
+    pass

yggdrasil/libs/databrickslib.py

@@ -1,24 +1,13 @@
 """Optional Databricks SDK dependency helpers."""
+from ..types.dummy_class import DummyModuleClass
 
-try:
-    import databricks
-    import databricks.sdk  # type: ignore
 
-    from databricks.sdk import WorkspaceClient
+class DatabricksDummyClass(DummyModuleClass):
+    """Placeholder object that raises if Databricks SDK is required."""
 
-    databricks = databricks
-    databricks_sdk = databricks.sdk
-except ImportError:
-    class _DatabricksDummy:
-        """Placeholder object that raises if Databricks SDK is required."""
-        def __getattr__(self, item):
-            """Raise an error when accessing missing Databricks SDK attributes."""
-            require_databricks_sdk()
-
-    databricks = _DatabricksDummy
-    databricks_sdk = _DatabricksDummy
-
-    WorkspaceClient = _DatabricksDummy
+    @classmethod
+    def module_name(cls) -> str:
+        return "databricks"
 
 
 def require_databricks_sdk():
@@ -34,9 +23,24 @@ def require_databricks_sdk():
         )
 
 
+try:
+    import databricks
+    import databricks.sdk
+
+    from databricks.sdk import WorkspaceClient
+
+    databricks = databricks
+    databricks_sdk = databricks.sdk
+except ImportError:
+    databricks = DatabricksDummyClass
+    databricks_sdk = DatabricksDummyClass
+    WorkspaceClient = DatabricksDummyClass
+
+
 __all__ = [
     "databricks",
     "databricks_sdk",
     "require_databricks_sdk",
-    "WorkspaceClient"
+    "WorkspaceClient",
+    "DatabricksDummyClass"
 ]

yggdrasil/libs/extensions/spark_extensions.py

@@ -7,7 +7,6 @@ from typing import List, Union, Optional, Iterable, Callable, TYPE_CHECKING, Map
 
 import pyarrow as pa
 
-from ..pandaslib import pandas
 from ..sparklib import (
     pyspark,
     SparkDataFrame,
@@ -18,6 +17,7 @@ from ..sparklib import (
 from ...types.cast.registry import convert
 
 if TYPE_CHECKING:  # pragma: no cover
+    from ..pandaslib import pandas
     from ...types.cast.cast_options import CastOptions
 
 if pyspark is not None:

yggdrasil/libs/pandaslib.py

@@ -1,15 +1,22 @@
 """Optional pandas dependency helpers."""
 
+from ..types.dummy_class import DummyModuleClass
+
+class DummyPandasClass(DummyModuleClass):
+
+    @classmethod
+    def module_name(cls) -> str:
+        return "pandas"
+
+
 try:
     import pandas  # type: ignore
-    pandas = pandas
-
     PandasDataFrame = pandas.DataFrame
+    PandasSeries = pandas.Series
 except ImportError:
     pandas = None
-
-    class PandasDataFrame:
-        pass
+    PandasDataFrame = DummyPandasClass
+    PandasSeries = DummyPandasClass
 
 
 def require_pandas():
@@ -28,5 +35,7 @@ def require_pandas():
 __all__ = [
     "pandas",
     "require_pandas",
-    "PandasDataFrame"
+    "PandasDataFrame",
+    "PandasSeries",
+    "DummyPandasClass",
 ]

yggdrasil/libs/polarslib.py

@@ -1,4 +1,13 @@
 """Optional Polars dependency helpers."""
+from yggdrasil.types.dummy_class import DummyModuleClass
+
+
+class DummyPolarsClass(DummyModuleClass):
+
+    @classmethod
+    def module_name(cls) -> str:
+        return "polars"
+
 
 try:
     import polars  # type: ignore
@@ -6,27 +15,54 @@ try:
     polars = polars
 
     PolarsDataFrame = polars.DataFrame
+    PolarsSeries = polars.Series
+    PolarsExpr = polars.Expr
+    PolarsDataFrame = polars.DataFrame
+    PolarsField = polars.Field
+    PolarsSchema = polars.Schema
+    PolarsDataType = polars.DataType
+
+
+    def require_polars():
+        """Ensure polars is available before using polars helpers.
+
+        Returns:
+            None.
+        """
+        return None
+
 except ImportError:
     polars = None
+    PolarsDataFrame = DummyPolarsClass
+    PolarsSeries = DummyPolarsClass
+    PolarsExpr = DummyPolarsClass
+    PolarsDataFrame = DummyPolarsClass
+    PolarsField = DummyPolarsClass
+    PolarsSchema = DummyPolarsClass
+    PolarsDataType = DummyPolarsClass
+
+
+    def require_polars():
+        """Ensure polars is available before using polars helpers.
+
+        Returns:
+            None.
+        """
+        import polars
 
-    class PolarsDataFrame:
-        pass
 
 __all__ = [
     "polars",
     "require_polars",
-    "PolarsDataFrame"
+    "PolarsDataFrame",
+    "PolarsSeries",
+    "PolarsExpr",
+    "PolarsDataFrame",
+    "PolarsField",
+    "PolarsSchema",
+    "PolarsDataType",
+    "DummyPolarsClass"
 ]
 
 
-def require_polars():
-    """Ensure polars is available before using polars helpers.
 
-    Returns:
-        None.
-    """
-    if polars is None:
-        raise ImportError(
-            "polars is required to use this function. "
-            "Install it with `pip install polars`."
-        )

yggdrasil/pyutils/__init__.py

@@ -1,8 +1,7 @@
 """Python utility helpers for retries, parallelism, and environment management."""
 
+from .mimetypes import *
 from .retry import retry
 from .parallel import parallelize
 from .python_env import PythonEnv
 from .callable_serde import CallableSerde
-
-__all__ = ["retry", "parallelize", "PythonEnv", "CallableSerde"]