Fix eventual consistency and introduce new flag

Author: pentest-gg

fastapi_users_db_dynamodb/__init__.py

@@ -225,12 +225,19 @@ def _ensure_email_lower(self, data: dict[str, Any]) -> None:
         if "email" in data and isinstance(data["email"], str):
             data["email"] = data["email"].lower()
 
-    async def get(self, id: ID | str) -> UP | None:
+    async def get(
+        self,
+        id: ID | str,
+        instant_update: bool = False,
+    ) -> UP | None:
         """Get a user by id and hydrate oauth_accounts if available."""
         id_str = self._ensure_id_str(id)
 
         async with self._table(self.user_table_name, self._resource_region) as table:
-            resp = await table.get_item(Key={self.primary_key: id_str})
+            resp = await table.get_item(
+                Key={self.primary_key: id_str},
+                ConsistentRead=instant_update,
+            )
             item = resp.get("Item")
             user = self._item_to_user(item)
 
@@ -335,11 +342,19 @@ async def create(self, create_dict: dict[str, Any]) -> UP:
             raise ValueError("Could not cast DB item to User model")
         return refreshed_user
 
-    async def update(self, user: UP, update_dict: dict[str, Any]) -> UP:
+    async def update(
+        self,
+        user: UP,
+        update_dict: dict[str, Any],
+        instant_update: bool = False,
+    ) -> UP:
         """Update a user with update_dict and return the updated UP instance."""
         user_id = self._extract_id_from_user(user)
         async with self._table(self.user_table_name, self._resource_region) as table:
-            resp = await table.get_item(Key={self.primary_key: user_id})
+            resp = await table.get_item(
+                Key={self.primary_key: user_id},
+                ConsistentRead=instant_update,
+            )
             current = resp.get("Item", None)
 
             if not current:
@@ -414,7 +429,7 @@ async def update_oauth_account(
         user: UP,
         oauth_account: OAP,  # type: ignore
         update_dict: dict[str, Any],
-    ) -> UP | None:
+    ) -> UP:
         """Update an OAuth account and return the refreshed user (UP)."""
         if self.oauth_account_table is None or self.oauth_account_table_name is None:
             raise NotImplementedError()

fastapi_users_db_dynamodb/access_token.py

@@ -119,12 +119,16 @@ def _ensure_token(self, token: Any) -> str:
         return str(token)
 
     async def get_by_token(
-        self, token: str, max_age: datetime | None = None
+        self,
+        token: str,
+        max_age: datetime | None = None,
+        instant_update: bool = False,
     ) -> AP | None:
         """Retrieve an access token by token string."""
         async with self._table(self.table_name, self._resource_region) as table:
             resp = await table.get_item(
-                Key={self.primary_key: self._ensure_token(token)}
+                Key={self.primary_key: self._ensure_token(token)},
+                ConsistentRead=instant_update,
             )
             item = resp.get("Item")
 

tests/test_users.py

@@ -199,8 +199,13 @@ async def test_queries_oauth(
     )
     assert user.oauth_accounts[0].access_token == "NEW_TOKEN"  # type: ignore
 
+    #! IMPORTANT: Since DynamoDB uses eventual consistency, we need a small delay (e.g. `time.sleep(0.01)`) \
+    #! to ensure the user was fully updated. In production, this should be negligible. \
+    #! Alternatively, the `get` and `update` methods of the `DynamoDBDatabase` class allow users \
+    #! to enable consistent reads via the `instant_update` argument.
+
     # Get by id
-    id_user = await dynamodb_user_db_oauth.get(user.id)
+    id_user = await dynamodb_user_db_oauth.get(user.id, instant_update=True)
     assert id_user is not None
     assert id_user.id == user.id
     assert id_user.oauth_accounts[0].access_token == "NEW_TOKEN"  # type: ignore