diff --git a/src/coldfront_plugin_cloud/management/commands/report_total_quotas.py b/src/coldfront_plugin_cloud/management/commands/report_total_quotas.py new file mode 100644 index 00000000..698f6b5f --- /dev/null +++ b/src/coldfront_plugin_cloud/management/commands/report_total_quotas.py @@ -0,0 +1,159 @@ +"""Report total cloud quota values across active allocations.""" + +import csv +import json +import logging +from django.core.exceptions import ObjectDoesNotExist +from django.core.management.base import BaseCommand + +from coldfront.core.allocation.models import ( + Allocation, + AllocationStatusChoice, +) +from coldfront.core.resource.models import Resource, ResourceType + +from coldfront_plugin_cloud import attributes +from coldfront_plugin_cloud.models.quota_models import QuotaSpecs + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +CLOUD_TYPES = ["OpenStack", "OpenShift", "OpenShift Virtualization", "ESI"] + +TOTAL_KEY = "all" + + +class Command(BaseCommand): + help = "Show total cloud quotas across active allocations" + + def add_arguments(self, parser): + parser.add_argument( + "--cloud-type", + choices=["all"] + CLOUD_TYPES, + default="all", + help="cloud type", + ) + parser.add_argument("--project-id", help="limit scope to project id") + parser.add_argument( + "--per-project", + action="store_true", + help="report a total per project instead of a single grand total", + ) + parser.add_argument( + "--format", choices=["json", "csv"], default="json", help="output format" + ) + + def get_quota_names(self, resource): + """Return the quota attribute names defined on a resource.""" + quota_specs = resource.get_attribute(attributes.RESOURCE_QUOTA_RESOURCES) + if not quota_specs: + logger.warning( + f"Resource {resource.name} has no " + f"'{attributes.RESOURCE_QUOTA_RESOURCES}' attribute, skipping. Run " + "register_default_quotas or add_quota_to_resource to define its quotas." + ) + return [] + return list(QuotaSpecs.model_validate(json.loads(quota_specs)).root.keys()) + + def get_totals(self, cloud_type, project_id=None, per_project=False): + try: + resource_type = ResourceType.objects.get(name=cloud_type) + except ObjectDoesNotExist: + logger.info(f"Skipping {cloud_type} - resource type does not exist") + return [] + + resources = Resource.objects.filter(resource_type=resource_type) + + # Each resource defines its own quotas; merge them into one sorted column set. + quota_names_by_resource = {r.pk: self.get_quota_names(r) for r in resources} + all_quota_names = sorted( + {name for names in quota_names_by_resource.values() for name in names} + ) + if not all_quota_names: + return [] + + filter_kwargs = {} + if project_id: + filter_kwargs["project_id"] = project_id + + # distinct() so an allocation on two resources isn't counted twice. + allocations = Allocation.objects.filter( + resources__in=resources, + status=AllocationStatusChoice.objects.get(name="Active"), + **filter_kwargs, + ).distinct() + + def new_group(project=None): + return { + "cloud_type": cloud_type, + "project_id": project.id if project else TOTAL_KEY, + "project_title": project.title if project else TOTAL_KEY, + "allocations": 0, + "quotas": {name: 0.0 for name in all_quota_names}, + } + + # Emit the total row even when it's zero. + groups = {} if per_project else {TOTAL_KEY: new_group()} + + for allocation in allocations: + # Only the key differs between the two modes. + key = allocation.project_id if per_project else TOTAL_KEY + group = groups.setdefault( + key, new_group(allocation.project if per_project else None) + ) + group["allocations"] += 1 + + resource = allocation.resources.filter(resource_type=resource_type).first() + for name in quota_names_by_resource.get(resource.pk, []): + try: + # Values are text; unset quotas come back as None. + group["quotas"][name] += float(allocation.get_attribute(name)) + except (TypeError, ValueError): + logger.debug( + f"No usable value for '{name}' on allocation {allocation.id}" + ) + + # Per-project keys are ints, the total key is a string; can't sort together. + if per_project: + return [groups[key] for key in sorted(groups)] + return list(groups.values()) + + def render_json(self, rows): + self.stdout.write(json.dumps(rows, indent=4)) + + def render_csv(self, rows): + quota_names = sorted({name for row in rows for name in row["quotas"]}) + writer = csv.writer(self.stdout) + writer.writerow( + ["cloud_type", "project_id", "project_title", "allocations"] + + [name.replace(" ", "_") for name in quota_names] + ) + for row in rows: + writer.writerow( + [ + row["cloud_type"], + row["project_id"], + row["project_title"], + row["allocations"], + ] + + [row["quotas"].get(name, 0.0) for name in quota_names] + ) + + def handle(self, *args, **options): + cloud_type = options["cloud_type"] + cloud_types = CLOUD_TYPES if cloud_type == "all" else [cloud_type] + + rows = [] + for ct in cloud_types: + rows.extend( + self.get_totals( + ct, + project_id=options.get("project_id"), + per_project=options["per_project"], + ) + ) + + if options["format"] == "json": + self.render_json(rows) + else: + self.render_csv(rows) diff --git a/src/coldfront_plugin_cloud/tests/unit/test_report_total_quotas.py b/src/coldfront_plugin_cloud/tests/unit/test_report_total_quotas.py new file mode 100644 index 00000000..eb381850 --- /dev/null +++ b/src/coldfront_plugin_cloud/tests/unit/test_report_total_quotas.py @@ -0,0 +1,48 @@ +import json +from io import StringIO + +from django.core.management import call_command + +from coldfront_plugin_cloud import attributes, utils +from coldfront_plugin_cloud.tests import base + + +class TestReportTotalQuotas(base.TestBase): + @classmethod + def setUpTestData(cls) -> None: + super().setUpTestData() + cls.resource = cls.new_openstack_resource(name="test-quota-totals") + call_command("register_default_quotas", apply=True) + + def report(self, **kwargs): + out = StringIO() + call_command("report_total_quotas", stdout=out, **kwargs) + return json.loads(out.getvalue()) + + def test_totals_across_allocations(self): + for vcpus in (2, 5): + project = self.new_project() + allocation = self.new_allocation(project, self.resource, 1) + utils.set_attribute_on_allocation(allocation, attributes.QUOTA_VCPU, vcpus) + + (row,) = [r for r in self.report(cloud_type="OpenStack")] + self.assertEqual(row["allocations"], 2) + self.assertEqual(row["quotas"][attributes.QUOTA_VCPU], 7.0) + # Quotas with no value set still report as zero. + self.assertEqual(row["quotas"][attributes.QUOTA_FLOATING_IPS], 0.0) + + def test_per_project_and_project_id_filter(self): + project_a = self.new_project() + project_b = self.new_project() + for project, vcpus in ((project_a, 3), (project_b, 4)): + allocation = self.new_allocation(project, self.resource, 1) + utils.set_attribute_on_allocation(allocation, attributes.QUOTA_VCPU, vcpus) + + rows = self.report(cloud_type="OpenStack", per_project=True) + totals = {r["project_id"]: r["quotas"][attributes.QUOTA_VCPU] for r in rows} + self.assertEqual(totals[project_a.id], 3.0) + self.assertEqual(totals[project_b.id], 4.0) + + (scoped,) = self.report(cloud_type="OpenStack", project_id=project_a.id) + self.assertEqual(scoped["quotas"][attributes.QUOTA_VCPU], 3.0) + self.assertEqual(scoped["allocations"], 1)