Modernize codebase to be fully async (#77)

* Modernize codebase to be fully async

* Make API poll timeout configurable

* Default to 5 seconds on API polling requests

5 seconds is also httpx's default.

Author: goncalossilva

requirements-dev.in

@@ -9,11 +9,4 @@ pytest
 pytest-asyncio
 pytest-mock
 reorder-python-imports
-httpx
 certifi
-
-# lock the version because `starlette`(from requirements.in) explicitly depends on it
-# but httpx tries to fetch the latest version causing conflict between
-# requirements.txt and requirements-dev.txt
-anyio==3.7.1
-

requirements-dev.txt

@@ -4,21 +4,14 @@
 #
 #    pip-compile requirements-dev.in
 #
-anyio==3.7.1
-    # via
-    #   -r requirements-dev.in
-    #   httpcore
-astroid==2.15.6
+astroid==2.15.8
     # via pylint
 autopep8==2.0.4
     # via -r requirements-dev.in
 black==23.9.1
     # via -r requirements-dev.in
 certifi==2023.7.22
-    # via
-    #   -r requirements-dev.in
-    #   httpcore
-    #   httpx
+    # via -r requirements-dev.in
 cfgv==3.4.0
     # via pre-commit
 classify-imports==4.2.0
@@ -33,18 +26,8 @@ filelock==3.12.4
     # via virtualenv
 flake8==6.1.0
     # via -r requirements-dev.in
-h11==0.14.0
-    # via httpcore
-httpcore==0.18.0
-    # via httpx
-httpx==0.25.0
-    # via -r requirements-dev.in
-identify==2.5.28
+identify==2.5.29
     # via pre-commit
-idna==3.4
-    # via
-    #   anyio
-    #   httpx
 iniconfig==2.0.0
     # via pytest
 isort==5.12.0
@@ -84,7 +67,7 @@ pycodestyle==2.11.0
     #   flake8
 pyflakes==3.1.0
     # via flake8
-pylint==2.17.5
+pylint==2.17.6
     # via -r requirements-dev.in
 pytest==7.4.2
     # via
@@ -97,13 +80,8 @@ pytest-mock==3.11.1
     # via -r requirements-dev.in
 pyyaml==6.0.1
     # via pre-commit
-reorder-python-imports==3.10.0
+reorder-python-imports==3.11.0
     # via -r requirements-dev.in
-sniffio==1.3.0
-    # via
-    #   anyio
-    #   httpcore
-    #   httpx
 tomlkit==0.12.1
     # via pylint
 virtualenv==20.24.5

requirements.in

@@ -1,14 +1,13 @@
 fastapi
 uvicorn
-requests
 marshmallow
 flagsmith-flag-engine
 python-decouple
 python-dotenv
 pydantic
 orjson
+httpx
 # sse-stuff
 sse-starlette
 asyncio
 redis
-

requirements.txt

@@ -2,37 +2,44 @@
 # This file is autogenerated by pip-compile with Python 3.11
 # by the following command:
 #
-#    pip-compile
+#    pip-compile requirements.in
 #
 anyio==3.7.1
     # via
     #   fastapi
+    #   httpcore
     #   starlette
 asyncio==3.4.3
     # via -r requirements.in
 certifi==2023.7.22
-    # via requests
-charset-normalizer==3.2.0
-    # via requests
+    # via
+    #   httpcore
+    #   httpx
 click==8.1.7
     # via uvicorn
-fastapi==0.103.1
+fastapi==0.103.2
     # via -r requirements.in
-flagsmith-flag-engine==4.0.4
+flagsmith-flag-engine==4.1.0
     # via -r requirements.in
 h11==0.14.0
-    # via uvicorn
+    # via
+    #   httpcore
+    #   uvicorn
+httpcore==0.18.0
+    # via httpx
+httpx==0.25.0
+    # via -r requirements.in
 idna==3.4
     # via
     #   anyio
-    #   requests
+    #   httpx
 marshmallow==3.20.1
     # via -r requirements.in
 orjson==3.9.7
     # via -r requirements.in
 packaging==23.1
     # via marshmallow
-pydantic==1.10.12
+pydantic==1.10.13
     # via
     #   -r requirements.in
     #   fastapi
@@ -44,26 +51,25 @@ python-decouple==3.8
     # via -r requirements.in
 python-dotenv==1.0.0
     # via -r requirements.in
-redis==4.6.0
-    # via -r requirements.in
-requests==2.31.0
+redis==5.0.1
     # via -r requirements.in
 semver==2.13.0
     # via flagsmith-flag-engine
 sniffio==1.3.0
-    # via anyio
+    # via
+    #   anyio
+    #   httpcore
+    #   httpx
 sse-starlette==1.6.5
     # via -r requirements.in
 starlette==0.27.0
     # via
     #   fastapi
     #   sse-starlette
-typing-extensions==4.7.1
+typing-extensions==4.8.0
     # via
     #   fastapi
     #   pydantic
     #   pydantic-collections
-urllib3==2.0.6
-    # via requests
 uvicorn==0.23.2
     # via -r requirements.in

src/cache.py

@@ -1,8 +1,8 @@
 import logging
 from datetime import datetime
 
+import httpx
 import orjson
-import requests
 
 from .exceptions import FlagsmithUnknownKeyError
 from .settings import Settings
@@ -14,25 +14,25 @@ class CacheService:
     def __init__(self, settings: Settings):
         self.settings = settings
         self.last_updated_at = None
-        self._session = requests.Session()
         self._cache = {}
+        self._client = httpx.AsyncClient(timeout=settings.api_poll_timeout)
 
-    def fetch_document(self, server_side_key):
-        url = f"{self.settings.api_url}/environment-document/"
-        response = self._session.get(
-            url, headers={"X-Environment-Key": server_side_key}
+    async def fetch_document(self, server_side_key):
+        response = await self._client.get(
+            url=f"{self.settings.api_url}/environment-document/",
+            headers={"X-Environment-Key": server_side_key},
         )
         response.raise_for_status()
         return orjson.loads(response.text)
 
-    def refresh(self):
+    async def refresh(self):
         received_error = False
         for key_pair in self.settings.environment_key_pairs:
             try:
-                self._cache[key_pair.client_side_key] = self.fetch_document(
+                self._cache[key_pair.client_side_key] = await self.fetch_document(
                     key_pair.server_side_key
                 )
-            except (requests.exceptions.HTTPError, orjson.JSONDecodeError):
+            except (httpx.HTTPError, orjson.JSONDecodeError):
                 received_error = True
                 logger.exception(
                     f"Failed to fetch document for {key_pair.client_side_key}"

src/main.py

@@ -1,3 +1,4 @@
+import logging
 from contextlib import suppress
 from datetime import datetime
 
@@ -44,7 +45,7 @@ async def unknown_key_error(request, exc):
 
 @app.get("/health", response_class=ORJSONResponse, deprecated=True)
 @app.get("/proxy/health", response_class=ORJSONResponse)
-def health_check():
+async def health_check():
     with suppress(TypeError):
         last_updated = datetime.now() - cache_service.last_updated_at
         buffer = 30 * len(settings.environment_key_pairs)  # 30s per environment
@@ -55,7 +56,7 @@ def health_check():
 
 
 @app.get("/api/v1/flags/", response_class=ORJSONResponse)
-def flags(feature: str = None, x_environment_key: str = Header(None)):
+async def flags(feature: str = None, x_environment_key: str = Header(None)):
     environment_document = cache_service.get_environment(x_environment_key)
     environment = build_environment_model(environment_document)
 
@@ -87,7 +88,7 @@ def flags(feature: str = None, x_environment_key: str = Header(None)):
 
 
 @app.post("/api/v1/identities/", response_class=ORJSONResponse)
-def identity(
+async def identity(
     input_data: IdentityWithTraits,
     x_environment_key: str = Header(None),
 ):
@@ -116,9 +117,13 @@ def identity(
 
 
 @app.on_event("startup")
-@repeat_every(seconds=settings.api_poll_frequency, raise_exceptions=True)
-def refresh_cache():
-    cache_service.refresh()
+@repeat_every(
+    seconds=settings.api_poll_frequency,
+    raise_exceptions=True,
+    logger=logging.getLogger(__name__),
+)
+async def refresh_cache():
+    await cache_service.refresh()
 
 
 app.add_middleware(

src/settings.py

@@ -24,7 +24,8 @@ class EnvironmentKeyPair(BaseModel):
 class Settings(BaseSettings):
     environment_key_pairs: List[EnvironmentKeyPair]
     api_url: HttpUrl = "https://edge.api.flagsmith.com/api/v1"
-    api_poll_frequency: int = 10
+    api_poll_frequency: int = 10  # minutes
+    api_poll_timeout: int = 5  # seconds
 
     # sse settings
     stream_delay: int = 1  # seconds

tests/test_cache.py

@@ -1,7 +1,7 @@
 import unittest.mock
 
+import httpx
 import pytest
-import requests
 
 from src.cache import CacheService
 from src.exceptions import FlagsmithUnknownKeyError
@@ -16,62 +16,64 @@
 )
 
 
-def test_refresh_makes_correct_http_call(mocker):
+@pytest.mark.asyncio
+async def test_refresh_makes_correct_http_call(mocker):
     # Given
-    mocked_get = mocker.patch("src.cache.requests.Session.get")
+    mocked_get = mocker.patch("src.cache.httpx.AsyncClient.get")
     mocked_get.side_effect = [
-        unittest.mock.AsyncMock(text='{"key1": "value1"}'),
-        unittest.mock.AsyncMock(text='{"key2": "value2"}'),
+        unittest.mock.Mock(text='{"key1": "value1"}'),
+        unittest.mock.Mock(text='{"key2": "value2"}'),
     ]
     mocked_datetime = mocker.patch("src.cache.datetime")
+
     cache_service = CacheService(settings)
 
     # When
-    cache_service.refresh()
+    await cache_service.refresh()
 
     # Then
     mocked_get.assert_has_calls(
         [
             mocker.call(
-                f"{settings.api_url}/environment-document/",
+                url=f"{settings.api_url}/environment-document/",
                 headers={
                     "X-Environment-Key": settings.environment_key_pairs[
                         0
                     ].server_side_key
                 },
-            )
-        ],
-        [
+            ),
             mocker.call(
-                f"{settings.api_url}/environment-document/",
+                url=f"{settings.api_url}/environment-document/",
                 headers={
                     "X-Environment-Key": settings.environment_key_pairs[
                         1
                     ].server_side_key
                 },
-            )
-        ],
+            ),
+        ]
     )
     assert cache_service.last_updated_at == mocked_datetime.now.return_value
 
 
-def test_refresh_does_not_update_last_updated_at_if_any_request_fails(mocker):
+@pytest.mark.asyncio
+async def test_refresh_does_not_update_last_updated_at_if_any_request_fails(mocker):
     # Given
-    mocked_session = mocker.patch("src.cache.requests.Session")
-    mocked_session.return_value.get.side_effect = [
-        mocker.MagicMock(),
-        requests.exceptions.HTTPError(),
+    mocked_get = mocker.patch("src.cache.httpx.AsyncClient.get")
+    mocked_get.side_effect = [
+        httpx.ConnectTimeout("timeout"),
+        unittest.mock.Mock(text='{"key2": "value2"}'),
     ]
     cache_service = CacheService(settings)
 
     # When
-    cache_service.refresh()
+    await cache_service.refresh()
 
     # Then
     assert cache_service.last_updated_at is None
 
 
-def test_get_environment_works_correctly(mocker):
+@pytest.mark.asyncio
+async def test_get_environment_works_correctly(mocker):
     # Given
     cache_service = CacheService(settings)
     doc_1 = {"key1": "value1"}
@@ -83,7 +85,7 @@ def test_get_environment_works_correctly(mocker):
     )
 
     # When
-    cache_service.refresh()
+    await cache_service.refresh()
 
     # Next, test that get environment return correct document
     assert (