Refactor credential refresh logic in IdentityClient

Simplifies and restructures the credential refresh flow to prioritize constructor/environment credentials, handle expired session tokens, and streamline VeFaaS IAM and AssumeRole fallback logic. Improves readability and maintainability by removing redundant checks and helper functions.

Author: loveyana

veadk/integrations/ve_identity/identity_client.py

@@ -56,76 +56,55 @@ def refresh_credentials(func):
     """
     import asyncio
 
+    def _try_get_vefaas_credentials():
+        """Attempt to retrieve credentials from VeFaaS IAM."""
+        try:
+            ve_iam_cred = get_credential_from_vefaas_iam()
+            return (
+                ve_iam_cred.access_key_id,
+                ve_iam_cred.secret_access_key,
+                ve_iam_cred.session_token,
+            )
+        except FileNotFoundError:
+            pass  # VeFaaS IAM file not found, ignore
+        except Exception as e:
+            logger.warning(f"Failed to retrieve credentials from VeFaaS IAM: {e}")
+        return None
+
     @wraps(func)
     def _refresh_creds(self: IdentityClient):
         """Helper to refresh credentials."""
-        # Try to get credentials from environment variables first
+        # Step 1: Get initial credentials from constructor or environment variables
         ak = self._initial_access_key or os.getenv("VOLCENGINE_ACCESS_KEY", "")
         sk = self._initial_secret_key or os.getenv("VOLCENGINE_SECRET_KEY", "")
         session_token = self._initial_session_token or os.getenv(
             "VOLCENGINE_SESSION_TOKEN", ""
         )
 
-        # Helper function to attempt VeFaaS IAM credential retrieval
-        def try_get_vefaas_credentials():
-            """Attempt to retrieve credentials from VeFaaS IAM."""
-            try:
-                ve_iam_cred = get_credential_from_vefaas_iam()
-                return (
-                    ve_iam_cred.access_key_id,
-                    ve_iam_cred.secret_access_key,
-                    ve_iam_cred.session_token,
-                )
-            except FileNotFoundError:
-                pass  # If VeFaaS IAM file not found, ignore
-            except Exception as e:
-                logger.warning(f"Failed to retrieve credentials from VeFaaS IAM: {e}")
-            return None
+        # Step 2: Clear expired session_token
+        if self._is_sts_credential_expired():
+            logger.info("STS credentials expired, clearing...")
+            session_token = ""
 
-        # If no AK/SK, try to get from VeFaaS IAM
-        if not (ak and sk):
-            logger.info(
-                "Credentials not found in environment, attempting to fetch from VeFaaS IAM..."
-            )
-            credentials = try_get_vefaas_credentials()
-            if credentials:
-                ak, sk, session_token = credentials
+        # Step 3: Try VeFaaS IAM if no credentials or no session_token
+        # VeFaaS IAM provides complete credentials (ak, sk, session_token)
+        if not (ak and sk) or (ak and sk and not session_token):
+            ak, sk, session_token = _try_get_vefaas_credentials()
 
-        # If we have AK/SK but no session token, or STS credentials are expired,
-        # try to get complete credentials
-        need_refresh = False
+        # Step 4: If still no session_token, try AssumeRole
         if ak and sk and not session_token:
-            need_refresh = True
-        elif ak and sk and session_token:
-            # Check if STS credentials are expired
-            if self._is_sts_credential_expired():
-                logger.info("STS credentials expired, refreshing...")
-                need_refresh = True
-                # Clear expired session token to force refresh
-                session_token = ""
-
-        if need_refresh:
-            # First attempt: try VeFaaS IAM
-            credentials = try_get_vefaas_credentials()
-            if credentials:
-                ak, sk, session_token = credentials
-
-            # Second attempt: if still no session token, try AssumeRole
-            if not session_token:
-                role_trn = self._get_iam_role_trn_from_vefaas_iam() or os.getenv(
-                    "RUNTIME_IAM_ROLE_TRN", ""
-                )
-
-                if role_trn:
-                    try:
-                        sts_credentials = self._assume_role(ak, sk, role_trn)
-                        ak = sts_credentials.access_key_id
-                        sk = sts_credentials.secret_access_key
-                        session_token = sts_credentials.session_token
-                    except Exception as e:
-                        logger.warning(f"Failed to assume role: {e}")
+            if role_trn := self._get_iam_role_trn_from_vefaas_iam() or os.getenv(
+                "RUNTIME_IAM_ROLE_TRN", ""
+            ):
+                try:
+                    sts_cred = self._assume_role(ak, sk, role_trn)
+                    ak = sts_cred.access_key_id
+                    sk = sts_cred.secret_access_key
+                    session_token = sts_cred.session_token
+                except Exception as e:
+                    logger.warning(f"Failed to assume role: {e}")
 
-        # Update configuration with the credentials
+        # Step 5: Update configuration with the credentials
         self._api_client.api_client.configuration.ak = ak
         self._api_client.api_client.configuration.sk = sk
         self._api_client.api_client.configuration.session_token = session_token