-
Notifications
You must be signed in to change notification settings - Fork 3.3k
[Core] Refactor MSAL HTTP cache to use JSON format #32511
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: dev
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,112 @@ | ||
| # -------------------------------------------------------------------------------------------- | ||
| # Copyright (c) Microsoft Corporation. All rights reserved. | ||
| # Licensed under the MIT License. See License.txt in the project root for license information. | ||
| # -------------------------------------------------------------------------------------------- | ||
| import json | ||
| from collections.abc import MutableMapping | ||
|
|
||
| from azure.cli.core.decorators import retry | ||
| from knack.log import get_logger | ||
| from msal.throttled_http_client import NormalizedResponse | ||
|
|
||
| logger = get_logger(__name__) | ||
|
|
||
|
|
||
| class NormalizedResponseJsonEncoder(json.JSONEncoder): | ||
| def default(self, o): | ||
| if isinstance(o, NormalizedResponse): | ||
| return { | ||
| "status_code": o.status_code, | ||
| "text": o.text, | ||
| "headers": o.headers, | ||
| } | ||
| return super().default(o) | ||
|
|
||
|
|
||
| class JsonCache(MutableMapping): | ||
| """ | ||
| A simple dict-like class that is backed by a json file. This is designed for the MSAL HTTP cache. | ||
|
|
||
| All direct modifications with `__setitem__` and `__delitem__` will save the file. | ||
| Indirect modifications should be followed by a call to `save`. | ||
| """ | ||
| def __init__(self, file_name): | ||
| super().__init__() | ||
| self.filename = file_name | ||
| self.data = {} | ||
| self.load() | ||
|
|
||
| @retry() | ||
| def _load(self): | ||
| """Load cache with retry. If it still fails at last, raise the original exception as-is.""" | ||
| try: | ||
| with open(self.filename, 'r', encoding='utf-8') as f: | ||
| data = json.load(f) | ||
| response_keys = [key for key in data if key != "_index_"] | ||
| for key in response_keys: | ||
| try: | ||
| response_dict = data[key] | ||
| # Reconstruct NormalizedResponse from the stored dict | ||
| response = NormalizedResponse.__new__(NormalizedResponse) | ||
| response.status_code = response_dict["status_code"] | ||
| response.text = response_dict["text"] | ||
| response.headers = response_dict["headers"] | ||
|
Comment on lines
+50
to
+53
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't feel it is appropriate for Azure CLI as a downstream application of MSAL to serialize MSAL's HTTP cache as JSON. Azure CLI doesn't know which fields of the MSAL itself should provide a serializable HTTP cache, similar to the token cache |
||
| data[key] = response | ||
| except KeyError as e: | ||
| logger.debug("Failed to reconstruct NormalizedResponse for key %s: %s", key, e) | ||
| # If reconstruction fails, remove the entry from cache | ||
| del data[key] | ||
| return data | ||
| except FileNotFoundError: | ||
| # The cache file has not been created. This is expected. No need to retry. | ||
| logger.debug("%s not found. Using a fresh one.", self.filename) | ||
| return {} | ||
|
|
||
| def load(self): | ||
| logger.debug("load: %s", self.filename) | ||
| try: | ||
| self.data = self._load() | ||
| except Exception as ex: # pylint: disable=broad-exception-caught | ||
| # If we still get exception after retry, ignore all types of exceptions and use a new cache. | ||
| # - EOFError is caused by empty cache file created by other az instance, but hasn't been filled yet. | ||
| # - KeyError is caused by reading cache generated by different MSAL versions. | ||
| logger.debug("Failed to load cache: %s. Using a fresh one.", ex) | ||
| self.data = {} # Ignore a non-existing or corrupted http_cache | ||
|
|
||
| @retry() | ||
| def _save(self): | ||
| with open(self.filename, 'w', encoding='utf-8') as f: | ||
| # At this point, an empty cache file will be created. Loading this cache file will | ||
| # raise EOFError. This can be simulated by adding time.sleep(30) here. | ||
| # So during loading, EOFError is ignored. | ||
| json.dump(self.data, f, cls=NormalizedResponseJsonEncoder) | ||
|
|
||
| def save(self): | ||
| logger.debug("save: %s", self.filename) | ||
| # If 2 processes write at the same time, the cache will be corrupted, | ||
| # but that is fine. Subsequent runs would reach eventual consistency. | ||
| try: | ||
| self._save() | ||
| except TypeError as e: | ||
| # If serialization fails, skip saving to avoid corrupting the cache file | ||
| logger.debug("Failed to save cache due to TypeError: %s", e) | ||
|
|
||
| def get(self, key, default=None): | ||
| return self.data.get(key, default) | ||
|
|
||
| def __getitem__(self, key): | ||
| return self.data[key] | ||
|
|
||
| def __setitem__(self, key, value): | ||
| self.data[key] = value | ||
| self.save() | ||
|
|
||
| def __delitem__(self, key): | ||
| del self.data[key] | ||
| self.save() | ||
|
|
||
| def __iter__(self): | ||
| return iter(self.data) | ||
|
|
||
| def __len__(self): | ||
| return len(self.data) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The
BinaryCacheis designed for saving anything, so the nameBinaryCacheis accurate.However, the new
JsonCachecan only handleNormalizedResponse, so the name is not accurate. It should be something likeNormalizedResponseJsonCache.