Add databricks.bundles.core

Author: kanterov

experimental/python/databricks/bundles/build.py

@@ -0,0 +1,46 @@
+__all__ = [
+    "Bundle",
+    "Diagnostic",
+    "Diagnostics",
+    "Location",
+    "Resource",
+    "ResourceMutator",
+    "Resources",
+    "Severity",
+    "Variable",
+    "VariableOr",
+    "VariableOrDict",
+    "VariableOrList",
+    "VariableOrOptional",
+    "job_mutator",
+    "load_resources_from_current_package_module",
+    "load_resources_from_module",
+    "load_resources_from_modules",
+    "load_resources_from_package_module",
+    "variables",
+]
+
+from databricks.bundles.core._bundle import Bundle
+from databricks.bundles.core._diagnostics import (
+    Diagnostic,
+    Diagnostics,
+    Severity,
+)
+from databricks.bundles.core._load import (
+    load_resources_from_current_package_module,
+    load_resources_from_module,
+    load_resources_from_modules,
+    load_resources_from_package_module,
+)
+from databricks.bundles.core._location import Location
+from databricks.bundles.core._resource import Resource
+from databricks.bundles.core._resource_mutator import ResourceMutator, job_mutator
+from databricks.bundles.core._resources import Resources
+from databricks.bundles.core._variable import (
+    Variable,
+    VariableOr,
+    VariableOrDict,
+    VariableOrList,
+    VariableOrOptional,
+    variables,
+)

experimental/python/databricks/bundles/core/__init__.py

@@ -0,0 +1,102 @@
+from dataclasses import dataclass, field
+from typing import Any, TypeVar, Union, get_origin
+
+from databricks.bundles.core._variable import Variable, VariableOr, VariableOrList
+
+__all__ = [
+    "Bundle",
+]
+
+_T = TypeVar("_T")
+
+_VAR_PREFIX = "var"
+
+
+@dataclass(frozen=True, kw_only=True)
+class Bundle:
+    """
+    Bundle contains information about a bundle accessible in functions
+    loading and mutating resources.
+    """
+
+    target: str
+    """
+    Selected target where the bundle is being loaded. E.g.: 'development', 'staging', or 'production'.
+    """
+
+    variables: dict[str, Any] = field(default_factory=dict)
+    """
+    Values of bundle variables resolved for selected target. Bundle variables are defined in databricks.yml.
+    For accessing variables as structured data, use :meth:`resolve_variable`.
+    
+    Example:
+    
+    .. code-block:: yaml
+    
+        variables:
+          default_dbr_version:
+            description: Default version of Databricks Runtime
+            default: "14.3.x-scala2.12"
+    """
+
+    def resolve_variable(self, variable: VariableOr[_T]) -> _T:
+        """
+        Resolve a variable to its value.
+
+        If the value is a variable, it will be resolved and returned.
+        Otherwise, the value will be returned as is.
+        """
+        if not isinstance(variable, Variable):
+            return variable
+
+        if not variable.path.startswith(_VAR_PREFIX + "."):
+            raise ValueError(
+                "You can only get values of variables starting with 'var.*'"
+            )
+        else:
+            variable_name = variable.path[len(_VAR_PREFIX + ".") :]
+
+        if variable_name not in self.variables:
+            raise ValueError(
+                f"Can't find '{variable_name}' variable. Did you define it in databricks.yml?"
+            )
+
+        value = self.variables.get(variable_name)
+
+        # avoid circular import
+        from databricks.bundles.core._transform import (
+            _display_type,
+            _find_union_arg,
+            _transform,
+            _unwrap_variable_path,
+        )
+
+        if nested := _unwrap_variable_path(value):
+            can_be_variable = get_origin(variable.type) == Union and _find_union_arg(
+                nested, variable.type
+            )
+            can_be_variable = can_be_variable or get_origin(variable.type) == Variable
+
+            if not can_be_variable:
+                display_type = _display_type(variable.type)
+
+                raise ValueError(
+                    f"Failed to resolve '{variable_name}' because refers to another "
+                    f"variable '{nested}'. Change variable type to "
+                    f"Variable[VariableOr[{display_type}]]"
+                )
+
+        try:
+            return _transform(variable.type, value)
+        except Exception as e:
+            raise ValueError(f"Failed to read '{variable_name}' variable value") from e
+
+    def resolve_variable_list(self, variable: VariableOrList[_T]) -> list[_T]:
+        """
+        Resolve a list variable to its value.
+
+        If the value is a variable, or the list item is a variable, it will be resolved and returned.
+        Otherwise, the value will be returned as is.
+        """
+
+        return [self.resolve_variable(item) for item in self.resolve_variable(variable)]

experimental/python/databricks/bundles/core/_bundle.py

@@ -0,0 +1,222 @@
+import traceback
+from dataclasses import dataclass, field, replace
+from enum import Enum
+from io import StringIO
+from typing import TYPE_CHECKING, Optional, TypeVar
+
+from databricks.bundles.core._location import Location
+
+_T = TypeVar("_T")
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+__all__ = [
+    "Diagnostic",
+    "Diagnostics",
+    "Severity",
+]
+
+
+class Severity(Enum):
+    WARNING = "warning"
+    ERROR = "error"
+
+
+@dataclass(kw_only=True, frozen=True)
+class Diagnostic:
+    severity: Severity
+    """
+    Severity of the diagnostics item.
+    """
+
+    summary: str
+    """
+    Short summary of the error or warning.
+    """
+
+    detail: Optional[str] = None
+    """
+    Explanation of the error or warning.
+    """
+
+    path: Optional[tuple[str, ...]] = None
+    """
+    Path in databricks.yml where the error or warning occurred.
+    """
+
+    location: Optional[Location] = None
+    """
+    Source code location where the error or warning occurred.
+    """
+
+    def as_dict(self) -> dict:
+        def omit_none(values: dict):
+            return {key: value for key, value in values.items() if value is not None}
+
+        if self.location:
+            location = self.location.as_dict()
+        else:
+            location = None
+
+        return omit_none(
+            {
+                "severity": self.severity.value,
+                "summary": self.summary,
+                "detail": self.detail,
+                "path": self.path,
+                "location": location,
+            }
+        )
+
+
+@dataclass(frozen=True)
+class Diagnostics:
+    """
+    Diagnostics is a collection of errors and warnings we print to users.
+
+    Each item can have source location or path associated, that is reported in output to
+    indicate where the error or warning occurred.
+    """
+
+    items: tuple[Diagnostic, ...] = field(default_factory=tuple, kw_only=False)
+
+    def extend(self, diagnostics: "Self") -> "Self":
+        """
+        Extend items with another diagnostics. This pattern allows
+        to accumulate errors and warnings.
+
+        Example:
+
+        .. code-block:: python
+
+            def foo() -> Diagnostics: ...
+            def bar() -> Diagnostics: ...
+
+            diagnostics = Diagnostics()
+            diagnostics = diagnostics.extend(foo())
+            diagnostics = diagnostics.extend(bar())
+        """
+
+        return replace(
+            self,
+            items=(*self.items, *diagnostics.items),
+        )
+
+    def extend_tuple(self, pair: tuple[_T, "Self"]) -> tuple[_T, "Self"]:
+        """
+        Extend items with another diagnostics. This variant is useful when
+        methods return a pair of value and diagnostics. This pattern allows
+        to accumulate errors and warnings.
+
+        Example:
+
+        .. code-block:: python
+
+            def foo() -> (int, Diagnostics): ...
+
+            diagnostics = Diagnostics()
+            value, diagnostics = diagnostics.extend_tuple(foo())
+        """
+
+        value, other_diagnostics = pair
+
+        return value, self.extend(other_diagnostics)
+
+    def has_error(self) -> bool:
+        """
+        Returns True if there is at least one error in diagnostics.
+        """
+
+        for item in self.items:
+            if item.severity == Severity.ERROR:
+                return True
+
+        return False
+
+    @classmethod
+    def create_error(
+        cls,
+        msg: str,
+        *,
+        detail: Optional[str] = None,
+        location: Optional[Location] = None,
+        path: Optional[tuple[str, ...]] = None,
+    ) -> "Self":
+        """
+        Create an error diagnostics.
+        """
+
+        return cls(
+            items=(
+                Diagnostic(
+                    severity=Severity.ERROR,
+                    summary=msg,
+                    detail=detail,
+                    location=location,
+                    path=path,
+                ),
+            ),
+        )
+
+    @classmethod
+    def create_warning(
+        cls,
+        msg: str,
+        *,
+        detail: Optional[str] = None,
+        location: Optional[Location] = None,
+        path: Optional[tuple[str, ...]] = None,
+    ) -> "Self":
+        """
+        Create a warning diagnostics.
+        """
+
+        return cls(
+            items=(
+                Diagnostic(
+                    severity=Severity.WARNING,
+                    summary=msg,
+                    detail=detail,
+                    location=location,
+                    path=path,
+                ),
+            )
+        )
+
+    @classmethod
+    def from_exception(
+        cls,
+        exc: Exception,
+        *,
+        summary: str,
+        location: Optional[Location] = None,
+        path: Optional[tuple[str, ...]] = None,
+        explanation: Optional[str] = None,
+    ) -> "Self":
+        """
+        Create diagnostics from an exception.
+
+        :param exc: exception to create diagnostics from
+        :param summary: short summary of the error
+        :param location: optional location in the source code where the error occurred
+        :param path: optional path to relevant property in databricks.yml
+        :param explanation: optional explanation to add to the details
+        """
+
+        detail_io = StringIO()
+        traceback.print_exception(exc, file=detail_io)
+
+        detail = detail_io.getvalue()
+        if explanation:
+            detail = f"{detail}\n\n\033[0;36mExplanation:\033[0m {explanation}"
+
+        diagnostic = Diagnostic(
+            severity=Severity.ERROR,
+            summary=summary,
+            location=location,
+            path=path,
+            detail=detail,
+        )
+
+        return cls(items=(diagnostic,))