diff --git a/requirements.txt b/requirements.txt index e518d261..064d671d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ maggma==0.29.4 requests==2.25.1 monty==2021.6.10 emmet-core==0.4.1 +ratelimit==2.2.1 \ No newline at end of file diff --git a/setup.py b/setup.py index 19028214..bfa5ba64 100644 --- a/setup.py +++ b/setup.py @@ -31,6 +31,7 @@ "monty", "emmet-core", "maggma", + "ratelimit", ], extras_require={ "server": [ diff --git a/src/mp_api/core/client.py b/src/mp_api/core/client.py index 7f3a9cb8..ba09db6f 100644 --- a/src/mp_api/core/client.py +++ b/src/mp_api/core/client.py @@ -23,6 +23,8 @@ from emmet.core.utils import jsanitize from maggma.api.utils import api_sanitize +from mp_api.core.ratelimit import check_limit + try: from pymatgen.core import __version__ as pmg_version # type: ignore except ImportError: # pragma: no cover @@ -148,6 +150,8 @@ def _post_resource( available. """ + check_limit() + payload = jsanitize(body) try: @@ -223,6 +227,8 @@ def _query_resource( available. """ + check_limit() + if criteria: criteria = {k: v for k, v in criteria.items() if v is not None} else: diff --git a/src/mp_api/core/ratelimit.py b/src/mp_api/core/ratelimit.py new file mode 100644 index 00000000..05fcacdf --- /dev/null +++ b/src/mp_api/core/ratelimit.py @@ -0,0 +1,11 @@ +from ratelimit import sleep_and_retry, limits +from mp_api.core.settings import MAPISettings + + +@sleep_and_retry +@limits(calls=MAPISettings().requests_per_min, period=60) +def check_limit(): + """ + Empty function for enabling global rate limiting. + """ + return diff --git a/src/mp_api/core/settings.py b/src/mp_api/core/settings.py index 729354cb..631b861a 100644 --- a/src/mp_api/core/settings.py +++ b/src/mp_api/core/settings.py @@ -22,5 +22,9 @@ class MAPISettings(BaseSettings): db_version: str = Field("2021_04_26", description="Database version") + requests_per_min: int = Field( + 60, description="Number of requests per minute to for rate limit." + ) + class Config: env_prefix = "mapi_"