stack/bundle/resources/__init__.py

"""Databricks Asset Bundle resources — Python-defined.

Called by the DAB CLI at deploy time via::

    # databricks.yml
    python:
      resources:
        - "resources:load_resources"

Generates Job, Schema, and Volume resources from the pipeline
registry, eliminating the need for YAML generation.
"""

from __future__ import annotations

from databricks.bundles.core import Resources, Variable
from databricks.bundles.jobs import (
    CronSchedule,
    Job,
    JobCluster,
    Library,
    PythonWheelTask,
    Task,
    TaskDependency,
)
from databricks.bundles.schemas import Schema
from databricks.bundles.volumes import Volume, VolumeType


def _pipeline_registry() -> dict:
    """Import the pipeline registry. Returns empty dict if unavailable."""
    try:
        from aco.pipe import registry

        return registry
    except ImportError:
        return {}


def _pipeline_schemas(registry: dict) -> list[str]:
    """Derive schema names from pipeline expression outputs."""
    schemas: set[str] = set()
    for pipe in registry.values():
        for expr in pipe.exprs:
            name = expr.name if isinstance(expr.name, str) else expr[0]
            if "." in name:
                schemas.add(name.split(".")[0])
    return sorted(schemas)


def _config() -> dict:
    """Read [databricks] section from stack.toml."""
    try:
        from conf import cfg

        return cfg.databricks.to_dict()
    except (ImportError, AttributeError):
        return {}


def load_resources() -> Resources:
    """Generate all DAB resources from the pipeline registry."""
    registry = _pipeline_registry()
    cfg = _config()
    schema_names = _pipeline_schemas(registry)

    catalog = Variable(path="catalog", type=str)

    # Schemas — one per pipeline-produced schema
    schemas = {
        name: Schema(
            name=name,
            catalog_name=catalog,
        )
        for name in schema_names
    }

    # Volume — staging area for data uploads
    volume = Volume(
        name="staging",
        catalog_name=catalog,
        schema_name="default",
        volume_type=VolumeType.MANAGED,
    )

    # Job cluster
    cluster_cfg = cfg.get("clusters", {}).get("default", {})
    job_cluster = JobCluster(
        job_cluster_key="default_cluster",
        new_cluster={
            "spark_version": cluster_cfg.get("spark_version", "15.4.x-scala2.12"),
            "node_type_id": cluster_cfg.get("node_type_id", "i3.xlarge"),
            "num_workers": cluster_cfg.get("num_workers", 1),
        },
    )

    # Tasks — one per pipeline
    package_name = cfg.get("package_name", "stack")
    entry_point = cfg.get("entry_point", "cli")
    tasks = []
    for pipe_name, pipe in sorted(registry.items()):
        task = Task(
            task_key=pipe_name,
            depends_on=[TaskDependency(task_key=dep) for dep in pipe.upstream] or None,
            job_cluster_key="default_cluster",
            python_wheel_task=PythonWheelTask(
                package_name=package_name,
                entry_point=entry_point,
                named_parameters={
                    "run": pipe_name,
                    "--target": "spark",
                    "--catalog": catalog,
                    "--save": "",
                },
            ),
            libraries=[Library(whl="dbfs:/FileStore/wheels/stack-latest.whl")],
        )
        tasks.append(task)

    # Schedule
    schedule_expr = cfg.get("schedule", "0 0 6 * * ?")
    timezone = cfg.get("timezone", "America/New_York")

    job = Job(
        name=f"{cfg.get('bundle_name', 'stack')}-pipelines",
        description=f"Run all {len(registry)} ACO pipelines in dependency order",
        schedule=CronSchedule(
            quartz_cron_expression=schedule_expr,
            timezone_id=timezone,
        ),
        job_clusters=[job_cluster],
        tasks=tasks,
    )

    resources = Resources()
    resources.add_job("stack_pipelines", job)
    for name, schema in schemas.items():
        resources.add_schema(name, schema)
    resources.add_volume("staging", volume)
    return resources