feat: add api keys rotating

aioetherscan/client.py

@@ -1,5 +1,5 @@
 from asyncio import AbstractEventLoop
-from typing import AsyncContextManager
+from typing import AsyncContextManager, Union
 
 from aiohttp import ClientTimeout
 from aiohttp_retry import RetryOptionsBase
@@ -20,7 +20,7 @@
 class Client:
     def __init__(
         self,
-        api_key: str,
+        api_key: Union[str, list[str]],
         api_kind: str = 'eth',
         network: str = 'main',
         loop: AbstractEventLoop = None,
@@ -29,7 +29,8 @@ def __init__(
         throttler: AsyncContextManager = None,
         retry_options: RetryOptionsBase = None,
     ) -> None:
-        self._url_builder = UrlBuilder(api_key, api_kind, network)
+        api_keys = [api_key] if isinstance(api_key, str) else api_key
+        self._url_builder = UrlBuilder(api_keys, api_kind, network)
         self._http = Network(self._url_builder, loop, timeout, proxy, throttler, retry_options)
 
         self.account = Account(self)

aioetherscan/exceptions.py

@@ -20,6 +20,10 @@ def __str__(self):
         return f'[{self.message}] {self.result}'
 
 
+class EtherscanClientApiRateLimitError(EtherscanClientApiError):
+    pass
+
+
 class EtherscanClientProxyError(EtherscanClientError):
     """JSON-RPC 2.0 Specification
 

aioetherscan/modules/extra/generators/generator_utils.py

@@ -84,7 +84,7 @@ async def _parse_by_pages(
         api_method: Callable, request_params: dict[str, Any]
     ) -> AsyncIterator[Transfer]:
         page = count(1)
-        while True:
+        while True:  # pragma: no cover
             request_params['page'] = next(page)
             try:
                 result = await api_method(**request_params)

aioetherscan/network.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 from asyncio import AbstractEventLoop
+from functools import wraps
 from typing import Union, AsyncContextManager, Optional
 
 import aiohttp
@@ -15,10 +16,30 @@
     EtherscanClientError,
     EtherscanClientApiError,
     EtherscanClientProxyError,
+    EtherscanClientApiRateLimitError,
 )
 from aioetherscan.url_builder import UrlBuilder
 
 
+def retry_limit_attempt(f):
+    @wraps(f)
+    async def inner(self, *args, **kwargs):
+        attempt = 1
+        max_attempts = self._url_builder.keys_count
+        while True:  # pragma: no cover
+            try:
+                return await f(self, *args, **kwargs)
+            except EtherscanClientApiRateLimitError as e:
+                self._logger.warning(f'Key daily limit exceeded, {attempt=}: {e}')
+                attempt += 1
+                if attempt > max_attempts:
+                    raise e
+                await asyncio.sleep(0.01)
+                self._url_builder.rotate_api_key()
+
+    return inner
+
+
 class Network:
     def __init__(
         self,
@@ -48,9 +69,11 @@ async def close(self):
         if self._retry_client is not None:
             await self._retry_client.close()
 
+    @retry_limit_attempt
     async def get(self, params: dict = None) -> Union[dict, list, str]:
         return await self._request(METH_GET, params=self._url_builder.filter_and_sign(params))
 
+    @retry_limit_attempt
     async def post(self, data: dict = None) -> Union[dict, list, str]:
         return await self._request(METH_POST, data=self._url_builder.filter_and_sign(data))
 
@@ -68,6 +91,7 @@ async def _request(
         if self._retry_client is None:
             self._retry_client = self._get_retry_client()
         session_method = getattr(self._retry_client, method.lower())
+
         async with self._throttler:
             async with session_method(
                 self._url_builder.API_URL, params=params, data=data, proxy=self._proxy
@@ -93,6 +117,10 @@ async def _handle_response(self, response: aiohttp.ClientResponse) -> Union[dict
     def _raise_if_error(response_json: dict):
         if 'status' in response_json and response_json['status'] != '1':
             message, result = response_json.get('message'), response_json.get('result')
+
+            if 'max daily rate limit reached' in result.lower():
+                raise EtherscanClientApiRateLimitError(message, result)
+
             raise EtherscanClientApiError(message, result)
 
         if 'error' in response_json:

aioetherscan/url_builder.py

@@ -1,3 +1,5 @@
+import logging
+from itertools import cycle
 from typing import Optional
 from urllib.parse import urlunsplit, urljoin
 
@@ -19,15 +21,19 @@ class UrlBuilder:
     BASE_URL: str = None
     API_URL: str = None
 
-    def __init__(self, api_key: str, api_kind: str, network: str) -> None:
-        self._API_KEY = api_key
+    def __init__(self, api_keys: list[str], api_kind: str, network: str) -> None:
+        self._api_keys = api_keys
+        self._api_keys_cycle = cycle(self._api_keys)
+        self._api_key = self._get_next_api_key()
 
         self._set_api_kind(api_kind)
         self._network = network.lower().strip()
 
         self.API_URL = self._get_api_url()
         self.BASE_URL = self._get_base_url()
 
+        self._logger = logging.getLogger(__name__)
+
     def _set_api_kind(self, api_kind: str) -> None:
         api_kind = api_kind.lower().strip()
         if api_kind not in self._API_KINDS:
@@ -87,9 +93,35 @@ def filter_and_sign(self, params: dict):
     def _sign(self, params: dict) -> dict:
         if not params:
             params = {}
-        params['apikey'] = self._API_KEY
+        params['apikey'] = self._api_key
         return params
 
     @staticmethod
     def _filter_params(params: dict) -> dict:
         return {k: v for k, v in params.items() if v is not None}
+
+    def _get_next_api_key(self) -> str:
+        return next(self._api_keys_cycle)
+
+    def rotate_api_key(self) -> None:
+        prev_api_key = self._api_key
+        next_api_key = self._get_next_api_key()
+
+        self._logger.info(
+            f'Rotating API key from {self._mask_api_key(prev_api_key)!r} to {self._mask_api_key(next_api_key)!r}'
+        )
+
+        self._api_key = next_api_key
+
+    @staticmethod
+    def _mask_api_key(api_key: str, masked_chars_count: int = 30, symbol: str = '*') -> str:
+        api_key_len = len(api_key)
+        if masked_chars_count >= api_key_len or masked_chars_count <= 0:
+            return symbol * len(api_key)
+
+        right_part = api_key[-(api_key_len - masked_chars_count) :]
+        return right_part.rjust(len(api_key), symbol)
+
+    @property
+    def keys_count(self) -> int:
+        return len(self._api_keys)

tests/test_network.py

@@ -1,7 +1,7 @@
 import asyncio
 import json
 import logging
-from unittest.mock import patch, AsyncMock, MagicMock, Mock
+from unittest.mock import patch, AsyncMock, MagicMock, Mock, call
 
 import aiohttp
 import aiohttp_retry
@@ -17,8 +17,9 @@
     EtherscanClientError,
     EtherscanClientApiError,
     EtherscanClientProxyError,
+    EtherscanClientApiRateLimitError,
 )
-from aioetherscan.network import Network
+from aioetherscan.network import Network, retry_limit_attempt
 from aioetherscan.url_builder import UrlBuilder
 
 
@@ -45,7 +46,7 @@ def get_loop():
 
 @pytest_asyncio.fixture
 async def ub():
-    ub = UrlBuilder('test_api_key', 'eth', 'main')
+    ub = UrlBuilder(['test_api_key'], 'eth', 'main')
     yield ub
 
 
@@ -86,25 +87,25 @@ def test_no_loop(ub):
 async def test_get(nw):
     with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock:
         await nw.get()
-        mock.assert_called_once_with(METH_GET, params={'apikey': nw._url_builder._API_KEY})
+        mock.assert_called_once_with(METH_GET, params={'apikey': nw._url_builder._api_key})
 
 
 @pytest.mark.asyncio
 async def test_post(nw):
     with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock:
         await nw.post()
-        mock.assert_called_once_with(METH_POST, data={'apikey': nw._url_builder._API_KEY})
+        mock.assert_called_once_with(METH_POST, data={'apikey': nw._url_builder._api_key})
 
     with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock:
         await nw.post({'some': 'data'})
         mock.assert_called_once_with(
-            METH_POST, data={'apikey': nw._url_builder._API_KEY, 'some': 'data'}
+            METH_POST, data={'apikey': nw._url_builder._api_key, 'some': 'data'}
         )
 
     with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock:
         await nw.post({'some': 'data', 'null': None})
         mock.assert_called_once_with(
-            METH_POST, data={'apikey': nw._url_builder._API_KEY, 'some': 'data'}
+            METH_POST, data={'apikey': nw._url_builder._api_key, 'some': 'data'}
         )
 
 
@@ -238,3 +239,77 @@ def test_get_retry_client(nw):
             retry_options=nw._retry_options,
         )
         assert result is m.return_value
+
+
+def test_raise_if_error_daily_limit_reached(nw):
+    data = dict(
+        status='0',
+        message='NOTOK',
+        result='Max daily rate limit reached. 110000 (100%) of 100000 day/limit',
+    )
+    with pytest.raises(EtherscanClientApiRateLimitError) as e:
+        nw._raise_if_error(data)
+
+    assert e.value.message == data['message']
+    assert e.value.result == data['result']
+
+
+class TestRetryClass:
+    def __init__(self, limit: int, keys_count: int):
+        self._url_builder = Mock()
+        self._url_builder.keys_count = keys_count
+        self._url_builder.rotate_api_key = Mock()
+
+        self._logger = Mock()
+        self._logger.warning = Mock()
+
+        self._count = 1
+        self._limit = limit
+
+    @retry_limit_attempt
+    async def some_method(self):
+        self._count += 1
+
+        if self._count > self._limit:
+            raise EtherscanClientApiRateLimitError(
+                'NOTOK',
+                'Max daily rate limit reached. 110000 (100%) of 100000 day/limit',
+            )
+
+
+@pytest.mark.asyncio
+async def test_retry_limit_attempt_error_limit_exceeded(nw):
+    c = TestRetryClass(1, 1)
+
+    with pytest.raises(EtherscanClientApiRateLimitError):
+        await c.some_method()
+    c._url_builder.rotate_api_key.assert_not_called()
+    c._logger.warning.assert_called_once_with(
+        'Key daily limit exceeded, attempt=1: [NOTOK] Max daily rate limit reached. 110000 (100%) of 100000 day/limit'
+    )
+
+
+@pytest.mark.asyncio
+async def test_retry_limit_attempt_error_limit_rotate(nw):
+    c = TestRetryClass(1, 2)
+
+    with pytest.raises(EtherscanClientApiRateLimitError):
+        await c.some_method()
+    c._url_builder.rotate_api_key.assert_called_once()
+    c._logger.warning.assert_has_calls(
+        [
+            call(
+                'Key daily limit exceeded, attempt=1: [NOTOK] Max daily rate limit reached. 110000 (100%) of 100000 day/limit'
+            ),
+            call(
+                'Key daily limit exceeded, attempt=2: [NOTOK] Max daily rate limit reached. 110000 (100%) of 100000 day/limit'
+            ),
+        ]
+    )
+
+
+@pytest.mark.asyncio
+async def test_retry_limit_attempt_ok(nw):
+    c = TestRetryClass(2, 1)
+
+    await c.some_method()

tests/test_url_builder.py

@@ -1,4 +1,5 @@
-from unittest.mock import patch
+from itertools import cycle
+from unittest.mock import patch, Mock
 
 import pytest
 import pytest_asyncio
@@ -7,7 +8,7 @@
 
 
 def apikey():
-    return 'test_api_key'
+    return ['test_api_key']
 
 
 @pytest_asyncio.fixture
@@ -17,8 +18,8 @@ async def ub():
 
 
 def test_sign(ub):
-    assert ub._sign({}) == {'apikey': ub._API_KEY}
-    assert ub._sign({'something': 'something'}) == {'something': 'something', 'apikey': ub._API_KEY}
+    assert ub._sign({}) == {'apikey': ub._api_key}
+    assert ub._sign({'something': 'something'}) == {'something': 'something', 'apikey': ub._api_key}
 
 
 def test_filter_params(ub):
@@ -130,3 +131,42 @@ def test_get_link(ub):
         path = 'some_path'
         ub.get_link(path)
         join_mock.assert_called_once_with(ub.BASE_URL, path)
+
+
+@pytest.mark.parametrize(
+    'api_key,masked_chars_count,expected',
+    [
+        ('abc', 4, '***'),
+        ('abc', 2, '**c'),
+        ('abcd', 4, '****'),
+        ('abcdef', 4, '****ef'),
+    ],
+)
+def test_mask_api_key(ub, api_key, masked_chars_count, expected):
+    assert ub._mask_api_key(api_key, masked_chars_count) == expected
+
+    assert ub._mask_api_key('qwe', 2, '#') == '##e'
+
+
+def test_keys_count(ub):
+    ub._api_keys = [1, 2]
+    assert ub.keys_count == 2
+
+
+def test_rotate_api_key(ub):
+    api_keys = ['one', 'two', 'three']
+    first, second, _ = api_keys
+
+    ub._api_keys = api_keys
+    ub._api_keys_cycle = cycle(ub._api_keys)
+    ub._api_key = ub._get_next_api_key()
+
+    assert ub._api_key == first
+
+    ub._logger = Mock()
+    ub._logger.info = Mock()
+
+    ub.rotate_api_key()
+
+    ub._logger.info.assert_called_once()
+    assert ub._api_key == second