Trigger the Globus Search Rate Limit

trigger-globus-search-rate-limit.py

#!/usr/bin/env python3
"""
Intentionally trigger rate limiting within Globus Search.

This script runs a request loop, up to a maximum of N, until a rate limiting
error is reached. If given an access token, it will make authenticated calls.
Without an access token, it makes unauthenticated calls.
"""
import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request

ENV = os.getenv("GLOBUS_SDK_ENVIRONMENT")
if ENV == "preview":
    URL = "https://search.api.preview.globus.org/"
elif ENV and ENV != "production":
    URL = f"https://search.api.{ENV}.globuscs.info/"
else:
    URL = "https://search.api.globus.org/"


def run(n, delay, token, abort_on_rate_limit):
    for i in range(1, n + 1):
        req = urllib.request.Request(URL)
        if token:
            req.add_header("Authorization", f"Bearer {token}")

        try:
            response = urllib.request.urlopen(req)  # nosec
        except urllib.error.URLError as e:
            body = json.loads(e.read())
            print(
                f"n={i} status={e.code} code={body.get('code')} "
                f'message="{body.get("message")}"'
            )
            if e.code == 429:
                if abort_on_rate_limit:
                    sys.exit(0)
        else:
            body = json.loads(response.read())
            print(
                f"n={i} status={response.code} "
                f'description="{body.get("description")}"'
            )
        if delay:
            time.sleep(delay)


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "-t",
        "--token",
        help="An access token for triggering rate limiting as an authenticated user.",
    )
    parser.add_argument(
        "-N",
        help="Number of requests to make to try to trigger an error. default=100",
        default=100,
        type=int,
    )
    parser.add_argument(
        "--no-stop",
        help="Don't stop when a rate limit error is triggered; keep making requests",
        default=False,
        action="store_true",
    )
    parser.add_argument(
        "--delay",
        help="A fixed delay between calls, in milliseconds. Defaults to 0.",
        default=0,
        type=int,
    )
    args = parser.parse_args()
    run(args.N, args.delay / 1000 if args.delay else None, args.token, not args.no_stop)


if __name__ == "__main__":
    main()