Allow setting default and individual SU budgets for project members

allocations/migrations/0009_chargebudget.py

@@ -0,0 +1,30 @@
+# Generated by Django 3.2 on 2024-03-11 21:47
+
+from django.conf import settings
+import django.core.validators
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('projects', '0021_project_default_su_budget'),
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('allocations', '0008_allocation_balance_service_version'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='ChargeBudget',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('su_budget', models.IntegerField(default=0, validators=[django.core.validators.MinValueValidator(0)])),
+                ('project', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='projects.project')),
+                ('user', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='projectbudgets', to=settings.AUTH_USER_MODEL)),
+            ],
+            options={
+                'unique_together': {('user', 'project')},
+            },
+        ),
+    ]

allocations/models.py

@@ -1,6 +1,8 @@
-from django.db import models
 import logging
+
 from django.conf import settings
+from django.db import models
+from django.core.validators import MinValueValidator
 from projects.models import Project
 from util.consts import allocation
 
@@ -100,3 +102,18 @@ class Charge(models.Model):
 
     def __str__(self):
         return f"{self.allocation.project}: {self.start_time}-{self.end_time}"
+
+
+class ChargeBudget(models.Model):
+    user = models.ForeignKey(
+        settings.AUTH_USER_MODEL,
+        related_name="projectbudgets",
+        on_delete=models.CASCADE,
+    )
+    project = models.ForeignKey(
+        Project, on_delete=models.CASCADE
+    )
+    su_budget = models.IntegerField(default=0, validators=[MinValueValidator(0)])
+
+    class Meta:
+        unique_together = ('user', 'project',)

balance_service/enforcement/usage_enforcement.py

@@ -21,10 +21,11 @@
 from django.contrib.auth import get_user_model
 from django.utils import timezone
 
-from allocations.models import Charge
+from allocations.models import Charge, ChargeBudget
 from balance_service.enforcement import exceptions
 from balance_service.utils import su_calculators
 from projects.models import Project
+from util.keycloak_client import KeycloakClient
 
 LOG = logging.getLogger(__name__)
 
@@ -132,6 +133,34 @@ def get_balance_service_version(self, data):
         else:
             return 2
 
+    def _check_usage_against_user_budget(self, user, project, new_charge):
+        """Raises error if user charges exceed allocated budget
+
+        Raises:
+            exceptions.BillingError:
+        """
+        try:
+            user_budget = ChargeBudget.objects.get(user=user, project=project)
+        except ChargeBudget.DoesNotExist:
+            # if the default SU budget for project is 0, then there are no limits
+            if project.default_su_budget == 0:
+                return
+            else:
+                user_budget = ChargeBudget(
+                    user=user, project=project, su_budget=project.default_su_budget
+                )
+                user_budget.save()
+        left = user_budget.su_budget - su_calculators.calculate_user_total_su_usage(user, project)
+        if left < new_charge:
+            raise exceptions.BillingError(
+                message=(
+                    "Reservation for user {} would spend {:.2f} SUs, "
+                    "only {:.2f} left in user budget".format(
+                        user.username, new_charge, left
+                    )
+                )
+            )
+
     def check_usage_against_allocation(self, data):
         """Check if we have enough available SUs for this reservation
 
@@ -156,6 +185,14 @@ def check_usage_against_allocation(self, data):
             )
 
         alloc = su_calculators.get_active_allocation(lease_eval.project)
+        keycloak_client = KeycloakClient()
+        role, scopes = keycloak_client.get_user_project_role_scopes(
+            lease_eval.user.username, lease_eval.project.charge_code
+        )
+        if role == 'member':
+            self._check_usage_against_user_budget(
+                lease_eval.user, lease_eval.project, lease_eval.amount
+            )
         approved_alloc = su_calculators.get_consecutive_approved_allocation(
             lease_eval.project, alloc
         )

balance_service/utils/su_calculators.py

@@ -1,6 +1,9 @@
 from django.utils import timezone
+from django.db import models
+from django.db.models import ExpressionWrapper, F, Sum, functions
 
 from projects.models import Project
+from allocations.models import Charge
 
 
 def get_used_sus(charge):
@@ -76,3 +79,37 @@ def project_balances(project_ids) -> "list[dict]":
             }
         )
     return project_balances
+
+
+def calculate_user_total_su_usage(user, project):
+    """
+    Calculate the current SU usage for the args:user in args:project
+    """
+    allocation = get_active_allocation(project)
+    if not allocation:
+        return 0
+
+    charges = Charge.objects.filter(allocation=allocation, user=user)
+
+    # Avoid doing the calculation if there are no charges
+    if not charges.exists():
+        return 0
+    microseconds_per_hour = 1_000_000 * 3600
+    # could've used output field as DurationField and further annotate with
+    # Extract('charge_duration', 'hour') to get the duration in hours but
+    # Extract requires native DurationField database support.
+    charges_with_duration_in_ms = charges.annotate(
+        charge_duration=ExpressionWrapper(
+            F('end_time') - F('start_time'), output_field=models.FloatField()
+        )
+    )
+    # Calculate cost of each charge in SUs by converting ms to hours
+    charges_with_actual_cost = charges_with_duration_in_ms.annotate(
+        charge_cost=F('charge_duration') / microseconds_per_hour * F('hourly_cost')
+    )
+    # calculates the total cost of charges for the user on the project
+    # by summing up the charge_cost values calculated for each charge.
+    # If there are no charges, it returns 0.0
+    return charges_with_actual_cost.aggregate(
+        total_cost=functions.Coalesce(Sum('charge_cost'), 0.0, output_field=models.IntegerField())
+    )['total_cost']

projects/membership.py

@@ -1,5 +1,7 @@
 from datetime import datetime
 from django.contrib.auth import get_user_model
+from allocations.models import ChargeBudget
+from projects.models import Project
 from util.keycloak_client import KeycloakClient, UserAttributes
 
 
@@ -27,6 +29,8 @@ def add_user_to_project(tas_project, user_ref):
     If this is the first time the user has been added to any project, also set an
     attribute with the current date indicating when the user joined their first project.
 
+    Also apply the project's default budget to the user
+
     Args:
         tas_project (pytas.Project): the TAS project representation.
         user_ref (str): the username or email address of the user.
@@ -41,6 +45,13 @@ def add_user_to_project(tas_project, user_ref):
         keycloak_client.update_user(
             user.username, lifecycle_allocation_joined=datetime.now()
         )
+    # add the project's default budget to the user
+    project = Project.objects.get(charge_code=tas_project.chargeCode)
+    if project.default_su_budget != 0:
+        user_budget = ChargeBudget(
+            user=user, project=project, su_budget=project.default_su_budget
+        )
+        user_budget.save()
 
     return True
 

projects/migrations/0021_project_default_su_budget.py

@@ -0,0 +1,19 @@
+# Generated by Django 3.2 on 2024-03-11 21:47
+
+import django.core.validators
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('projects', '0020_duplicate_publications'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='project',
+            name='default_su_budget',
+            field=models.IntegerField(default=0, validators=[django.core.validators.MinValueValidator(0)]),
+        ),
+    ]

projects/models.py

@@ -12,6 +12,7 @@
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from django.contrib.sites.models import Site
+from django.core.validators import MinValueValidator
 
 from projects.user_publication.utils import PublicationUtils
 
@@ -47,6 +48,7 @@ class Project(models.Model):
     title = models.TextField(blank=False)
     nickname = models.CharField(max_length=255, blank=False, unique=True)
     charge_code = models.CharField(max_length=50, blank=False, unique=True)
+    default_su_budget = models.IntegerField(default=0, validators=[MinValueValidator(0)])
 
     def __str__(self) -> str:
         return self.charge_code