Refactor user_id_index

Author: pentest-gg

fastapi_users_db_dynamodb/__init__.py

@@ -142,21 +142,6 @@ class Meta:
 
         oauth_name = UnicodeAttribute(hash_key=True)
 
-    class UserIdIndex(GlobalSecondaryIndex):
-        """Enable the `user_id` attribute to be a Global Secondary Index.
-
-        Args:
-            GlobalSecondaryIndex (_type_): The Global Secondary Index base class.
-        """
-
-        class Meta:
-            """The metadata for the Global Secondary Index."""
-
-            index_name = "user_id-index"
-            projection = AllProjection()
-
-        user_id = GUID(hash_key=True)
-
     if TYPE_CHECKING:  # pragma: no cover
         id: ID
         oauth_name: str
@@ -176,7 +161,6 @@ class Meta:
     # Global Secondary Index
     account_id_index = AccountIdIndex()
     oauth_name_index = OAuthNameIndex()
-    user_id_index = UserIdIndex()
 
 
 class DynamoDBBaseOAuthAccountTableUUID(DynamoDBBaseOAuthAccountTable[UUID_ID]):
@@ -186,13 +170,32 @@ class DynamoDBBaseOAuthAccountTableUUID(DynamoDBBaseOAuthAccountTable[UUID_ID]):
         DynamoDBBaseOAuthAccountTable (_type_): The underlying table object.
     """
 
+    # MANDATORY GSI (MUST BE IMPLEMENTED)
+    class UserIdIndex(GlobalSecondaryIndex):
+        """Enable the `user_id` attribute to be a Global Secondary Index.
+
+        Args:
+            GlobalSecondaryIndex (_type_): The Global Secondary Index base class.
+        """
+
+        class Meta:
+            """The metadata for the Global Secondary Index."""
+
+            index_name = "user_id-index"
+            projection = AllProjection()
+
+        user_id = GUID(hash_key=True)
+
     if TYPE_CHECKING:  # pragma: no cover
         id: UUID_ID
         user_id: UUID_ID
     else:
         id: GUID = GUID(hash_key=True, default=uuid.uuid4)
         user_id: GUID = GUID(null=False)
 
+    # Global Secondary Index
+    user_id_index = UserIdIndex()
+
 
 class DynamoDBUserDatabase(Generic[UP, ID], BaseUserDatabase[UP, ID]):
     """
@@ -240,6 +243,13 @@ async def _hydrate_oauth_accounts(
 
         user.oauth_accounts = []  # type: ignore
 
+        if not hasattr(self.oauth_account_table, "user_id_index") or not isinstance(
+            self.oauth_account_table.user_id_index,  # type: ignore
+            GlobalSecondaryIndex,
+        ):
+            raise ValueError(
+                "Attribute 'user_id_index' not found: OAuthAccount table scheme must implement a Global Secondary Index for attribute 'user_id'."
+            )
         async for oauth_acc in self.oauth_account_table.user_id_index.query(  # type: ignore
             user.id,
             consistent_read=instant_update,

fastapi_users_db_dynamodb/access_token.py

@@ -69,11 +69,30 @@ class DynamoDBBaseAccessTokenTableUUID(DynamoDBBaseAccessTokenTable[UUID_ID]):
         DynamoDBBaseAccessTokenTable (_type_): The underlying table object.
     """
 
+    # OPTIONAL GSI
+    class UserIdIndex(GlobalSecondaryIndex):
+        """Enable the `user_id` attribute to be a Global Secondary Index.
+
+        Args:
+            GlobalSecondaryIndex (_type_): The Global Secondary Index base class.
+        """
+
+        class Meta:
+            """The metadata for the Global Secondary Index."""
+
+            index_name = "user_id-index"
+            projection = AllProjection()
+
+        user_id = GUID(hash_key=True)
+
     if TYPE_CHECKING:  # pragma: no cover
         user_id: UUID_ID
     else:
         user_id: GUID = GUID(null=False)
 
+    # Global Secondary Index
+    user_id_index = UserIdIndex()
+
 
 class DynamoDBAccessTokenDatabase(Generic[AP], AccessTokenDatabase[AP]):
     """

fastapi_users_db_dynamodb/config.py

@@ -1,7 +1,7 @@
 from enum import StrEnum
 from typing import Any, Literal, TypedDict
 
-__version__ = "1.0.1"
+__version__ = "1.0.2"
 
 
 # Right now, only ON-DEMAND/PAY_PER_REQUEST mode is supported!

tests/test_users.py

@@ -258,6 +258,7 @@ async def test_queries_oauth(
     oauth_account1: dict[str, Any],
     oauth_account2: dict[str, Any],
     user_id: UUID_ID,
+    monkeypatch,
 ):
     """Test `OAuth` implemenatation and basic **CRUD** operations.
 
@@ -338,3 +339,14 @@ def _get_account(_user: UserOAuth):
             oauth_account,
             {"access_token": "NEW_TOKEN"},
         )
+
+    with pytest.raises(
+        ValueError,
+        match="OAuthAccount table scheme must implement a Global Secondary Index",
+    ):
+        monkeypatch.setattr(
+            dynamodb_user_db_oauth.oauth_account_table,
+            "user_id_index",
+            None,
+        )
+        await dynamodb_user_db_oauth._hydrate_oauth_accounts(oauth_user)