Implement aiobotocore patch

Author: pentest-gg

fastapi_users_db_dynamodb/__init__.py

@@ -26,6 +26,7 @@
 from fastapi_users.db.base import BaseUserDatabase
 from fastapi_users.models import ID, OAP, UP
 
+from fastapi_users_db_dynamodb._aioboto3_patch import *  # noqa: F403
 from fastapi_users_db_dynamodb.generics import GUID
 
 __version__ = "1.0.0"

fastapi_users_db_dynamodb/_aioboto3_patch.py

@@ -0,0 +1,87 @@
+import inspect
+from binascii import crc32
+
+import aiobotocore.endpoint
+import aiobotocore.retryhandler
+from aiobotocore.endpoint import HttpxStreamingBody, StreamingBody
+from aiobotocore.retryhandler import ChecksumError, logger
+
+try:
+    import httpx
+except ImportError:
+    httpx = None
+
+
+async def _fixed_check_response(self, attempt_number, response):
+    http_response = response[0]
+    expected_crc = http_response.headers.get(self._header_name)
+    if expected_crc is None:
+        logger.debug(
+            "crc32 check skipped, the %s header is not in the http response.",
+            self._header_name,
+        )
+    else:
+        if inspect.isawaitable(http_response.content):
+            data_buf = await http_response.content
+        else:
+            data_buf = http_response.content
+
+        actual_crc32 = crc32(data_buf) & 0xFFFFFFFF
+        if not actual_crc32 == int(expected_crc):
+            logger.debug(
+                "retry needed: crc32 check failed, expected != actual: %s != %s",
+                int(expected_crc),
+                actual_crc32,
+            )
+            raise ChecksumError(
+                checksum_type="crc32",
+                expected_checksum=int(expected_crc),
+                actual_checksum=actual_crc32,
+            )
+
+
+async def convert_to_response_dict(http_response, operation_model):
+    """Convert an HTTP response object to a request dict.
+
+    This converts the HTTP response object to a dictionary.
+
+    :type http_response: botocore.awsrequest.AWSResponse
+    :param http_response: The HTTP response from an AWS service request.
+
+    :rtype: dict
+    :return: A response dictionary which will contain the following keys:
+        * headers (dict)
+        * status_code (int)
+        * body (string or file-like object)
+
+    """
+    response_dict = {
+        "headers": http_response.headers,
+        "status_code": http_response.status_code,
+        "context": {
+            "operation_name": operation_model.name,
+        },
+    }
+    if response_dict["status_code"] >= 300:
+        if inspect.isawaitable(http_response.content):
+            response_dict["body"] = await http_response.content
+        else:
+            response_dict["body"] = http_response.content
+    elif operation_model.has_event_stream_output:
+        response_dict["body"] = http_response.raw
+    elif operation_model.has_streaming_output:
+        if httpx and isinstance(http_response.raw, httpx.Response):
+            response_dict["body"] = HttpxStreamingBody(http_response.raw)
+        else:
+            length = response_dict["headers"].get("content-length")
+            response_dict["body"] = StreamingBody(http_response.raw, length)
+    else:
+        if inspect.isawaitable(http_response.content):
+            response_dict["body"] = await http_response.content
+        else:
+            response_dict["body"] = http_response.content
+    return response_dict
+
+
+aiobotocore.retryhandler.AioCRC32Checker._check_response = _fixed_check_response
+aiobotocore.endpoint.convert_to_response_dict = convert_to_response_dict

fastapi_users_db_dynamodb/access_token.py

@@ -15,6 +15,8 @@
 from fastapi_users.authentication.strategy.db import AP, AccessTokenDatabase
 from fastapi_users.models import ID
 
+from fastapi_users_db_dynamodb._aioboto3_patch import *  # noqa: F403
+
 
 class DynamoDBBaseAccessTokenTable(Generic[ID]):
     """Base access token table schema for DynamoDB."""

pyproject.toml

@@ -37,6 +37,8 @@ dependencies = [
     "asgi_lifespan",
     "ruff",
     "moto[all]",
+    "types-aioboto3",
+    "types-aiobotocore",
 ]
 
 [tool.hatch.envs.default.scripts]
@@ -84,8 +86,8 @@ classifiers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "fastapi-users >= 10.0.0",
-    "aioboto3 >= 11.0.0",
+    "fastapi-users >= 14.0.0",
+    "aioboto3 >= 15.0.0",
 ]
 
 [project.urls]

tests/tables.py

@@ -0,0 +1,21 @@
+import aioboto3
+
+
+async def ensure_table_exists(session: aioboto3.Session, table_name: str, region: str):
+    async with session.resource("dynamodb", region_name=region) as client:
+        try:
+            await client.describe_table(TableName=table_name)
+        except client.exceptions.ResourceNotFoundException:
+            await client.create_table(
+                TableName=table_name,
+                KeySchema=[
+                    {"AttributeName": "id", "KeyType": "HASH"},
+                ],
+                AttributeDefinitions=[
+                    {"AttributeName": "id", "AttributeType": "S"},
+                ],
+                ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5},
+            )
+
+            waiter = client.get_waiter("table_exists")
+            await waiter.wait(TableName=table_name)

tests/test_access_token.py

@@ -9,11 +9,13 @@
 from pydantic import UUID4
 
 from fastapi_users_db_dynamodb import DynamoDBBaseUserTableUUID, DynamoDBUserDatabase
+from fastapi_users_db_dynamodb._aioboto3_patch import *  # noqa: F403
 from fastapi_users_db_dynamodb.access_token import (
     DynamoDBAccessTokenDatabase,
     DynamoDBBaseAccessTokenTableUUID,
 )
 from tests.conftest import DATABASE_REGION
+from tests.tables import ensure_table_exists
 
 
 class Base:
@@ -41,6 +43,8 @@ async def dynamodb_access_token_db(
         session = aioboto3.Session()
         user_table_name = "users_test"
         token_table_name = "access_tokens_test"
+        await ensure_table_exists(session, user_table_name, DATABASE_REGION)
+        await ensure_table_exists(session, token_table_name, DATABASE_REGION)
 
         user_db = DynamoDBUserDatabase(
             session,

tests/test_users.py

@@ -11,7 +11,9 @@
     DynamoDBBaseUserTableUUID,
     DynamoDBUserDatabase,
 )
+from fastapi_users_db_dynamodb._aioboto3_patch import *  # noqa: F403
 from tests.conftest import DATABASE_REGION
+from tests.tables import ensure_table_exists
 
 
 class Base:
@@ -40,6 +42,7 @@ async def dynamodb_user_db() -> AsyncGenerator[DynamoDBUserDatabase, None]:
     with mock_aws():
         session = aioboto3.Session()
         table_name = "users_test"
+        await ensure_table_exists(session, table_name, DATABASE_REGION)
 
         db = DynamoDBUserDatabase(
             session,
@@ -56,6 +59,8 @@ async def dynamodb_user_db_oauth() -> AsyncGenerator[DynamoDBUserDatabase, None]
         session = aioboto3.Session()
         user_table_name = "users_test_oauth"
         oauth_table_name = "oauth_accounts_test"
+        await ensure_table_exists(session, user_table_name, DATABASE_REGION)
+        await ensure_table_exists(session, oauth_table_name, DATABASE_REGION)
 
         db = DynamoDBUserDatabase(
             session,