Fix race condition in protected decorator (#14)

* Fix race condition in protected decorator

* Bump project version

* Fix len8 py312 issue in noxfile

* Update changelog

CHANGELOG.md

@@ -1,5 +1,23 @@
 # Changelog
 
+## v0.7.2 (May 2024)
+
+### Fixes
+
+- Fix race condition for concurrent requests using the `protected` decorator.
+
+### Additions
+
+- The `Client` can now be used as an async context manager which starts
+  and closes the client automatically.
+
+### Changes
+
+- The `InvalidKeyHandlerT` and `ExcHandlerT` types no longer include `Optional`,
+  and instead are wrapped in `Optional` in the function signature.
+
+---
+
 ## v0.7.1 (Feb 2024)
 
 ### Fixes

noxfile.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import functools
+import platform
 from typing import Callable
 from pathlib import Path
 
@@ -75,10 +76,15 @@ def types(session: nox.Session) -> None:
 
 
 @nox.session(reuse_venv=True)
-@install("black", "len8")
+@install("black")
 def formatting(session: nox.Session) -> None:
     session.run("black", ".", "--check")
-    session.run("len8")
+
+    major, minor, *_ = platform.python_version_tuple()
+    if major == "3" and int(minor) < 12:
+        # This is a hack but it doesnt support python 3.12
+        session.install(DEPS["len8"])
+        session.run("len8")
 
 
 @nox.session(reuse_venv=True)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unkey.py"
-version = "0.7.1"
+version = "0.7.2"
 description = "An asynchronous Python SDK for unkey.dev."
 authors = ["Jonxslays"]
 license = "GPL-3.0-only"

unkey/__init__.py

@@ -3,7 +3,7 @@
 from typing import Final
 
 __packagename__: Final[str] = "unkey.py"
-__version__: Final[str] = "0.7.1"
+__version__: Final[str] = "0.7.2"
 __author__: Final[str] = "Jonxslays"
 __copyright__: Final[str] = "2023-present Jonxslays"
 __description__: Final[str] = "An asynchronous Python SDK for unkey.dev."

unkey/client.py

@@ -51,6 +51,13 @@ def __init_service(self, service: t.Type[ServiceT]) -> ServiceT:
 
         return service(self._http, self._serializer)  # type: ignore[return-value]
 
+    async def __aenter__(self) -> Client:
+        await self.start()
+        return self
+
+    async def __aexit__(self, *_args: t.Any, **_kwargs: t.Any) -> None:
+        await self.close()
+
     @property
     def keys(self) -> services.KeyService:
         """The key service used to make key related requests."""

unkey/decorators.py

@@ -26,18 +26,18 @@
 functions `*args` and `**kwargs`.
 """
 
-InvalidKeyHandlerT = Optional[Callable[[Dict[str, Any], Optional[models.ApiKeyVerification]], Any]]
+InvalidKeyHandlerT = Callable[[Dict[str, Any], Optional[models.ApiKeyVerification]], Any]
 """The type of a callback used to handle cases where the key was invalid."""
 
-ExcHandlerT = Optional[Callable[[Exception], Any]]
+ExcHandlerT = Callable[[Exception], Any]
 """The type of a callback used to handle exceptions during verification."""
 
 
 def protected(
     api_id: str,
     key_extractor: ExtractorT,
-    on_invalid_key: InvalidKeyHandlerT = None,
-    on_exc: ExcHandlerT = None,
+    on_invalid_key: Optional[InvalidKeyHandlerT] = None,
+    on_exc: Optional[ExcHandlerT] = None,
 ) -> DecoratorT:
     """A framework agnostic second order decorator that is used to protect
     api routes with Unkey key verification.
@@ -82,7 +82,6 @@ def protected(
             altered by `on_invalid_key` if it was passed. If verification
             succeeds the original functions return value is returned.
     """
-    _client = client.Client()
 
     def _on_invalid_key(
         data: Dict[str, Any], verification: Optional[models.ApiKeyVerification] = None
@@ -108,9 +107,8 @@ async def inner(*args: Any, **kwargs: Any) -> VerificationResponseT[T]:
                     message = "Failed to extract API key"
                     return _on_invalid_key({"code": None, "message": message})
 
-                await _client.start()
-                result = await _client.keys.verify_key(key, api_id)
-                await _client.close()
+                async with client.Client() as c:
+                    result = await c.keys.verify_key(key, api_id)
 
                 if result.is_err:
                     err = result.unwrap_err()