Add support for refill and update remaining (#9)

* Add support for refill

* Add update remaining

* Update tests

* Bump project version

* Update changelog

* Add py3.12 to ci

.github/workflows/ci.yml

@@ -13,7 +13,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
     runs-on: ${{ matrix.os }}
 
     steps:

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## v0.6.0 (Dec 2023)
+
+### Additions
+
+- Add `Refill`, `RefillInterval`, and `UpdateOp` models/enums.
+- Add `id` property onto `ApiKeyVerification`.
+- Add `refill` property onto `ApiKeyMeta` and `ApiKeyVerification`.
+- Add serialization methods for new properties and models.
+- Add support for `refill` when creating and updating a key.
+- Add `update_remaining` method to `KeyService` and corresponding `Route`.
+
+---
+
 ## v0.5.0 (Dec 2023)
 
 ### Breaking Changes

pyproject.toml

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

tests/test_serializer.py

@@ -243,6 +243,7 @@ def test_to_ratelimit_state(
 
 def _raw_api_key_verification() -> DictT:
     return {
+        "keyId": "key_uuuuuu",
         "valid": False,
         "owner_id": None,
         "meta": None,
@@ -251,6 +252,11 @@ def _raw_api_key_verification() -> DictT:
         "expires": 12345,
         "error": "some error",
         "code": "NOT_FOUND",
+        "refill": {
+            "amount": 100,
+            "interval": "daily",
+            "lastRefilledAt": 12345,
+        },
     }
 
 
@@ -261,6 +267,7 @@ def raw_api_key_verification() -> DictT:
 
 def _full_api_key_verification() -> models.ApiKeyVerification:
     model = models.ApiKeyVerification()
+    model.id = "key_uuuuuu"
     model.valid = False
     model.owner_id = None
     model.meta = None
@@ -269,6 +276,7 @@ def _full_api_key_verification() -> models.ApiKeyVerification:
     model.expires = 12345
     model.error = "some error"
     model.code = models.ErrorCode.NotFound
+    model.refill = models.Refill(100, models.RefillInterval.Daily, 12345)
     return model
 
 
@@ -350,6 +358,11 @@ def _raw_api_key_meta() -> DictT:
             "refillRate": 2,
             "refillInterval": 3,
         },
+        "refill": {
+            "amount": 100,
+            "interval": "daily",
+            "lastRefilledAt": 12345,
+        },
     }
 
 
@@ -369,6 +382,7 @@ def _full_api_key_meta() -> models.ApiKeyMeta:
     model.owner_id = "jonxslays"
     model.created_at = 456
     model.workspace_id = "ws_GGG"
+    model.refill = models.Refill(100, models.RefillInterval.Daily, 12345)
     model.ratelimit = models.Ratelimit(
         models.RatelimitType.Fast,
         limit=1,

unkey/__init__.py

@@ -3,7 +3,7 @@
 from typing import Final
 
 __packagename__: Final[str] = "unkey.py"
-__version__: Final[str] = "0.5.0"
+__version__: Final[str] = "0.6.0"
 __author__: Final[str] = "Jonxslays"
 __copyright__: Final[str] = "2023-present Jonxslays"
 __description__: Final[str] = "An asynchronous Python SDK for unkey.dev."
@@ -63,11 +63,14 @@
     "Ratelimit",
     "RatelimitState",
     "RatelimitType",
+    "Refill",
+    "RefillInterval",
     "Result",
     "Route",
     "Serializer",
     "UndefinedNoneOr",
     "UndefinedOr",
     "UnwrapError",
     "UNDEFINED",
+    "UpdateOp",
 )

unkey/models/__init__.py

@@ -17,4 +17,7 @@
     "Ratelimit",
     "RatelimitState",
     "RatelimitType",
+    "Refill",
+    "RefillInterval",
+    "UpdateOp",
 )

unkey/models/base.py

@@ -52,7 +52,7 @@ def from_str(cls: t.Type[T], value: str) -> T:
             ) from None
 
     @classmethod
-    def from_str_maybe(cls: t.Type[T], value: str) -> t.Optional[T]:
+    def from_str_maybe(cls: t.Type[T], value: t.Optional[str]) -> t.Optional[T]:
         """Attempt to generate this enum from the given value.
 
         Args:

unkey/models/keys.py

@@ -15,6 +15,9 @@
     "Ratelimit",
     "RatelimitState",
     "RatelimitType",
+    "Refill",
+    "RefillInterval",
+    "UpdateOp",
 )
 
 
@@ -23,6 +26,17 @@ class RatelimitType(BaseEnum):
     Consistent = "consistent"
 
 
+class RefillInterval(BaseEnum):
+    Daily = "daily"
+    Monthly = "monthly"
+
+
+class UpdateOp(BaseEnum):
+    Increment = "increment"
+    Decrement = "decrement"
+    Set = "set"
+
+
 @attrs.define(weakref_slot=False)
 class Ratelimit(BaseModel):
     """Data representing a particular ratelimit."""
@@ -92,10 +106,16 @@ class ApiKeyMeta(BaseModel):
     be ignored.
     """
 
+    refill: t.Optional[Refill]
+    """The keys refill state, if any."""
+
 
 @attrs.define(init=False, weakref_slot=False)
 class ApiKeyVerification(BaseModel):
-    """Data about whether this api key is valid."""
+    """Data about whether this api key and its validity."""
+
+    id: t.Optional[str]
+    """The id of this key."""
 
     valid: bool
     """Whether or not this key is valid and passes ratelimit."""
@@ -121,6 +141,9 @@ class ApiKeyVerification(BaseModel):
     ratelimit: t.Optional[RatelimitState]
     """The state of the ratelimit set on this key, if any."""
 
+    refill: t.Optional[Refill]
+    """The keys refill state, if any."""
+
     code: t.Optional[ErrorCode]
     """The optional error code returned by the unkey api."""
 
@@ -140,3 +163,19 @@ class RatelimitState(BaseModel):
 
     reset: int
     """The unix timestamp in milliseconds until the next window."""
+
+
+@attrs.define(weakref_slot=False)
+class Refill(BaseModel):
+    """Data regarding how a key's verifications should be refilled."""
+
+    amount: int
+    """The number of verifications to refill."""
+
+    interval: RefillInterval
+    """The interval at which to refill the verifications."""
+
+    last_refilled_at: t.Optional[int] = None
+    """The UNIX timestamp in milliseconds indicating when the key was
+    las refilled, if it has been.
+    """

unkey/routes.py

@@ -84,6 +84,7 @@ def compile(self, *args: t.Union[str, int]) -> CompiledRoute:
 VERIFY_KEY: t.Final[Route] = Route(c.POST, "/keys.verifyKey")
 REVOKE_KEY: t.Final[Route] = Route(c.POST, "/keys.deleteKey")
 UPDATE_KEY: t.Final[Route] = Route(c.POST, "/keys.updateKey")
+UPDATE_REMAINING: t.Final[Route] = Route(c.POST, "/keys.updateRemaining")
 GET_KEY: t.Final[Route] = Route(c.GET, "/keys.getKey")
 
 # Apis

unkey/serializer.py

@@ -68,8 +68,14 @@ def to_api_key(self, data: DictT) -> models.ApiKey:
 
     def to_api_key_verification(self, data: DictT) -> models.ApiKeyVerification:
         model = models.ApiKeyVerification()
+        model.id = data.get("keyId")
+
         ratelimit = data.get("ratelimit")
         model.ratelimit = self.to_ratelimit_state(ratelimit) if ratelimit else ratelimit
+
+        refill = data.get("refill")
+        model.refill = self.to_refill(refill) if refill else refill
+
         model.code = models.ErrorCode.from_str_maybe(data.get("code", ""))
         self._set_attrs_cased(
             model, data, "valid", "owner_id", "meta", "remaining", "error", "expires", maybe=True
@@ -97,8 +103,13 @@ def to_ratelimit(self, data: DictT) -> models.Ratelimit:
 
     def to_api_key_meta(self, data: DictT) -> models.ApiKeyMeta:
         model = models.ApiKeyMeta()
+
         ratelimit = data.get("ratelimit")
         model.ratelimit = self.to_ratelimit(ratelimit) if ratelimit else ratelimit
+
+        refill = data.get("refill")
+        model.refill = self.to_refill(refill) if refill else refill
+
         self._set_attrs_cased(
             model,
             data,
@@ -122,3 +133,12 @@ def to_api_key_list(self, data: DictT) -> models.ApiKeyList:
         model.total = data["total"]
         model.keys = [self.to_api_key_meta(key) for key in data["keys"]]
         return model
+
+    def to_refill(self, data: DictT) -> models.Refill:
+        interval = models.RefillInterval.from_str(data["interval"])
+        amount = data["amount"]
+
+        model = models.Refill(amount, interval)
+        model.last_refilled_at = data.get("lastRefilledAt")
+
+        return model

unkey/services/apis.py

@@ -37,15 +37,6 @@ async def get_api(self, api_id: str) -> ResultT[models.Api]:
         if isinstance(data, models.HttpResponse):
             return result.Err(data)
 
-        if "error" in data:
-            return result.Err(
-                models.HttpResponse(
-                    404,
-                    data["error"].get("message", "Unknown error"),
-                    models.ErrorCode.from_str_maybe(data["error"].get("code", "UNKNOWN")),
-                )
-            )
-
         return result.Ok(self._serializer.to_api(data))
 
     async def list_keys(
@@ -78,13 +69,4 @@ async def list_keys(
         if isinstance(data, models.HttpResponse):
             return result.Err(data)
 
-        if "error" in data:
-            return result.Err(
-                models.HttpResponse(
-                    404,
-                    data["error"].get("message", "Unknown error"),
-                    models.ErrorCode.from_str_maybe(data["error"].get("code", "UNKNOWN")),
-                )
-            )
-
         return result.Ok(self._serializer.to_api_key_list(data))

unkey/services/http.py

@@ -71,14 +71,15 @@ async def _request(
         if isinstance(data, models.HttpResponse):
             return data
 
-        # Skipping 404's seems hacky but whatever
-        if response.status not in (*self._ok_responses, 404):
-            return models.HttpResponse(
-                response.status,
-                data.get("error")
-                or data.get("message")
-                or "An unexpected error occurred while making the request.",
-            )
+        if response.status not in self._ok_responses:
+            error: t.Union[t.Any, t.Dict[str, t.Any]] = data.get("error")
+            is_dict = isinstance(error, dict)
+
+            message = error.get("message") if is_dict else error
+            code = models.ErrorCode.from_str_maybe(error.get("code") if is_dict else "UNKNOWN")
+            message = message or "An unexpected error occurred while making the request."
+
+            return models.HttpResponse(response.status, str(message), code=code)
 
         return data