sentry-infra-tools 0.0.3__py3-none-any.whl

assistant/prdocs.py

@@ -0,0 +1,147 @@
+import argparse
+import sys
+from dataclasses import dataclass
+from json import loads
+from pathlib import Path
+from typing import MutableMapping, MutableSet, Sequence
+
+COMMENT_TEMPLATE = """# Ops Assistant
+
+The ops repo is a large monorepo and there are multiple processes
+to deploy its content depending on what is changed.
+This message is here to provide context.
+
+{content}
+"""
+
+INSTRUCTION_TEMPLATE = """
+{content}
+
+Files changed:
+```
+{files}
+```
+---
+"""
+
+INSTRUCTIONS_FILE = "deploy_instructions.md"
+INSTRUCTIONS_CONF_FILE = "deploy_instructions.json"
+
+
+@dataclass
+class Instruction:
+    paths: MutableSet[Path]
+    text: str
+
+
+class InstructionsMessage:
+    """
+    Accumulates instructions to add to a comment in a PR.
+
+    Given a list of paths (touched by the PR) this scans all its
+    parent directories looking for instructions on how to deploy
+    the change contained in the file.
+
+    It then merges all instructions into one message.
+    """
+
+    def __init__(self, root: Path) -> None:
+        # The root directory where the repo is.
+        self.__root = root.resolve()
+        self.__instructions: MutableMapping[Path, Instruction] = {}
+
+    def __fetch_instruction(self, path: Path, updated_file: Path) -> None:
+        """
+        Looks for a `deploy_instructions.md` file to add to the comment.
+
+        Many directories should share the same instructions (think about
+        the k8s services deployable with GoCD). For that there is the
+        `deploy_instructions.json` file that contains a reference to a
+        markdown instructions file.
+
+        This file will also be used to provide more sophisticated
+        directives on how to compose the message.
+        """
+        config_path = path / INSTRUCTIONS_CONF_FILE
+        instructions_path = path / INSTRUCTIONS_FILE
+
+        if config_path.exists() and config_path.is_file():
+            conf = loads(config_path.read_text())
+            assert isinstance(conf, MutableMapping), f"Invalid content of {INSTRUCTIONS_CONF_FILE}"
+            ref_path = conf.get("ref")
+            if ref_path:
+                instructions_path = path / ref_path
+
+        if instructions_path.exists() and instructions_path.is_file():
+            content = instructions_path.read_text()
+
+            instruction_obj = self.__instructions.get(instructions_path.resolve())
+            if instruction_obj is None:
+                instruction_obj = Instruction({updated_file}, content)
+                self.__instructions[instructions_path] = instruction_obj
+            else:
+                instruction_obj.paths.add(updated_file)
+
+    def add_path(self, path: Path) -> None:
+        path = path.resolve()
+        if self.__root not in path.parents:
+            return
+
+        updated_file = path
+        if not path.is_dir():
+            path = path.parent
+
+        while path != self.__root:
+            self.__fetch_instruction(path, updated_file)
+            path = path.parent
+
+    def produce_message(self) -> str | None:
+        if not self.__instructions:
+            return None
+
+        instructions_content = ""
+        for instruction in self.__instructions.values():
+            files = "".join(sorted(str(p) for p in instruction.paths))
+            instructions_content += INSTRUCTION_TEMPLATE.format(
+                **{"files": files, "content": instruction.text}
+            )
+
+        return COMMENT_TEMPLATE.format(**{"content": instructions_content})
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        description="""
+            Collects all the relevant documentation to deploy the changes
+            in a pr given the list of involved files. It, then, produces a single
+            message to attach to the PR itself.
+        """
+    )
+    parser.add_argument(
+        "-r",
+        "--root",
+        action="store",
+        help="""
+            The root directory of all changes
+        """,
+    )
+
+    args = parser.parse_args()
+
+    root = args.root
+    assert (
+        root is not None and Path(root).exists() and Path(root).is_dir()
+    ), "The root path does not exists or is not a directory."
+    message = InstructionsMessage(Path(root))
+
+    for path in sys.stdin:
+        print(f"Processing {path}", file=sys.stderr)
+        message.add_path(Path(path))
+
+    response = message.produce_message()
+    if response is not None:
+        print(response)
+
+
+if __name__ == "__main__":
+    main()

assistant/test_prdocs.py

@@ -0,0 +1,88 @@
+import os
+import tempfile
+from json import dumps
+from pathlib import Path
+from typing import Generator
+
+import pytest
+
+from assistant.prdocs import InstructionsMessage
+
+
+@pytest.fixture
+def valid_structure() -> Generator[str, None, None]:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        full_path1 = Path(temp_dir) / "shared_config/kafka/generated_files"
+        full_path2 = Path(temp_dir) / "shared_config/kafka/other_files"
+        full_path3 = Path(temp_dir) / "shared_config/kafka/more_files"
+        os.makedirs(full_path1)
+        os.makedirs(full_path2)
+        os.makedirs(full_path3)
+
+        with open(full_path1 / "topic1.yaml", "w") as file:
+            file.write("something")
+
+        with open(full_path1 / "topic2.yaml", "w") as file:
+            file.write("something")
+
+        with open(full_path1 / "deploy_instructions.md", "w") as file:
+            file.write("Kafka content")
+
+        with open(full_path2 / "topic3.yaml", "w") as file:
+            file.write("something")
+
+        with open(full_path2 / "deploy_instructions.json", "w") as file:
+            file.write(dumps({"ref": str(full_path1 / "deploy_instructions.md")}))
+
+        yield temp_dir
+
+
+def test_empty(valid_structure) -> None:
+    message = InstructionsMessage(Path(valid_structure))
+    message.add_path(Path(valid_structure) / "shared_config/kafka/more_files/something.yaml")
+
+    assert message.produce_message() is None
+
+
+def test_full_message(valid_structure) -> None:
+    message = InstructionsMessage(Path(valid_structure))
+    message.add_path(Path(valid_structure) / "shared_config/kafka/generated_files/topic1.yaml")
+    message.add_path(Path(valid_structure) / "shared_config/kafka/generated_files/topic2.yaml")
+
+    two_files_comment = f"""# Ops Assistant
+
+The ops repo is a large monorepo and there are multiple processes
+to deploy its content depending on what is changed.
+This message is here to provide context.
+
+
+Kafka content
+
+Files changed:
+```
+{Path(valid_structure).resolve() / "shared_config/kafka/generated_files/topic1.yaml"}{Path(valid_structure).resolve() / "shared_config/kafka/generated_files/topic2.yaml"}
+```
+---
+
+"""
+    assert message.produce_message() == two_files_comment
+
+    message.add_path(Path(valid_structure) / "shared_config/kafka/other_files/topic3.yaml")
+
+    three_files_comment = f"""# Ops Assistant
+
+The ops repo is a large monorepo and there are multiple processes
+to deploy its content depending on what is changed.
+This message is here to provide context.
+
+
+Kafka content
+
+Files changed:
+```
+{Path(valid_structure).resolve() / "shared_config/kafka/generated_files/topic1.yaml"}{Path(valid_structure).resolve() / "shared_config/kafka/generated_files/topic2.yaml"}{Path(valid_structure).resolve() / "shared_config/kafka/other_files/topic3.yaml"}
+```
+---
+
+"""
+    assert message.produce_message() == three_files_comment

config_builder/clean_generated.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Sequence
+
+from config_builder.combined_generator import clean_all as clean_generated
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        description="""
+        Clean up all generated and materialized files.
+        """
+    )
+    parser.add_argument(
+        "root_dir",
+        type=str,
+        action="store",
+        help="""
+            The root directory of the shared configs.
+        """,
+    )
+
+    args = parser.parse_args(argv)
+    materialized_dir = Path(args.root_dir)
+    clean_generated(materialized_dir)
+
+
+if __name__ == "__main__":
+    main()

config_builder/combined_generator.py

@@ -0,0 +1,170 @@
+import os
+from enum import Enum
+from pathlib import Path
+from shutil import rmtree
+from typing import Generator, Sequence, Tuple
+
+from config_builder.loaders import YamlFileLoader
+from config_builder.merger import FileMerger
+from config_builder.merger.libsonnet import LibsonnetMerger
+from config_builder.merger.yamljson import YamlMerger
+
+from config_builder.json_schema_validator import JsonSchemaValidator
+
+DEFAULT_LIBSONNET_OUTPUT_FILE_NAME = "_generated.libsonnet"
+DEFAULT_JSON_OUTPUT_FILE_NAME = "_generated.json"
+CONFIG_GENERATOR_SETTINGS = "_config_generator.json"
+GENERATED_DIR = "generated"
+
+
+class Outcome(Enum):
+    NEW = "new"
+    UNCHANGED = "unchanged"
+    UPDATED = "updated"
+
+
+def combined_file_name(root_directory: Path, output_file_name: str) -> Path:
+    return root_directory / GENERATED_DIR / output_file_name
+
+
+def combine_files(
+    merger: FileMerger,
+    root_directory: Path,
+) -> str:
+    """
+    Combine all the files in `root_directory` using the FileMerger
+    provided.
+    """
+    assert root_directory.is_dir(), "Root directory must be a directory"
+
+    input_files = [x for x in Path(root_directory).iterdir() if x.is_file()]
+    for f in input_files:
+        merger.add_file(f)
+
+    return merger.serialize_content()
+
+
+def validate_schema(validator: JsonSchemaValidator, root_directory: Path) -> None:
+    input_files = [x for x in Path(root_directory).iterdir() if x.is_file()]
+    for f in input_files:
+        if str(f).endswith(("yaml", "yml")):
+            validator.validate_yaml(f)
+
+
+def combine_and_write(
+    merger: FileMerger,
+    root_directory: Path,
+    output_file_name: str,
+) -> Tuple[Outcome, Path]:
+    """
+    Generates a unified file from the all the yaml/json/libsonnet files
+    contained in a directory.
+
+    A FileMerger is provided that is able to merge files of different types.
+    There are two implementations: one for yaml/json and the other for
+    libsonnet files.
+
+    The output is written in a predefined location `root_dir/generated` if
+    the generated file is in a ny way different from the pre-existing
+    combined import file in such directory.
+
+    It return whether the file was unchanged, new or updated.
+    """
+    generated_file = combine_files(merger, root_directory)
+    if not (root_directory / GENERATED_DIR).exists():
+        os.makedirs(root_directory / GENERATED_DIR)
+
+    generated_file_path = combined_file_name(root_directory, output_file_name)
+    if not generated_file_path.exists():
+        outcome = Outcome.NEW
+    else:
+        with open(generated_file_path) as f:
+            existing_content = f.read()
+            outcome = Outcome.UNCHANGED if existing_content == generated_file else Outcome.UPDATED
+
+    if outcome != Outcome.UNCHANGED:
+        with open(generated_file_path, "w") as f:
+            f.write(generated_file)
+
+    return (outcome, generated_file_path)
+
+
+def iterate_config_directories(root_dir: Path) -> Generator[Path, None, None]:
+    """
+    Finds all the subdirectories of `root_dir` that are suitable for jsonnet
+    import file generation.
+
+    Such directories contain:
+    - a `_config_generator.json` file to mark it
+    - a `generated` subdirectory where the generated file will be (if it doesn't exsit, it'll get created in the upcoming combine_and_write step)
+    """
+    for item in root_dir.rglob("*"):
+        if item.is_dir() and (item / CONFIG_GENERATOR_SETTINGS).exists():
+            yield item
+
+
+def iterate_roots_and_regions(
+    root_dir: Path,
+) -> Generator[Tuple[Path, bool | None], None, None]:
+    """
+    Extends `iterate_config_directories` by managing properly region
+    subdirectories.
+
+    If a root_directory (a directory with containing a CONFIG_GENERATOR_SETTINGS
+    file) contains a `regional_overrides` subdirectory, the root_directory
+    is also returned, as well as each dir inside 'regional_overrides'.
+    """
+
+    for config_root in iterate_config_directories(root_dir):
+        yield (config_root, False)
+        regions_path = Path(config_root / "regional_overrides")
+        if regions_path.exists() and regions_path.is_dir():
+            for region_root in regions_path.iterdir():
+                if region_root.is_dir():
+                    yield (region_root, True)
+
+
+def validate_all_files(root_directory: Path) -> None:
+    validator = JsonSchemaValidator()
+    for config_root, _ in iterate_roots_and_regions(root_directory):
+        input_files = [
+            x
+            for x in Path(config_root).iterdir()
+            if x.is_file() and str(x).endswith(("yaml", "yml"))
+        ]
+        for f in input_files:
+            validator.validate_yaml(f)
+
+
+def generate_all_files(root_directory: Path) -> Sequence[Tuple[Outcome, Path]]:
+    """
+    Scan all the directories under `root_directory` to find directories
+    that has to be treated as root configs, thus requiring us to combine
+    multiple files into one to be imported by a jsonnet script.
+    """
+    ret = []
+
+    for config_root, _ in iterate_roots_and_regions(root_directory):
+        outcome, output_file_name = combine_and_write(
+            LibsonnetMerger(), config_root, DEFAULT_LIBSONNET_OUTPUT_FILE_NAME
+        )
+        ret.append((outcome, output_file_name))
+
+        content_loader = YamlFileLoader(config_root)
+        outcome, output_file_name = combine_and_write(
+            YamlMerger(CONFIG_GENERATOR_SETTINGS, content_loader),
+            config_root,
+            DEFAULT_JSON_OUTPUT_FILE_NAME,
+        )
+        ret.append((outcome, output_file_name))
+    return ret
+
+
+def clean_all(root_directory: Path) -> None:
+    """
+    Deletes all the generated files.
+    """
+    for config_root, _ in iterate_roots_and_regions(root_directory):
+        if (config_root / GENERATED_DIR).exists():
+            assert (config_root / GENERATED_DIR).is_dir()
+            rmtree(config_root / GENERATED_DIR)

config_builder/generate_topic_data.py

@@ -0,0 +1,26 @@
+import os
+import json
+from sys import stderr
+from sentry_kafka_schemas import get_topic, list_topics
+from pathlib import Path
+
+
+def main() -> None:
+    root = Path(os.path.abspath(__file__)).parent.parent.parent.parent
+    output = f"{root}/shared_config/kafka/topics/generated/_generated_raw_topic_data.json"
+
+    print("Generating raw topic data", file=stderr)
+    file_path = Path(output)
+
+    topic_data = {
+        name: {"topic_creation_config": get_topic(name)["topic_creation_config"]}
+        for name in list_topics()
+    }
+
+    os.makedirs(file_path.parent, exist_ok=True)
+    with open(file_path, "w") as f:
+        f.write(json.dumps(topic_data, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

config_builder/json_schema_validator.py

@@ -0,0 +1,81 @@
+import json
+from pathlib import Path
+from typing import Any, Mapping, Optional
+
+import jsonschema
+import yaml
+from jsonschema.exceptions import ValidationError
+from referencing import Registry, Resource
+
+# TODO: this should not be hardcoded. Find a dynamic way to specify schema files, maybe in _config_generator.json?
+SCHEMAS = {
+    "kafka/topics/regional_overrides/*/*.yaml": "kafka_topic_overrides.schema.json",
+    "kafka/consumer_groups/regional_overrides/*/*.yaml": "consumer_groups/regional_override_consumer_group.schema.json",
+    "kafka/consumer_groups/*.yaml": "consumer_groups/default_consumer_group.schema.json",
+    "kafka/clusters/*.yaml": "kafka_clusters.schema.json",
+    "kafka/consumer_types/*.yaml": "consumer_groups/consumer_type.schema.json",
+}
+
+ROOT = Path("shared_config")
+SCHEMAS_DIR = "schemas"
+REGIONAL_OVERRIDE_DIR = "regional_overrides"
+
+
+class ValidationException(Exception):
+    def __init__(self, file: str, schema: str):
+        self.file = file
+        self.schema = schema
+
+
+class JsonSchemaValidator:
+    def __init__(
+        self,
+        root: Optional[Path] = None,
+        schemas: Optional[Mapping[str, str]] = None,
+        schemas_dir: Optional[str] = None,
+    ) -> None:
+        self.root = ROOT if root is None else root
+        self.schema_root = (
+            (self.root / SCHEMAS_DIR) if schemas_dir is None else (self.root / schemas_dir)
+        )
+        self.schemas = SCHEMAS if schemas is None else schemas
+        self.registry = self.__build_registry()
+
+    def __build_registry(self) -> Registry[Mapping[str, Resource[Mapping[str, Any]]]]:
+        """
+        The jsonschema's concept of registry is for $ref keyword to
+        resolve file paths and find the referenced object
+        """
+
+        def retrieve_ref(ref: str) -> Resource[Mapping[str, Any]]:
+            ref_file = self.schema_root / Path(ref)
+            contents = json.loads(ref_file.read_text())
+            return Resource.from_contents(contents)
+
+        # library bug, typing of this argument is wrong
+        registry: Registry = Registry(retrieve=retrieve_ref)  # type: ignore
+        return registry
+
+    def __get_schema(self, file: Path) -> Path | None:
+        """
+        Finds which schema to use for a yaml file
+        """
+        for glob, schema in self.schemas.items():
+            if file.match(str(self.root / glob)):
+                return self.schema_root / schema
+        return None
+
+    def validate_yaml(self, file: Path) -> int | None:
+        schema = self.__get_schema(file)
+        if schema:
+            file_content = yaml.safe_load(file.read_text())
+            try:
+                jsonschema.validate(
+                    schema=json.load(open(schema)),
+                    instance=file_content,
+                    registry=self.registry,
+                )
+                return 0
+            except ValidationError as e:
+                raise ValidationException(str(file), str(schema)) from e
+        return None

config_builder/loaders.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Mapping, cast
+
+from yaml import safe_load
+
+
+class ContentLoader(ABC):
+    """
+    Abstractions that allow the config materializers to load content
+    from different sources, being them directories or external
+    packages.
+
+    Instances of this class can be reused.
+    """
+
+    @abstractmethod
+    def load_dict(self, file_name: str) -> Mapping[str, Any] | None:
+        """
+        Loads a dictionary from a a source. The file name can be an actual
+        file on disk or something logical representing a file.
+        `file_name` should contain the extension.
+
+        If a non existing file is asked for, this returns None. It does not
+        fail.
+        """
+        raise NotImplementedError
+
+
+class YamlFileLoader(ContentLoader):
+    """
+    Loader that loads simple yaml file from a directory
+    """
+
+    def __init__(self, directory: Path) -> None:
+        assert (
+            directory.exists() and directory.is_dir()
+        ), f"The provided path {directory} does not exists or it is not a directory"
+        self.__directory = directory
+
+    def load_dict(self, file_name: str) -> Mapping[str, Any] | None:
+        if (self.__directory / file_name).exists():
+            with open(self.__directory / file_name) as content:
+                return cast(Mapping[str, Any], safe_load(content))
+
+        return None