Skip to content

Commit

Permalink
feat: add client cached_session
Browse files Browse the repository at this point in the history
Use this cached session in short lived scopes. Prevent calling the API too often when iterating over a list of resource that contains the same objects.
  • Loading branch information
jooola committed Nov 21, 2023
1 parent 155a565 commit afc2d3a
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 1 deletion.
19 changes: 19 additions & 0 deletions examples/use_cached_session.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
from __future__ import annotations

from os import environ

from hcloud import Client

assert (
"HCLOUD_TOKEN" in environ
), "Please export your API token in the HCLOUD_TOKEN environment variable"
token = environ["HCLOUD_TOKEN"]

client = Client(token=token)

with client.cached_session() as session:
# This will query the API only once
for i in range(100):
servers = session.locations.get_all()

print(servers)
45 changes: 44 additions & 1 deletion hcloud/_client.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
from __future__ import annotations

import copy
import time
from typing import NoReturn
from contextlib import contextmanager
from typing import Generator, NoReturn

import requests

Expand Down Expand Up @@ -231,3 +233,44 @@ def request( # type: ignore[no-untyped-def]

# TODO: return an empty dict instead of an empty string when content == "".
return content # type: ignore[return-value]

def session(self, session: requests.Session) -> None:
"""
Configure a custom :class:`Session <requests.Session>` to use when calling the API.
:param session: The session to use when making API requests.
"""
self._requests_session = session

@contextmanager
def cached_session(self) -> Generator[Client, None, None]:
"""
Provide a copy of the :class:`Client <hcloud.Client>` as context manager that
will cache all GET requests.
Cached response will not expire automatically, therefor the cached client must
not be used for long living scopes.
"""
client = copy.deepcopy(self)
client.session(CachedSession())
yield client


class CachedSession(requests.Session):
cache: dict[str, requests.Response] = {}

def send(self, request: requests.PreparedRequest, **kwargs) -> requests.Response: # type: ignore[no-untyped-def]
"""
Send a given PreparedRequest.
"""
if request.method != "GET" or request.url is None:
return super().send(request, **kwargs)

if request.url in self.cache:
return self.cache[request.url]

response = super().send(request, **kwargs)
if response.ok:
self.cache[request.url] = response

return response
24 changes: 24 additions & 0 deletions tests/unit/test_client.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
from __future__ import annotations

import json
from io import BytesIO
from unittest.mock import MagicMock

import pytest
import requests

from hcloud import APIException, Client
from hcloud._client import CachedSession


class TestHetznerClient:
Expand Down Expand Up @@ -182,3 +184,25 @@ def test_request_limit_then_success(self, client, rate_limit_response):
"POST", "http://url.com", params={"argument": "value"}, timeout=2
)
assert client._requests_session.request.call_count == 2


class TestCachedSession:
def test_cache(self):
response = requests.Response()
response.status_code = 200
response.raw = BytesIO(json.dumps({"result": "data"}).encode("utf-8"))

adapter = MagicMock()
adapter.send.return_value = response

session = CachedSession()
session.get_adapter = MagicMock(return_value=adapter)

resp1 = session.request("GET", "https://url.com", params={"argument": "value"})
resp2 = session.request("GET", "https://url.com", params={"argument": "value"})

assert adapter.send.call_count == 1
assert resp1 is session.cache["https://url.com/?argument=value"]
assert resp2 is session.cache["https://url.com/?argument=value"]
assert resp1.json() == {"result": "data"}
assert resp2.json() == {"result": "data"}

0 comments on commit afc2d3a

Please sign in to comment.