Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 64 additions & 0 deletions google/auth/_cache.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import OrderedDict


class LRUCache(dict):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a benefit to using a custom implementation over functools.lru_cache?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Daniel:

Thanks for asking the question. The short answer is that functools.lru_cache is not the right tool for this job, despite a similarity in names.

Incompatible Introspection: The consuming code that uses LRUCache (in google.auth.jwt.OnDemandCredentials) uses the cache as a data store where it retrieves a token and then checks the expiration to confirm that the expiry < helpers.utcnow(). functools.lru_cache does not support retrieving a value without "using" it in a way that affects eviction, nor does it support inspecting the value (expiry) to decide if it is valid.

Explicit Invalidation: If the token is expired, the code needs to explicitly replace it. functools.lru_cache does not expose a way to invalidate or update a specific key; you can only clear the entire cache.

Usage Pattern: functools.lru_cache is a function decorator intended to hold the arguments for a decorated function so that the function does not need to be called again. Here, an object-based cache is needed to store state associated with the OnDemandCredentials instance.

def __init__(self, maxsize):
super().__init__()
self._order = OrderedDict()
self.maxsize = maxsize

def clear(self):
super().clear()
self._order.clear()

def get(self, key, default=None):
try:
value = super().__getitem__(key)
self._update(key)
return value
except KeyError:
return default

def __getitem__(self, key):
value = super().__getitem__(key)
self._update(key)
return value

def __setitem__(self, key, value):
maxsize = self.maxsize
if maxsize <= 0:
return
if key not in self:
while len(self) >= maxsize:
self.popitem()
super().__setitem__(key, value)
self._update(key)

def __delitem__(self, key):
super().__delitem__(key)
del self._order[key]

def popitem(self):
"""Remove and return the least recently used key-value pair."""
key, _ = self._order.popitem(last=False)
return key, super().pop(key)

def _update(self, key):
try:
self._order.move_to_end(key)
except KeyError:
self._order[key] = None
7 changes: 3 additions & 4 deletions google/auth/jwt.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,7 @@
import json
import urllib

import cachetools

from google.auth import _cache
from google.auth import _helpers
from google.auth import _service_account_info
from google.auth import crypt
Expand Down Expand Up @@ -630,7 +629,7 @@ def __init__(
token_lifetime (int): The amount of time in seconds for
which the token is valid. Defaults to 1 hour.
max_cache_size (int): The maximum number of JWT tokens to keep in
cache. Tokens are cached using :class:`cachetools.LRUCache`.
cache. Tokens are cached using :class:`google.auth._cache.LRUCache`.
quota_project_id (Optional[str]): The project ID used for quota
and billing.

Expand All @@ -646,7 +645,7 @@ def __init__(
additional_claims = {}

self._additional_claims = additional_claims
self._cache = cachetools.LRUCache(maxsize=max_cache_size)
self._cache = _cache.LRUCache(maxsize=max_cache_size)

@classmethod
def _from_signer_and_info(cls, signer, info, **kwargs):
Expand Down
1 change: 0 additions & 1 deletion noxfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,6 @@ def mypy(session):
session.install("-e", ".")
session.install(
"mypy",
"types-cachetools",
"types-certifi",
"types-freezegun",
"types-pyOpenSSL",
Expand Down
1 change: 0 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@


DEPENDENCIES = (
"cachetools>=2.0.0,<7.0",
"pyasn1-modules>=0.2.1",
# rsa==4.5 is the last version to support 2.7
# https://github.com/sybrenstuvel/python-rsa/issues/152#issuecomment-643470233
Expand Down
1 change: 0 additions & 1 deletion testing/constraints-3.7.txt
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
#
# e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev",
# Then this file should have foo==1.14.0
cachetools==2.0.0
pyasn1-modules==0.2.1
setuptools==40.3.0
rsa==3.1.4
Expand Down
82 changes: 82 additions & 0 deletions tests/test__cache.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
from google.auth._cache import LRUCache


def test_lru_cache():
"""Test the LRUCache for generally expected functionality and ordering."""
lru_cache = LRUCache(2)
lru_cache["a"] = 1
lru_cache["b"] = 2
assert lru_cache["a"] == 1
lru_cache["c"] = 3
assert "b" not in lru_cache
assert lru_cache["a"] == 1
assert lru_cache["c"] == 3
lru_cache["d"] = 4
assert "a" not in lru_cache
assert lru_cache["c"] == 3
assert lru_cache["d"] == 4


def test_zero_size_lru_cache():
"""Confirm the LRUCache handles zero-size correctly."""
lru_cache = LRUCache(0)
lru_cache["a"] = 1
assert "a" not in lru_cache


def test_lru_cache_get_updates_lru():
"""Confirm the LRUCache handles get calls correctly."""
lru_cache = LRUCache(2)
lru_cache["a"] = 1
lru_cache["b"] = 2

# Access "a" via get(), making it MRU.
assert lru_cache.get("a") == 1

# Add "c", which should evict "b" (LRU), not "a".
lru_cache["c"] = 3

assert "a" in lru_cache
assert "b" not in lru_cache
assert "c" in lru_cache


def test_lru_cache_get_missing():
"""Confirm the LRUCache handles missing keys correctly."""
lru_cache = LRUCache(2)
assert lru_cache.get("missing") is None
assert lru_cache.get("missing", "default") == "default"


def test_lru_cache_clear():
"""Confirm the LRUCache clears the cache properly."""
lru_cache = LRUCache(2)
lru_cache["a"] = 1
lru_cache["b"] = 2
assert len(lru_cache) == 2

lru_cache.clear()
assert len(lru_cache) == 0
assert "a" not in lru_cache
assert "b" not in lru_cache
# Ensure internal order is also cleared
assert len(lru_cache._order) == 0


def test_lru_cache_delitem():
"""Confirm the LRUCache deletes individual items properly."""
lru_cache = LRUCache(2)
lru_cache["a"] = 1
lru_cache["b"] = 2

del lru_cache["a"]
assert "a" not in lru_cache
assert len(lru_cache) == 1
# Ensure it's removed from internal order
assert "a" not in lru_cache._order

# Test that we can continue using the cache
lru_cache["c"] = 3
assert "c" in lru_cache
assert "b" in lru_cache
assert len(lru_cache) == 2