databricks-bundle-decorators 0.1.2__py3-none-any.whl

databricks_bundle_decorators/__init__.py

@@ -0,0 +1,55 @@
+"""databricks-bundle-decorators – decorator-based Databricks job/task framework.
+
+Public API
+----------
+Decorators:
+    ``@task``, ``@job``
+
+Cluster configuration:
+    ``job_cluster()``
+
+Data management:
+    ``IoManager``, ``OutputContext``, ``InputContext``
+
+Task values (small scalars):
+    ``set_task_value``, ``get_task_value``
+
+Job parameters:
+    ``params``
+"""
+
+from databricks_bundle_decorators.context import params as params
+from databricks_bundle_decorators.decorators import job as job
+from databricks_bundle_decorators.decorators import job_cluster as job_cluster
+from databricks_bundle_decorators.decorators import task as task
+from databricks_bundle_decorators.discovery import (
+    discover_pipelines as discover_pipelines,
+)
+from databricks_bundle_decorators.io_manager import InputContext as InputContext
+from databricks_bundle_decorators.io_manager import IoManager as IoManager
+from databricks_bundle_decorators.io_manager import OutputContext as OutputContext
+from databricks_bundle_decorators.registry import (
+    DuplicateResourceError as DuplicateResourceError,
+)
+from databricks_bundle_decorators.sdk_types import ClusterConfig as ClusterConfig
+from databricks_bundle_decorators.sdk_types import JobConfig as JobConfig
+from databricks_bundle_decorators.sdk_types import TaskConfig as TaskConfig
+from databricks_bundle_decorators.task_values import get_task_value as get_task_value
+from databricks_bundle_decorators.task_values import set_task_value as set_task_value
+
+__all__ = [
+    "task",
+    "job",
+    "job_cluster",
+    "discover_pipelines",
+    "IoManager",
+    "OutputContext",
+    "InputContext",
+    "DuplicateResourceError",
+    "ClusterConfig",
+    "JobConfig",
+    "TaskConfig",
+    "set_task_value",
+    "get_task_value",
+    "params",
+]

databricks_bundle_decorators/cli.py

@@ -0,0 +1,253 @@
+"""CLI for databricks-bundle-decorators.
+
+Provides scaffolding commands for pipeline repositories.
+
+Usage::
+
+    uv run dbxdec init
+"""
+
+import argparse
+import sys
+import textwrap
+from pathlib import Path
+
+try:
+    import tomllib
+except ModuleNotFoundError:  # Python < 3.11
+    import tomli as tomllib  # type: ignore[no-redef]
+
+
+def _read_pyproject(cwd: Path) -> dict:
+    """Read and parse pyproject.toml from *cwd*."""
+    path = cwd / "pyproject.toml"
+    if not path.exists():
+        print(
+            "Error: No pyproject.toml found in the current directory.", file=sys.stderr
+        )
+        print(
+            "Run this command from the root of your pipeline project.", file=sys.stderr
+        )
+        sys.exit(1)
+    return tomllib.loads(path.read_text())
+
+
+def _detect_package_name(pyproject: dict) -> str:
+    """Derive the Python import name from the project name in pyproject.toml."""
+    name = pyproject.get("project", {}).get("name")
+    if not name:
+        print("Error: No [project].name found in pyproject.toml.", file=sys.stderr)
+        sys.exit(1)
+    return name.replace("-", "_")
+
+
+def _detect_src_layout(cwd: Path, package_name: str) -> Path:
+    """Return the package directory, detecting flat or src layout."""
+    src_path = cwd / "src" / package_name
+    flat_path = cwd / package_name
+    if src_path.exists():
+        return src_path
+    if flat_path.exists():
+        return flat_path
+    # Default to src layout (will be created)
+    return src_path
+
+
+# --- File templates -------------------------------------------------------
+
+_RESOURCES_INIT = '''\
+"""Resource loader for ``databricks bundle deploy``.
+
+Referenced from ``python.resources`` in ``databricks.yaml``::
+
+    python:
+      venv_path: .venv
+      resources:
+        - 'resources:load_resources'
+"""
+
+from databricks.bundles.core import Bundle, Resources
+
+
+def load_resources(bundle: Bundle) -> Resources:
+    """Entry-point called by ``databricks bundle deploy``."""
+    from databricks_bundle_decorators.discovery import discover_pipelines
+    from databricks_bundle_decorators.codegen import generate_resources
+
+    discover_pipelines()
+
+    resources = Resources()
+    for key, job_resource in generate_resources().items():
+        resources.add_resource(key, job_resource)
+    return resources
+'''
+
+_PIPELINES_INIT = '''\
+"""Pipeline auto-discovery.
+
+Every .py module in this package is imported automatically, triggering
+@task / @job / @job_cluster decorator registration.
+"""
+
+
+
+import importlib
+import pkgutil
+
+for _loader, _module_name, _is_pkg in pkgutil.walk_packages(__path__):
+    importlib.import_module(f"{__name__}.{_module_name}")
+'''
+
+_DATABRICKS_YAML = """\
+bundle:
+  name: {project_name}
+
+artifacts:
+  {package_name}:
+    type: whl
+    build: uv build --wheel
+    path: .
+
+python:
+  venv_path: .venv
+  resources:
+    - 'resources:load_resources'
+
+targets:
+  dev:
+    mode: development
+    workspace:
+      host: https://<your-workspace>.azuredatabricks.net/
+"""
+
+_EXAMPLE_PIPELINE = '''\
+"""Example pipeline – replace with your own tasks and jobs."""
+
+from databricks_bundle_decorators import job, job_cluster, task, params
+
+
+default_cluster = job_cluster(
+    name="default_cluster",
+    spark_version="14.3.x-scala2.12",
+    node_type_id="Standard_DS3_v2",
+    num_workers=2,
+)
+
+
+@task
+def hello():
+    print(f"Hello from databricks-bundle-decorators! url={{params.get('url', 'N/A')}}")
+
+
+@job(
+    params={{"url": "https://example.com"}},
+    cluster="default_cluster",
+)
+def example_job():
+    hello()
+'''
+
+
+# --- Init command ----------------------------------------------------------
+
+
+def _cmd_init(args: argparse.Namespace) -> None:
+    """Scaffold a new databricks-bundle-decorators pipeline project."""
+    cwd = Path.cwd()
+    pyproject = _read_pyproject(cwd)
+    package_name = _detect_package_name(pyproject)
+    project_name = pyproject["project"]["name"]
+    pkg_dir = _detect_src_layout(cwd, package_name)
+
+    created: list[str] = []
+    skipped: list[str] = []
+
+    def _write(path: Path, content: str) -> None:
+        if path.exists():
+            skipped.append(str(path.relative_to(cwd)))
+            return
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(content)
+        created.append(str(path.relative_to(cwd)))
+
+    # 1. resources/__init__.py
+    _write(cwd / "resources" / "__init__.py", _RESOURCES_INIT)
+
+    # 2. pipelines/__init__.py  (auto-discovery)
+    _write(pkg_dir / "pipelines" / "__init__.py", _PIPELINES_INIT)
+
+    # 3. Example pipeline
+    _write(
+        pkg_dir / "pipelines" / "example.py",
+        _EXAMPLE_PIPELINE,
+    )
+
+    # 4. databricks.yaml
+    _write(
+        cwd / "databricks.yaml",
+        _DATABRICKS_YAML.format(
+            project_name=project_name,
+            package_name=package_name,
+        ),
+    )
+
+    # 5. Ensure package __init__.py exists
+    _write(pkg_dir / "__init__.py", "")
+
+    # --- Summary -----------------------------------------------------------
+    print()
+    print("databricks-bundle-decorators project initialized!")
+    print()
+
+    if created:
+        print("Created:")
+        for f in created:
+            print(f"  {f}")
+
+    if skipped:
+        print("Skipped (already exist):")
+        for f in skipped:
+            print(f"  {f}")
+
+    # --- Check for entry point in pyproject.toml ---------------------------
+    entry_points = (
+        pyproject.get("project", {})
+        .get("entry-points", {})
+        .get("databricks_bundle_decorators.pipelines", {})
+    )
+    if not entry_points:
+        print()
+        print("Next step: add the pipeline entry point to your pyproject.toml:")
+        print()
+        print(
+            textwrap.dedent(f"""\
+            [project.entry-points."databricks_bundle_decorators.pipelines"]
+            {package_name} = "{package_name}.pipelines"
+            """)
+        )
+
+    print("Done! Define your @task and @job functions in the pipelines/ directory.")
+
+
+# --- Main ------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="dbxdec",
+        description="databricks-bundle-decorators CLI",
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    subparsers.add_parser(
+        "init",
+        help="Scaffold a new databricks-bundle-decorators pipeline project",
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "init":
+        _cmd_init(args)
+    else:
+        parser.print_help()
+        sys.exit(1)

databricks_bundle_decorators/codegen.py

@@ -0,0 +1,107 @@
+"""Convert registries into ``databricks.bundles.jobs`` resource objects.
+
+Called at deploy time by the resource loader.  Reads the global
+registries populated by ``@task``, ``@job_cluster``, and ``@job`` decorators
+and produces ``Job`` dataclass instances that the Databricks CLI
+serialises into the bundle configuration.
+"""
+
+from databricks_bundle_decorators.registry import (
+    _CLUSTER_REGISTRY,
+    _JOB_REGISTRY,
+    _TASK_REGISTRY,
+)
+
+
+def generate_resources(package_name: str = "databricks_bundle_decorators") -> dict:
+    """Build ``{resource_key: Job}`` from the global registries.
+
+    Parameters
+    ----------
+    package_name:
+        The Python package name used in ``PythonWheelTask``.  Must match
+        the ``[project.name]`` in *pyproject.toml*.
+    """
+    from databricks.bundles.jobs import (
+        ClusterSpec,
+        Job,
+        JobCluster,
+        JobParameterDefinition,
+        Library,
+        PythonWheelTask,
+        Task,
+        TaskDependency,
+    )
+
+    jobs: dict[str, Job] = {}
+
+    for job_name, job_meta in _JOB_REGISTRY.items():
+        tasks: list[Task] = []
+
+        for task_key, upstream_keys in job_meta.dag.items():
+            depends_on = [TaskDependency(task_key=uk) for uk in upstream_keys]
+
+            # ----- named_parameters sent to the wheel entry-point ----------
+            named_params: dict[str, str] = {
+                "__job_name__": job_name,
+                "__task_key__": task_key,
+                "__run_id__": "{{job.run_id}}",
+            }
+
+            # Upstream edge info so the runtime can invoke IoManager.load()
+            edges = job_meta.dag_edges.get(task_key, {})
+            for param_name, upstream_task in edges.items():
+                named_params[f"__upstream__{param_name}"] = upstream_task
+
+            # Forward every job-level parameter to the task CLI
+            for param_name in job_meta.params:
+                named_params[param_name] = (
+                    "{{" + f'job.parameters["{param_name}"]' + "}}"
+                )
+
+            # ----- per-task SDK config (max_retries, timeout, etc.) -----
+            qualified_key = f"{job_name}.{task_key}"
+            task_meta = _TASK_REGISTRY.get(qualified_key)
+            task_sdk_config = task_meta.sdk_config if task_meta else {}
+
+            task_obj = Task(
+                task_key=task_key,
+                depends_on=depends_on,
+                job_cluster_key=job_meta.cluster,
+                python_wheel_task=PythonWheelTask(
+                    package_name=package_name,
+                    entry_point="dbxdec-run",
+                    named_parameters=named_params,  # type: ignore[arg-type]  # SDK Variable wrappers
+                ),
+                libraries=[Library(whl="dist/*.whl")],
+                **task_sdk_config,
+            )
+            tasks.append(task_obj)
+
+        # ----- job clusters -----------------------------------------------
+        job_clusters: list[JobCluster] = []
+        if job_meta.cluster and job_meta.cluster in _CLUSTER_REGISTRY:
+            cluster_meta = _CLUSTER_REGISTRY[job_meta.cluster]
+            job_clusters.append(
+                JobCluster(
+                    job_cluster_key=cluster_meta.name,
+                    new_cluster=ClusterSpec.from_dict(cluster_meta.spec),  # type: ignore[arg-type]  # typed as ClusterSpecDict
+                )
+            )
+
+        # ----- parameters -------------------------------------------------
+        parameters = [
+            JobParameterDefinition(name=k, default=v)
+            for k, v in job_meta.params.items()
+        ]
+
+        job_obj = Job(
+            name=job_name,
+            tasks=tasks,  # type: ignore[arg-type]  # SDK Variable wrappers
+            parameters=parameters,
+            job_clusters=job_clusters,  # type: ignore[arg-type]  # SDK Variable wrappers
+            **job_meta.sdk_config,
+        )
+        jobs[job_name] = job_obj
+
+    return jobs

databricks_bundle_decorators/context.py

@@ -0,0 +1,31 @@
+"""Runtime context – provides job parameters to task functions.
+
+At runtime the entry point populates the global ``params`` dict from CLI
+arguments (parsed via ``argparse``).  Task code imports and reads it::
+
+    from databricks_bundle_decorators import params
+
+    @task
+    def my_task():
+        url = params["url"]
+"""
+
+from typing import Any
+
+
+class _Params(dict[str, Any]):
+    """Dict subclass that holds job parameters.
+
+    An instance lives at module level and is populated by the runtime
+    runner before the task function is invoked.
+    """
+
+
+# Global params instance – importable by user code.
+params: _Params = _Params()
+
+
+def _populate_params(values: dict[str, Any]) -> None:
+    """Replace the contents of the global *params* dict."""
+    params.clear()
+    params.update(values)