OpenConceptLab/ocl_online#517 | capacity tier per group from the groups config, output language permission

core/capabilities/management/commands/seed_group_capabilities.py

@@ -5,12 +5,17 @@
 from django.core.management import BaseCommand, CommandError
 from django.db import transaction
 
+from core.capacity.constants import TIERS
+
 GROUPS_CONFIG_ENV = 'OCL_GROUPS_CONFIG'
 
-# Plan groups are deployment config: {"groups": {"<name>": {"permissions": [...], "capabilities": {...}}}}
+# Groups are deployment config:
+# {"groups": {"<name>": {"permissions": [...], "capabilities": {...}, "capacity_tier": "<tier>"}}}
 # Sync - the config is the source of truth for every group it declares: users.* permissions are set exactly,
-# declared limits are upserted and undeclared ones deleted. Groups it doesn't declare are left alone.
+# declared limits are upserted and undeclared ones deleted, and so is the capacity tier. Groups it doesn't declare
+# are left alone.
 # Limits: 0 = unlimited, -1 = blocked (the kill switch), N > 0 = the limit; no row = not entitled (blocked).
+# capacity_tier: one of core.capacity.constants.TIERS; a user gets the highest of their groups' (core.capacity.limiter).
 
 
 def seed_groups(config, stdout=None):
@@ -47,6 +52,17 @@ def seed_group(group_name, group_config, warn):
         GroupCapability.objects.update_or_create(group=group, capability=capability, defaults={'limit': limit})
         declared_ids.append(capability.id)
     GroupCapability.objects.filter(group=group).exclude(capability_id__in=declared_ids).delete()
+    seed_capacity_tier(group, group_config.get('capacity_tier'))
+
+
+def seed_capacity_tier(group, tier):
+    from core.capacity.models import GroupCapacityTier
+
+    # sync: the declared capacity tier is upserted, or the group's is deleted
+    if tier:
+        GroupCapacityTier.objects.update_or_create(group=group, defaults={'tier': tier})
+    else:
+        GroupCapacityTier.objects.filter(group=group).delete()
 
 
 def validate_config(config):
@@ -66,6 +82,9 @@ def validate_config(config):
                 not all(isinstance(limit, int) and not isinstance(limit, bool) for limit in capabilities.values())
         ):
             raise CommandError(f'{group_name}: capabilities must map names to integer limits')
+        tier = group_config.get('capacity_tier')
+        if tier is not None and tier not in TIERS:
+            raise CommandError(f'{group_name}: capacity_tier must be one of {", ".join(TIERS)}')
     return config['groups']
 
 

core/capabilities/tests/tests.py

@@ -410,6 +410,23 @@ def test_creates_group_with_permissions_and_limits(self):
         self.assertEqual(self.codenames('pro'), {'mapper_use', 'mapper_custom_algorithms', 'mapper_scispacy'})
         self.assertEqual(self.limits('pro'), {'mapper.projects': 20, 'mapper.match_operations': 0})
 
+    def test_seeds_and_syncs_the_capacity_tier(self):
+        from core.capacity.models import GroupCapacityTier
+        config = {'groups': {'group_a': {'capacity_tier': 'early_access'}, 'group_b': {}}}
+        self.seed(config)
+        self.assertEqual(GroupCapacityTier.objects.get(group__name='group_a').tier, 'early_access')
+        self.assertFalse(GroupCapacityTier.objects.filter(group__name='group_b').exists())
+
+        self.seed({'groups': {'group_a': {'capacity_tier': 'core'}}})
+        self.assertEqual(GroupCapacityTier.objects.get(group__name='group_a').tier, 'core')
+
+        self.seed({'groups': {'group_a': {'permissions': ['mapper_use']}}})
+        self.assertFalse(GroupCapacityTier.objects.filter(group__name='group_a').exists())
+
+    def test_rejects_an_unknown_capacity_tier(self):
+        with self.assertRaisesMessage(CommandError, 'group_a: capacity_tier must be one of'):
+            self.seed({'groups': {'group_a': {'capacity_tier': 'gold'}}})
+
     def test_reseed_syncs_permissions_and_limits(self):
         self.seed(self.CONFIG)
         self.seed({

core/capacity/constants.py

@@ -9,12 +9,14 @@
 ENFORCE_FOR = (ENFORCE_FOR_AWARE, ENFORCE_FOR_ALL)
 CAPACITY_AWARE_METADATA_KEY = 'capacity_aware'
 
-# A user's tier is their highest plan group, highest first.
+# Capacity tiers, highest first. Staff and superusers get the staff tier; everyone else the highest tier their groups
+# are set to (`capacity_tier` in the groups config, core.capacity.models.GroupCapacityTier), else the lowest.
 TIER_STAFF = 'staff'
 TIER_CORE = 'core'
 TIER_EARLY_ACCESS = 'early_access'
 TIER_PREVIEW = 'preview'
 TIERS = (TIER_STAFF, TIER_CORE, TIER_EARLY_ACCESS, TIER_PREVIEW)
+DEFAULT_TIER = TIERS[-1]
 
 # Lanes: each counts the heavy calls in flight in one scope.
 LANE_API_HEAVY = 'api_heavy'  # every heavy call, cluster-wide

core/capacity/limiter.py

@@ -34,15 +34,14 @@
 
 from core.capacity.config import get_config
 from core.capacity.constants import (
-    TIERS, MODE_OFF, MODE_ENFORCE, ENFORCE_FOR_ALL, CAPACITY_AWARE_METADATA_KEY, TIER_STAFF, TIER_CORE,
-    TIER_EARLY_ACCESS, TIER_PREVIEW, LANE_API_HEAVY, LANE_API_HEAVY_TASK, LANE_ES2_KNN, LANE_TIER, LANE_USER,
+    TIERS, DEFAULT_TIER, MODE_OFF, MODE_ENFORCE, ENFORCE_FOR_ALL, CAPACITY_AWARE_METADATA_KEY, TIER_STAFF,
+    LANE_API_HEAVY, LANE_API_HEAVY_TASK, LANE_ES2_KNN, LANE_TIER, LANE_USER,
     DECISION_ADMITTED, DECISION_SHADOW_REFUSED, DECISION_REFUSED, DECISION_UNAVAILABLE, ENDPOINT_MATCH,
     HEADER_DECISION, HEADER_LIMIT, HEADER_IN_FLIGHT, HEADER_TIER, HEADER_TIER_LIMIT, HEADER_TIER_IN_FLIGHT,
     HEADER_SUGGESTED_CONCURRENCY, CAPACITY_EXCEEDED_ERROR_CODE, LOG_EVENT, REDIS_KEY_PREFIX)
 from core.capacity.logs import emit
 from core.capacity.threads import RedisThread
 from core.common.utils import get_event_metadata
-from core.users.constants import CORE_USER_GROUP, EARLY_ACCESS_GROUP
 
 _NOW_MS = """
 if redis.replicate_commands then redis.replicate_commands() end
@@ -248,15 +247,15 @@ def get_not_after_ms():
 
 
 def get_tier(user):
-    """The user's highest plan tier, as capabilities resolve it: staff > core > early_access > preview."""
+    """
+    The user's capacity tier: staff for staff and superusers, else the highest tier their groups are set to
+    (`capacity_tier` in the groups config, OpenConceptLab/ocl_online#517), else the lowest.
+    """
+    from core.capacity.models import GroupCapacityTier
     if user.is_staff or user.is_superuser:
         return TIER_STAFF
-    groups = set(user.groups.values_list('name', flat=True))
-    if CORE_USER_GROUP in groups:
-        return TIER_CORE
-    if EARLY_ACCESS_GROUP in groups:
-        return TIER_EARLY_ACCESS
-    return TIER_PREVIEW
+    tiers = set(GroupCapacityTier.objects.filter(group__in=user.groups.all()).values_list('tier', flat=True))
+    return next((tier for tier in TIERS if tier in tiers), DEFAULT_TIER)
 
 
 def get_queue_ms(request, now=None):

core/capacity/migrations/0002_groupcapacitytier.py

@@ -0,0 +1,24 @@
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('auth', '0012_alter_user_first_name_max_length'),
+        ('capacity', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='GroupCapacityTier',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('tier', models.CharField(max_length=32)),
+                ('group', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, related_name='capacity_tier', to='auth.group')),
+            ],
+            options={
+                'db_table': 'group_capacity_tiers',
+            },
+        ),
+    ]

core/capacity/models.py

@@ -20,3 +20,16 @@ class Meta:
     @classmethod
     def get_latest(cls):
         return cls.objects.order_by('-id').first()
+
+
+class GroupCapacityTier(models.Model):
+    """
+    The capacity tier a group's members get (OpenConceptLab/ocl_online#517), from `capacity_tier` in the groups
+    config, which `seed_group_capabilities` syncs. A user gets the highest tier of their groups; core.capacity.limiter
+    never reads group names.
+    """
+    class Meta:
+        db_table = 'group_capacity_tiers'
+
+    group = models.OneToOneField('auth.Group', on_delete=models.CASCADE, related_name='capacity_tier')
+    tier = models.CharField(max_length=32)