From e2a00a6852f29a0b07f2cca7803a49025880bf6e Mon Sep 17 00:00:00 2001 From: Aleksander Wieckowski Date: Thu, 30 Jul 2026 14:34:36 +0200 Subject: [PATCH] Add oauth token caching to DNAAS client --- src/ralph/dns/dnsaas.py | 44 +++++++++++++++++++++--------- src/ralph/dns/tests.py | 59 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 13 deletions(-) diff --git a/src/ralph/dns/dnsaas.py b/src/ralph/dns/dnsaas.py index 2adcbddc53..277aea0d8e 100644 --- a/src/ralph/dns/dnsaas.py +++ b/src/ralph/dns/dnsaas.py @@ -3,6 +3,7 @@ import logging from datetime import datetime, timedelta from functools import wraps +from threading import Lock from typing import List, Optional, Tuple, Union from urllib.parse import parse_qs, urlencode, urljoin, urlsplit @@ -41,6 +42,9 @@ def wrapper(self, *args, **kwargs): class DNSaaS: + _oauth_tokens = {} + _oauth_token_lock = Lock() + def __init__(self, headers: dict = None): self.session = requests.Session() _headers = { @@ -52,25 +56,39 @@ def __init__(self, headers: dict = None): _headers.update(headers) self.session.headers.update(_headers) - def _get_oauth_token(self): + def _get_oauth_token(self, force_refresh=False): client_id = settings.OAUTH_CLIENT_ID secret = settings.OAUTH_SECRET token_url = settings.OAUTH_TOKEN_URL - client = BackendApplicationClient(client_id=client_id) - oauth = OAuth2Session(client=client) - try: - token = oauth.fetch_token( - token_url=token_url, client_id=client_id, client_secret=secret - ) - except CustomOAuth2Error as e: - logger.error(str(e)) + cache_key = (token_url, client_id) + + with self._oauth_token_lock: + now = datetime.now() + cached_token = self._oauth_tokens.get(cache_key) + if not force_refresh and cached_token and now < cached_token["expiration"]: + self.token_expiration = cached_token["expiration"] + return cached_token["access_token"] - expire_in = token.get("expires_in") - self.token_expiration = datetime.now() + timedelta(0, expire_in - 60) - return token.get("access_token") + client = BackendApplicationClient(client_id=client_id) + oauth = OAuth2Session(client=client) + try: + token = oauth.fetch_token( + token_url=token_url, client_id=client_id, client_secret=secret + ) + except CustomOAuth2Error: + logger.exception("Fetching an OAuth token for DNSaaS failed") + raise + + expires_in = token.get("expires_in") + self.token_expiration = now + timedelta(seconds=max(expires_in - 60, 0)) + self._oauth_tokens[cache_key] = { + "access_token": token.get("access_token"), + "expiration": self.token_expiration, + } + return token.get("access_token") def _update_oauth_token(self): - token = self._get_oauth_token() + token = self._get_oauth_token(force_refresh=True) self.session.headers["Authorization"] = "Bearer {}".format(token) def _verify_oauth_token_validity(self): diff --git a/src/ralph/dns/tests.py b/src/ralph/dns/tests.py index 504e3cf7a3..27195c7aba 100644 --- a/src/ralph/dns/tests.py +++ b/src/ralph/dns/tests.py @@ -1,8 +1,10 @@ # -*- coding: utf-8 -*- +from datetime import datetime, timedelta from unittest.mock import patch from django.db import transaction from django.test import override_settings, TestCase, TransactionTestCase +from oauthlib.oauth2.rfc6749.errors import CustomOAuth2Error from ralph.assets.tests.factories import ConfigurationClassFactory, EthernetFactory from ralph.data_center.models import BaseObjectCluster, DataCenterAsset @@ -15,6 +17,63 @@ from ralph.virtual.tests.factories import VirtualServerFactory +@override_settings( + OAUTH_CLIENT_ID="ralph", + OAUTH_SECRET="secret", + OAUTH_TOKEN_URL="https://oauth.example.com/token", +) +class TestOAuthToken(TestCase): + def setUp(self): + DNSaaS._oauth_tokens.clear() + + def tearDown(self): + DNSaaS._oauth_tokens.clear() + + @patch("ralph.dns.dnsaas.OAuth2Session") + def test_reuses_token_between_clients_until_it_expires(self, oauth_session_mock): + oauth_session_mock.return_value.fetch_token.return_value = { + "access_token": "token", + "expires_in": 3600, + } + + first_client = DNSaaS() + second_client = DNSaaS() + + oauth_session_mock.return_value.fetch_token.assert_called_once() + self.assertEqual(first_client.session.headers["Authorization"], "Bearer token") + self.assertEqual(second_client.session.headers["Authorization"], "Bearer token") + + @patch("ralph.dns.dnsaas.OAuth2Session") + def test_fetches_new_token_when_cached_token_has_expired(self, oauth_session_mock): + oauth_session_mock.return_value.fetch_token.side_effect = [ + {"access_token": "old-token", "expires_in": 3600}, + {"access_token": "new-token", "expires_in": 3600}, + ] + first_client = DNSaaS() + cached_token = next(iter(DNSaaS._oauth_tokens.values())) + cached_token["expiration"] = datetime.now() - timedelta(seconds=1) + + second_client = DNSaaS() + + self.assertEqual( + first_client.session.headers["Authorization"], "Bearer old-token" + ) + self.assertEqual( + second_client.session.headers["Authorization"], "Bearer new-token" + ) + self.assertEqual(oauth_session_mock.return_value.fetch_token.call_count, 2) + + @patch("ralph.dns.dnsaas.OAuth2Session") + def test_preserves_oauth_error_when_fetching_token_fails(self, oauth_session_mock): + oauth_error = CustomOAuth2Error(description="rate limited") + oauth_session_mock.return_value.fetch_token.side_effect = oauth_error + + with self.assertRaises(CustomOAuth2Error) as raised: + DNSaaS() + + self.assertIs(raised.exception, oauth_error) + + class TestGetDnsRecords(TestCase): @patch.object(DNSaaS, "_get_oauth_token") def setUp(self, mocked):