schemautil

#!/usr/bin/env python3

"""schemautil -- a hacky tool for managing our schema repository
"""

import click
import json
import logging
import subprocess
import sys

from pathlib import Path
from typing import Any, Self

LOG = logging.getLogger(__name__)


class Schema:
    """Utility class for represenating an openapi schema definition. Has
    methods for reading in and writing out the faux-manifest documents
    we use so that we can patch schemas using kustomize."""

    def __init__(
        self,
        name: str,
        definition: dict[str, Any],
        labels: None | dict[str, str] = None,
        annotations: None | dict[str, str] = None,
    ):
        self.name = name
        self.definition = definition
        self.annotations = annotations if annotations else {}
        self.labels = labels if labels else {}

    def __str__(self) -> str:
        return f"<Schema {self.name}>"

    @classmethod
    def from_manifest(cls, manifest: dict[str, Any]) -> Self:
        name = manifest["metadata"]["name"]
        labels = manifest["metadata"].get("labels")
        annotations = manifest["metadata"].get("annotations")

        return cls(
            name, manifest["definitions"][name], labels=labels, annotations=annotations
        )

    @classmethod
    def from_manifest_file(cls, path: str | Path) -> Self:
        with Path(path).open("r") as fd:
            manifest = json.load(fd)
            return cls.from_manifest(manifest)

    def as_dict(self) -> dict[str, Any]:
        """Return a dictionary representation of this Schema. Suitable
        for writing out to a file."""

        metadata = {"name": self.name}
        if self.labels:
            metadata["labels"] = self.labels
        if self.annotations:
            metadata["annotations"] = self.annotations

        return {
            "apiVersion": "org.openapis.spec/v3",
            "kind": "OpenAPISchema",
            "metadata": metadata,
            "definitions": {
                self.name: self.definition,
            },
        }

    def update_labels(self, labels: list[str]) -> None:
        self.update_kvlist("labels", labels)

    def update_annotations(self, annotations: list[str]) -> None:
        self.update_kvlist("annotations", annotations)

    def update_kvlist(self, name: str, kvlist: list[str]) -> None:
        kvdict = dict(x.split("=", 1) for x in kvlist)
        setattr(
            self,
            name,
            {k: v for k, v in (getattr(self, name, {}) | kvdict).items() if v != "-"},
        )


@click.group()
@click.option("--verbose", "-v", count=True)
def main(verbose):
    init_logging(verbose)


@main.command("kustomize")
@click.argument("schema_directory", type=Path)
def generate_kustomization(schema_directory):
    """Generate `<schema_directory>/kustomization.yaml` by running
    `kustomize create --autodetect --recursive` in
    `schema_directory`."""

    LOG.info("generating %s/kustomization.yaml", schema_directory)

    (schema_directory / "kustomization.yaml").unlink(missing_ok=True)
    subprocess.run(
        ["kustomize", "create", "--autodetect", "--recursive"],
        cwd=schema_directory,
        check=True,
    )


@main.command("regenerate")
@click.option("--label", "-l", "labels", multiple=True)
@click.option("--annotation", "-a", "annotations", multiple=True)
@click.argument("schema_directory", type=Path)
def regenerate_schemas(schema_directory, labels, annotations):
    """Regenerate all the schema files in `schema_directory`
    by reading them in and writing them back out. Use this to add
    labels or annotations or to apply changes to the schema
    template."""

    for schemafile in schema_directory.rglob("*.json"):
        schema = Schema.from_manifest_file(schemafile)
        schema.update_labels(labels)
        schema.update_annotations(annotations)
        LOG.debug("processing %s", schema.name)

        with schemafile.open("w") as fd:
            json.dump(
                schema.as_dict(),
                fd,
                indent=2,
            )
            fd.write("\n")


@main.command("fetch")
@click.option("-o", "--output", default=sys.stdout, type=click.File(mode="w"))
def fetch_schemas(output):
    """Fetch openapi schemas from Kubernetes by running `kustomize openapi fetch`."""

    LOG.info("fetching schemas")
    with output:
        subprocess.run(["kustomize", "openapi", "fetch"], stdout=output, check=True)


@main.command("split")
@click.option("-d", "--schema-directory", default="schemas", type=Path)
@click.option("--label", "-l", "labels", multiple=True)
@click.option("--annotation", "-a", "annotations", multiple=True)
@click.argument("schemas", type=click.File(mode="r"))
def split_schemas(schema_directory, schemas, labels, annotations):
    """Split an openapi schema into one-file-per-definition manifests.
    Apply labels and/or annotations as requested."""

    LOG.info("splitting schemas")

    with schemas as fd:
        all_schemas = json.load(fd)

    for name, definition in all_schemas["definitions"].items():
        LOG.debug("processing %s", name)
        apiversion, kind = name.rsplit(".", 1)
        path = schema_directory / apiversion
        path.mkdir(parents=True, exist_ok=True)

        schema = Schema(name, definition)
        schema.update_labels(labels)
        schema.update_annotations(annotations)

        with (path / f"{kind}.json").open("w") as fd:
            json.dump(schema.as_dict(), fd, indent=2)
            fd.write("\n")


def init_logging(verbose):
    levels = [logging.WARNING, logging.INFO, logging.DEBUG]
    loglevel = levels[min(verbose, len(levels))]
    logging.basicConfig(level=loglevel)


if __name__ == "__main__":
    main()