update focus

id.saladeaula.digital/app/integrations/apigateway_oauth2/authorization_server.py

@@ -28,6 +28,12 @@ GRANT_TYPES_EXPIRES_IN = {
 logger = Logger(__name__)
 
 
+def expires_in(client, grant_type: str) -> int:
+    return GRANT_TYPES_EXPIRES_IN.get(
+        grant_type, BearerTokenGenerator.DEFAULT_EXPIRES_IN
+    )
+
+
 class AuthorizationServer(oauth2.AuthorizationServer):
     def __init__(
         self,
@@ -36,15 +42,9 @@ class AuthorizationServer(oauth2.AuthorizationServer):
     ) -> None:
         self._persistence_layer = persistence_layer
 
-        super().__init__(
-            scopes_supported=(
-                set(OAUTH2_SCOPES_SUPPORTED.split()) if OAUTH2_SCOPES_SUPPORTED else []
-            )
-        )
-
-        def expires_in(client, grant_type: str) -> int:
-            return GRANT_TYPES_EXPIRES_IN.get(
-                grant_type, BearerTokenGenerator.DEFAULT_EXPIRES_IN
+        if OAUTH2_SCOPES_SUPPORTED:
+            super().__init__(
+                scopes_supported=set(OAUTH2_SCOPES_SUPPORTED.split()),
             )
 
         self.register_token_generator(
@@ -61,6 +61,7 @@ class AuthorizationServer(oauth2.AuthorizationServer):
         token: dict,
         request: requests.OAuth2Request,
     ) -> None:
+        """Define function to save the generated token into database."""
         if not request.payload:
             raise ValueError('Missing request payload')
 
@@ -88,6 +89,7 @@ class AuthorizationServer(oauth2.AuthorizationServer):
                 item={
                     'id': user_id,
                     'sk': f'SESSION#ACCESS_TOKEN#{access_token}',
+                    'client_id': client_id,
                     'ttl': access_token_ttl,
                     'created_at': now_,
                 }
@@ -112,6 +114,7 @@ class AuthorizationServer(oauth2.AuthorizationServer):
                     item={
                         'id': user_id,
                         'sk': f'SESSION#REFRESH_TOKEN#{refresh_token}',
+                        'client_id': client_id,
                         'ttl': access_token_ttl,
                         'created_at': now_,
                     }
@@ -133,7 +136,11 @@ class AuthorizationServer(oauth2.AuthorizationServer):
 
         return None
 
-    def query_client(self, client_id: str):
+    def query_client(
+        self,
+        client_id: str,
+    ):
+        """Query OAuth client by client_id."""
         client = self._persistence_layer.collection.get_item(
             KeyPair(
                 pk='OAUTH2',