Add pickling support for SharedTokenCacheCredential (#36404)

Author: xiangyan99

sdk/identity/azure-identity/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed the issue that `SharedTokenCacheCredential` was not picklable.
+
 ### Other Changes
 
 ## 1.17.1 (2024-06-21)

sdk/identity/azure-identity/azure/identity/_credentials/silent.py

@@ -38,6 +38,10 @@ def __init__(
         validate_tenant_id(self._tenant_id)
         self._cache = kwargs.pop("_cache", None)
         self._cae_cache = kwargs.pop("_cae_cache", None)
+        if self._cache or self._cae_cache:
+            self._custom_cache = True
+        else:
+            self._custom_cache = False
 
         self._cache_persistence_options = kwargs.pop("cache_persistence_options", None)
 
@@ -162,3 +166,18 @@ def _acquire_token_silent(self, *scopes: str, **kwargs: Any) -> AccessToken:
 
         # cache doesn't contain a matching refresh (or access) token
         raise CredentialUnavailableError(message=NO_TOKEN.format(self._auth_record.username))
+
+    def __getstate__(self) -> Dict[str, Any]:
+        state = self.__dict__.copy()
+        # Remove the non-picklable entries
+        if not self._custom_cache:
+            del state["_cache"]
+            del state["_cae_cache"]
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        # Re-create the unpickable entries
+        if not self._custom_cache:
+            self._cache = None
+            self._cae_cache = None

sdk/identity/azure-identity/azure/identity/_internal/shared_token_cache.py

@@ -5,7 +5,7 @@
 import abc
 import platform
 import time
-from typing import Any, Iterable, List, Mapping, Optional, cast
+from typing import Any, Iterable, List, Mapping, Optional, cast, Dict
 from urllib.parse import urlparse
 import msal
 
@@ -96,6 +96,10 @@ def __init__(
         self._tenant_id = tenant_id
         self._cache = kwargs.pop("_cache", None)
         self._cae_cache = kwargs.pop("_cae_cache", None)
+        if self._cache or self._cae_cache:
+            self._custom_cache = True
+        else:
+            self._custom_cache = False
         self._cache_persistence_options = kwargs.pop("cache_persistence_options", None)
         self._client_kwargs = kwargs
         self._client_kwargs["tenant_id"] = "organizations"
@@ -267,3 +271,18 @@ def supported() -> bool:
         :rtype: bool
         """
         return platform.system() in {"Darwin", "Linux", "Windows"}
+
+    def __getstate__(self) -> Dict[str, Any]:
+        state = self.__dict__.copy()
+        # Remove the non-picklable entries
+        if not self._custom_cache:
+            del state["_cache"]
+            del state["_cae_cache"]
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        # Re-create the unpickable entries
+        if not self._custom_cache:
+            self._cache = None
+            self._cae_cache = None

sdk/identity/azure-identity/tests/test_pickling.py

@@ -3,7 +3,7 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import pickle
-from azure.identity import DefaultAzureCredential
+from azure.identity import DefaultAzureCredential, SharedTokenCacheCredential
 from azure.identity._internal.msal_credentials import MsalCredential
 
 
@@ -15,6 +15,15 @@ def test_pickle_dac():
         data_loaded = pickle.load(infile)
 
 
+def test_pickle_shared_token_cache():
+    cred = SharedTokenCacheCredential()
+    cred._credential._initialize_cache()
+    with open("data.pkl", "wb") as outfile:
+        pickle.dump(cred, outfile)
+    with open("data.pkl", "rb") as infile:
+        data_loaded = pickle.load(infile)
+
+
 def test_pickle_msal_credential():
     cred = MsalCredential(client_id="CLIENT_ID")
     app = cred._get_app()