fix region param

Author: pentest-gg

fastapi_users_db_dynamodb/__init__.py

@@ -90,28 +90,31 @@ class DynamoDBUserDatabase(Generic[UP, ID], BaseUserDatabase[UP, ID]):
     user_table: type[UP]
     oauth_account_table: type[DynamoDBBaseOAuthAccountTable] | None
     user_table_name: str
-    oauth_table_name: str | None
+    oauth_account_table_name: str | None
     _resource: Any | None
+    _resource_region: str | None
 
     def __init__(
         self,
         session: aioboto3.Session,
         user_table: type[UP],
         user_table_name: str,
         oauth_account_table: type[DynamoDBBaseOAuthAccountTable] | None = None,
-        oauth_table_name: str | None = None,
+        oauth_account_table_name: str | None = None,
         dynamodb_resource: Any | None = None,
+        dynamodb_resource_region: str | None = None,
     ):
         self.session = session
         self.user_table = user_table
         self.oauth_account_table = oauth_account_table
         self.user_table_name = user_table_name
-        self.oauth_table_name = oauth_table_name
+        self.oauth_account_table_name = oauth_account_table_name
 
         self._resource = dynamodb_resource
+        self._resource_region = dynamodb_resource_region
 
     @asynccontextmanager
-    async def _table(self, table_name: str):
+    async def _table(self, table_name: str, region: str | None = None):
         """Async context manager that yields a Table object.
 
         If a long-lived resource was provided at init, it's reused (no enter/exit).
@@ -121,7 +124,13 @@ async def _table(self, table_name: str):
             table = await self._resource.Table(table_name)
             yield table
         else:
-            async with self.session.resource("dynamodb") as dynamodb:
+            if region is None:
+                raise ValueError(
+                    "Parameter `region` must be specified when `dynamodb_resource` is omitted"
+                )
+            async with self.session.resource(
+                "dynamodb", region_name=region
+            ) as dynamodb:
                 table = await dynamodb.Table(table_name)
                 yield table
 
@@ -178,15 +187,15 @@ def _ensure_email_lower(self, data: dict[str, Any]) -> None:
     async def get(self, id: ID | str) -> UP | None:
         """Get a user by id."""
         id_str = self._ensure_id_str(id)
-        async with self._table(self.user_table_name) as table:
+        async with self._table(self.user_table_name, self._resource_region) as table:
             resp = await table.get_item(Key={"id": id_str})
             item = resp.get("Item")
             return self._item_to_user(item)
 
     async def get_by_email(self, email: str) -> UP | None:
         """Get a user by email (case-insensitive: emails are stored lowercased)."""
         email_norm = email.lower()
-        async with self._table(self.user_table_name) as table:
+        async with self._table(self.user_table_name, self._resource_region) as table:
             resp = await table.scan(
                 FilterExpression=Attr("email").eq(email_norm),
                 Limit=1,
@@ -198,10 +207,12 @@ async def get_by_email(self, email: str) -> UP | None:
 
     async def get_by_oauth_account(self, oauth: str, account_id: str) -> UP | None:
         """Find a user by oauth provider and provider account id."""
-        if self.oauth_account_table is None or self.oauth_table_name is None:
+        if self.oauth_account_table is None or self.oauth_account_table_name is None:
             raise NotImplementedError()
 
-        async with self._table(self.oauth_table_name) as oauth_table:
+        async with self._table(
+            self.oauth_account_table_name, self._resource_region
+        ) as oauth_table:
             resp = await oauth_table.scan(
                 FilterExpression=Attr("oauth_name").eq(oauth)
                 & Attr("account_id").eq(account_id),
@@ -227,7 +238,7 @@ async def create(self, create_dict: dict[str, Any]) -> UP:
 
         self._ensure_email_lower(item)
 
-        async with self._table(self.user_table_name) as table:
+        async with self._table(self.user_table_name, self._resource_region) as table:
             await table.put_item(Item=item)
 
             resp = await table.get_item(Key={"id": item["id"]})
@@ -241,7 +252,7 @@ async def create(self, create_dict: dict[str, Any]) -> UP:
     async def update(self, user: UP, update_dict: dict[str, Any]) -> 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) as table:
+        async with self._table(self.user_table_name, self._resource_region) as table:
             resp = await table.get_item(Key={"id": user_id})
             current = resp.get("Item", {})
             if not current:
@@ -263,12 +274,12 @@ async def update(self, user: UP, update_dict: dict[str, Any]) -> UP:
     async def delete(self, user: UP) -> None:
         """Delete a user."""
         user_id = self._extract_id_from_user(user)
-        async with self._table(self.user_table_name) as table:
+        async with self._table(self.user_table_name, self._resource_region) as table:
             await table.delete_item(Key={"id": user_id})
 
     async def add_oauth_account(self, user: UP, create_dict: dict[str, Any]) -> UP:
         """Add an OAuth account for `user`. Returns the refreshed user (UP)."""
-        if self.oauth_account_table is None or self.oauth_table_name is None:
+        if self.oauth_account_table is None or self.oauth_account_table_name is None:
             raise NotImplementedError()
 
         oauth_item = dict(create_dict)
@@ -278,7 +289,9 @@ async def add_oauth_account(self, user: UP, create_dict: dict[str, Any]) -> UP:
         user_id = self._extract_id_from_user(user)
         oauth_item["user_id"] = user_id
 
-        async with self._table(self.oauth_table_name) as oauth_table:
+        async with self._table(
+            self.oauth_account_table_name, self._resource_region
+        ) as oauth_table:
             await oauth_table.put_item(Item=oauth_item)
 
         refreshed_user = await self.get(user_id)
@@ -305,7 +318,7 @@ async def update_oauth_account(
         update_dict: dict[str, Any],
     ) -> UP:
         """Update an OAuth account and return the refreshed user (UP)."""
-        if self.oauth_account_table is None or self.oauth_table_name is None:
+        if self.oauth_account_table is None or self.oauth_account_table_name is None:
             raise NotImplementedError()
 
         oauth_id = None
@@ -328,7 +341,9 @@ async def update_oauth_account(
 
         oauth_id_str = self._ensure_id_str(oauth_id)
 
-        async with self._table(self.oauth_table_name) as oauth_table:
+        async with self._table(
+            self.oauth_account_table_name, self._resource_region
+        ) as oauth_table:
             resp = await oauth_table.get_item(Key={"id": oauth_id_str})
             current = resp.get("Item", {})
             if not current:

fastapi_users_db_dynamodb/access_token.py

@@ -8,7 +8,7 @@
 
 import uuid
 from contextlib import asynccontextmanager
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Generic, get_type_hints
 
 import aioboto3
@@ -45,27 +45,36 @@ class DynamoDBAccessTokenDatabase(Generic[AP], AccessTokenDatabase[AP]):
     access_token_table: type[AP]
     table_name: str
     _resource: Any | None
+    _resource_region: str | None
 
     def __init__(
         self,
         session: aioboto3.Session,
         access_token_table: type[AP],
         table_name: str,
         dynamodb_resource: Any | None = None,
+        dynamodb_resource_region: Any | None = None,
     ):
         self.session = session
         self.access_token_table = access_token_table
         self.table_name = table_name
         self._resource = dynamodb_resource
+        self._resource_region = dynamodb_resource_region
 
     @asynccontextmanager
-    async def _table(self, table_name: str):
+    async def _table(self, table_name: str, region: str | None = None):
         """Async context manager that yields a Table object."""
         if self._resource is not None:
             table = await self._resource.Table(table_name)
             yield table
         else:
-            async with self.session.resource("dynamodb") as dynamodb:
+            if region is None:
+                raise ValueError(
+                    "Parameter `region` must be specified when `dynamodb_resource` is omitted"
+                )
+            async with self.session.resource(
+                "dynamodb", region_name=region
+            ) as dynamodb:
                 table = await dynamodb.Table(table_name)
                 yield table
 
@@ -98,7 +107,7 @@ async def get_by_token(
         self, token: str, max_age: datetime | None = None
     ) -> AP | None:
         """Retrieve an access token by token string."""
-        async with self._table(self.table_name) as table:
+        async with self._table(self.table_name, self._resource_region) as table:
             resp = await table.get_item(Key={"token": self._ensure_token(token)})
             item = resp.get("Item")
 
@@ -119,11 +128,11 @@ async def create(self, create_dict: dict[str, Any]) -> AP:
         if "token" not in item or item["token"] is None:
             item["token"] = uuid.uuid4().hex[:43]
         if "created_at" not in item or not isinstance(item["created_at"], str):
-            item["created_at"] = datetime.utcnow().isoformat()
+            item["created_at"] = datetime.now(timezone.utc).isoformat()
         if isinstance(item.get("user_id"), uuid.UUID):
             item["user_id"] = str(item["user_id"])
 
-        async with self._table(self.table_name) as table:
+        async with self._table(self.table_name, self._resource_region) as table:
             await table.put_item(Item=item)
 
             resp = await table.get_item(Key={"token": item["token"]})
@@ -151,7 +160,7 @@ async def update(self, access_token: AP, update_dict: dict[str, Any]) -> AP:
         if isinstance(token_dict.get("created_at"), datetime):
             token_dict["created_at"] = token_dict["created_at"].isoformat()
 
-        async with self._table(self.table_name) as table:
+        async with self._table(self.table_name, self._resource_region) as table:
             await table.put_item(Item=token_dict)
 
             resp = await table.get_item(Key={"token": token_dict["token"]})
@@ -170,5 +179,5 @@ async def delete(self, access_token: AP) -> None:
         if token is None:
             raise ValueError("access_token has no 'token' field")
 
-        async with self._table(self.table_name) as table:
+        async with self._table(self.table_name, self._resource_region) as table:
             await table.delete_item(Key={"token": self._ensure_token(token)})

tests/conftest.py

@@ -4,6 +4,8 @@
 import pytest
 from fastapi_users import schemas
 
+DATABASE_REGION: str = "eu-central-1"
+
 
 class User(schemas.BaseUser):
     first_name: Optional[str]

tests/test_access_token.py

@@ -13,6 +13,7 @@
     DynamoDBAccessTokenDatabase,
     DynamoDBBaseAccessTokenTableUUID,
 )
+from tests.conftest import DATABASE_REGION
 
 
 class Base:
@@ -42,7 +43,10 @@ async def dynamodb_access_token_db(
         token_table_name = "access_tokens_test"
 
         user_db = DynamoDBUserDatabase(
-            session, DynamoDBBaseUserTableUUID, user_table_name
+            session,
+            DynamoDBBaseUserTableUUID,
+            user_table_name,
+            dynamodb_resource_region=DATABASE_REGION,
         )
         user = await user_db.create(
             {
@@ -52,7 +56,12 @@ async def dynamodb_access_token_db(
             }
         )
 
-        token_db = DynamoDBAccessTokenDatabase(session, AccessToken, token_table_name)
+        token_db = DynamoDBAccessTokenDatabase(
+            session,
+            AccessToken,
+            token_table_name,
+            dynamodb_resource_region=DATABASE_REGION,
+        )
 
         # Vorherigen Token löschen, falls er existiert
         token_obj = await token_db.get_by_token("TOKEN")

tests/test_users.py

@@ -11,6 +11,7 @@
     DynamoDBBaseUserTableUUID,
     DynamoDBUserDatabase,
 )
+from tests.conftest import DATABASE_REGION
 
 
 class Base:
@@ -40,7 +41,12 @@ async def dynamodb_user_db() -> AsyncGenerator[DynamoDBUserDatabase, None]:
         session = aioboto3.Session()
         table_name = "users_test"
 
-        db = DynamoDBUserDatabase(session, DynamoDBBaseUserTableUUID, table_name)
+        db = DynamoDBUserDatabase(
+            session,
+            DynamoDBBaseUserTableUUID,
+            table_name,
+            dynamodb_resource_region=DATABASE_REGION,
+        )
         yield db
 
 
@@ -57,6 +63,7 @@ async def dynamodb_user_db_oauth() -> AsyncGenerator[DynamoDBUserDatabase, None]
             user_table_name,
             OAuthAccount,
             oauth_table_name,
+            dynamodb_resource_region=DATABASE_REGION,
         )
         yield db