feat: Add more functions related to session claims

Author: KShivendu

supertokens_python/recipe/session/asyncio/__init__.py

@@ -11,7 +11,7 @@
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations
 # under the License.
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, TypeVar
 
 from supertokens_python.recipe.openid.interfaces import (
     GetOpenIdDiscoveryConfigurationResult,
@@ -20,16 +20,18 @@
     RegenerateAccessTokenOkResult,
     SessionContainer,
     SessionInformationResult,
+    SessionClaim,
 )
 from supertokens_python.recipe.session.recipe import SessionRecipe
 from supertokens_python.utils import FRAMEWORKS
-
 from ...jwt.interfaces import (
     CreateJwtOkResult,
     CreateJwtResultUnsupportedAlgorithm,
     GetJWKSResult,
 )
 
+_T = TypeVar("_T")
+
 
 async def create_new_session(
     request: Any,
@@ -49,6 +51,43 @@ async def create_new_session(
     )
 
 
+async def get_claim_value(
+    session_handle: str,
+    claim: SessionClaim[_T],
+    user_context: Union[None, Dict[str, Any]] = None,
+) -> Union[_T, None]:
+    if user_context is None:
+        user_context = {}
+    return await SessionRecipe.get_instance().recipe_implementation.get_claim_value(
+        session_handle, claim, user_context
+    )
+
+
+async def set_claim_value(
+    session_handle: str,
+    claim: SessionClaim[_T],
+    value: _T,
+    user_context: Union[None, Dict[str, Any]] = None,
+) -> bool:
+    if user_context is None:
+        user_context = {}
+    return await SessionRecipe.get_instance().recipe_implementation.set_claim_value(
+        session_handle, claim, value, user_context
+    )
+
+
+async def remove_claim(
+    session_handle: str,
+    claim: SessionClaim[Any],
+    user_context: Union[None, Dict[str, Any]] = None,
+) -> bool:
+    if user_context is None:
+        user_context = {}
+    return await SessionRecipe.get_instance().recipe_implementation.remove_claim(
+        session_handle, claim, user_context
+    )
+
+
 async def get_session(
     request: Any,
     anti_csrf_check: Union[bool, None] = None,

supertokens_python/recipe/session/interfaces.py

@@ -20,6 +20,7 @@
 from supertokens_python.async_to_sync_wrapper import sync
 from supertokens_python.types import APIResponse, GeneralErrorResponse
 
+from ...utils import Promise
 from .utils import SessionConfig
 
 if TYPE_CHECKING:
@@ -79,6 +80,15 @@ async def create_new_session(
     ) -> SessionContainer:
         pass
 
+    @abstractmethod
+    async def get_global_claim_validators(
+        self,
+        user_id: str,
+        claim_validators_added_by_other_recipes: List[SessionClaimValidator],
+        user_context: Dict[str, Any],
+    ) -> Union[SessionClaimValidator, Promise[SessionClaimValidator]]:
+        pass
+
     @abstractmethod
     async def get_session(
         self,
@@ -151,6 +161,43 @@ async def get_access_token_lifetime_ms(self, user_context: Dict[str, Any]) -> in
     async def get_refresh_token_lifetime_ms(self, user_context: Dict[str, Any]) -> int:
         pass
 
+    @abstractmethod
+    async def fetch_and_set_claim(
+        self,
+        session_handle: str,
+        claim: SessionClaim[Any],
+        user_context: Dict[str, Any],
+    ) -> bool:
+        pass
+
+    @abstractmethod
+    async def set_claim_value(
+        self,
+        session_handle: str,
+        claim: SessionClaim[_T],
+        value: _T,
+        user_context: Dict[str, Any],
+    ) -> bool:
+        pass
+
+    @abstractmethod
+    async def get_claim_value(
+        self,
+        session_handle: str,
+        claim: SessionClaim[Any],
+        user_context: Dict[str, Any],
+    ):
+        pass
+
+    @abstractmethod
+    async def remove_claim(
+        self,
+        session_handle: str,
+        claim: SessionClaim[Any],
+        user_context: Dict[str, Any],
+    ) -> bool:
+        pass
+
     @abstractmethod
     async def regenerate_access_token(
         self,
@@ -217,7 +264,7 @@ async def verify_session(
         pass
 
 
-class SessionContainer(ABC):
+class SessionContainer(ABC):  # pylint: disable=too-many-public-methods
     def __init__(
         self,
         recipe_implementation: RecipeInterface,
@@ -260,6 +307,12 @@ async def update_access_token_payload(
         self,
         new_access_token_payload: Dict[str, Any],
         user_context: Union[Dict[str, Any], None] = None,
+    ) -> None:
+        """DEPRECATED: Use merge_into_access_token_payload instead"""
+
+    @abstractmethod
+    async def merge_into_access_token_payload(
+        self, access_token_payload_update: Dict[str, Any], user_context: Any
     ) -> None:
         pass
 
@@ -291,6 +344,41 @@ async def get_time_created(
     async def get_expiry(self, user_context: Union[Dict[str, Any], None] = None) -> int:
         pass
 
+    @abstractmethod
+    async def assert_claims(
+        self,
+        claim_validators: List[SessionClaimValidator],
+        user_context: Union[Dict[str, Any], None] = None,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    async def fetch_and_set_claim(
+        self, claim: SessionClaim[Any], user_context: Union[Dict[str, Any], None] = None
+    ) -> None:
+        pass
+
+    @abstractmethod
+    async def set_claim_value(
+        self,
+        claim: SessionClaim[_T],
+        value: _T,
+        user_context: Union[Dict[str, Any], None] = None,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    async def get_claim_value(
+        self, claim: SessionClaim[_T], user_context: Union[Dict[str, Any], None] = None
+    ) -> Union[_T, None]:
+        pass
+
+    @abstractmethod
+    async def remove_claim(
+        self, claim: SessionClaim[Any], user_context: Union[Dict[str, Any], None] = None
+    ) -> None:
+        pass
+
     def sync_get_expiry(self, user_context: Union[Dict[str, Any], None] = None) -> int:
         return sync(self.get_expiry(user_context))
 
@@ -334,7 +422,6 @@ def __getitem__(self, item: str):
 _T = TypeVar("_T")
 JSONObject = Dict[str, Any]
 
-
 JSONPrimitive = Union[str, int, bool, None, Dict[str, Any]]
 
 FetchValueReturnType = Union[_T, None]

supertokens_python/recipe/session/recipe.py

@@ -39,7 +39,13 @@
 
 from .api.implementation import APIImplementation
 from .constants import SESSION_REFRESH, SIGNOUT
-from .interfaces import APIInterface, APIOptions, RecipeInterface
+from .interfaces import (
+    APIInterface,
+    APIOptions,
+    RecipeInterface,
+    SessionClaim,
+    SessionClaimValidator,
+)
 from .recipe_implementation import RecipeImplementation
 from .utils import (
     InputErrorHandlers,
@@ -138,6 +144,9 @@ def __init__(
             else self.config.override.apis(api_implementation)
         )
 
+        self.claims_added_by_other_recipes: List[SessionClaim[Any]] = []
+        self.claim_validators_added_by_other_recipes: List[SessionClaimValidator] = []
+
     def is_error_from_this_recipe_based_on_instance(self, err: Exception) -> bool:
         return isinstance(err, SuperTokensError) and (
             isinstance(err, SuperTokensSessionError)
@@ -278,6 +287,22 @@ def reset():
             raise_general_exception("calling testing function in non testing env")
         SessionRecipe.__instance = None
 
+    def add_claim_from_other_recipe(self, claim: SessionClaim[Any]):
+        self.claims_added_by_other_recipes.append(claim)
+
+    def get_claims_added_by_other_recipes(self) -> List[SessionClaim[Any]]:
+        return self.claims_added_by_other_recipes
+
+    def add_claim_validator_from_other_recipe(
+        self, claim_validator: SessionClaimValidator
+    ):
+        self.claim_validators_added_by_other_recipes.append(claim_validator)
+
+    def get_claim_validators_added_by_other_recipes(
+        self,
+    ) -> List[SessionClaimValidator]:
+        return self.claim_validators_added_by_other_recipes
+
     async def verify_session(
         self,
         request: BaseRequest,

supertokens_python/recipe/session/recipe_implementation.py

@@ -13,13 +13,14 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, TypeVar
 
 from supertokens_python.framework.request import BaseRequest
 from supertokens_python.logger import log_debug_message
 from supertokens_python.normalised_url_path import NormalisedURLPath
 from supertokens_python.process_state import AllowedProcessStates, ProcessState
 from supertokens_python.utils import (
+    Promise,
     execute_async,
     frontend_has_interceptor,
     get_timestamp_ms,
@@ -39,6 +40,8 @@
     AccessTokenObj,
     RecipeInterface,
     RegenerateAccessTokenOkResult,
+    SessionClaim,
+    SessionClaimValidator,
     SessionInformationResult,
     SessionObj,
 )
@@ -53,6 +56,8 @@
 
 from .interfaces import SessionContainer
 
+_T = TypeVar("_T")
+
 
 class HandshakeInfo:
     def __init__(self, info: Dict[str, Any]):
@@ -74,7 +79,7 @@ def get_jwt_signing_public_key_list(self) -> List[Dict[str, Any]]:
         ]
 
 
-class RecipeImplementation(RecipeInterface):
+class RecipeImplementation(RecipeInterface):  # pylint: disable=too-many-public-methods
     def __init__(self, querier: Querier, config: SessionConfig):
         super().__init__()
         self.querier = querier
@@ -338,6 +343,91 @@ async def get_access_token_lifetime_ms(self, user_context: Dict[str, Any]) -> in
     async def get_refresh_token_lifetime_ms(self, user_context: Dict[str, Any]) -> int:
         return (await self.get_handshake_info()).refresh_token_validity
 
+    async def merge_into_access_token_payload(
+        self,
+        session_handle: str,
+        access_token_payload_update: Dict[str, Any],
+        user_context: Dict[str, Any],
+    ) -> bool:
+        session_info = await self.get_session_information(session_handle, user_context)
+        if session_info is None:
+            return False
+
+        new_access_token_payload = {
+            **session_info.access_token_payload,
+            **access_token_payload_update,
+        }
+        for k in access_token_payload_update.keys():
+            if new_access_token_payload[k] is None:
+                del new_access_token_payload[k]
+
+        return await self.update_access_token_payload(
+            session_handle, new_access_token_payload, user_context
+        )
+
+    async def fetch_and_set_claim(
+        self,
+        session_handle: str,
+        claim: SessionClaim[Any],
+        user_context: Dict[str, Any],
+    ) -> bool:
+        session_info = await self.get_session_information(session_handle, user_context)
+        if session_info is None:
+            return False
+
+        access_token_payload_update = await claim.build(
+            session_info.user_id, user_context
+        )
+        return await self.merge_into_access_token_payload(
+            session_handle, access_token_payload_update, user_context
+        )
+
+    async def set_claim_value(
+        self,
+        session_handle: str,
+        claim: SessionClaim[Any],
+        value: Any,
+        user_context: Dict[str, Any],
+    ):
+        access_token_payload_update = claim.add_to_payload_({}, value, user_context)
+        return await self.merge_into_access_token_payload(
+            session_handle, access_token_payload_update, user_context
+        )
+
+    async def get_claim_value(
+        self,
+        session_handle: str,
+        claim: SessionClaim[Any],
+        user_context: Dict[str, Any],
+    ):
+        session_info = await self.get_session_information(session_handle, user_context)
+        if session_info is None:
+            raise Exception("Session does not exist")
+
+        return claim.get_value_from_payload(
+            session_info.access_token_payload, user_context
+        )
+
+    async def get_global_claim_validators(
+        self,
+        user_id: str,
+        claim_validators_added_by_other_recipes: List[SessionClaimValidator],
+        user_context: Dict[str, Any],
+    ) -> Union[SessionClaimValidator, Promise[SessionClaimValidator]]:
+        # TODO: Implement this
+        return claim_validators_added_by_other_recipes[0]
+
+    async def remove_claim(
+        self,
+        session_handle: str,
+        claim: SessionClaim[Any],
+        user_context: Dict[str, Any],
+    ) -> bool:
+        access_token_payload = claim.remove_from_payload_by_merge_({}, user_context)
+        return await self.merge_into_access_token_payload(
+            session_handle, access_token_payload, user_context
+        )
+
     async def regenerate_access_token(
         self,
         access_token: str,