Support custom hooks for resource context (#64)

* add option for custom hook --resource-context-hook

* provide example hook

* no need to pass default args

* cache objects from Kubernetes API (pods for PVC)

* test get_hook_function

Author: hjacobs

README.rst

@@ -111,6 +111,8 @@ Available command line options:
     Optional: filename pointing to a YAML file with a list of rules to apply TTL values to arbitrary Kubernetes objects, e.g. to delete all deployments without a certain label automatically after N days. See Rules File configuration section below.
 ``--deployment-time-annotation``
     Optional: name of the annotation that would be used instead of the creation timestamp of the resource. This option should be used if you want the resources to not be cleaned up if they've been recently redeployed, and your deployment tooling can set this annotation.
+``--resource-context-hook``
+    Optional: string pointing to a Python function to populate the ``_context`` object with additional information, e.g. by calling external services. Built-in example to set ``_context.random_dice`` to a random dice value (1-6): ``--resource-context-hook=kube_janitor.example_hooks.random_dice``.
 
 Example flags:
 

kube_janitor/cmd.py

@@ -1,10 +1,21 @@
 import argparse
+import importlib
 import os
+from typing import Callable
 
 DEFAULT_EXCLUDE_RESOURCES = "events,controllerrevisions"
 DEFAULT_EXCLUDE_NAMESPACES = "kube-system"
 
 
+def get_hook_function(value: str) -> Callable:
+    module_name, attr_path = value.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    function = getattr(module, attr_path)
+    if not callable(function):
+        raise ValueError(f"Not a callable function: {value}")
+    return function
+
+
 def get_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -55,4 +66,9 @@ def get_parser():
         "--deployment-time-annotation",
         help="Annotation that contains a resource's last deployment time, overrides creationTime",
     )
+    parser.add_argument(
+        "--resource-context-hook",
+        type=get_hook_function,
+        help="Optional hook to extend the '_context' object with custom information, e.g. to base decisions on external systems",
+    )
     return parser

kube_janitor/example_hooks.py

@@ -0,0 +1,30 @@
+"""
+Example resource context hooks for Kubernetes Janitor.
+
+Usage: --resource-context-hook=kube_janitor.example_hooks.random_dice
+"""
+import logging
+import random
+from typing import Any
+from typing import Dict
+
+from pykube.objects import APIObject
+
+logger = logging.getLogger(__name__)
+
+CACHE_KEY = "random_dice"
+
+
+def random_dice(resource: APIObject, cache: Dict[str, Any]) -> Dict[str, Any]:
+    """Built-in example resource context hook to set ``_context.random_dice`` to a random dice value (1-6)."""
+
+    # re-use any value from the cache to have only one dice roll per janitor run
+    dice_value = cache.get(CACHE_KEY)
+
+    if dice_value is None:
+        # roll the dice
+        dice_value = random.randint(1, 6)
+        logger.debug(f"The random dice value is {dice_value}!")
+        cache[CACHE_KEY] = dice_value
+
+    return {"random_dice": dice_value}

kube_janitor/janitor.py

@@ -1,11 +1,15 @@
 import datetime
 import logging
 from collections import Counter
+from typing import Any
+from typing import Callable
+from typing import Dict
 from typing import Optional
 
 import pykube
 from pykube import Event
 from pykube import Namespace
+from pykube.objects import APIObject
 
 from .helper import format_duration
 from .helper import parse_expiry
@@ -168,16 +172,18 @@ def handle_resource_on_ttl(
     resource,
     rules,
     delete_notification: int,
-    deployment_time_annotation: Optional[str],
-    dry_run: bool,
+    deployment_time_annotation: Optional[str] = None,
+    resource_context_hook: Optional[Callable[[APIObject, dict], Dict[str, Any]]] = None,
+    cache: Dict[str, Any] = None,
+    dry_run: bool = False,
 ):
     counter = {"resources-processed": 1}
 
     ttl = resource.annotations.get(TTL_ANNOTATION)
     if ttl:
         reason = f"annotation {TTL_ANNOTATION} is set"
     else:
-        context = get_resource_context(resource)
+        context = get_resource_context(resource, resource_context_hook, cache)
         for rule in rules:
             if rule.matches(resource, context):
                 logger.debug(
@@ -272,11 +278,13 @@ def clean_up(
     exclude_namespaces: frozenset,
     rules: list,
     delete_notification: int,
-    deployment_time_annotation: Optional[str],
-    dry_run: bool,
+    deployment_time_annotation: Optional[str] = None,
+    resource_context_hook: Optional[Callable[[APIObject, dict], Dict[str, Any]]] = None,
+    dry_run: bool = False,
 ):
 
     counter: Counter = Counter()
+    cache: Dict[str, Any] = {}
 
     for namespace in Namespace.objects(api):
         if matches_resource_filter(
@@ -292,6 +300,8 @@ def clean_up(
                     rules,
                     delete_notification,
                     deployment_time_annotation,
+                    resource_context_hook,
+                    cache,
                     dry_run,
                 )
             )
@@ -340,6 +350,8 @@ def clean_up(
                 rules,
                 delete_notification,
                 deployment_time_annotation,
+                resource_context_hook,
+                cache,
                 dry_run,
             )
         )

kube_janitor/main.py

@@ -1,6 +1,8 @@
 #!/usr/bin/env python3
 import logging
 import time
+from typing import Callable
+from typing import Optional
 
 from kube_janitor import __version__
 from kube_janitor import cmd
@@ -43,6 +45,7 @@ def main(args=None):
         args.interval,
         args.delete_notification,
         args.deployment_time_annotation,
+        args.resource_context_hook,
         args.dry_run,
     )
 
@@ -56,8 +59,9 @@ def run_loop(
     rules,
     interval,
     delete_notification,
-    deployment_time_annotation,
-    dry_run,
+    deployment_time_annotation: Optional[str],
+    resource_context_hook: Optional[Callable],
+    dry_run: bool,
 ):
     handler = shutdown.GracefulShutdown()
     while True:
@@ -72,6 +76,7 @@ def run_loop(
                 rules=rules,
                 delete_notification=delete_notification,
                 deployment_time_annotation=deployment_time_annotation,
+                resource_context_hook=resource_context_hook,
                 dry_run=dry_run,
             )
         except Exception as e:

kube_janitor/resource_context.py

@@ -1,8 +1,11 @@
 import logging
 import re
 from typing import Any
+from typing import Callable
 from typing import Dict
+from typing import Optional
 
+from pykube import HTTPClient
 from pykube.objects import APIObject
 from pykube.objects import NamespacedAPIObject
 from pykube.objects import Pod
@@ -11,13 +14,28 @@
 logger = logging.getLogger(__name__)
 
 
-def get_persistent_volume_claim_context(pvc: NamespacedAPIObject):
+def get_objects_in_namespace(
+    clazz, api: HTTPClient, namespace: str, cache: Dict[str, Any]
+):
+    """Get (cached) objects from the Kubernetes API."""
+    cache_key = f"{namespace}/{clazz.endpoint}"
+    objects = cache.get(cache_key)
+    if objects is None:
+        objects = list(clazz.objects(api, namespace=namespace))
+        cache[cache_key] = objects
+
+    return objects
+
+
+def get_persistent_volume_claim_context(
+    pvc: NamespacedAPIObject, cache: Dict[str, Any]
+):
     """Get context for PersistentVolumeClaim: whether it's mounted by a Pod and whether it's referenced by a StatefulSet."""
     pvc_is_mounted = False
     pvc_is_referenced = False
 
     # find out whether a Pod mounts the PVC
-    for pod in Pod.objects(pvc.api, namespace=pvc.namespace):
+    for pod in get_objects_in_namespace(Pod, pvc.api, pvc.namespace, cache):
         for volume in pod.obj.get("spec", {}).get("volumes", []):
             if "persistentVolumeClaim" in volume:
                 if volume["persistentVolumeClaim"].get("claimName") == pvc.name:
@@ -28,7 +46,7 @@ def get_persistent_volume_claim_context(pvc: NamespacedAPIObject):
                     break
 
     # find out whether the PVC is still referenced somewhere
-    for sts in StatefulSet.objects(pvc.api, namespace=pvc.namespace):
+    for sts in get_objects_in_namespace(StatefulSet, pvc.api, pvc.namespace, cache):
         # see https://kubernetes.io/docs/concepts/workloads/controllers/statefulset/
         for claim_template in sts.obj.get("spec", {}).get("volumeClaimTemplates", []):
             claim_prefix = claim_template.get("metadata", {}).get("name")
@@ -47,12 +65,27 @@ def get_persistent_volume_claim_context(pvc: NamespacedAPIObject):
     }
 
 
-def get_resource_context(resource: APIObject):
+def get_resource_context(
+    resource: APIObject,
+    hook: Optional[Callable[[APIObject, dict], Dict[str, Any]]] = None,
+    cache: Optional[Dict[str, Any]] = None,
+):
     """Get additional context information for a single resource, e.g. whether a PVC is mounted/used or not."""
 
     context: Dict[str, Any] = {}
 
+    if cache is None:
+        cache = {}
+
     if resource.kind == "PersistentVolumeClaim":
-        context.update(get_persistent_volume_claim_context(resource))
+        context.update(get_persistent_volume_claim_context(resource, cache))
+
+    if hook:
+        try:
+            context.update(hook(resource, cache))
+        except Exception as e:
+            logger.exception(
+                f"Failed populating _context from resource context hook: {e}"
+            )
 
     return context

tests/test_cmd.py

@@ -1,6 +1,13 @@
+import kube_janitor.example_hooks
+from kube_janitor.cmd import get_hook_function
 from kube_janitor.cmd import get_parser
 
 
 def test_parse_args():
     parser = get_parser()
     parser.parse_args(["--dry-run", "--rules-file=/config/rules.yaml"])
+
+
+def test_get_example_hook_function():
+    func = get_hook_function("kube_janitor.example_hooks.random_dice")
+    assert func == kube_janitor.example_hooks.random_dice

tests/test_resource_context.py

@@ -1,7 +1,9 @@
 from unittest.mock import MagicMock
 
+from pykube.objects import Namespace
 from pykube.objects import PersistentVolumeClaim
 
+import kube_janitor.example_hooks
 from kube_janitor.resource_context import get_resource_context
 
 
@@ -83,3 +85,16 @@ def get(**kwargs):
 
     context = get_resource_context(pvc)
     assert not context["pvc_is_not_referenced"]
+
+
+def test_example_hook():
+    namespace = Namespace(None, {"metadata": {"name": "my-ns"}})
+    hook = kube_janitor.example_hooks.random_dice
+    cache = {}
+    context = get_resource_context(namespace, hook, cache)
+    value = context["random_dice"]
+    assert 1 <= value <= 6
+
+    # check that cache is used
+    new_context = get_resource_context(namespace, hook, cache)
+    assert new_context["random_dice"] == value