view roundup/rate_limit.py @ 8180:d02ce1d14acd

feat: issue2551068 - Provide way to retrieve file/msg data via rest endpoint. Use Allow header to change format of /binary_content endpoint. If Allow header for endpoint is not application/json, it will be matched against the mime type for the file. */*, text/* are supported and will return the native mime type if present. Changes: move */* mime type from static dict of supported types. It was hardcoded to return json only. Now it can return a matching non-json mime type for the /binary_content endpoint. Edited some errors to explicitly add */* mime type. Cleanups to use ', ' separation in lists of valid mime types rather than just space separated. Remove ETag header when sending raw content. See issue 2551375 for background. Doc added to rest.txt. Small format fix up (add dash) in CHANGES.txt. Make passing an unset/None/False accept_mime_type to format_dispatch_output a 500 error. This used to be the fallback to produce a 406 error after all processing had happened. It should no longer be possible to take that code path as all 406 errors (with valid accept_mime_types) are generated before processing takes place. Make format_dispatch_output handle output other than json/xml so it can send back binary_content data. Removed a spurious client.response_code = 400 that seems to not be used. Tests added for all code paths. Database setup for tests msg and file entry. This required a file upload test to change so it doesn't look for file1 as the link returned by the upload. Download the link and verify the data rather than verifying the link. Multiple formatting changes to error messages to make all lists of valid mime types ', ' an not just space separated.
author John Rouillard <rouilj@ieee.org>
date Sun, 08 Dec 2024 17:22:33 -0500
parents 5fbd3af526bd
children 224ccb8b49ca
line wrap: on
line source

# Originaly from
# https://smarketshq.com/implementing-gcra-in-python-5df1f11aaa96?gi=4b9725f99bfa
# with imports, modifications for python 2, implementation of
# set/get_tat and marshaling as string, support for testonly
# and status method.

from datetime import timedelta, datetime

try:
    # used by python 3.11 and newer use tz aware dates
    from datetime import UTC
    dt_min = datetime.min.replace(tzinfo=UTC)
    # start of unix epoch
    dt_epoch = datetime(1970, 1, 1, tzinfo=UTC)
    fromisoformat = datetime.fromisoformat
except ImportError:
    # python 2.7 and older than 3.11 - use naive dates
    dt_min = datetime.min
    dt_epoch = datetime(1970, 1, 1)
    def fromisoformat(date):
        # only for naive dates
        return datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%f")        

from roundup.anypy.datetime_ import utcnow


class RateLimit:  # pylint: disable=too-few-public-methods
    def __init__(self, count, period):
        self.count = count
        self.period = period

    @property
    def inverse(self):
        return self.period.total_seconds() / self.count


class Gcra:
    def __init__(self):
        self.memory = {}

    def get_tat(self, key):
        # This should return a previous tat for the key or the current time.
        if key in self.memory:
            return self.memory[key]
        else:
            return dt_min

    def set_tat(self, key, tat):
        self.memory[key] = tat

    def get_tat_as_string(self, key):
        # get value as string:
        #  YYYY-MM-DDTHH:MM:SS.mmmmmm
        # to allow it to be marshalled/unmarshaled
        if key in self.memory:
            return self.memory[key].isoformat()
        else:
            return dt_min.isoformat()

    def set_tat_as_string(self, key, tat):
        # Take value as string and unmarshall:
        #  YYYY-MM-DDTHH:MM:SS.mmmmmm
        # to datetime
        self.memory[key] = fromisoformat(tat)

    def update(self, key, limit, testonly=False):
        '''Determine if the item associated with the key should be
           rejected given the RateLimit limit.
        '''
        now = utcnow()
        tat = max(self.get_tat(key), now)
        separation = (tat - now).total_seconds()
        max_interval = limit.period.total_seconds() - limit.inverse
        if separation > max_interval:
            reject = True
        else:
            reject = False
            if not testonly:
                new_tat = max(tat, now) + timedelta(seconds=limit.inverse)
                self.set_tat(key, new_tat)
        return reject

    def status(self, key, limit):
        '''Return status suitable for displaying as headers:
            X-RateLimit-Limit: calls allowed per period. Period/window
                is not specified in any api I found.
            X-RateLimit-Limit-Period: Non standard. Defines period in
                seconds for RateLimit-Limit.
            X-RateLimit-Remaining: How many calls are left in this window.
            X-RateLimit-Reset: window ends in this many seconds (not an
                 epoch timestamp) and all RateLimit-Limit calls are
                 available again.
            Retry-After: if user's request fails, this is the next time there
                 will be at least 1 available call to be consumed.
        '''

        ret = {}
        tat = self.get_tat(key)
        # static defined headers according to limit
        # all values are strings as that is required when used as headers
        ret['X-RateLimit-Limit'] = str(limit.count)
        ret['X-RateLimit-Limit-Period'] = str(
                                           int(
                                            limit.period.total_seconds())
                                          )

        # status of current limit as of now
        now = utcnow()

        current_count = int((limit.period - (tat - now)).total_seconds() /
                            limit.inverse)
        ret['X-RateLimit-Remaining'] = str(min(current_count, limit.count))

        # tat_in_epochsec = (tat - datetime(1970, 1, 1)).total_seconds()
        seconds_to_tat = (tat - now).total_seconds()
        ret['X-RateLimit-Reset'] = str(max(seconds_to_tat, 0))
        ret['X-RateLimit-Reset-date'] = "%s" % tat
        ret['Now'] = str((now - dt_epoch).total_seconds())
        ret['Now-date'] = "%s" % now

        if self.update(key, limit, testonly=True):
            # A new request would be rejected if it was processes.
            # The user has to wait until an item is dequeued.
            # One item is dequeued every limit.inverse seconds.
            ret['Retry-After'] = str(int(limit.inverse))
            ret['Retry-After-Timestamp'] = "%s" % \
                    (now + timedelta(seconds=limit.inverse))  # noqa: E127
        else:
            # if we are not rejected, the user can post another
            # attempt immediately.
            # Do we even need this header if not rejected?
            # RFC implies this is used with a 503 (or presumably
            # 429 which may postdate the rfc). So if no error, no header?
            # ret['Retry-After'] = '0'
            # ret['Retry-After-Timestamp'] = str(ret['Now-date'])
            pass

        return ret

Roundup Issue Tracker: http://roundup-tracker.org/