c7a7a3f81a
# What this PR does Follow up PR for https://github.com/grafana/oncall/pull/5004 Tests haven’t caught a bug, so the method and the tests are fixed ## Which issue(s) this PR closes Related to [issue link here] <!-- *Note*: If you want the issue to be auto-closed once the PR is merged, change "Related to" to "Closes" in the line above. If you have more than one GitHub issue that this PR closes, be sure to preface each issue link with a [closing keyword](https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/using-keywords-in-issues-and-pull-requests#linking-a-pull-request-to-an-issue). This ensures that the issue(s) are auto-closed once the PR has been merged. --> ## Checklist - [ ] Unit, integration, and e2e (if applicable) tests updated - [ ] Documentation added (or `pr:no public docs` PR label added if not required) - [ ] Added the relevant release notes label (see labels prefixed w/ `release:`). These labels dictate how your PR will show up in the autogenerated release notes.
33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from django.conf import settings
|
|
from ratelimit.utils import _split_rate
|
|
from rest_framework.throttling import UserRateThrottle
|
|
|
|
|
|
class CustomRateUserThrottler(UserRateThrottle):
|
|
""" """
|
|
|
|
def parse_rate(self, rate):
|
|
"Use django ratelimit format to parse rate, i.e. '30/1m', instead of '30/m'"
|
|
return _split_rate(rate)
|
|
|
|
def allow_request(self, request, view):
|
|
# Override default rate limit, if organization id is specified in CUSTOM_RATELIMITS
|
|
custom_ratelimits = settings.CUSTOM_RATELIMITS
|
|
organization_id = str(request.user.organization_id)
|
|
if organization_id in custom_ratelimits:
|
|
self.rate = custom_ratelimits[organization_id].public_api
|
|
self.num_requests, self.duration = self.parse_rate(self.rate)
|
|
|
|
return super().allow_request(request, view)
|
|
|
|
|
|
class CustomRateOrganizationThrottler(CustomRateUserThrottler):
|
|
scope = "organization"
|
|
|
|
def get_cache_key(self, request, view):
|
|
if request.user.is_authenticated:
|
|
ident = request.user.organization.pk
|
|
else:
|
|
ident = self.get_ident(request)
|
|
|
|
return self.cache_format % {"scope": self.scope, "ident": ident}
|