[Comp-792] fix caching versions for multiple pods (#712)

Author: tolkamps1

compliance-api/src/compliance_api/services/authorize_service/auth_service.py

@@ -14,6 +14,21 @@
 class AuthService:
     """Handle service request for epic.authorize with integrated cache management."""
 
+    AUTH_CACHE_VERSION_KEY = "auth_cache_version"
+
+    @classmethod
+    def _get_auth_cache_version(cls):
+        """
+        Get current auth cache version from shared cache.
+
+        This ensures all pods see the same version number.
+        """
+        version = cache.get(cls.AUTH_CACHE_VERSION_KEY)
+        if version is None:
+            version = 0
+            cache.set(cls.AUTH_CACHE_VERSION_KEY, version)
+        return version
+
     @staticmethod
     def get_epic_user_by_guid(auth_user_guid: str):
         """
@@ -24,9 +39,10 @@ def get_epic_user_by_guid(auth_user_guid: str):
         """
         from compliance_api.services.cached_staff_user import CachedStaffUserService
 
-        # Include token hash in cache key for security
+        # Include version AND token hash in cache key
+        version = AuthService._get_auth_cache_version()
         token_hash = CachedStaffUserService._get_token_hash()
-        cache_key = f"auth_user:{auth_user_guid}:{token_hash}"
+        cache_key = f"auth_user:{auth_user_guid}:{token_hash}:v{version}"
 
         cached_result = cache.get(cache_key)
         if cached_result is not None:
@@ -56,9 +72,10 @@ def get_epic_users_by_app():
         from compliance_api.utils.constant import AUTH_APP
         from compliance_api.exceptions import BusinessError
 
-        # Include token hash in cache key for security
+        # Include version AND token hash in cache key
+        version = AuthService._get_auth_cache_version()
         token_hash = CachedStaffUserService._get_token_hash()
-        cache_key = f"auth_users_app:{AUTH_APP}:{token_hash}"
+        cache_key = f"auth_users_app:{AUTH_APP}:{token_hash}:v{version}"
 
         cached_result = cache.get(cache_key)
         if cached_result is not None:
@@ -91,11 +108,8 @@ def update_user_group(auth_user_guid: str, payload: dict):
                 f"Update group in the auth server failed for user : {auth_user_guid}"
             )
 
-        # Invalidate auth caches for this specific user
-        AuthService._invalidate_auth_user_cache(auth_user_guid)
-
-        # Invalidate the "all users by app" cache since this user's groups changed
-        AuthService._invalidate_auth_users_by_app_cache()
+        # Invalidate all auth caches by bumping version
+        AuthService._invalidate_all_auth_cache()
 
         # Invalidate ALL staff caches since permissions changed
         CachedStaffUserService.invalidate_staff_cache(auth_user_guid)
@@ -116,50 +130,30 @@ def delete_user_group(auth_user_guid: str, group: str, del_sub_group_mappings=Tr
         if delete_response.status_code != 204:
             raise BusinessError("Delete group mapping failed")
 
-        # Invalidate auth caches for this specific user
-        AuthService._invalidate_auth_user_cache(auth_user_guid)
-
-        # Invalidate the "all users by app" cache since this user's groups changed
-        AuthService._invalidate_auth_users_by_app_cache()
+        # Invalidate all auth caches by bumping version
+        AuthService._invalidate_all_auth_cache()
 
-        # Invalidate ALL staff caches since permissions changed
+        # Invalidate all staff caches since permissions changed
         CachedStaffUserService.invalidate_staff_cache(auth_user_guid)
 
         return delete_response
 
     @staticmethod
-    def _invalidate_auth_user_cache(auth_user_guid: str):
-        """
-        Invalidate the individual auth user cache.
-
-        Since cache keys include token hashes, we can't delete all variations.
-        Instead, we use a pattern-based approach or accept that cache will expire naturally.
-        """
-        from compliance_api.services.cached_staff_user import CachedStaffUserService
-
-        # Get current token hash
-        token_hash = CachedStaffUserService._get_token_hash()
-        cache_key = f"auth_user:{auth_user_guid}:{token_hash}"
-
-        cache.delete(cache_key)
-        current_app.logger.info(f"Invalidated auth user cache for {auth_user_guid}")
-
-    @staticmethod
-    def _invalidate_auth_users_by_app_cache():
-        """
-        Invalidate the "all users by app" cache.
-
-        Since cache keys include token hashes, we can't delete all variations.
-        Instead, we delete for the current token.
-        """
-        from compliance_api.services.cached_staff_user import CachedStaffUserService
-
-        # Get current token hash
-        token_hash = CachedStaffUserService._get_token_hash()
-        cache_key = f"auth_users_app:{AUTH_APP}:{token_hash}"
-
-        cache.delete(cache_key)
-        current_app.logger.info("Invalidated auth users by app cache")
+    def _invalidate_all_auth_cache():
+        """Invalidate all auth cache by bumping version number."""
+        try:
+            current_version = cache.get(AuthService.AUTH_CACHE_VERSION_KEY)
+            if current_version is None:
+                current_version = 0
+
+            new_version = current_version + 1
+            cache.set(AuthService.AUTH_CACHE_VERSION_KEY, new_version)
+
+            current_app.logger.debug(
+                f"Bumped auth cache version from {current_version} to {new_version}"
+            )
+        except (AttributeError, RuntimeError) as e:
+            current_app.logger.error(f"Error invalidating auth cache: {e}")
 
 
 def _request_auth_service(

compliance-api/src/compliance_api/services/cached_staff_user.py

@@ -16,7 +16,20 @@ class CachedStaffUserService:
     STAFF_CACHE_KEY_PREFIX = "staff_user:"
     ALL_STAFF_CACHE_KEY = "all_staff_users"
     ALL_STAFF_WITH_AUTH_PREFIX = "all_staff_with_auth:"
-    CACHE_VERSION_KEY = "all_staff_with_auth_version"
+    CACHE_VERSION_KEY = "staff_cache_version"
+
+    @classmethod
+    def _get_cache_version(cls):
+        """
+        Get current cache version from shared cache.
+
+        This ensures all pods see the same version number.
+        """
+        version = cache.get(cls.CACHE_VERSION_KEY)
+        if version is None:
+            version = 0
+            cache.set(cls.CACHE_VERSION_KEY, version)
+        return version
 
     @classmethod
     def exists_staff_by_auth_guid(cls, auth_guid: str) -> bool:
@@ -35,7 +48,9 @@ def exists_staff_by_auth_guid(cls, auth_guid: str) -> bool:
         if not auth_guid:
             return False
 
-        cache_key = f"{cls.STAFF_CACHE_KEY_PREFIX}{auth_guid}"
+        # Include version in cache key
+        version = cls._get_cache_version()
+        cache_key = f"{cls.STAFF_CACHE_KEY_PREFIX}{auth_guid}:v{version}"
 
         cached_staff_exists = cache.get(cache_key)
         if cached_staff_exists is not None:
@@ -63,10 +78,8 @@ def get_all_staff_users_with_auth():
         # pylint: disable=import-outside-toplevel
         from compliance_api.services.staff_user import _set_permission_level_in_compliance_user_obj
 
-        # Get current version
-        version = cache.get(CachedStaffUserService.CACHE_VERSION_KEY)
-        if version is None:
-            version = 0
+        # Get current version from shared cache
+        version = CachedStaffUserService._get_cache_version()
 
         # Create cache key that includes token hash AND version
         token_hash = CachedStaffUserService._get_token_hash()
@@ -111,20 +124,18 @@ def invalidate_staff_cache(auth_guid: str = None):
         Args:
             auth_guid: Specific auth_guid to invalidate, or None to clear all staff cache
         """
+        # Bump version to invalidate all caches across all pods
+        CachedStaffUserService._invalidate_all_staff_cache()
+
         if auth_guid:
-            cache_key = f"{CachedStaffUserService.STAFF_CACHE_KEY_PREFIX}{auth_guid}"
-            cache.delete(cache_key)
             current_app.logger.info(f"Invalidated cache for staff user: {auth_guid}")
 
-        # Always invalidate the aggregate cache when any staff user changes
-        CachedStaffUserService._invalidate_all_staff_with_auth_cache()
-
     @staticmethod
-    def _invalidate_all_staff_with_auth_cache():
+    def _invalidate_all_staff_cache():
         """
-        Invalidate all staff+auth cache by bumping version number.
+        Invalidate all staff cache by bumping version number.
 
-        This invalidates cache for all tokens/users since the version
+        This invalidates cache for across pods since the version
         is part of every cache key.
         """
         try:
@@ -135,10 +146,10 @@ def _invalidate_all_staff_with_auth_cache():
             new_version = current_version + 1
             cache.set(CachedStaffUserService.CACHE_VERSION_KEY, new_version)
             current_app.logger.info(
-                f"Bumped staff+auth cache version from {current_version} to {new_version}"
+                f"Bumped staff cache version from {current_version} to {new_version}"
             )
         except (AttributeError, RuntimeError) as e:
-            current_app.logger.error(f"Error invalidating staff+auth cache: {e}")
+            current_app.logger.error(f"Error invalidating staff cache: {e}")
 
     @classmethod
     def _clear_all_staff_cache(cls):

compliance-api/src/compliance_api/services/staff_user.py

@@ -146,6 +146,10 @@ def update_user(user_id, user_data):
             if new_auth_guid and new_auth_guid != existing_user.auth_user_guid:
                 CachedStaffUserService.invalidate_staff_cache(new_auth_guid)
 
+            # Return the latest user object
+            final_auth_guid = new_auth_guid if new_auth_guid else existing_user.auth_user_guid
+            return StaffUserService.get_user_by_auth_guid(final_auth_guid)
+
         return updated_user
 
     @staticmethod

compliance-api/tests/unit/services/test_cached_staff_user.py

@@ -39,7 +39,8 @@ def test_exists_staff_by_auth_guid_cache_miss(self):
         assert exists is True
 
         # Verify cache was populated
-        cache_key = f"{CachedStaffUserService.STAFF_CACHE_KEY_PREFIX}{auth_guid}"
+        version = CachedStaffUserService._get_cache_version()
+        cache_key = f"{CachedStaffUserService.STAFF_CACHE_KEY_PREFIX}{auth_guid}:v{version}"
         cached_value = cache.get(cache_key)
         assert cached_value is True
 
@@ -86,20 +87,32 @@ def test_invalidate_staff_cache_specific_user(self):
         exists1 = CachedStaffUserService.exists_staff_by_auth_guid(auth_guid)
         assert exists1 is True
 
-        # Verify cache exists
-        cache_key = f"{CachedStaffUserService.STAFF_CACHE_KEY_PREFIX}{auth_guid}"
-        assert cache.get(cache_key) is True
+        # Get the version and verify cache exists
+        version_before = CachedStaffUserService._get_cache_version()
+        cache_key_before = f"{CachedStaffUserService.STAFF_CACHE_KEY_PREFIX}{auth_guid}:v{version_before}"
+        assert cache.get(cache_key_before) is True
 
-        # Invalidate specific user cache
+        # Bumps the version
         CachedStaffUserService.invalidate_staff_cache(auth_guid)
 
-        # Verify cache was cleared
-        assert cache.get(cache_key) is None
+        # Verify version was bumped
+        version_after = CachedStaffUserService._get_cache_version()
+        assert version_after == version_before + 1
 
-        # Should still work (will fetch from database)
+        # Verify old cache key still exists (not deleted, just obsolete)
+        assert cache.get(cache_key_before) is True
+
+        # Verify new cache key doesn't exist yet
+        cache_key_after = f"{CachedStaffUserService.STAFF_CACHE_KEY_PREFIX}{auth_guid}:v{version_after}"
+        assert cache.get(cache_key_after) is None
+
+        # Should still fetch from database with new version
         exists2 = CachedStaffUserService.exists_staff_by_auth_guid(auth_guid)
         assert exists2 is True
 
+        # Verify new cache was populated
+        assert cache.get(cache_key_after) is True
+
     def test_invalidate_staff_cache_all_users(self):
         """Test invalidating all staff cache."""
         # Populate cache for both users