redispython
Ben Gorman

Ben Gorman

Life's a garden. Dig it.

Challenge

You run a company called Ticket King 👑 that sells tickets to concerts and events. You have an API that lets people check ticket prices, but now you need to lock it down to prevent your servers from being overloaded.

Your company offers three plans:

  1. Peasant: 10 API calls per day
  2. Noble: 20 API calls per day
  3. Royal: 30 API calls per day

Quotas are tracked per user_id, and are reset at the start of each UTC day. Implement this fixed-window api rate limiter.

Starter Code

def get_plan_limit(plan):
    """
    Identify the daily API limit for this plan
 
    :param plan: The plan name
    :return: The daily limit. 0 if plan is not recognized
    """
 
    plan_limits = {
        'peasant': 10,
        'noble': 20,
        'royal': 30,
    }
 
    return plan_limits.get(plan, 0)
 
 
def get_user_plan(user_id: int):
    """
    Fetch this user's plan from the SQL database.
    (This function is slow and expensive..)
    """
 
    users = {
        22912157: 'peasant',
        64792475: 'noble',
        56488868: 'royal',
        92899704: 'noble',
        73532154: 'peasant',
        68472103: 'peasant'
    }
 
    plan = users.get(user_id)
    return plan
 
 
def allow_request(user_id: int):
    """
    Check if this user is allowed to make a request right now.
    :return: True or False
    """
 
    pass

Solution

import redis
import datetime
 
 
def get_plan_limit(plan):
    """
    Identify the daily API limit for this plan
 
    :param plan: The plan name
    :return: The daily limit. 0 if plan is not recognized
    """
 
    plan_limits = {
        "peasant": 10,
        "noble": 20,
        "royal": 30,
    }
 
    return plan_limits.get(plan, 0)
 
 
def get_user_plan(user_id: int):
    """
    Fetch this user's plan from the SQL database.
    (This function is slow and expensive..)
    """
 
    users = {
        22912157: "peasant",
        64792475: "noble",
        56488868: "royal",
        92899704: "noble",
        73532154: "peasant",
        68472103: "peasant",
    }
 
    plan = users.get(user_id)
    return plan
 
 
def allow_request(user_id: int):
    """
    Check if this user is allowed to make a request right now.
    :return: True or False
    """
 
    # Connect to redis
    r = redis.Redis(host="localhost", port=6379, db=0, decode_responses=True)
 
    # Define redis keys
    current_date = datetime.datetime.now(tz=datetime.timezone.utc).strftime("%Y-%m-%d")
    key_limit = f"limit:{user_id}:{current_date}"
    key_usage = f"usage:{user_id}:{current_date}"
 
    # Determine the user's limit
    # If it's not already stored on redis, fetch and store it.
    if (limit := r.get(name=key_limit)) is None:
        plan = get_user_plan(user_id)
        limit = get_plan_limit(plan)
        r.set(name=key_limit, value=limit)
    else:
        limit = int(limit)
 
    # Determine the user's usage (before this api call)
    if (usage := r.get(name=key_usage)) is None:
        usage = 0
    else:
        usage = int(usage)
 
    # If usage is below the limit, update its value on redis
    if usage < limit:
        r.set(name=key_usage, value=usage + 1)
 
    # Return True if usage is below limit, False otherwise
    return usage < limit

Tests

# user 73532154: peasant plan. Should be allowed 10 requests per day.
for i in range(11):
    print(i, allow_request(user_id=73532154))
 
# 0 True
# 1 True
# 2 True
# ...
# 8 True
# 9 True
# 10 False
# user 92899704: noble plan. Should be allowed 20 requests per day.
for i in range(21):
    print(i, allow_request(user_id=92899704))
 
# 0 True
# 1 True
# 2 True
# ...
# 18 True
# 19 True
# 20 False
# user 56488868: Non-existent user. Should be allowed 30 requests per day.
for i in range(31):
    print(i, allow_request(user_id=56488868))
 
# 0 True
# 1 True
# 2 True
# ...
# 28 True
# 29 True
# 30 False
# user 123: Non-existent user. Should be allowed 0 requests per day.
for i in range(1):
    print(i, allow_request(user_id=123))
 
# 0 False

How to clear the database

You might want to flush your database before running the tests above. You can do so with flushdb().

r = redis.Redis(host="localhost", port=6379, db=0, decode_responses=True)
r.flushdb()

Explanation

  1. The first step is to connect to Redis.

    import redis
     
    # Connect to redis
    r = redis.Redis(host="localhost", port=6379, db=0, decode_responses=True) 
    type(r)  # redis.client.Redis
  2. Define redis keys. We'll keep track of two key-values per user:

    • key_limit:
      Stores the limit for a particular (user, date).
      Example: (key = limit:123456:2022-05-15, value = 20)
    • key_usage:
      Stores the number of API calls made for a particular (user, date).
      Example: (key = usage:123456:2022-05-15, value = 12)
    import datetime
     
    # Define redis keys
    current_date = datetime.datetime.now(tz=datetime.timezone.utc).strftime("%Y-%m-%d")
    key_limit = f"limit:{user_id}:{current_date}"
    key_usage = f"usage:{user_id}:{current_date}"
  3. Determine the user's limit (for today).

    # Determine the user's limit
    # If it's not already stored on redis, fetch and store it.
    if (limit := r.get(name=key_limit)) is None: 
        plan = get_user_plan(user_id)
        limit = get_plan_limit(plan)
        r.set(name=key_limit, value=limit)
    else:
        limit = int(limit) 
  4. Determine how many API calls this user made prior to the current one.

    # Determine the user's usage (before this api call)
    if (usage := r.get(name=key_usage)) is None:
        usage = 0
    else:
        usage = int(usage)
  5. If the user's usage is below the user's limit, increment usage by one and save the new value on Redis.

    # If usage is below the limit, update its value on redis
    if usage < limit:
        r.set(name=key_usage, value=usage + 1)

    Why don't we always increment usage by 1?

    You certainly could. But it's important to make the distinction between a successful API call and an unsuccessful API call. If the user's usage equals the limit, this call is going to get denied, making it an unsuccessful call. We've made the design decision to store successful API calls only.

    One benefit of this is, if the user upgrades their plan midday, their unsuccessful API calls from before the upgrade will not be counted against their new quota.

  6. Return True or False.

    # Return True if usage is below limit, False otherwise
    return usage < limit