databricks-labs-blueprint 0.0.2__py3-none-any.whl

databricks/__init__.py

@@ -0,0 +1,3 @@
+# DO NOT ADD ANYTHING ELSE TO THIS FILE FOR COMPATIBILITY WITH OTHER databricks.* PACKAGES
+# SEE https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

databricks/labs/__init__.py

@@ -0,0 +1,3 @@
+# DO NOT ADD ANYTHING ELSE TO THIS FILE FOR COMPATIBILITY WITH OTHER databricks.labs.* PACKAGES
+# SEE https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

databricks/labs/blueprint/__about__.py

@@ -0,0 +1 @@
+__version__ = "0.0.2"

databricks/labs/blueprint/__main__.py

@@ -0,0 +1,122 @@
+from pathlib import Path
+
+from databricks.labs.blueprint.cli import App
+from databricks.labs.blueprint.entrypoint import (
+    find_project_root,
+    get_logger,
+    relative_paths,
+)
+from databricks.labs.blueprint.tui import Prompts
+
+blueprint = App(__file__)
+logger = get_logger(__file__)
+
+main_py_file = '''from databricks.sdk import AccountClient, WorkspaceClient
+from databricks.labs.blueprint.entrypoint import get_logger
+from databricks.labs.blueprint.cli import App
+
+__app__ = App(__file__)
+logger = get_logger(__file__)
+
+
+@__app__.command
+def me(w: WorkspaceClient, greeting: str):
+    """Shows current username"""
+    logger.info(f"{greeting}, {w.current_user.me().user_name}!")
+
+
+@__app__.command(is_account=True)
+def workspaces(a: AccountClient):
+    """Shows workspaces"""
+    for ws in a.workspaces.list():
+        logger.info(f"Workspace: {ws.workspace_name} ({ws.workspace_id})")
+
+
+if "__main__" == __name__:
+    __app__()
+
+'''
+
+labs_yml_file = """---
+name: __app__
+description: Common libraries for Databricks Labs
+install:
+  script: src/databricks/labs/__app__/__init__.py
+entrypoint: src/databricks/labs/__app__/__main__.py
+min_python: 3.10
+  - name: me
+    description: shows current username
+    flags:
+     - name: greeting
+       default: Hello
+       description: Greeting prefix
+  - name: workspaces
+    is_account: true
+    description: shows current workspaces
+"""
+
+
+@blueprint.command(is_unauthenticated=True)
+def init_project(target):
+    """Creates the required boilerplate structure"""
+    prompts = Prompts()
+
+    project_root = find_project_root()
+    target_folder = Path(target)
+
+    project_name = prompts.question("Name of the project", default=target_folder.name)
+    src_dir, dst_dir = relative_paths(project_root, target_folder.absolute())
+
+    ignore_names = {
+        ".git",
+        ".venv",
+        ".databricks",
+        ".mypy_cache",
+        ".idea",
+        ".coverage",
+        "htmlcov",
+        "__pycache__",
+        "tests",
+        ".databricks-login.json",
+        "coverage.xml",
+        "dist",
+    }
+    queue: list[Path] = [src_dir]  # type: ignore[annotation-unchecked]
+    while queue:
+        current = queue.pop(0)
+        if current.name in ignore_names:
+            continue
+        if current.is_file():
+            relative_file_name = current.as_posix().replace("blueprint", project_name)
+            dst_file = dst_dir / relative_file_name
+            dst_file.parent.mkdir(exist_ok=True, parents=True)
+            with current.open("r") as r, dst_file.open("w") as w:
+                content = r.read().replace("blueprint", project_name)
+                content = content.replace("databricks-sdk", "databricks-labs-blueprint")
+                w.write(content)
+            continue
+        virtual_env_marker = current / "pyvenv.cfg"
+        if virtual_env_marker.exists():
+            continue
+        for file in current.iterdir():
+            if file.as_posix() == "src/databricks/labs/blueprint":
+                continue
+            queue.append(file)
+    inner_package_dir = dst_dir / "src" / "databricks" / "labs" / project_name
+    inner_package_dir.mkdir(parents=True, exist_ok=True)
+    with (inner_package_dir / "__main__.py").open("w") as f:
+        f.write(main_py_file.replace("__app__", project_name))
+    with (inner_package_dir / "__init__.py").open("w") as f:
+        f.write(f"from databricks.labs.{project_name}.__about__ import __version__")
+    with (inner_package_dir / "__about__.py").open("w") as f:
+        f.write('__version__ = "0.0.0"\n')
+    with (dst_dir / "labs.yml").open("w") as f:
+        f.write(labs_yml_file.replace("__app__", project_name))
+    with (dst_dir / "CODEOWNERS").open("w") as f:
+        f.write(f"* @nfx\n/src @databrickslabs/{project_name}-write\n/tests @databrickslabs/{project_name}-write\n")
+    with (dst_dir / "CHANGELOG.md").open("w") as f:
+        f.write(f"# Version changelog\n\n## 0.0.0\n\nInitial {project_name} commit\n")
+
+
+if "__main__" == __name__:
+    blueprint()

databricks/labs/blueprint/cli.py

@@ -0,0 +1,79 @@
+import json
+import logging
+from dataclasses import dataclass
+from typing import Callable
+
+from databricks.sdk import AccountClient, WorkspaceClient
+
+from databricks.labs.blueprint.entrypoint import get_logger, run_main
+from databricks.labs.blueprint.wheels import ProductInfo
+
+
+@dataclass
+class Command:
+    name: str
+    description: str
+    fn: Callable[..., None]
+    is_account: bool = False
+    is_unauthenticated: bool = False
+
+    def needs_workspace_client(self):
+        if self.is_unauthenticated:
+            return False
+        if self.is_account:
+            return False
+        return True
+
+
+class App:
+    def __init__(self, __file: str):
+        self._mapping: dict[str, Command] = {}
+        self._logger = get_logger(__file)
+        self._product_info = ProductInfo()
+
+    def command(self, is_account: bool = False, is_unauthenticated: bool = False):
+        def decorator(func):
+            command_name = func.__name__.replace("_", "-")
+            if not func.__doc__:
+                raise SyntaxError(f"{func.__name__} must have some doc comment")
+            self._mapping[command_name] = Command(
+                name=command_name,
+                description=func.__doc__,
+                fn=func,
+                is_account=is_account,
+                is_unauthenticated=is_unauthenticated,
+            )
+            return func
+
+        return decorator
+
+    def _route(self, raw):
+        payload = json.loads(raw)
+        command = payload["command"]
+        if command not in self._mapping:
+            msg = f"cannot find command: {command}"
+            raise KeyError(msg)
+        flags = payload["flags"]
+        log_level = flags.pop("log_level")
+        if log_level == "disabled":
+            log_level = "info"
+        databricks_logger = logging.getLogger("databricks")
+        databricks_logger.setLevel(log_level.upper())
+        kwargs = {k.replace("-", "_"): v for k, v in flags.items()}
+        try:
+            product_name = self._product_info.product_name()
+            product_version = self._product_info.version()
+            if self._mapping[command].needs_workspace_client():
+                kwargs["w"] = WorkspaceClient(product=product_name, product_version=product_version)
+            elif self._mapping[command].is_account:
+                kwargs["a"] = AccountClient(product=product_name, product_version=product_version)
+            self._mapping[command].fn(**kwargs)
+        except Exception as err:
+            logger = self._logger.getChild(command)
+            if log_level.lower() in ("debug", "trace"):
+                logger.error(f"Failed to call {command}", exc_info=err)
+            else:
+                logger.error(f"{err.__class__.__name__}: {err}")
+
+    def __call__(self):
+        run_main(self._route)

databricks/labs/blueprint/commands.py

@@ -0,0 +1,227 @@
+import ast
+import html
+import json
+import logging
+import platform
+import re
+import sys
+import threading
+from typing import Any, Callable
+
+from databricks.sdk.core import DatabricksError
+from databricks.sdk.mixins.compute import ClustersExt
+from databricks.sdk.service import compute
+
+_out_re = re.compile(r"Out\[[\d\s]+]:\s")
+_tag_re = re.compile(r"<[^>]*>")
+_exception_re = re.compile(r".*Exception:\s+(.*)")
+_execution_error_re = re.compile(r"ExecutionError: ([\s\S]*)\n(StatusCode=[0-9]*)\n(StatusDescription=.*)\n")
+_error_message_re = re.compile(r"ErrorMessage=(.+)\n")
+_ascii_escape_re = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]")
+_LOG = logging.getLogger("databricks.sdk")
+
+
+class _ReturnToPrintJson(ast.NodeTransformer):
+    def __init__(self) -> None:
+        self._has_json_import = False
+        self._has_print = False
+        self.has_return = False
+
+    @staticmethod
+    def transform(code: str) -> tuple[str, bool]:
+        major, minor, _ = platform.python_version_tuple()
+        unsupported_version = (int(major), int(minor)) < (3, 8)
+        has_return_in_last_line = code.splitlines()[-1].startswith("return ")
+
+        if unsupported_version and not has_return_in_last_line:
+            return code, False
+
+        if unsupported_version and has_return_in_last_line:
+            raise ValueError(
+                "dynamic conversion of return .. to print(json.dumps(..)) " "is only possible starting from Python 3.8"
+            )
+
+        # perform AST transformations for very repetitive tasks, like JSON serialization
+        code_tree = ast.parse(code)
+        transform = _ReturnToPrintJson()
+        new_tree = transform.apply(code_tree)
+        code = ast.unparse(new_tree)
+        return code, transform.has_return
+
+    def apply(self, raw_node: ast.AST) -> ast.AST:
+        node: ast.stmt = self.visit(raw_node)
+        if self.has_return and self._has_print:
+            msg = "cannot have print() call, return .. is converted to print(json.dumps(..))"
+            raise ValueError(msg)
+        if self.has_return and not self._has_json_import:
+            new_import: ast.stmt = ast.parse("import json").body[0]
+            node.body.insert(0, new_import)  # type: ignore[attr-defined]
+        return node
+
+    def visit_Import(self, node: ast.Import) -> ast.Import:  # noqa: N802
+        for name in node.names:
+            if "json" != ast.unparse(name):
+                continue
+            self._has_json_import = True
+            break
+        return node
+
+    def visit_Call(self, node: ast.Call):  # noqa: N802
+        if ast.unparse(node.func) == "print":
+            self._has_print = True
+        return node
+
+    def visit_Return(self, node):  # noqa: N802
+        value = node.value
+        if not value:
+            # Remove the original return statement
+            return None
+        return_expr = ast.unparse(value)
+        replaced_code = f"print(json.dumps({return_expr}))"
+        print_call = ast.parse(replaced_code).body[0]
+        self.has_return = True
+        return print_call
+
+
+class CommandExecutor:
+    def __init__(
+        self,
+        clusters: ClustersExt,
+        command_execution: compute.CommandExecutionAPI,
+        cluster_id_provider: Callable[[], str],
+        *,
+        language: compute.Language = compute.Language.PYTHON,
+    ):
+        self._cluster_id_provider = cluster_id_provider
+        self._language = language
+        self._clusters = clusters
+        self._commands = command_execution
+        self._lock = threading.Lock()
+        self._ctx: compute.ContextStatusResponse | None = None
+
+    def run(self, code: str, *, result_as_json=False, detect_return=True) -> Any:
+        code = self._trim_leading_whitespace(code)
+
+        if self._language == compute.Language.PYTHON and detect_return and not result_as_json:
+            code, result_as_json = _ReturnToPrintJson.transform(code)
+
+        ctx = self._running_command_context()
+        cluster_id = self._cluster_id_provider()
+        command_status_response = self._commands.execute(
+            cluster_id=cluster_id, language=self._language, context_id=ctx.id, command=code
+        ).result()
+
+        results = command_status_response.results
+        assert results is not None
+        if command_status_response.status == compute.CommandStatus.FINISHED:
+            self._raise_if_failed(results)
+            if results.result_type == compute.ResultType.TEXT and result_as_json:
+                try:
+                    # parse json from converted return statement
+                    assert results.data is not None
+                    return json.loads(results.data)
+                except json.JSONDecodeError as e:
+                    _LOG.warning("cannot parse converted return statement. Just returning text", exc_info=e)
+                    return results.data
+            return results.data
+        else:
+            # there might be an opportunity to convert builtin exceptions
+            assert results.summary is not None
+            raise DatabricksError(results.summary)
+
+    def install_notebook_library(self, library: str):
+        return self.run(
+            f"""
+            get_ipython().run_line_magic('pip', 'install {library}')
+            dbutils.library.restartPython()
+            """
+        )
+
+    def _running_command_context(self) -> compute.ContextStatusResponse:
+        if self._ctx:
+            return self._ctx
+        with self._lock:
+            if self._ctx:
+                return self._ctx
+            cluster_id = self._cluster_id_provider()
+            self._clusters.ensure_cluster_is_running(cluster_id)
+            self._ctx = self._commands.create(cluster_id=cluster_id, language=self._language).result()
+        return self._ctx
+
+    @staticmethod
+    def _is_failed(results: compute.Results) -> bool:
+        return results.result_type == compute.ResultType.ERROR
+
+    @staticmethod
+    def _text(results: compute.Results) -> str:
+        if results.result_type != compute.ResultType.TEXT:
+            return ""
+        return _out_re.sub("", str(results.data))
+
+    def _raise_if_failed(self, results: compute.Results):
+        if not self._is_failed(results):
+            return
+        raise DatabricksError(self._error_from_results(results))
+
+    def _error_from_results(self, results: compute.Results):
+        """Converts results into an error
+
+        :param results: compute.Results:
+
+        """
+        if not self._is_failed(results):
+            return
+        results_cause = results.cause
+        if results_cause:
+            sys.stderr.write(_ascii_escape_re.sub("", results_cause))
+        else:
+            results_cause = ""
+
+        summary = ""
+        if results.summary:
+            summary = results.summary
+        summary = _tag_re.sub("", summary)
+        summary = html.unescape(summary)
+
+        exception_matches = _exception_re.findall(summary)
+        if len(exception_matches) == 1:
+            summary = exception_matches[0].replace("; nested exception is:", "")
+            summary = summary.rstrip(" ")
+            return summary
+
+        execution_error_matches = _execution_error_re.findall(results_cause)
+        if len(execution_error_matches) == 1:
+            return "\n".join(execution_error_matches[0])
+
+        error_message_matches = _error_message_re.findall(results_cause)
+        if len(error_message_matches) == 1:
+            return error_message_matches[0]
+
+        return summary
+
+    @staticmethod
+    def _trim_leading_whitespace(command_str: str) -> str:
+        """Removes leading whitespace, so that Python code blocks that
+        are embedded into Python code still could be interpreted properly."""
+        lines = command_str.replace("\t", "    ").split("\n")
+        leading_whitespace = sys.maxsize
+        if lines[0] == "":
+            lines = lines[1:]
+        for line in lines:
+            pos = 0
+            for char in line:
+                if char in (" ", "\t"):
+                    pos += 1
+                else:
+                    break
+            if pos < leading_whitespace:
+                leading_whitespace = pos
+        new_command = ""
+        for line in lines:
+            if line == "" or line.strip(" \t") == "":
+                continue
+            if len(line) < leading_whitespace:
+                new_command += line + "\n"
+            else:
+                new_command += line[leading_whitespace:] + "\n"
+        return new_command.strip()

databricks/labs/blueprint/entrypoint.py

@@ -0,0 +1,88 @@
+import logging
+import os
+import sys
+from pathlib import Path
+
+from databricks.labs.blueprint.logger import install_logger
+
+
+def get_logger(file_name: str):
+    """Used as `get_logger(__file__)` to return a relevant logger for a file
+
+    :param file_name: str: use __file__ special constant
+
+    """
+    project_root = find_project_root().absolute()
+    entrypoint = Path(file_name).absolute()
+
+    relative = entrypoint.relative_to(project_root).as_posix()
+    relative = relative.removeprefix("src" + os.sep)
+    relative = relative.removesuffix("/__main__.py")
+    relative = relative.removesuffix("/__init__.py")
+    relative = relative.removesuffix("/cli.py")
+    relative = relative.removesuffix(".py")
+    module_name = relative.replace(os.sep, ".")
+
+    logger = logging.getLogger(module_name)
+
+    level = "INFO"
+    if is_in_debug():
+        level = "DEBUG"
+    logger.setLevel(level)
+
+    return logger
+
+
+def run_main(main):
+    """Runs main function with a logger
+
+    :param main: function that takes command-line arguments
+
+    """
+    install_logger()
+    main(*sys.argv[1:])
+
+
+def find_project_root() -> Path:
+    """Returns pathlib.Path for the nearest folder with pyproject.toml or setup.py file"""
+    this_path = Path(__file__)  # TODO: check if this works in downstreams
+    # TODO: detect when in wheel
+    for leaf in ["pyproject.toml", "setup.py"]:
+        root = find_dir_with_leaf(this_path, leaf)
+        if root is not None:
+            return root
+    msg = "Cannot find project root"
+    raise NotADirectoryError(msg)
+
+
+def find_dir_with_leaf(folder: Path, leaf: str) -> Path | None:
+    """Returns path object for the nearest folder with a leaf file
+
+    :param folder: Path: starting location
+    :param leaf: str: name of the file or folder
+
+    """
+    root = folder.root
+    while str(folder.absolute()) != root:
+        if (folder / leaf).exists():
+            return folder
+        folder = folder.parent
+    return None
+
+
+def is_in_debug() -> bool:
+    """Returns true if run from VSCode or IntelliJ"""
+    if "IDE_PROJECT_ROOTS" in os.environ:
+        return True
+    return os.path.basename(sys.argv[0]) in ["_jb_pytest_runner.py", "testlauncher.py"]
+
+
+def relative_paths(*maybe_paths) -> list[Path]:
+    """Converts list of paths to relative path objects
+
+    :param *maybe_paths: string-like arguments
+
+    """
+    all_paths = [Path(str(_)) for _ in maybe_paths]
+    common_path = Path(os.path.commonpath([_.as_posix() for _ in all_paths]))
+    return [_.relative_to(common_path) for _ in all_paths]

databricks/labs/blueprint/installer.py

@@ -0,0 +1,88 @@
+import json
+import logging
+import threading
+from json import JSONDecodeError
+from typing import TypedDict
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import NotFound
+from databricks.sdk.service.workspace import ImportFormat
+
+logger = logging.getLogger(__name__)
+
+Resources = dict[str, str]
+
+
+class RawState(TypedDict):
+    resources: dict[str, Resources]
+
+
+class IllegalState(ValueError):
+    pass
+
+
+class InstallState:
+    """Manages ~/.{product}/state.json file on WorkspaceFS to track installations"""
+
+    _state: RawState | None = None
+
+    def __init__(
+        self, ws: WorkspaceClient, product: str, config_version: int = 1, *, install_folder: str | None = None
+    ):
+        self._ws = ws
+        self._product = product
+        self._install_folder = install_folder
+        self._config_version = config_version
+        self._lock = threading.Lock()
+
+    def product(self) -> str:
+        return self._product
+
+    def install_folder(self) -> str:
+        if self._install_folder:
+            return self._install_folder
+        me = self._ws.current_user.me()
+        self._install_folder = f"/Users/{me.user_name}/.{self._product}"
+        return self._install_folder
+
+    def __getattr__(self, item: str) -> Resources:
+        with self._lock:
+            if not self._state:
+                self._state = self._load()
+            if item not in self._state["resources"]:
+                self._state["resources"][item] = {}
+            return self._state["resources"][item]
+
+    def _state_file(self) -> str:
+        return f"{self.install_folder()}/state.json"
+
+    def _load(self) -> RawState:
+        """Loads remote state"""
+        default_state: RawState = {"resources": {}}
+        try:
+            raw = json.load(self._ws.workspace.download(self._state_file()))
+            version = raw.pop("$version", None)
+            if version != self._config_version:
+                msg = f"expected state $version={self._config_version}, got={version}"
+                raise IllegalState(msg)
+            return raw
+        except NotFound:
+            return default_state
+        except JSONDecodeError:
+            logger.warning(f"JSON state file corrupt: {self._state_file}")
+            return default_state
+
+    def save(self) -> None:
+        """Saves remote state"""
+        with self._lock:
+            state: dict = {}
+            if self._state:
+                state = self._state.copy()  # type: ignore[assignment]
+            state["$version"] = self._config_version
+            state_dump = json.dumps(state, indent=2).encode("utf8")
+            self._ws.workspace.upload(
+                self._state_file(),
+                state_dump,  # type: ignore[arg-type]
+                format=ImportFormat.AUTO,
+                overwrite=True,
+            )

databricks/labs/blueprint/limiter.py

@@ -0,0 +1,67 @@
+import logging
+import threading
+import time
+from functools import wraps
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimiter:
+    """Used to limit the rate of requests to a certain maximum rate within a specified burst period,
+    defaulting to one second. It is nice to control and manage the flow of requests to prevent
+    overloading a service.
+
+    It uses a variation of the Leaky Bucket algorithm (https://en.wikipedia.org/wiki/Leaky_bucket)
+    for its simplicity."""
+
+    def __init__(self, *, max_requests: int = 30, burst_period_seconds: int = 1):
+        self._capacity = max_requests
+        self._burst_period_seconds = burst_period_seconds
+        self._lock = threading.RLock()
+        self._last = time.time()
+        self._bucket = 0
+
+    def throttle(self):
+        """The throttle method is used to check and control the rate of requests. It's intended to be called
+        before making a new request and is thread-safe.
+
+        It calculates the time elapsed since the last request and subtracts this from the _burst_period_seconds
+        to determine how much time remains in the current burst period. If the calculated delay is less than zero,
+        it means that the burst period has elapsed, so it resets the request count and updates the last request
+        timestamp. If the request count exceeds the allowed maximum (_max_requests), it sleeps for the remaining
+        time in the burst period (delay) to enforce the rate limit."""
+        with self._lock:
+            now = time.time()
+            delay = self._burst_period_seconds - (now - self._last)
+            if delay < 0:
+                # If the bucket is empty, it stops leaking.
+                self._bucket = 0
+                self._last = now
+            self._bucket += 1
+            if self._bucket > self._capacity:
+                # If the bucket is over capacity, start leaking
+                # by sleeping on the current thread.
+                logger.debug(f"Throttled for {delay}s")
+                time.sleep(delay)
+
+
+def rate_limited(*, max_requests: int = 30, burst_period_seconds: int = 1):
+    """Limits execution frequency of decorated function
+
+    :param *: this function accepts keyword-only arguments
+    :param max_requests: int:  (Default value = 30)
+    :param burst_period_seconds: int:  (Default value = 1)
+
+    """
+
+    def decorator(func):
+        rate_limiter = RateLimiter(max_requests=max_requests, burst_period_seconds=burst_period_seconds)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            rate_limiter.throttle()
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator