Merge pull request #1173 from kapicorp/test/omegaconf

feat: omegaconf inventory by @MatteoVoges

Dockerfile

@@ -13,7 +13,7 @@ RUN apt-get update \
         git \
         default-jre
 
-ENV POETRY_VERSION=1.7.1
+ENV POETRY_VERSION=1.8.3
 ENV VIRTUAL_ENV=/opt/venv
 ENV PATH="$VIRTUAL_ENV/bin:/usr/local/go/bin:${PATH}"
 RUN python -m venv $VIRTUAL_ENV \

examples/kubernetes/compiled/jsonnet-env/jsonnet-env/env.yml

@@ -8,7 +8,6 @@ classes:
 exports: {}
 parameters:
   _reclass_:
-    environment: base
     name:
       full: jsonnet-env
       parts:

examples/kubernetes/components/jsonnet-env/env.jsonnet

@@ -1,11 +1,18 @@
 local kap = import 'lib/kapitan.libjsonnet';
 local inventory = kap.inventory();
 
+
+
 {
   env: {
     applications: inventory.applications,
     classes: inventory.classes,
-    parameters: inventory.parameters,
+    parameters: inventory.parameters {
+      ["_kapitan_"]:: std.get(inventory.parameters, "_kapitan_"),  // Ignore this in compile tests because reclass doesn't support it
+      ["_reclass_"]: std.get(inventory.parameters, "_reclass_") {
+        ["environment"]:: "base" // ignore because unused
+      }
+    },
     exports: inventory.exports,
   },
 }

examples/kubernetes/inventory/targets/jsonnet-env.yml

@@ -7,6 +7,7 @@ parameters:
     compile:
       - output_path: jsonnet-env
         input_type: jsonnet
+        input_params: {}
         input_paths:
           - components/jsonnet-env/env.jsonnet
         output_type: yml

kapitan/cached.py

@@ -9,6 +9,8 @@
 from argparse import Namespace
 
 inv = {}
+global_inv = {}
+inventory_global_kadet = {}
 inv_cache = {}
 gpg_obj = None
 gkms_obj = None
@@ -17,14 +19,15 @@
 dot_kapitan = {}
 ref_controller_obj = None
 revealer_obj = None
-args = args = Namespace() # args won't need resetting
+args = Namespace() # args won't need resetting
 inv_sources = set()
 
 
 def reset_cache():
-    global inv, inv_cache, gpg_obj, gkms_obj, awskms_obj, azkms_obj, dot_kapitan, ref_controller_obj, revealer_obj, inv_sources
+    global inv, global_inv, inv_cache, gpg_obj, gkms_obj, awskms_obj, azkms_obj, dot_kapitan, ref_controller_obj, revealer_obj, inv_sources
 
     inv = {}
+    global_inv = {}
     inv_cache = {}
     inv_sources = set()
     gpg_obj = None
@@ -37,9 +40,10 @@ def reset_cache():
 
 
 def from_dict(cache_dict):
-    global inv, inv_cache, gpg_obj, gkms_obj, awskms_obj, azkms_obj, dot_kapitan, ref_controller_obj, revealer_obj, inv_sources, args
+    global inv, global_inv, inv_cache, gpg_obj, gkms_obj, awskms_obj, azkms_obj, dot_kapitan, ref_controller_obj, revealer_obj, inv_sources, args
 
     inv = cache_dict["inv"]
+    global_inv = cache_dict["global_inv"]
     inv_cache = cache_dict["inv_cache"]
     inv_sources = cache_dict["inv_sources"]
     gpg_obj = cache_dict["gpg_obj"]
@@ -55,6 +59,7 @@ def from_dict(cache_dict):
 def as_dict():
     return {
         "inv": inv,
+        "global_inv": global_inv,
         "inv_cache": inv_cache,
         "inv_sources": inv_sources,
         "gpg_obj": gpg_obj,

kapitan/cli.py

@@ -601,7 +601,6 @@ def build_parser():
 
 def main():
     """main function for command line usage"""
-
     try:
         multiprocessing.set_start_method("spawn")
     # main() is explicitly multiple times in tests

kapitan/inputs/base.py

@@ -127,7 +127,7 @@ def write_yaml(self, obj):
         target_name = self.kwargs.get("target_name", None)
 
         # TODO(ademaria): make it configurable per input type
-        style_selection = cached.inv[target_name].parameters.get("multiline_string_style", None)
+        style_selection = cached.inv[target_name]["parameters"].get("multiline_string_style", None)
 
         if not style_selection: 
             if hasattr(cached.args, "multiline_string_style"):

kapitan/inputs/jinja2.py

@@ -11,7 +11,7 @@
 from kapitan.inputs.base import CompiledFile, InputType
 from kapitan.resources import inventory
 from kapitan.utils import render_jinja2
-
+from kapitan import cached
 logger = logging.getLogger(__name__)
 
 
@@ -41,8 +41,9 @@ def compile_file(self, file_path, compile_path, ext_vars, **kwargs):
 
         # set ext_vars and inventory for jinja2 context
         context = ext_vars.copy()
-        context["inventory"] = inventory(self.search_paths, target_name)
-        context["inventory_global"] = inventory(self.search_paths, None)
+
+        context["inventory_global"] = cached.inv.inventory
+        context["inventory"] = cached.inv.inventory[target_name]
         context["input_params"] = input_params
 
         jinja2_filters = kwargs.get("jinja2_filters")

kapitan/inputs/kadet.py

@@ -11,32 +11,34 @@
 import os
 import sys
 from importlib.util import module_from_spec, spec_from_file_location
-
+import time
 import kadet
 from kadet import BaseModel, BaseObj, Dict
 
-from kapitan import cached
 from kapitan.errors import CompileError
+from functools import lru_cache
 from kapitan.inputs.base import CompiledFile, InputType
-from kapitan.resources import inventory as inventory_func
 from kapitan.utils import prune_empty
+from kapitan import cached
 
 # Set external kadet exception to kapitan.error.CompileError
 kadet.ABORT_EXCEPTION_TYPE = CompileError
 
 logger = logging.getLogger(__name__)
-inventory_path = vars(cached.args).get("inventory_path")
-# XXX think about this as it probably breaks usage as library
 search_paths = contextvars.ContextVar("current search_paths in thread")
 current_target = contextvars.ContextVar("current target in thread")
 
 
-def inventory(lazy=False):
-    return Dict(inventory_func(search_paths.get(), current_target.get(), inventory_path), default_box=lazy)
+@lru_cache(maxsize=None)
+def inventory_global(lazy=False):
+    # At hoc inventory for kadet
+    if not cached.inventory_global_kadet:
+        cached.inventory_global_kadet = Dict(cached.global_inv, default_box=lazy)
+    return cached.inventory_global_kadet
 
 
-def inventory_global(lazy=False):
-    return Dict(inventory_func(search_paths.get(), None, inventory_path), default_box=lazy)
+def inventory(lazy=False):
+    return inventory_global(lazy)[current_target.get()]
 
 
 def module_from_path(path, check_name=None):
@@ -63,7 +65,6 @@ def module_from_path(path, check_name=None):
 
     return mod, spec
 
-
 def load_from_search_paths(module_name):
     """
     loads and executes python module with module_name from search paths
@@ -100,7 +101,6 @@ def compile_file(self, file_path, compile_path, ext_vars, **kwargs):
         prune_output = kwargs.get("prune_output", False)
         reveal = kwargs.get("reveal", False)
         target_name = kwargs.get("target_name", None)
-        # inventory_path = kwargs.get("inventory_path", None)
         indent = kwargs.get("indent", 2)
 
         current_target.set(target_name)

kapitan/inventory/__init__.py

@@ -2,11 +2,13 @@
 
 from .inv_reclass import ReclassInventory
 from .inv_reclass_rs import ReclassRsInventory
+from .inv_omegaconf.inv_omegaconf import OmegaConfInventory
 from .inventory import Inventory
 
 # Dict mapping values for command line flag `--inventory-backend` to the
 # associated `Inventory` subclass.
 AVAILABLE_BACKENDS: dict[str, Type[Inventory]] = {
     "reclass": ReclassInventory,
     "reclass-rs": ReclassRsInventory,
+    "omegaconf": OmegaConfInventory,
 }

kapitan/inventory/inv_omegaconf/inv_omegaconf.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+
+# Copyright 2019 The Kapitan Authors
+# SPDX-FileCopyrightText: 2020 The Kapitan Authors <kapitan-admins@googlegroups.com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import multiprocessing as mp
+import os
+import time
+from functools import singledispatch
+from cachetools import cached, LRUCache
+from omegaconf import OmegaConf, ListMergeMode, DictConfig
+from pydantic import ConfigDict
+from ..inventory import InventoryError, Inventory, InventoryTarget
+from .resolvers import register_resolvers, register_user_resolvers
+from kadet import Dict
+from .migrate import migrate
+import yaml
+
+logger = logging.getLogger(__name__)
+
+@singledispatch
+def keys_to_strings(ob):
+    return ob
+
+
+@keys_to_strings.register
+def _handle_dict(ob: dict):
+    return {str(k): keys_to_strings(v) for k, v in ob.items()}
+
+
+@keys_to_strings.register
+def _handle_list(ob: list):
+    return [keys_to_strings(v) for v in ob]
+
+
+class OmegaConfTarget(InventoryTarget):
+    resolved: bool = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__add_metadata()
+
+    def __add_metadata(self):
+        metadata = {
+            "name": {
+                "short": self.name.split(".")[-1],
+                "full": self.name,
+                "path": os.path.splitext(self.path)[0],
+                "parts": self.name.split("."),
+            }
+        }
+        self.parameters["_kapitan_"] = metadata
+        self.parameters["_reclass_"] = metadata
+
+
+class OmegaConfInventory(Inventory):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, target_class=OmegaConfTarget)
+
+    def render_targets(self, targets: list[OmegaConfTarget] = None, ignore_class_not_found: bool = False) -> None:
+        if not self.initialised:
+            manager = mp.Manager()
+            shared_targets = manager.dict()
+            with mp.Pool(min(len(targets), os.cpu_count())) as pool:
+                r = pool.map_async(self.inventory_worker, [(self, target, shared_targets) for target in targets.values()])
+                r.wait()
+
+            for target in shared_targets.values():
+                self.targets[target.name] = target
+
+    @staticmethod
+    def inventory_worker(zipped_args):
+        self, target, shared_targets = zipped_args
+        try:
+            register_resolvers()
+            self.load_target(target)
+            shared_targets[target.name] = target
+
+        except Exception as e:
+            logger.error(f"{target.name}: could not render due to error {e}")
+            raise
+
+    @cached(cache=LRUCache(maxsize=1024))
+    def resolve_class_file_path(self, class_name: str, class_parent_dir: str = None, class_parent_name: str = None):
+        class_file = None
+
+        # Finds relative paths based on the parent directory 
+        if class_name.startswith(".") and class_parent_dir:
+            class_path_base = os.path.join(self.classes_path, class_parent_dir)
+        else:
+            class_path_base = self.classes_path
+
+        # Now try to find the class file
+        extension = ".yml"
+
+        cases = [
+            # case components.kapicorp is absolute and a directory, look for components/kapicorp/init.yml
+            # case .components.kapicorp is relative and a directory, look for <class_parent_dir>/components/kapicorp/init.yml
+            os.path.join(class_path_base, *class_name.split("."), "init" + extension),
+
+            # case components.kapicorp is absolute and a file, look for components/kapicorp.yml
+            # case components.kapicorp is relative and a file, look for <class_parent_dir>/components/kapicorp.yml
+            os.path.join(class_path_base, *class_name.split(".")) + extension,
+
+            # Reclass compatibility mode
+            # case .components.kapicorp  points to <class_parent_dir>/kapicorp.yml
+            os.path.join(class_path_base, *class_name.split(".")[2:]) + extension,
+
+            # case components.kapicorp points to components/kapicorp/init.yml
+            os.path.join(class_path_base, *class_name.split(".")[2:], "init" + extension),
+        ]
+
+        for case in cases:
+            if os.path.isfile(case):
+                class_file = case
+                return class_file
+
+        logger.error(f"class file not found for class {class_name}, tried {cases}")
+        return None
+
+    @cached(cache=LRUCache(maxsize=1024))
+    def load_file(self, filename):
+        with open(filename, "r") as f:
+            return yaml.safe_load(f)
+
+    def load_parameters_from_file(self, filename, parameters={}) -> Dict:
+        parameters = OmegaConf.create(parameters)
+        applications = []
+        classes = []
+        exports = Dict()
+
+        content = self.load_file(filename)
+
+        _classes = content.get("classes", [])
+        _parameters = keys_to_strings(content.get("parameters", {}))
+        _applications = content.get("applications", [])
+        _exports = content.get("exports", {})
+
+        # first processes all classes 
+        for class_name in _classes:
+
+            class_parent_dir = os.path.dirname(filename.removeprefix(self.classes_path).removeprefix("/"))
+            class_parent_name = os.path.basename(filename)
+            class_file = self.resolve_class_file_path(class_name, class_parent_dir=class_parent_dir, class_parent_name=class_parent_name)
+            if not class_file:
+                if self.ignore_class_not_found:
+                    continue
+                raise InventoryError(f"Class {class_name} not found")
+            p, c, a, e = self.load_parameters_from_file(class_file)
+            if p:
+                parameters = OmegaConf.unsafe_merge(parameters, p, list_merge_mode=ListMergeMode.EXTEND_UNIQUE)
+            classes.extend(c)
+            classes.append(class_name)
+            applications.extend(a)
+            exports.merge_update(e, box_merge_lists="unique")
+
+        # finally merges the parameters from the current file
+        if _parameters:
+            parameters = OmegaConf.unsafe_merge(parameters, _parameters, list_merge_mode=ListMergeMode.EXTEND_UNIQUE)
+
+        exports.merge_update(_exports, box_merge_lists="unique")
+        applications.extend(_applications)
+        return parameters, classes, applications, exports
+
+    def load_target(self, target: OmegaConfTarget):
+        full_target_path = os.path.join(self.targets_path, target.path)
+
+        start = time.perf_counter()
+        parameters = OmegaConf.create(keys_to_strings(target.parameters))
+        p, c, a, e = self.load_parameters_from_file(full_target_path, parameters=parameters)
+        load_parameters = time.perf_counter() - start
+        target.parameters = OmegaConf.to_container(p, resolve=True)
+        to_container = time.perf_counter() - load_parameters
+        target.classes = c
+        target.applications = a
+        target.exports = e
+        finish_loading = time.perf_counter() - start
+        logger.debug(f"{target.name}: Config loaded in {load_parameters:.2f}s, resolved in {to_container:.2f}s. Total {finish_loading:.2f}s")
+
+    def migrate(self):
+        migrate(self.inventory_path)
+
+    def resolve_targets(self, targets: list[OmegaConfTarget] = None) -> None:
+        if not targets:
+            targets = self.targets.values()
+        map(lambda target: target.resolve(), targets)