[Identity] Ensure DAC contextvar is reset (#42413)

If ChainedTokenCredential raises an exception, then the DAC contextvar
is never reset. This updates to ensure that it is reset.

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

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

@@ -332,8 +332,10 @@ def get_token(
             )
             return token
         within_dac.set(True)
-        token = super().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
-        within_dac.set(False)
+        try:
+            token = super().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
+        finally:
+            within_dac.set(False)
         return token
 
     def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
@@ -361,6 +363,8 @@ def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] =
             return token_info
 
         within_dac.set(True)
-        token_info = cast(SupportsTokenInfo, super()).get_token_info(*scopes, options=options)
-        within_dac.set(False)
+        try:
+            token_info = cast(SupportsTokenInfo, super()).get_token_info(*scopes, options=options)
+        finally:
+            within_dac.set(False)
         return token_info

sdk/identity/azure-identity/azure/identity/aio/_credentials/default.py

@@ -278,8 +278,10 @@ async def get_token(
             return token
 
         within_dac.set(True)
-        token = await super().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
-        within_dac.set(False)
+        try:
+            token = await super().get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
+        finally:
+            within_dac.set(False)
         return token
 
     async def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
@@ -309,6 +311,8 @@ async def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptio
             return token_info
 
         within_dac.set(True)
-        token_info = await cast(AsyncSupportsTokenInfo, super()).get_token_info(*scopes, options=options)
-        within_dac.set(False)
+        try:
+            token_info = await cast(AsyncSupportsTokenInfo, super()).get_token_info(*scopes, options=options)
+        finally:
+            within_dac.set(False)
         return token_info