diff --git a/core/capabilities/apps.py b/core/capabilities/apps.py new file mode 100644 index 000000000..fb05b224b --- /dev/null +++ b/core/capabilities/apps.py @@ -0,0 +1,9 @@ +from django.apps import AppConfig + + +class CapabilitiesConfig(AppConfig): + name = 'core.capabilities' + verbose_name = "Capabilities" + + def ready(self): + from core.capabilities import checks # pylint: disable=unused-variable, unused-import diff --git a/core/capabilities/checks.py b/core/capabilities/checks.py new file mode 100644 index 000000000..522c9f93f --- /dev/null +++ b/core/capabilities/checks.py @@ -0,0 +1,60 @@ +import inspect + +from django.conf import settings +from django.core.checks import Error, register +from django.utils.module_loading import import_string + +from core.capabilities.policy import CapabilityPolicy + + +def takes(function, *args, **kwargs): + try: + inspect.signature(function).bind(*args, **kwargs) + except TypeError: + return False + return True + + +def get_arguments(interface): + """What a call of this CapabilityPolicy method passes: required arguments by position, the others by name.""" + args, kwargs = [], {} + for name, param in list(inspect.signature(interface).parameters.items())[1:]: # after `self` + if param.kind == param.VAR_KEYWORD: + kwargs['any_other_keyword'] = None + elif param.default is param.empty: + args.append(None) + else: + kwargs[name] = None + return args, kwargs + + +def get_misfits(policy): + """What can't take the arguments it's called with. Reads signatures only: a policy's __init__ doesn't run here.""" + bare = object.__new__(policy) + misfits = [] if takes(policy) else ['__init__'] + for name, interface in inspect.getmembers(CapabilityPolicy, inspect.isfunction): + args, kwargs = get_arguments(interface) + if not takes(getattr(bare, name), *args, **kwargs): + misfits.append(name) + return misfits + + +@register() +def check_capability_policy(**_): + """At start-up (`migrate` runs the checks): a policy that can't be used would fail every use of a capability.""" + path = settings.CAPABILITY_POLICY + try: + policy = import_string(path) + except Exception as ex: # a module that imports but fails (a RuntimeError from a model, say) counts too + return [Error(f'CAPABILITY_POLICY {path!r} does not import: {ex}', id='capabilities.E001')] + if not (isinstance(policy, type) and issubclass(policy, CapabilityPolicy)): + return [Error(f'CAPABILITY_POLICY {path!r} is not a subclass of CapabilityPolicy', id='capabilities.E002')] + if inspect.isabstract(policy): + missing = ', '.join(sorted(policy.__abstractmethods__)) + return [Error(f'CAPABILITY_POLICY {path!r} does not implement {missing}', id='capabilities.E003')] + misfits = get_misfits(policy) + if misfits: + return [Error( + f'CAPABILITY_POLICY {path!r} can\'t take the arguments of {", ".join(misfits)}', id='capabilities.E004' + )] + return [] diff --git a/core/capabilities/exceptions.py b/core/capabilities/exceptions.py index a9ee30138..ed66cb8a9 100644 --- a/core/capabilities/exceptions.py +++ b/core/capabilities/exceptions.py @@ -6,29 +6,35 @@ ) -class CapabilityExceeded(APIException): +class UsageRefused(APIException): """ - An APIException (not a plain Exception) so that a call site which forgets to - catch this - there is more than one check_and_consume_capability() caller - - fails as a 403 under DRF's default exception handler instead of an uncaught 500. - Existing callers still catch it explicitly to build their own structured - {detail, error_code, limit, used} response; this only changes the failure mode - for whichever call site doesn't. + A capability policy refused a use. `detail` is the refusal body, sent as the policy built it: a call site adds + only its own sentence as `detail`, when the policy gave none. + An APIException, so a call site that doesn't catch it fails as a 403 rather than a 500. """ status_code = status.HTTP_403_FORBIDDEN default_code = 'capability_limit_reached' - def __init__(self, capability_name, limit, used, requested): # pylint: disable=super-init-not-called - self.capability_name = capability_name - self.limit = limit - self.used = used - self.requested = requested - self.detail = f'{capability_name}: {used}+{requested} > {limit}' + def __init__(self, detail, no_access=False): # pylint: disable=super-init-not-called + # set directly (as MapProjectCapacityExceeded does), so numbers stay numbers in the JSON + self.detail = detail + self.no_access = no_access - @property - def not_entitled(self): + +class CapabilityExceeded(UsageRefused): + """The refusal by limit: no limit or a blocked one, or not enough left.""" + + def __init__(self, capability_name, limit, used): # None = no row anywhere, -1 = blocked: both are "not available", never a count - return self.limit is None or self.limit < 0 + no_access = limit is None or limit < 0 + super().__init__( + { + 'error_code': CAPABILITY_NOT_ENTITLED_ERROR_CODE.get(capability_name, 'capability_not_entitled') + if no_access else CAPABILITY_EXCEEDED_ERROR_CODE.get(capability_name, 'capability_limit_reached'), + 'limit': None if no_access else limit, 'used': used, + }, + no_access=no_access + ) class MapProjectCapacityExceeded(APIException): diff --git a/core/capabilities/migrations/0003_usage_window_start_and_reversed_units.py b/core/capabilities/migrations/0003_usage_window_start_and_reversed_units.py new file mode 100644 index 000000000..d0a86c931 --- /dev/null +++ b/core/capabilities/migrations/0003_usage_window_start_and_reversed_units.py @@ -0,0 +1,52 @@ +# Generated by Django 5.2.17 on 2026-10-24 14:03 + +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('capabilities', '0002_usage_event_idempotency_key'), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.AlterUniqueTogether( + name='usagecounter', + unique_together=set(), + ), + migrations.AddField( + model_name='usagecounter', + name='window_start', + field=models.DateTimeField(blank=True, null=True), + ), + migrations.AddField( + model_name='usageevent', + name='reversed_units', + field=models.BigIntegerField(db_default=0, default=0), + ), + migrations.AddField( + model_name='usageevent', + name='window_start', + field=models.DateTimeField(blank=True, null=True), + ), + migrations.AlterField( + model_name='usagecounter', + name='used', + field=models.BigIntegerField(default=0), + ), + migrations.AlterField( + model_name='usageevent', + name='units', + field=models.BigIntegerField(default=1), + ), + migrations.AddConstraint( + model_name='usagecounter', + constraint=models.UniqueConstraint(fields=('user', 'capability', 'window_start'), name='usage_counters_unique_window'), + ), + migrations.AddConstraint( + model_name='usagecounter', + constraint=models.UniqueConstraint(condition=models.Q(('window_start__isnull', True)), fields=('user', 'capability'), name='usage_counters_unique_unwindowed'), + ), + ] diff --git a/core/capabilities/models.py b/core/capabilities/models.py index a18ff1f5d..c5cf609d2 100644 --- a/core/capabilities/models.py +++ b/core/capabilities/models.py @@ -1,4 +1,5 @@ from django.db import models, transaction +from django.db.models import F class Capability(models.Model): @@ -42,34 +43,62 @@ class Meta: class UsageCounter(models.Model): - """Fast, concurrency-safe running total for one (user, capability) pair.""" + """ + Fast, concurrency-safe running total for one (user, capability) pair in one usage window. + `window_start` None is the window that never ends. + """ class Meta: db_table = 'usage_counters' - unique_together = ('user', 'capability') + constraints = [ + models.UniqueConstraint( + fields=['user', 'capability', 'window_start'], name='usage_counters_unique_window' + ), + # PostgreSQL 14 treats nulls as distinct in the constraint above + models.UniqueConstraint( + fields=['user', 'capability'], condition=models.Q(window_start__isnull=True), + name='usage_counters_unique_unwindowed' + ), + ] user = models.ForeignKey('users.UserProfile', on_delete=models.CASCADE, related_name='usage_counters') capability = models.ForeignKey(Capability, on_delete=models.CASCADE, related_name='usage_counters') - used = models.IntegerField(default=0) + used = models.BigIntegerField(default=0) + window_start = models.DateTimeField(null=True, blank=True) @classmethod - def refund(cls, user, capability_id, units=1): + def get_used(cls, user, capability_id, window_start=None): + return cls.objects.filter( + user=user, capability_id=capability_id, window_start=window_start + ).values_list('used', flat=True).first() or 0 + + @classmethod + def get_locked(cls, user, capability_id, window_start=None): """ - Manual ledger correction (e.g. staff fixing an over-count after a bug). - Not part of the normal consumption flow: callers should consume a - capability only after the action it gates has actually succeeded (see - UserProfile.check_and_consume_capability), not reserve-then-refund - - that would need a failure path reachable by the same caller that - consumed, and CapabilityRefundView is staff-only precisely so a - cross-service caller (e.g. ocl-ai-assistant) never needs to reach it. - Floors at 0; never goes negative. `capability_id` must be an id of a - capability already seeded (Keycloak/fixtures) - a bad id fails with an + The window's counter, created at 0 on first use, locked until the caller's transaction ends. + `capability_id` must be an id of a capability already seeded (Keycloak/fixtures) - a bad id fails with an IntegrityError rather than silently creating a Capability row. """ + cls.objects.get_or_create(user=user, capability_id=capability_id, window_start=window_start) + return cls.objects.select_for_update().get(user=user, capability_id=capability_id, window_start=window_start) + + def consume(self, units, idempotency_key=None, **attribution): + """Adds `units` and logs the event. The caller holds this counter's lock.""" + UsageCounter.objects.filter(pk=self.pk).update(used=F('used') + units) + return UsageEvent.objects.create( + user_id=self.user_id, capability_id=self.capability_id, window_start=self.window_start, units=units, + idempotency_key=idempotency_key or None, **attribution + ) + + @classmethod + def refund(cls, user, capability_id, units=1, window_start=None): + """ + Manual correction (e.g. staff fixing an over-count after a bug) where there's no event to reverse. + CapabilityRefundView is staff-only, so a cross-service caller (e.g. ocl-ai-assistant) never reaches it. + Floors at 0; never goes negative. + """ with transaction.atomic(): - cls.objects.get_or_create(user=user, capability_id=capability_id) - counter = cls.objects.select_for_update().get(user=user, capability_id=capability_id) - new_used = max(counter.used - units, 0) - cls.objects.filter(pk=counter.pk).update(used=new_used) + counter = cls.get_locked(user, capability_id, window_start) + cls.objects.filter(pk=counter.pk).update(used=max(counter.used - units, 0)) class UsageEvent(models.Model): @@ -86,7 +115,12 @@ class Meta: user = models.ForeignKey('users.UserProfile', on_delete=models.CASCADE, related_name='usage_events') capability = models.ForeignKey(Capability, on_delete=models.CASCADE, related_name='usage_events') - units = models.IntegerField(default=1) + units = models.BigIntegerField(default=1) + # taken back after consumption (e.g. the search they were consumed for failed); `units` stays as consumed + # db_default: code that predates this column, still serving during a deploy, inserts events without it + reversed_units = models.BigIntegerField(default=0, db_default=0) + # the counter window the units counted in + window_start = models.DateTimeField(null=True, blank=True) action = models.CharField(max_length=100, blank=True, default='') algorithm = models.CharField(max_length=100, null=True, blank=True) map_project = models.ForeignKey( @@ -100,3 +134,24 @@ class Meta: # client-sent dedup key (e.g. per run/row): a repeat consume with the same key charges nothing idempotency_key = models.CharField(max_length=255, null=True, blank=True) created_at = models.DateTimeField(auto_now_add=True) + + @classmethod + def find(cls, user, capability_id, idempotency_key): + if not idempotency_key: + return None + return cls.objects.filter(user=user, capability_id=capability_id, idempotency_key=idempotency_key).first() + + def reverse(self, units): + """ + Takes back up to `units` of this event's units, from the counter of the window they counted in, even after + that window has ended. Returns the units taken back. + """ + with transaction.atomic(): + counter = UsageCounter.get_locked(self.user, self.capability_id, self.window_start) + reversed_units = UsageEvent.objects.filter(pk=self.pk).values_list('reversed_units', flat=True).get() + units = max(min(units, self.units - reversed_units), 0) + if units: + UsageCounter.objects.filter(pk=counter.pk).update(used=max(counter.used - units, 0)) + UsageEvent.objects.filter(pk=self.pk).update(reversed_units=F('reversed_units') + units) + self.reversed_units = reversed_units + units + return units diff --git a/core/capabilities/policy.py b/core/capabilities/policy.py new file mode 100644 index 000000000..7a4781e0c --- /dev/null +++ b/core/capabilities/policy.py @@ -0,0 +1,166 @@ +import logging +from abc import ABC, abstractmethod + +from django.conf import settings +from django.contrib.auth import get_user_model +from django.db import transaction +from django.utils.module_loading import import_string + +from core.capabilities.constants import ( + AUTHORING_CAPABILITY_DEFAULT_LIMITS, AUTHORING_CAPABILITY_IDS, CAPABILITY_NAME_BY_ID, + MAPPER_PROJECTS_CAPABILITY_ID, MAPPER_ROWS_PER_PROJECT_CAPABILITY_ID, +) +from core.capabilities.exceptions import CapabilityExceeded +from core.capabilities.models import GroupCapability, UsageCounter, UsageEvent +from core.users.constants import PREVIEW_GROUP + +logger = logging.getLogger('oclapi') + + +def get_capability_policy(): + return import_string(settings.CAPABILITY_POLICY)() + + +class CapabilityPolicy(ABC): + """ + Decides and records the use of capabilities. settings.CAPABILITY_POLICY names the class in use, + get_capability_policy() returns one, and core.capabilities.checks checks it at start-up. + + `user` is who is acting. `owner`, optional, is the user or organization whose content the use is for. A call site + takes it from that content, never from a request field. A request can name content that isn't the user's, so a + policy that decides by `owner` first checks that `user` may act for it. Windows, counters, events and reversals + stay per user. + """ + + @abstractmethod + def get_limit(self, user, capability_id, owner=None): + """None = no access, -1 = blocked, 0 = unlimited, N > 0 = the limit.""" + raise NotImplementedError + + @abstractmethod + def get_usage(self, user, capability_id, owner=None): + """`used` (None where nothing accumulates), plus any other fields the policy reports.""" + raise NotImplementedError + + @abstractmethod + def consume( # pylint: disable=too-many-arguments + self, user, capability_id, units=1, *, idempotency_key=None, owner=None, **attribution + ): + """ + Checks and records `units` in one transaction. Returns (usage_event, already_consumed). Raises UsageRefused, + whose `detail` is the refusal body. + """ + raise NotImplementedError + + @abstractmethod + def reverse(self, usage_event, units): + """Takes back up to `units` of an event's use. Returns the units taken back.""" + raise NotImplementedError + + @abstractmethod + def refund(self, user, capability_id, units=1): + """A staff correction, where there's no event to reverse: takes up to `units` off what `user` has used now.""" + raise NotImplementedError + + def describe(self, user, capability_id, owner=None): + """What usage reports send: `limit`, then every usage field.""" + return { + 'limit': self.get_limit(user, capability_id, owner=owner), + **self.get_usage(user, capability_id, owner=owner) + } + + +class GroupLimitPolicy(CapabilityPolicy): + """Limits from a per-user override, else from the user's groups. Ignores `owner`: all of it is the user's.""" + + def get_window_start(self, user, capability_id): # pylint: disable=unused-argument + """When the usage window counting now began. None is one window that never ends.""" + return None + + def get_limit(self, user, capability_id, owner=None): + """ + There is no implicit-unlimited fallback: None means no override and no group row for this user, and + consume() treats it as blocked. A new capability, or a new group, needs its limit configured (even to 0). + Staff and superusers are the deliberate exception. + + A per-user override wins over every group; else the highest of the groups' limits, where a 0 from any one + group wins outright. + + For the authoring capabilities (bulk import, $clone), which every account has always had, no row means the + `preview` group's value, not blocked (ocl_online#230). + """ + if user.is_superuser or user.is_staff: + return 0 + override = self._get_capability_limit(user, capability_id) + if override is not None: + return override + group_limits = self._get_group_capability_limit(user, capability_id) + if not group_limits: + return self._get_default_authoring_limit(capability_id) + if 0 in group_limits: + return 0 + return max(group_limits) + + @staticmethod + def _get_group_capability_limit(user, capability_id): + return list(GroupCapability.objects.filter( + group__in=user.groups.all(), capability_id=capability_id + ).values_list('limit', flat=True)) + + @staticmethod + def _get_capability_limit(user, capability_id): + return user.capability_overrides.filter(capability_id=capability_id).values_list('limit', flat=True).first() + + @staticmethod + def _get_default_authoring_limit(capability_id): + """None (blocked) unless `capability_id` is an authoring capability, which falls back to `preview`.""" + if capability_id not in AUTHORING_CAPABILITY_IDS: + return None + limit = GroupCapability.objects.filter( + group__name=PREVIEW_GROUP, capability_id=capability_id).values_list('limit', flat=True).first() + if limit is None: + logger.warning( + 'The %s group has no limit for capability %s; using the built-in default', PREVIEW_GROUP, capability_id) + return AUTHORING_CAPABILITY_DEFAULT_LIMITS[capability_id] + return limit + + def get_usage(self, user, capability_id, owner=None): + if capability_id in AUTHORING_CAPABILITY_IDS or capability_id == MAPPER_ROWS_PER_PROJECT_CAPABILITY_ID: + # per-request and per-project limits: nothing accumulates per user + return {'used': None} + if capability_id == MAPPER_PROJECTS_CAPABILITY_ID: + # enforced against the live count (HasMapProjectCapacity), so deleting a project frees its slot + return {'used': user.map_projects_used} + return {'used': UsageCounter.get_used(user, capability_id, self.get_window_start(user, capability_id))} + + def consume( # pylint: disable=too-many-arguments + self, user, capability_id, units=1, *, idempotency_key=None, owner=None, **attribution + ): + """ + 0 is unlimited and always succeeds; -1 is the kill switch (a per-user override blocks just that user). + With an `idempotency_key` already consumed, consumes nothing and returns that event, even past the limit. + Counts on `user`'s counter, even when a subclass's get_limit() answers by `owner`. + """ + limit = self.get_limit(user, capability_id, owner=owner) + + with transaction.atomic(): + # the account first, so a user's requests serialize even when they count in different windows. No key + # update, so a foreign-key check on the user (another transaction inserting a counter) isn't blocked. + get_user_model().objects.select_for_update(no_key=True).only('id').get(pk=user.pk) + counter = UsageCounter.get_locked(user, capability_id, self.get_window_start(user, capability_id)) + + # under the lock and before the limit check, so a retry after the last unit still succeeds + event = UsageEvent.find(user, capability_id, idempotency_key) + if event: + return event, True + + if limit is None or (limit and counter.used + units > limit): + raise CapabilityExceeded(CAPABILITY_NAME_BY_ID.get(capability_id, capability_id), limit, counter.used) + + return counter.consume(units, idempotency_key=idempotency_key, **attribution), False + + def reverse(self, usage_event, units): + return usage_event.reverse(units) + + def refund(self, user, capability_id, units=1): + UsageCounter.refund(user, capability_id, units=units, window_start=self.get_window_start(user, capability_id)) diff --git a/core/capabilities/serializers.py b/core/capabilities/serializers.py index 89d0de13c..a92cb3ac9 100644 --- a/core/capabilities/serializers.py +++ b/core/capabilities/serializers.py @@ -1,16 +1,22 @@ +from django.utils.functional import cached_property from rest_framework import serializers +from core.capabilities.policy import get_capability_policy + class CapabilitySerializer(serializers.Serializer): # pylint: disable=abstract-method name = serializers.CharField() - limit = serializers.SerializerMethodField() - used = serializers.SerializerMethodField() + limit = serializers.IntegerField(read_only=True, allow_null=True) + used = serializers.IntegerField(read_only=True, allow_null=True) - def get_limit(self, obj): - return self.context['user'].get_capability_limit(obj.id) + @cached_property + def policy(self): + # one for a whole list (many=True serializes every item with one of these), so it's one snapshot + return get_capability_policy() - def get_used(self, obj): - return self.context['user'].get_capability_usage(obj.id) + def to_representation(self, instance): + # every usage field the policy reports, not only `used` + return {'name': instance.name, **self.policy.describe(self.context['user'], instance.id)} class UserCapabilityOverrideSerializer(serializers.Serializer): # pylint: disable=abstract-method diff --git a/core/capabilities/tests/tests.py b/core/capabilities/tests/tests.py index 43e25b497..0a91ac798 100644 --- a/core/capabilities/tests/tests.py +++ b/core/capabilities/tests/tests.py @@ -1,18 +1,26 @@ import json +from datetime import datetime, timezone as dt_timezone from io import StringIO from unittest.mock import patch from django.contrib.auth.models import Group, Permission +from django.core.checks.registry import registry from django.core.management import call_command, CommandError +from django.db import connection +from django.test import SimpleTestCase, override_settings +from django.test.utils import CaptureQueriesContext +from core.capabilities.checks import check_capability_policy from core.capabilities.constants import ( AI_ASSISTANT_CALLS_CAPABILITY_ID, AI_ASSISTANT_CHANGE_COMMENTS_CAPABILITY_ID, CLONE_RESOURCES_PER_CALL_CAPABILITY_ID, IMPORTS_FILE_SIZE_CAPABILITY_ID, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, MAPPER_PROJECTS_CAPABILITY_ID, ) -from core.capabilities.exceptions import CapabilityExceeded -from core.capabilities.models import Capability, GroupCapability, UsageEvent, UserCapabilityOverride +from core.capabilities.exceptions import CapabilityExceeded, UsageRefused +from core.capabilities.models import Capability, GroupCapability, UsageCounter, UsageEvent, UserCapabilityOverride +from core.capabilities.policy import CapabilityPolicy, GroupLimitPolicy, get_capability_policy from core.common.tests import OCLAPITestCase, PREVIEW_GROUP_NAME +from core.orgs.tests.factories import OrganizationFactory from core.users.constants import PREVIEW_GROUP, PREVIEW_GRANDFATHERED_GROUP from core.users.tests.factories import UserProfileFactory @@ -111,7 +119,7 @@ def test_put_minus_one_blocks_the_user(self): self.assertEqual(self.target.get_capability_limit(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID), -1) with self.assertRaises(CapabilityExceeded) as ctx: self.target.check_and_consume_capability(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID) - self.assertTrue(ctx.exception.not_entitled) + self.assertTrue(ctx.exception.no_access) def test_list_reflects_created_overrides(self): self.client.put( @@ -534,3 +542,588 @@ def test_per_request_limits_report_no_usage(self): user = UserProfileFactory() self.assertIsNone(user.get_capability_usage(IMPORTS_FILE_SIZE_CAPABILITY_ID)) self.assertIsNone(user.get_capability_usage(CLONE_RESOURCES_PER_CALL_CAPABILITY_ID)) + + +class FakePolicy(CapabilityPolicy): + """Records every call, and answers with values GroupLimitPolicy never gives.""" + calls = [] + refusal = None + made = 0 + + def __init__(self): + type(self).made += 1 + + def get_limit(self, user, capability_id, owner=None): + self.calls.append(('get_limit', user, capability_id, owner)) + return 7 + + def get_usage(self, user, capability_id, owner=None): + self.calls.append(('get_usage', user, capability_id, owner)) + return {'used': 3, 'window': 'fake'} + + def consume( # pylint: disable=too-many-arguments + self, user, capability_id, units=1, *, idempotency_key=None, owner=None, **attribution + ): + self.calls.append(('consume', user, capability_id, units, idempotency_key, owner, attribution)) + if self.refusal: + raise UsageRefused(dict(self.refusal), no_access=self.refusal.get('error_code') == 'fake_no_access') + return UsageEvent.objects.create(user=user, capability_id=capability_id, units=units), False + + def reverse(self, usage_event, units): + self.calls.append(('reverse', usage_event, units)) + return units + + def refund(self, user, capability_id, units=1): + self.calls.append(('refund', user, capability_id, units)) + + +@override_settings(CAPABILITY_POLICY='core.capabilities.tests.tests.FakePolicy') +class CapabilityPolicyContractTest(OCLAPITestCase): + """Every call site goes through CAPABILITY_POLICY, and sends what the policy reports.""" + + def setUp(self): + super().setUp() + FakePolicy.calls = [] + FakePolicy.refusal = None + FakePolicy.made = 0 + self.user = UserProfileFactory() + self.user.groups.add(Group.objects.get(name=PREVIEW_GROUP_NAME)) # mapper.use, for $match + + def calls(self, name): + return [call for call in FakePolicy.calls if call[0] == name] + + def consume(self, **data): + return self.client.post( + '/capabilities/consume/', data={'capability': 'ai_assistant.calls', **data}, format='json', + HTTP_AUTHORIZATION='Token ' + self.user.get_token()) + + def match(self, rows=2): + return self.client.post( + '/concepts/$match/', {'rows': [{'name': 'Malaria'}] * rows, 'target_repo_url': '/orgs/OCL/sources/X/'}, + format='json', HTTP_AUTHORIZATION='Token ' + self.user.get_token()) + + def test_user_profile_delegates(self): + self.assertEqual(self.user.get_capability_limit(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID), 7) + self.assertEqual(self.user.get_capability_usage(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID), 3) + event, already_consumed = self.user.check_and_consume_capability( + MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=2, action='act', algorithm='alg', idempotency_key='key') + + self.assertFalse(already_consumed) + self.assertEqual(event.units, 2) + self.assertEqual( + self.calls('consume'), + [('consume', self.user, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, 2, 'key', None, + {'action': 'act', 'algorithm': 'alg', 'map_project': None, 'run': None})]) + + def test_the_owner_reaches_the_policy(self): + owner = OrganizationFactory() + capability_id = MAPPER_MATCH_OPERATIONS_CAPABILITY_ID + + self.user.get_capability_limit(capability_id, owner=owner) + self.user.get_capability_usage(capability_id, owner=owner) + self.user.check_and_consume_capability(capability_id, owner=owner) + self.assertEqual(self.calls('get_limit'), [('get_limit', self.user, capability_id, owner)]) + self.assertEqual(self.calls('get_usage'), [('get_usage', self.user, capability_id, owner)]) + self.assertEqual(self.calls('consume')[0][5], owner) + + FakePolicy.calls = [] + self.assertEqual( + FakePolicy().describe(self.user, capability_id, owner=owner), {'limit': 7, 'used': 3, 'window': 'fake'}) + self.assertEqual( + FakePolicy.calls, + [('get_limit', self.user, capability_id, owner), ('get_usage', self.user, capability_id, owner)]) + + def test_a_consume_request_cannot_name_the_owner(self): + response = self.consume(owner=OrganizationFactory().mnemonic) + + self.assertEqual(response.status_code, 200) + self.assertIsNone(self.calls('consume')[0][5]) + self.assertNotIn('owner', self.calls('consume')[0][6]) + + def test_user_capabilities_report_every_usage_field(self): + response = self.client.get( + '/user/?includeCapabilities=true', HTTP_AUTHORIZATION='Token ' + self.user.get_token()) + + self.assertEqual(response.status_code, 200) + capabilities_by_name = {c['name']: c for c in response.data['capabilities']} + self.assertEqual( + capabilities_by_name['ai_assistant.calls'], + {'name': 'ai_assistant.calls', 'limit': 7, 'used': 3, 'window': 'fake'}) + + def test_consume_reports_every_usage_field(self): + response = self.consume(units=2, action='act', idempotency_key='key') + + self.assertEqual(response.status_code, 200) + self.assertEqual( + {key: response.data[key] for key in ['capability', 'limit', 'used', 'window', 'units', 'already_consumed']}, + {'capability': 'ai_assistant.calls', 'limit': 7, 'used': 3, 'window': 'fake', 'units': 2, + 'already_consumed': False}) + self.assertEqual(self.calls('consume')[0][3:5], (2, 'key')) + + def test_consume_refusal_is_sent_as_the_policy_built_it(self): + FakePolicy.refusal = {'error_code': 'fake_limit_reached', 'limit': 5, 'used': 5, 'window': 'fake'} + + response = self.consume() + + self.assertEqual(response.status_code, 403) + self.assertEqual( + dict(response.data), + {'detail': 'ai_assistant.calls limit reached.', 'error_code': 'fake_limit_reached', 'limit': 5, + 'used': 5, 'window': 'fake'}) + + FakePolicy.refusal = {'detail': 'From the policy.', 'error_code': 'fake_no_access'} + + response = self.consume() + + self.assertEqual(response.status_code, 403) + self.assertEqual(dict(response.data), {'detail': 'From the policy.', 'error_code': 'fake_no_access'}) + + def test_match_refusal_is_sent_as_the_policy_built_it(self): + FakePolicy.refusal = {'error_code': 'fake_no_access', 'limit': None, 'used': 0, 'window': 'fake'} + + response = self.match() + + self.assertEqual(response.status_code, 403) + self.assertEqual( + dict(response.data), + {'detail': 'You do not have Mapper match access.', 'error_code': 'fake_no_access', 'limit': None, + 'used': 0, 'window': 'fake'}) + + def test_match_reverses_its_event_when_the_search_fails(self): + from core.concepts.views import MetadataToConceptsListView + with patch.object(MetadataToConceptsListView, 'filter_queryset', side_effect=ValueError('es down')): + with self.assertRaises(ValueError): + self.match(rows=3) + + consume, = self.calls('consume') + self.assertEqual(consume[1:4], (self.user, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, 3)) + reverse, = self.calls('reverse') + self.assertEqual((reverse[1].user, reverse[1].units, reverse[2]), (self.user, 3, 3)) + + def test_match_that_succeeds_is_not_reversed(self): + from core.concepts.views import MetadataToConceptsListView + with patch.object(MetadataToConceptsListView, 'filter_queryset', return_value=[]): + response = self.match() + + self.assertEqual(response.status_code, 200) + self.assertEqual(len(self.calls('consume')), 1) + self.assertEqual(self.calls('reverse'), []) + + def test_staff_refund_goes_through_the_policy(self): + staff = UserProfileFactory(is_staff=True) + UsageCounter.objects.create(user=self.user, capability_id=AI_ASSISTANT_CALLS_CAPABILITY_ID, used=5) + + response = self.client.post( + '/capabilities/refund/', data={'user': self.user.username, 'capability': 'ai_assistant.calls', 'units': 2}, + format='json', HTTP_AUTHORIZATION='Token ' + staff.get_token()) + + self.assertEqual(response.status_code, 200) + self.assertEqual( + dict(response.data), + {'user': self.user.username, 'capability': 'ai_assistant.calls', 'limit': 7, 'used': 3, 'window': 'fake'}) + self.assertEqual(self.calls('refund'), [('refund', self.user, AI_ASSISTANT_CALLS_CAPABILITY_ID, 2)]) + # the correction is the policy's to make + self.assertEqual(UsageCounter.objects.get(user=self.user).used, 5) + + def test_a_consume_response_comes_from_the_policy_that_consumed(self): + response = self.consume() + + self.assertEqual(response.status_code, 200) + self.assertEqual(len(self.calls('consume')), 1) + self.assertEqual(len(self.calls('get_usage')), 1) + self.assertEqual(FakePolicy.made, 1) + + def test_a_match_that_fails_is_reversed_by_the_policy_that_consumed(self): + from core.concepts.views import MetadataToConceptsListView + with patch.object(MetadataToConceptsListView, 'filter_queryset', side_effect=ValueError('es down')): + with self.assertRaises(ValueError): + self.match() + + self.assertEqual(len(self.calls('reverse')), 1) + self.assertEqual(FakePolicy.made, 1) + + def test_a_staff_refund_response_comes_from_the_policy_that_refunded(self): + response = self.client.post( + '/capabilities/refund/', data={'user': self.user.username, 'capability': 'ai_assistant.calls'}, + format='json', HTTP_AUTHORIZATION='Token ' + UserProfileFactory(is_staff=True).get_token()) + + self.assertEqual(response.status_code, 200) + self.assertEqual(FakePolicy.made, 1) + + def test_a_list_of_capabilities_comes_from_one_policy(self): + response = self.client.get( + '/user/?includeCapabilities=true', HTTP_AUTHORIZATION='Token ' + self.user.get_token()) + + self.assertEqual(response.status_code, 200) + self.assertGreater(len(response.json()['capabilities']), 1) + self.assertEqual(len(self.calls('get_usage')), len(response.json()['capabilities'])) + self.assertEqual(FakePolicy.made, 1) + + def test_a_list_of_overrides_comes_from_one_policy(self): + for capability_id in (AI_ASSISTANT_CALLS_CAPABILITY_ID, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID): + UserCapabilityOverride.objects.create(user=self.user, capability_id=capability_id, limit=5) + FakePolicy.made = 0 + + response = self.client.get( + f'/users/{self.user.username}/capabilities/overrides/', + HTTP_AUTHORIZATION='Token ' + UserProfileFactory(is_staff=True).get_token()) + + self.assertEqual(response.status_code, 200) + self.assertEqual(len(response.json()), 2) + self.assertEqual(FakePolicy.made, 1) + + def test_refund_is_for_staff(self): + response = self.client.post( + '/capabilities/refund/', data={'user': self.user.username, 'capability': 'ai_assistant.calls'}, + format='json', HTTP_AUTHORIZATION='Token ' + self.user.get_token()) + + self.assertEqual(response.status_code, 403) + self.assertEqual(self.calls('refund'), []) + + +class WindowedPolicy(GroupLimitPolicy): + window_start = None + + def get_window_start(self, user, capability_id): + return self.window_start + + +@override_settings(CAPABILITY_POLICY='core.capabilities.tests.tests.WindowedPolicy') +class GroupLimitPolicyWindowTest(OCLAPITestCase): + first_window = datetime(2030, 1, 1, tzinfo=dt_timezone.utc) + second_window = datetime(2030, 2, 1, tzinfo=dt_timezone.utc) + + def setUp(self): + super().setUp() + WindowedPolicy.window_start = self.first_window + self.user = UserProfileFactory() + UserCapabilityOverride.objects.create( + user=self.user, capability_id=MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, limit=3) + + def tearDown(self): + WindowedPolicy.window_start = None + super().tearDown() + + def consume(self, units=1, idempotency_key=None): + return self.user.check_and_consume_capability( + MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=units, idempotency_key=idempotency_key) + + def used_by_window(self): + return dict(UsageCounter.objects.filter(user=self.user).values_list('window_start', 'used')) + + def test_each_window_counts_from_zero(self): + event, _ = self.consume(units=3) + with self.assertRaises(CapabilityExceeded): + self.consume() + + WindowedPolicy.window_start = self.second_window + self.consume() + + self.assertEqual(event.window_start, self.first_window) + self.assertEqual(self.user.get_capability_usage(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID), 1) + self.assertEqual(self.used_by_window(), {self.first_window: 3, self.second_window: 1}) + + def test_a_key_replays_in_a_later_window(self): + event, _ = self.consume(units=2, idempotency_key='key') + + WindowedPolicy.window_start = self.second_window + replay, already_consumed = self.consume(units=2, idempotency_key='key') + + self.assertTrue(already_consumed) + self.assertEqual(replay.id, event.id) + self.assertEqual(self.used_by_window(), {self.first_window: 2, self.second_window: 0}) + + def test_reverse_lowers_the_window_the_units_counted_in(self): + event, _ = self.consume(units=2) + WindowedPolicy.window_start = self.second_window + self.consume() + + self.assertEqual(WindowedPolicy().reverse(event, 2), 2) + self.assertEqual(WindowedPolicy().reverse(event, 2), 0) + + event.refresh_from_db() + self.assertEqual((event.units, event.reversed_units), (2, 2)) + self.assertEqual(self.used_by_window(), {self.first_window: 0, self.second_window: 1}) + + def test_reverse_never_takes_back_more_than_the_event_or_the_counter(self): + event, _ = self.consume(units=3) + UsageCounter.refund(self.user, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=2, window_start=self.first_window) + + self.assertEqual(WindowedPolicy().reverse(event, 5), 3) + + event.refresh_from_db() + self.assertEqual(event.reversed_units, 3) + self.assertEqual(self.used_by_window(), {self.first_window: 0}) + + def test_refund_corrects_the_window_counting_now(self): + self.consume(units=3) + WindowedPolicy.window_start = self.second_window + self.consume(units=3) + + WindowedPolicy().refund(self.user, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=2) + + self.assertEqual(self.used_by_window(), {self.first_window: 3, self.second_window: 1}) + + +class GroupLimitPolicyDefaultWindowTest(OCLAPITestCase): + def test_consume_locks_the_account_without_blocking_foreign_key_checks(self): + user = UserProfileFactory() + UserCapabilityOverride.objects.create(user=user, capability_id=MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, limit=0) + + with CaptureQueriesContext(connection) as context: + user.check_and_consume_capability(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID) + + self.assertTrue(any( + 'FROM "user_profiles"' in query['sql'] and query['sql'].endswith('FOR NO KEY UPDATE') + for query in context.captured_queries)) + + def test_an_event_inserted_without_the_newer_columns_gets_their_defaults(self): + # code that predates reversed_units and window_start keeps serving while a deploy migrates + user = UserProfileFactory() + with connection.cursor() as cursor: + cursor.execute( + 'INSERT INTO usage_events (user_id, capability_id, units, action, created_at) ' + 'VALUES (%s, %s, 1, %s, now()) RETURNING id', [user.id, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, 'match']) + event = UsageEvent.objects.get(id=cursor.fetchone()[0]) + + self.assertEqual((event.reversed_units, event.window_start), (0, None)) + + def test_usage_counts_in_the_window_that_never_ends(self): + user = UserProfileFactory() + UserCapabilityOverride.objects.create(user=user, capability_id=MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, limit=0) + + event, _ = user.check_and_consume_capability(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=2) + + self.assertIsNone(event.window_start) + self.assertEqual( + list(UsageCounter.objects.filter(user=user).values_list('window_start', 'used')), [(None, 2)]) + self.assertEqual(get_capability_policy().reverse(event, 2), 2) + self.assertEqual(user.get_capability_usage(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID), 0) + + def test_staff_refund_lowers_the_counter_and_never_below_zero(self): + user = UserProfileFactory() + UsageCounter.objects.create(user=user, capability_id=AI_ASSISTANT_CALLS_CAPABILITY_ID, used=3) + staff_token = 'Token ' + UserProfileFactory(is_staff=True).get_token() + + for units, used in ((2, 1), (5, 0)): + response = self.client.post( + '/capabilities/refund/', + data={'user': user.username, 'capability': 'ai_assistant.calls', 'units': units}, + format='json', HTTP_AUTHORIZATION=staff_token) + + self.assertEqual(response.status_code, 200) + self.assertEqual(response.data['used'], used) + + +class GroupLimitPolicyOwnerTest(OCLAPITestCase): + """The default policy decides and counts by the user, whoever owns the content.""" + capability_id = MAPPER_MATCH_OPERATIONS_CAPABILITY_ID + + def setUp(self): + super().setUp() + self.owner = OrganizationFactory() + self.user = UserProfileFactory() + UserCapabilityOverride.objects.create(user=self.user, capability_id=self.capability_id, limit=3) + + def test_limit_and_usage_ignore_the_owner(self): + self.user.check_and_consume_capability(self.capability_id, units=2) + + self.assertEqual(self.user.get_capability_limit(self.capability_id, owner=self.owner), 3) + self.assertEqual(self.user.get_capability_usage(self.capability_id, owner=self.owner), 2) + + def test_consume_charges_the_user_at_the_users_limit(self): + event, already_consumed = self.user.check_and_consume_capability(self.capability_id, units=3, owner=self.owner) + + self.assertFalse(already_consumed) + self.assertEqual((event.user, event.units), (self.user, 3)) + self.assertEqual( + list(UsageCounter.objects.filter(capability_id=self.capability_id).values_list('user', 'used')), + [(self.user.id, 3)]) + with self.assertRaises(CapabilityExceeded): + self.user.check_and_consume_capability(self.capability_id, owner=self.owner) + + +class OwnerLimitPolicy(GroupLimitPolicy): + """Replaces only get_limit: with an owner, a positive limit becomes this one. Blocked and unlimited stay.""" + owner_limit = 5 + + def get_limit(self, user, capability_id, owner=None): + limit = super().get_limit(user, capability_id, owner=owner) + return self.owner_limit if owner and limit and limit > 0 else limit + + +@override_settings(CAPABILITY_POLICY='core.capabilities.tests.tests.OwnerLimitPolicy') +class ReplacedPolicyOwnerTest(OCLAPITestCase): + capability_id = MAPPER_MATCH_OPERATIONS_CAPABILITY_ID + + def setUp(self): + super().setUp() + self.owner = OrganizationFactory() + self.user = UserProfileFactory() + UserCapabilityOverride.objects.create(user=self.user, capability_id=self.capability_id, limit=1) + + def test_limit_follows_the_owner(self): + self.assertEqual(self.user.get_capability_limit(self.capability_id, owner=self.owner), 5) + self.assertEqual(self.user.get_capability_limit(self.capability_id), 1) + + def test_inherited_consume_enforces_the_limit_the_owner_gets(self): + self.user.check_and_consume_capability(self.capability_id, units=4, owner=self.owner) + + with self.assertRaises(CapabilityExceeded): + self.user.check_and_consume_capability(self.capability_id) + with self.assertRaises(CapabilityExceeded): + self.user.check_and_consume_capability(self.capability_id, units=2, owner=self.owner) + + self.user.check_and_consume_capability(self.capability_id, owner=self.owner) + self.assertEqual(self.user.get_capability_usage(self.capability_id), 5) + + +class PartialPolicy(CapabilityPolicy): + """Implements everything except reverse() and refund().""" + + def get_limit(self, user, capability_id, owner=None): + return 0 + + def get_usage(self, user, capability_id, owner=None): + return {'used': None} + + def consume( # pylint: disable=too-many-arguments + self, user, capability_id, units=1, *, idempotency_key=None, owner=None, **attribution + ): + return None, False + + +class NarrowRefundPolicy(GroupLimitPolicy): + """refund() can't take `units`.""" + + def refund(self, user, capability_id): # pylint: disable=arguments-differ + return None + + +class NoOwnerPolicy(GroupLimitPolicy): + """get_limit() can't take `owner`, as a policy written before it couldn't.""" + + def get_limit(self, user, capability_id): # pylint: disable=arguments-differ + return 0 + + +class RequiredInitPolicy(GroupLimitPolicy): + def __init__(self, setting): + self.setting = setting + + +class OptionalInitPolicy(GroupLimitPolicy): + def __init__(self, setting=None, **options): + self.setting = setting + self.options = options + + +class AnyArgumentsPolicy(GroupLimitPolicy): + """Every call fits, and `owner` is taken only by name, as call sites pass it.""" + + def get_limit(self, user, capability_id, *, owner=None): # pylint: disable=arguments-differ + return 0 + + def get_usage(self, *args, **kwargs): # pylint: disable=arguments-differ,unused-argument + return {'used': None} + + +class KeywordUnitsPolicy(GroupLimitPolicy): + """Takes `units` only by name, which the check allows: call sites must pass it that way.""" + + def consume( # pylint: disable=too-many-arguments,arguments-differ + self, user, capability_id, *, units=1, idempotency_key=None, owner=None, **attribution + ): + return super().consume(user, capability_id, units, idempotency_key=idempotency_key, owner=owner, **attribution) + + +@override_settings(CAPABILITY_POLICY='core.capabilities.tests.tests.KeywordUnitsPolicy') +class CallSitesFitTheCheckTest(OCLAPITestCase): + def test_a_policy_the_check_passes_can_be_called(self): + user = UserProfileFactory() + UserCapabilityOverride.objects.create(user=user, capability_id=MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, limit=0) + + self.assertEqual(check_capability_policy(), []) + + event, _ = user.check_and_consume_capability(MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=2) + self.assertEqual(event.units, 2) + + +class CapabilityPolicyInterfaceTest(SimpleTestCase): + def test_a_method_left_to_the_interface_says_so(self): + for call in ( + lambda: CapabilityPolicy.get_limit(None, None, None), + lambda: CapabilityPolicy.get_usage(None, None, None), + lambda: CapabilityPolicy.consume(None, None, None), + lambda: CapabilityPolicy.reverse(None, None, None), + lambda: CapabilityPolicy.refund(None, None, None), + ): + with self.assertRaises(NotImplementedError): + call() + + def test_a_policy_missing_a_method_cannot_be_made(self): + for policy in (CapabilityPolicy, PartialPolicy): + with self.subTest(policy=policy), self.assertRaises(TypeError): + policy() # pylint: disable=abstract-class-instantiated + + @override_settings(CAPABILITY_POLICY='core.capabilities.policy.GroupLimitPolicy') + def test_the_setting_names_the_policy_in_use(self): + self.assertIsInstance(get_capability_policy(), GroupLimitPolicy) + + +class CapabilityPolicyCheckTest(SimpleTestCase): + def errors(self, path): + with override_settings(CAPABILITY_POLICY=path): + return check_capability_policy() + + def test_it_runs_at_start_up(self): + self.assertIn(check_capability_policy, registry.get_checks()) + + def test_the_default_passes(self): + self.assertEqual(self.errors('core.capabilities.policy.GroupLimitPolicy'), []) + + def test_a_policy_that_does_not_import_fails(self): + for path in ('core.capabilities.policy.NoSuchPolicy', 'no_such_package.policy.Policy'): + with self.subTest(path): + self.assertEqual([error.id for error in self.errors(path)], ['capabilities.E001']) + + def test_a_policy_whose_module_fails_to_load_fails(self): + with patch('core.capabilities.checks.import_string', side_effect=RuntimeError('not in INSTALLED_APPS')): + self.assertEqual( + [error.id for error in self.errors('core.capabilities.policy.GroupLimitPolicy')], ['capabilities.E001']) + + def test_something_other_than_a_policy_fails(self): + for path in ('django.contrib.auth.models.Group', 'core.capabilities.policy.logger'): + with self.subTest(path): + self.assertEqual([error.id for error in self.errors(path)], ['capabilities.E002']) + + def test_a_policy_missing_a_method_fails_and_names_it(self): + error, = self.errors('core.capabilities.tests.tests.PartialPolicy') + + self.assertEqual(error.id, 'capabilities.E003') + self.assertIn('refund, reverse', error.msg) + + error, = self.errors('core.capabilities.policy.CapabilityPolicy') + self.assertEqual(error.id, 'capabilities.E003') + + def test_a_method_that_cannot_take_the_interface_s_arguments_fails_and_is_named(self): + for path, misfits in ( + ('core.capabilities.tests.tests.NarrowRefundPolicy', 'refund'), + ('core.capabilities.tests.tests.NoOwnerPolicy', 'get_limit'), + ('core.capabilities.tests.tests.RequiredInitPolicy', '__init__'), + ): + with self.subTest(path): + error, = self.errors(path) + self.assertEqual(error.id, 'capabilities.E004') + self.assertIn(misfits, error.msg) + + def test_a_policy_is_not_made_by_the_check(self): + with patch.object(GroupLimitPolicy, '__init__', side_effect=RuntimeError('read at start-up')) as init: + self.assertEqual(check_capability_policy(), []) + init.assert_not_called() + + def test_a_policy_that_takes_the_arguments_another_way_passes(self): + for path in ('core.capabilities.tests.tests.AnyArgumentsPolicy', + 'core.capabilities.tests.tests.OptionalInitPolicy'): + with self.subTest(path): + self.assertEqual(self.errors(path), []) diff --git a/core/capabilities/views.py b/core/capabilities/views.py index 321cc420b..5d22ef86f 100644 --- a/core/capabilities/views.py +++ b/core/capabilities/views.py @@ -5,11 +5,10 @@ from rest_framework.response import Response from rest_framework.views import APIView -from core.capabilities.constants import ( - CAPABILITY_EXCEEDED_ERROR_CODE, CAPABILITY_ID_BY_NAME, CAPABILITY_NOT_ENTITLED_ERROR_CODE, -) -from core.capabilities.exceptions import CapabilityExceeded -from core.capabilities.models import UsageCounter, UserCapabilityOverride +from core.capabilities.constants import CAPABILITY_ID_BY_NAME +from core.capabilities.exceptions import UsageRefused +from core.capabilities.models import UserCapabilityOverride +from core.capabilities.policy import get_capability_policy from core.capabilities.serializers import UserCapabilityOverrideSerializer from core.common.mixins import ListWithHeadersMixin from core.common.utils import parse_id @@ -93,22 +92,20 @@ def post(self, request): return error_response map_project, run = self.get_owned_map_project_and_run() + # one policy for the whole response, so what it reports is what it counted + policy = get_capability_policy() try: - event, already_consumed = request.user.check_and_consume_capability( - capability_id, units=units, + event, already_consumed = policy.consume( + request.user, capability_id, units=units, idempotency_key=idempotency_key, action=request.data.get('action', ''), algorithm=request.data.get('algorithm'), - map_project=map_project, run=run, idempotency_key=idempotency_key, + map_project=map_project, run=run, ) - except CapabilityExceeded as ex: - not_entitled = ex.not_entitled + except UsageRefused as ex: return Response( { - 'detail': f'You do not have access to {capability_name}.' if not_entitled else + 'detail': f'You do not have access to {capability_name}.' if ex.no_access else f'{capability_name} limit reached.', - 'error_code': CAPABILITY_NOT_ENTITLED_ERROR_CODE.get(capability_name, 'capability_not_entitled') - if not_entitled else - CAPABILITY_EXCEEDED_ERROR_CODE.get(capability_name, 'capability_limit_reached'), - 'limit': None if not_entitled else ex.limit, 'used': ex.used, + **ex.detail, }, status=status.HTTP_403_FORBIDDEN ) @@ -116,8 +113,7 @@ def post(self, request): return Response( { 'capability': capability_name, - 'limit': request.user.get_capability_limit(capability_id), - 'used': request.user.get_capability_usage(capability_id), + **policy.describe(request.user, capability_id), 'usage_event_id': event.id, 'units': event.units, 'already_consumed': already_consumed, @@ -140,7 +136,7 @@ def get_idempotency_key(self): class CapabilityRefundView(CapabilityBaseView): """ - Manual ledger correction (see UsageCounter.refund) - e.g. staff fixing an + Manual ledger correction (see CapabilityPolicy.refund) - e.g. staff fixing an over-count after a bug. Staff-only, and operates on a user given explicitly in the request body: this is an admin tool for correcting ANY user's ledger, not a per-request refund a service would issue against its own @@ -168,14 +164,14 @@ def post(self, request): if error_response: return error_response - UsageCounter.refund(user, capability_id, units=units) + policy = get_capability_policy() + policy.refund(user, capability_id, units=units) return Response( { 'user': username, 'capability': capability_name, - 'limit': user.get_capability_limit(capability_id), - 'used': user.get_capability_usage(capability_id), + **policy.describe(user, capability_id), }, status=status.HTTP_200_OK ) diff --git a/core/common/tests.py b/core/common/tests.py index 832fc6a77..9a8bd2057 100644 --- a/core/common/tests.py +++ b/core/common/tests.py @@ -1,4 +1,5 @@ import datetime +import importlib import json import os import sys @@ -19,6 +20,7 @@ from django.http import HttpResponse from django.test import TestCase, override_settings, RequestFactory, tag from django.test.runner import DiscoverRunner +from django.urls import clear_url_caches, resolve from django.utils import timezone from elastic_transport import Transport from mock.mock import call @@ -3647,3 +3649,22 @@ def assert_rolled_back(client): target = client.indices.split.call_args[1]['target'] client.indices.delete.assert_called_once_with(index=target, ignore_unavailable=True) client.indices.put_settings.assert_called_once_with(index='concepts', settings={'index.blocks.write': None}) + + +class ExtraURLConfsTest(OCLTestCase): + def tearDown(self): + from core import urls + importlib.reload(urls) + clear_url_caches() + super().tearDown() + + def test_extra_urlconfs_are_included_after_the_core_routes(self): + from core import urls + from core.capabilities.views import CapabilityConsumeView + with patch('core.settings.OCL_EXTRA_URLCONFS', ['core.capabilities.urls']): + importlib.reload(urls) + clear_url_caches() + + self.assertEqual(resolve('/consume/').func.view_class, CapabilityConsumeView) + self.assertEqual(resolve('/capabilities/consume/').func.view_class, CapabilityConsumeView) + self.assertEqual(urls.urlpatterns[-1].urlconf_name.__name__, 'core.capabilities.urls') diff --git a/core/concepts/views.py b/core/concepts/views.py index 64fceff91..a0c472ef2 100644 --- a/core/concepts/views.py +++ b/core/concepts/views.py @@ -27,9 +27,9 @@ HEAD, INCLUDE_INVERSE_MAPPINGS_PARAM, INCLUDE_RETIRED_PARAM, ACCESS_TYPE_NONE, LIST_DEFAULT_LIMIT) from core.common.exceptions import Http400, Http403, Http409 from core.common.mixins import ListWithHeadersMixin, ConceptDictionaryMixin -from core.capabilities.constants import CAPABILITY_EXCEEDED_ERROR_CODE, CAPABILITY_NOT_ENTITLED_ERROR_CODE, \ - MAPPER_MATCH_OPERATIONS_CAPABILITY, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID -from core.capabilities.exceptions import CapabilityExceeded +from core.capabilities.constants import MAPPER_MATCH_OPERATIONS_CAPABILITY_ID +from core.capabilities.exceptions import UsageRefused +from core.capabilities.policy import get_capability_policy from core.capacity.constants import ENDPOINT_MATCH, ENDPOINT_RERANK from core.capacity.limiter import CapacityLimitMixin from core.common.permissions import CanUseMapper @@ -1265,38 +1265,34 @@ def post(self, request, **kwargs): # pylint: disable=unused-argument def match(self, request): rows = request.data.get('rows') - consumed_units = 0 + usage_event, already_consumed = None, False + # one policy for the whole operation, so a failed search is reversed by the policy that consumed + policy = get_capability_policy() if isinstance(rows, list) and rows: algorithm_id, map_project = get_match_operations_attribution(request) try: - request.user.check_and_consume_capability( - MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=len(rows), action='match_concepts', + usage_event, already_consumed = policy.consume( + request.user, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=len(rows), action='match_concepts', algorithm=algorithm_id, map_project=map_project ) - consumed_units = len(rows) - except CapabilityExceeded as ex: - not_entitled = ex.not_entitled + except UsageRefused as ex: return Response( { - 'detail': 'You do not have Mapper match access.' if not_entitled else + 'detail': 'You do not have Mapper match access.' if ex.no_access else 'Match operation limit reached.', - 'error_code': CAPABILITY_NOT_ENTITLED_ERROR_CODE[MAPPER_MATCH_OPERATIONS_CAPABILITY] - if not_entitled else CAPABILITY_EXCEEDED_ERROR_CODE[MAPPER_MATCH_OPERATIONS_CAPABILITY], - 'limit': None if not_entitled else ex.limit, 'used': ex.used, + **ex.detail, }, status=status.HTTP_403_FORBIDDEN ) - # Consumption happens before the search so an over-quota caller is rejected - # without paying for an ES query it was never going to get results from. If the - # search itself then fails, that's not a successful match - refund the units so - # a flaky ES call doesn't silently burn an unrecoverable chunk of a capped, - # never-reset allowance. + # Consumption happens before the search so a caller over its limit is rejected + # without an ES query it was never going to get results from. If the search + # itself then fails, that's not a successful match - reverse the units so a + # flaky ES call doesn't use them up. try: results = self.filter_queryset() except Exception: - if consumed_units: - from core.capabilities.models import UsageCounter - UsageCounter.refund(request.user, MAPPER_MATCH_OPERATIONS_CAPABILITY_ID, units=consumed_units) + if usage_event and not already_consumed: # a replay consumed nothing + policy.reverse(usage_event, usage_event.units) raise response = Response(results) # num_returned is picked up by the analytics middleware as diff --git a/core/integration_tests/tests_concepts.py b/core/integration_tests/tests_concepts.py index 54038581e..a57045132 100644 --- a/core/integration_tests/tests_concepts.py +++ b/core/integration_tests/tests_concepts.py @@ -19,6 +19,7 @@ from core.bundles.models import Bundle from core.capabilities.exceptions import CapabilityExceeded +from core.capabilities.policy import GroupLimitPolicy from core.capacity.config import clear_cache, save_config from core.capacity.constants import HEADER_DECISION from core.capacity.limiter import RedisLanes @@ -4525,8 +4526,7 @@ def test_post_429_heavy_call_refused_by_capacity_limit(self): def test_post_403_match_operation_limit_reached(self): with patch.object( - UserProfile, 'check_and_consume_capability', - side_effect=CapabilityExceeded('mapper.match_operations', 10, 10, 1)): + GroupLimitPolicy, 'consume', side_effect=CapabilityExceeded('mapper.match_operations', 10, 10)): response = self.match() self.assertEqual(response.status_code, 403) @@ -4535,11 +4535,12 @@ def test_post_403_match_operation_limit_reached(self): def test_post_refunds_consumed_units_when_search_fails(self): with patch.object(MetadataToConceptsListView, 'filter_queryset', side_effect=ValueError('es down')), \ - patch('core.capabilities.models.UsageCounter.refund') as refund_mock: + patch.object(GroupLimitPolicy, 'reverse') as reverse_mock: with self.assertRaises(ValueError): self.match({'rows': [{'name': 'Malaria'}, {'name': 'Fever'}]}) - refund_mock.assert_called_once_with(self.user, ANY, units=2) + reverse_mock.assert_called_once_with(ANY, 2) + self.assertEqual(reverse_mock.call_args.args[0].user, self.user) def test_variants_are_only_for_core_users(self): with patch.object(ConceptFuzzySearch, 'search', wraps=ConceptFuzzySearch.search) as search_mock: diff --git a/core/settings.py b/core/settings.py index e5e8ad3cf..f3e8de165 100644 --- a/core/settings.py +++ b/core/settings.py @@ -667,6 +667,16 @@ def get_set_from_env(name): # SILKY_PYTHON_PROFILER = True # SILKY_PYTHON_PROFILER_RESULT_PATH = '/code/core/' +# Django apps and URLconfs a deployment adds to this build, as comma-separated dotted paths. Extra URLconfs are +# included after core.urls' own routes, so those keep precedence. +OCL_EXTRA_APPS = [app.strip() for app in os.environ.get('OCL_EXTRA_APPS', '').split(',') if app.strip()] +INSTALLED_APPS = [*INSTALLED_APPS, *OCL_EXTRA_APPS] +OCL_EXTRA_URLCONFS = [ + urlconf.strip() for urlconf in os.environ.get('OCL_EXTRA_URLCONFS', '').split(',') if urlconf.strip() +] +# The class that decides and records the use of capabilities (core.capabilities.policy) +CAPABILITY_POLICY = os.environ.get('CAPABILITY_POLICY', 'core.capabilities.policy.GroupLimitPolicy') + # MINIO storage settings MINIO_ENDPOINT = os.environ.get('MINIO_ENDPOINT', '') MINIO_EXTERNAL_ENDPOINT = os.environ.get('MINIO_EXTERNAL_ENDPOINT', '') diff --git a/core/urls.py b/core/urls.py index 781477677..2c0d47514 100644 --- a/core/urls.py +++ b/core/urls.py @@ -32,7 +32,7 @@ from core.fhir.swagger import ApiSchemaGenerator from core.importers.views import BulkImportView from core.map_projects.views import AutomatchRunView -from core.settings import ENV +from core.settings import ENV, OCL_EXTRA_URLCONFS api_info = openapi.Info( title="OCL API", @@ -127,6 +127,8 @@ *urlpatterns ] +urlpatterns += [path('', include(urlconf)) for urlconf in OCL_EXTRA_URLCONFS] + handler500 = 'rest_framework.exceptions.server_error' handler400 = 'rest_framework.exceptions.bad_request' diff --git a/core/users/models.py b/core/users/models.py index 5d58a14ab..947561ad4 100644 --- a/core/users/models.py +++ b/core/users/models.py @@ -1,7 +1,5 @@ -import logging import uuid from datetime import datetime -from typing import Any from dirtyfields import DirtyFieldsMixin from django.contrib.auth.models import AbstractUser, Group @@ -10,19 +8,17 @@ from django.contrib.contenttypes.models import ContentType from django.core.exceptions import ValidationError from django.db import models, transaction -from django.db.models import F from rest_framework.authtoken.models import Token +from core.capabilities.policy import get_capability_policy from core.common.mixins import SourceContainerMixin from core.common.models import BaseModel, CommonLogoModel from core.common.tasks import send_user_verification_email, send_user_reset_password_email from core.common.utils import web_url -from core.users.constants import STAFF_GROUP, SUPERADMIN_GROUP, GUEST_GROUP, CORE_USER_GROUP, PREVIEW_GROUP +from core.users.constants import STAFF_GROUP, SUPERADMIN_GROUP, GUEST_GROUP, CORE_USER_GROUP from .constants import USER_OBJECT_TYPE from ..common.checksums import ChecksumModel -logger = logging.getLogger('oclapi') - class Follow(models.Model): class Meta: @@ -248,127 +244,21 @@ def capabilities(self): from core.capabilities.models import Capability return Capability.objects.all() - def get_capability_limit(self, capability_id): - """ - Returns the configured limit as-is. 0 means unlimited. - Set explicitly via a per-user override or a group's own row. - There is no implicit-unlimited fallback: - None means this capability has no override and no group row for this user at - all, and callers must treat that as BLOCKED, not unlimited (see - check_and_consume_capability) - a new capability, or a new group that grants - Mapper access, must have its limit configured explicitly (even to 0) before - real usage is allowed against it. Silently failing open on missing - configuration was the previous behavior; don't reintroduce it here. - - Staff and superusers are the deliberate exception. - - Per-user override wins over every group; else the highest of the user's - groups' limits, where an explicit 0 from any one group wins outright (it - isn't just "the lowest number" - max() alone would let a capped group beat an - unlimited one). - - Exception: for the authoring capabilities in AUTHORING_CAPABILITY_IDS (bulk - import, $clone), which every account has always had, no row means the - `preview` group's value - the lowest tier - not blocked (ocl_online#230). - """ - if self.is_superuser or self.is_staff: - return 0 - override = self._get_capability_limit(capability_id) - if override is not None: - return override - group_limits = self._get_group_capability_limit(capability_id) - if not group_limits: - return self._get_default_authoring_limit(capability_id) - if 0 in group_limits: - return 0 - return max(group_limits) - - def _get_group_capability_limit(self, capability_id) -> list[Any]: - from core.capabilities.models import GroupCapability - return list(GroupCapability.objects.filter( - group__in=self.groups.all(), capability_id=capability_id - ).values_list('limit', flat=True)) - - def _get_capability_limit(self, capability_id): - return self.capability_overrides.filter(capability_id=capability_id).values_list('limit', flat=True).first() + # Limits and usage are decided by settings.CAPABILITY_POLICY (core.capabilities.policy) + def get_capability_limit(self, capability_id, owner=None): + return get_capability_policy().get_limit(self, capability_id, owner=owner) - @staticmethod - def _get_default_authoring_limit(capability_id): - """None (blocked) unless `capability_id` is an authoring capability, which falls back to `preview`.""" - from core.capabilities.constants import AUTHORING_CAPABILITY_IDS, AUTHORING_CAPABILITY_DEFAULT_LIMITS - from core.capabilities.models import GroupCapability - if capability_id not in AUTHORING_CAPABILITY_IDS: - return None - limit = GroupCapability.objects.filter( - group__name=PREVIEW_GROUP, capability_id=capability_id).values_list('limit', flat=True).first() - if limit is None: - logger.warning( - 'The %s group has no limit for capability %s; using the built-in default', PREVIEW_GROUP, capability_id) - return AUTHORING_CAPABILITY_DEFAULT_LIMITS[capability_id] - return limit - - def get_capability_usage(self, capability_id): - from core.capabilities.constants import ( - AUTHORING_CAPABILITY_IDS, MAPPER_PROJECTS_CAPABILITY_ID, MAPPER_ROWS_PER_PROJECT_CAPABILITY_ID - ) - if capability_id in AUTHORING_CAPABILITY_IDS: - return None # per-request limits: nothing accumulates - if capability_id == MAPPER_ROWS_PER_PROJECT_CAPABILITY_ID: - # A per-project cap has no per-user usage; clients show the cap itself. - return None - if capability_id == MAPPER_PROJECTS_CAPABILITY_ID: - # mapper.projects is enforced against the live count (HasMapProjectCapacity), - # not the monotonic UsageCounter - report the same number here, or a user who - # deletes their one project sees "used" stay at the old total while creation - # (correctly) succeeds again. - return self.map_projects_used - return self.usage_counters.filter(capability_id=capability_id).values_list('used', flat=True).first() or 0 + def get_capability_usage(self, capability_id, owner=None): + return get_capability_policy().get_usage(self, capability_id, owner=owner)['used'] def check_and_consume_capability( # pylint: disable=too-many-arguments self, capability_id, units=1, action='', algorithm=None, map_project=None, run=None, - idempotency_key=None + idempotency_key=None, owner=None ): - """ - Atomically checks this user's remaining `capability_id` quota and, if `units` - fits, consumes it and logs a `UsageEvent`. Raises `CapabilityExceeded` (nothing - consumed, nothing logged) if it doesn't fit. An explicit limit of 0 (only) means - unlimited and always succeeds; usage is still logged. A capability with no - override and no group row for this user at all (get_capability_limit() returns - None) is BLOCKED, not unlimited - see get_capability_limit's docstring. - kill switch = set the limit to -1 in groups.yaml (a deleted GroupCapability row is - recreated on restart); -1 on a per-user override blocks just that user. - (0 is NOT the kill switch - it means unlimited). - `capability_id` must be an id of a capability already seeded (Keycloak/fixtures) - - a bad id fails with an IntegrityError on insert rather than silently creating a - Capability row. - With an `idempotency_key` already consumed, charges nothing and returns that event, - even past the limit. Returns (usage_event, already_consumed). - """ - from core.capabilities.constants import CAPABILITY_NAME_BY_ID - from core.capabilities.exceptions import CapabilityExceeded - from core.capabilities.models import UsageCounter - limit = self.get_capability_limit(capability_id) - - with transaction.atomic(): - UsageCounter.objects.get_or_create(user=self, capability_id=capability_id) - counter = self.usage_counters.select_for_update().get(capability_id=capability_id) - - # under the counter lock and before the limit check, so a retry after the last unit still succeeds - if idempotency_key: - event = self.usage_events.filter(capability_id=capability_id, idempotency_key=idempotency_key).first() - if event: - return event, True - - if limit is None or (limit and counter.used + units > limit): - raise CapabilityExceeded( - CAPABILITY_NAME_BY_ID.get(capability_id, capability_id), limit, counter.used, units) - - UsageCounter.objects.filter(pk=counter.pk).update(used=F('used') + units) - event = self.usage_events.create( - capability_id=capability_id, units=units, action=action, algorithm=algorithm, - map_project=map_project, run=run, idempotency_key=idempotency_key or None - ) - return event, False + return get_capability_policy().consume( + self, capability_id, units=units, idempotency_key=idempotency_key, owner=owner, action=action, + algorithm=algorithm, map_project=map_project, run=run + ) def log_capability_event( # pylint: disable=too-many-arguments self, capability_id, units=1, action='', algorithm=None, map_project=None, run=None diff --git a/core/users/tests/tests.py b/core/users/tests/tests.py index b1e551823..2557bffce 100644 --- a/core/users/tests/tests.py +++ b/core/users/tests/tests.py @@ -1087,7 +1087,7 @@ def test_user_subscribed_orgs_and_pins_only_as_visible_to_requester(self): self.assertNotIn('pins', user_data) @patch('core.users.models.UserProfile.get_all_permissions', return_value=set()) - @patch('core.users.models.UserProfile.get_capability_usage', return_value=0) + @patch('core.capabilities.policy.GroupLimitPolicy.get_usage', return_value={'used': 0}) def test_user_detail_does_not_compute_plan_for_other_users(self, usage_mock, permissions_mock): user = UserProfileFactory(username='planusageuser') other_user = UserProfileFactory(username='planusageother')