From c89b050679e1cb16b923112d3bc72de1308eb830 Mon Sep 17 00:00:00 2001
From: Leo Kirchner <leo.kirchner98@gmail.com>
Date: Fri, 23 Dec 2022 14:53:46 +0100
Subject: [PATCH] Allow for filtering of model count metric by different
 fields.

---
 development/nautobot_config.py                |  6 +--
 nautobot_capacity_metrics/__init__.py         |  6 +--
 nautobot_capacity_metrics/metrics.py          | 40 +++++++++++++++++--
 .../tests/test_endpoints.py                   |  2 +-
 4 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/development/nautobot_config.py b/development/nautobot_config.py
index 80cb216..eef44c6 100644
--- a/development/nautobot_config.py
+++ b/development/nautobot_config.py
@@ -260,9 +260,9 @@
             "jobs": True,
             "models": {
                 "dcim": {
-                    "Site": True,
-                    "Rack": True,
-                    "Device": True,
+                    "Location": ["status__name", "parent__name"],
+                    "Rack": ["location__name", "location__parent__name"],
+                    "Device": ["location__parent__name", "device_type__manufacturer__name"],
                     "Interface": True,
                     "Cable": True,
                 },
diff --git a/nautobot_capacity_metrics/__init__.py b/nautobot_capacity_metrics/__init__.py
index 13e09bb..358b6cf 100644
--- a/nautobot_capacity_metrics/__init__.py
+++ b/nautobot_capacity_metrics/__init__.py
@@ -41,9 +41,9 @@ class MetricsExtConfig(NautobotAppConfig):
         "app_metrics": {
             "models": {
                 "dcim": {
-                    "Site": True,
-                    "Rack": True,
-                    "Device": True,
+                    "Location": ["status__name", "parent__name"],
+                    "Rack": ["location__name", "location__parent__name"],
+                    "Device": ["location__parent__name", "device_type__manufacturer__name"],
                 },
                 "ipam": {"IPAddress": True, "Prefix": True},
             },
diff --git a/nautobot_capacity_metrics/metrics.py b/nautobot_capacity_metrics/metrics.py
index a23dac7..c012122 100644
--- a/nautobot_capacity_metrics/metrics.py
+++ b/nautobot_capacity_metrics/metrics.py
@@ -88,21 +88,53 @@ def metric_models(params):
         Iterator[GaugeMetricFamily]
             nautobot_model_count: with model name and application name as labels
     """
-    gauge = GaugeMetricFamily("nautobot_model_count", "Per Nautobot Model count", labels=["app", "name"])
     for app in params:
         app_config = deepcopy(params[app])  # Avoid changing the dictionary we are iterating over
         module = app_config.pop("_module", "nautobot")
-        for model in app_config:
+        for model, additional_fields in app_config.items():
             try:
                 models = importlib.import_module(f"{module}.{app}.models")
                 model_class = getattr(models, model)
-                gauge.add_metric([app, model], model_class.objects.count())
             except ModuleNotFoundError:
                 logger.warning("Unable to find the python library %s.models", app)
+                continue
             except AttributeError:
                 logger.warning("Unable to load the module %s from the python library %s.models", model, app)
+                continue
 
-    yield gauge
+            # Prepare the queryset, prefetching any related fields that are necessary
+            if not isinstance(additional_fields, list):
+                additional_fields = []
+            prefetch_related = [field.split("__", maxsplit=1)[0] for field in additional_fields]
+            queryset = model_class.objects.prefetch_related(*prefetch_related).order_by()
+
+            # Default without any filtering by addition fields
+            gauge = GaugeMetricFamily(
+                f"nautobot_model_count_{model.lower()}_total", "Nautobot model count", labels=["app"]
+            )
+            gauge.add_metric([app], queryset.count())
+            yield gauge
+
+            for additional_field in additional_fields:
+                yield from _individual_model_metric(additional_field, app, model, queryset)
+
+
+def _individual_model_metric(by_field, app, model, queryset):
+    """Yield an individual metric for the metric_models generator."""
+    path = by_field.split("__")
+    related_model_name = path[-2]
+    gauge = GaugeMetricFamily(
+        f"nautobot_model_count_{model.lower()}_by_{related_model_name}_total",
+        f"Nautobot {model.lower()} count per {related_model_name}",
+        labels=["app", related_model_name],
+    )
+    related_fields = queryset.values_list(by_field, flat=True).distinct()
+    for field in related_fields:
+        if field is None:
+            continue
+        filtered_queryset = queryset.filter(**{by_field: field})
+        gauge.add_metric([app, field], filtered_queryset.count())
+        yield gauge
 
 
 def metric_versions():
diff --git a/nautobot_capacity_metrics/tests/test_endpoints.py b/nautobot_capacity_metrics/tests/test_endpoints.py
index 44377da..52b6126 100644
--- a/nautobot_capacity_metrics/tests/test_endpoints.py
+++ b/nautobot_capacity_metrics/tests/test_endpoints.py
@@ -20,5 +20,5 @@ def test_endpoint(self):
     def test_model_count_metrics(self):
         """Ensure that the model count metrics work correctly."""
         resp = self.client.get(self.app_metric_url)
-        if "TestModel" not in resp.content.decode("utf-8"):
+        if "nautobot_model_count_testmodel_total" not in resp.content.decode("utf-8"):
             self.fail("nautobot_capacity_metrics.test_models.models.TestModel does not report its count.")