stack/dev/scripts/backends/databricks.py

"""Databricks Asset Bundle (DAB) emitter.

Generates databricks.yml + transpiled SQL files from the pipeline
registry, table catalog, and stack.toml [databricks] config.

Each pipeline expression is transpiled to Databricks SQL via
``aco.lake.transpile`` and written to ``bundle/sql/<pipeline>/<expr>.sql``.
Table DDL is generated from ``aco.table`` Pydantic models.
The bundle YAML references these SQL files as ``sql_task`` definitions —
Databricks executes native SQL, not a Python wheel wrapper.

Follows the same emit() → dict[str, str] pattern as gitea.py / github.py.
"""

from __future__ import annotations

from collections import OrderedDict

import yaml

# ── YAML helpers ──────────────────────────────────────────────────

yaml.add_representer(
    OrderedDict,
    lambda dumper, data: dumper.represent_mapping(
        "tag:yaml.org,2002:map", data.items()
    ),
)


def _dump(data: dict) -> str:
    return yaml.dump(
        data,
        default_flow_style=False,
        sort_keys=False,
        width=120,
        allow_unicode=True,
    )


# ── Plugin protocol ──────────────────────────────────────────────


class DabPlugin:
    """Protocol for DAB resource plugins."""

    def resources(self, cfg: dict, registry: dict, catalog_schemas: list[str]) -> dict:
        return {}

    def files(self, cfg: dict, registry: dict) -> dict[str, str]:
        """Return {relative_path: content} for extra files."""
        return {}


# ── Built-in plugins ─────────────────────────────────────────────


class SQLPlugin(DabPlugin):
    """Transpile every pipeline expression to Databricks SQL files."""

    def files(self, cfg: dict, registry: dict) -> dict[str, str]:
        from aco.lake.transpile import transpile

        result: dict[str, str] = {}

        for pipe_name, pipe in sorted(registry.items()):
            try:
                sql_map = transpile(
                    pipe,
                    target_dialect="databricks",
                    catalog="${var.catalog}",
                    output_mode="ctas",
                )
            except Exception as e:
                result[f"bundle/sql/{pipe_name}/_error.txt"] = f"Transpile failed: {e}"
                continue

            for expr_name, sql in sql_map.items():
                # expr_name is like "core._stg_claims_member_months"
                # filename: bundle/sql/core/_stg_claims_member_months.sql
                schema, table = expr_name.split(".", 1)
                path = f"bundle/sql/{pipe_name}/{table}.sql"
                header = (
                    f"-- {expr_name}\n-- Generated by gen_config.py — DO NOT EDIT\n\n"
                )
                result[path] = header + sql + "\n"

        return result


class JobsPlugin(DabPlugin):
    """Generate jobs with sql_task references to transpiled SQL files."""

    def resources(self, cfg: dict, registry: dict, catalog_schemas: list) -> dict:
        # Collect which cluster profiles are actually used
        used_profiles: set[str] = set()
        for pipe in registry.values():
            used_profiles.add(getattr(pipe, "cluster_profile", "default"))

        # Build job_clusters
        clusters_cfg = cfg.get("clusters", {"default": {}})
        job_clusters = []
        for profile_name in sorted(used_profiles):
            cluster_spec = dict(clusters_cfg.get(profile_name, {}))
            job_clusters.append(
                OrderedDict(
                    [
                        ("job_cluster_key", f"{profile_name}_cluster"),
                        (
                            "new_cluster",
                            OrderedDict(
                                [
                                    (
                                        "spark_version",
                                        cluster_spec.get(
                                            "spark_version", "15.4.x-scala2.12"
                                        ),
                                    ),
                                    (
                                        "node_type_id",
                                        cluster_spec.get("node_type_id", "i3.xlarge"),
                                    ),
                                    ("num_workers", cluster_spec.get("num_workers", 1)),
                                ]
                            ),
                        ),
                    ]
                )
            )

        # Pipeline tasks — each pipeline runs as a Python wheel task
        pipeline_tasks = []
        for pipe_name, pipe in sorted(registry.items()):
            profile = getattr(pipe, "cluster_profile", "default")

            depends_on = []
            for dep in pipe.upstream:
                depends_on.append(OrderedDict([("task_key", dep)]))

            task = OrderedDict()
            task["task_key"] = pipe_name
            if depends_on:
                task["depends_on"] = depends_on
            task["job_cluster_key"] = f"{profile}_cluster"

            task["python_wheel_task"] = OrderedDict(
                [
                    ("package_name", cfg.get("package_name", "stack")),
                    ("entry_point", cfg.get("entry_point", "cli")),
                    (
                        "named_parameters",
                        OrderedDict(
                            [
                                ("run", pipe_name),
                                ("--target", "spark"),
                                ("--catalog", "${var.catalog}"),
                                ("--save", ""),
                            ]
                        ),
                    ),
                ]
            )
            task["libraries"] = [
                OrderedDict([("whl", "dbfs:/FileStore/wheels/stack-latest.whl")])
            ]
            pipeline_tasks.append(task)

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

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

        return {"jobs": OrderedDict([("stack_pipelines", job)])}


class SchemasPlugin(DabPlugin):
    """Generate Unity Catalog schema resources."""

    def resources(self, cfg: dict, registry: dict, catalog_schemas: list) -> dict:
        if not catalog_schemas:
            return {}
        schemas = OrderedDict()
        for schema_name in sorted(catalog_schemas):
            schemas[schema_name] = OrderedDict(
                [
                    ("name", schema_name),
                    ("catalog_name", "${var.catalog}"),
                ]
            )
        return {"schemas": schemas}


class VolumesPlugin(DabPlugin):
    """Generate a staging volume for data upload."""

    def resources(self, cfg: dict, registry: dict, catalog_schemas: list) -> dict:
        return {
            "volumes": OrderedDict(
                [
                    (
                        "staging",
                        OrderedDict(
                            [
                                ("name", "staging"),
                                ("catalog_name", "${var.catalog}"),
                                ("schema_name", "default"),
                                ("volume_type", "MANAGED"),
                            ]
                        ),
                    ),
                ]
            )
        }


# ── Core generator ───────────────────────────────────────────────

_HEADER = (
    "# DO NOT EDIT — generated by gen_config.py from stack.toml + pipeline registry\n"
    "# Re-generate: uv run python dev/scripts/gen_config.py\n"
)

_DEFAULT_PLUGINS: list[DabPlugin] = [
    JobsPlugin(),
    SchemasPlugin(),
    VolumesPlugin(),
]

# Plugins used only for SQL generation (--dab-sql flag)
_SQL_PLUGINS: list[DabPlugin] = [SQLPlugin()]


def _get_catalog_schemas() -> list[str]:
    """Return only schemas produced by registered pipelines."""
    try:
        from aco.lake.catalog import Catalog

        return Catalog.pipeline_schemas()
    except Exception:
        return []


def emit(cfg_data: dict, *, sql: bool = False) -> dict[str, str]:
    """Generate databricks.yml from config + pipeline registry.

    The generated databricks.yml delegates resource definitions to
    ``bundle/resources/__init__.py`` via the ``python.resources`` key.
    The YAML only contains bundle metadata, targets, and the Python
    entry point — all resources (jobs, schemas, volumes) are defined
    in typed Python code.

    When ``sql=True``, also generates transpiled SQL files via SQLPlugin
    (for audit/validation, not deployment).

    Returns {relative_path: content} for all generated files.
    """
    dab_cfg = cfg_data.get("databricks", {})
    bundle_name = dab_cfg.get("bundle_name", "stack")
    targets_cfg = dab_cfg.get("targets", {})

    all_files: dict[str, str] = {}

    # SQL generation (on demand only)
    if sql:
        from aco.pipe import registry

        for plugin in _SQL_PLUGINS:
            plugin_files = plugin.files(dab_cfg, registry)
            all_files.update(plugin_files)

    # Build minimal bundle YAML — resources come from Python
    bundle = OrderedDict()
    bundle["bundle"] = OrderedDict([("name", bundle_name)])

    bundle["python"] = OrderedDict(
        [
            ("venv_path", ".venv"),
            (
                "resources",
                ["resources:load_resources"],
            ),
        ]
    )

    bundle["workspace"] = OrderedDict([("host", "${DATABRICKS_HOST}")])

    # Targets
    targets = OrderedDict()
    for tname, tcfg in sorted(targets_cfg.items()):
        target = OrderedDict()
        if tcfg.get("mode"):
            target["mode"] = tcfg["mode"]
        if tcfg.get("default"):
            target["default"] = True
        target["variables"] = OrderedDict(
            [
                ("catalog", tcfg.get("catalog", "aco")),
            ]
        )
        if tcfg.get("mode") == "production":
            run_as = dab_cfg.get("run_as", {})
            if run_as.get("service_principal_name"):
                target["run_as"] = OrderedDict(
                    [
                        ("service_principal_name", run_as["service_principal_name"]),
                    ]
                )
        targets[tname] = target
    bundle["targets"] = targets

    all_files["databricks.yml"] = _HEADER + "\n" + _dump(dict(bundle))
    return all_files