Ticket #33882: failure_server.py

import asyncio
import os
import random

from aiohttp import web
from asgiref.sync import sync_to_async
from django.conf import settings
from django.contrib.auth import get_user_model
from django.core import management
from django.core.wsgi import get_wsgi_application
from django.db import close_old_connections
from django.db.transaction import Atomic
from functools import wraps, WRAPPER_ASSIGNMENTS


# https://stackoverflow.com/questions/74575922/how-to-use-transaction-with-async-functions-in-django
class AsyncAtomicContextManager(Atomic):
    def __init__(self, using=None, savepoint=True, durable=False):
        super().__init__(using, savepoint, durable)

    async def __aenter__(self):
        await sync_to_async(super().__enter__, thread_sensitive=True)()
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        await sync_to_async(super().__exit__, thread_sensitive=True)(exc_type, exc_value, traceback)


def aatomic(using=None, savepoint=True, durable=False):

    def decorator(func):
        @wraps(func, assigned=WRAPPER_ASSIGNMENTS)
        async def _wrapped(*args, **kwargs):
            async with AsyncAtomicContextManager():
                return await func(*args, **kwargs)

        return _wrapped

    return decorator


# Configure Django so the ORM works
if not settings.configured:
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "failure_settings")

application = get_wsgi_application()
management.call_command("migrate")


@aatomic()
async def i_am_atomic(num):
    user = await User.objects.aget(username=f"user{num}")
    return user


async def handle(request):
    await sync_to_async(close_old_connections)()
    num = request.query["num"]
    user = await i_am_atomic(num)
    return web.Response(text=user.username)


app = web.Application()
app.add_routes([web.get("/", handle),
                web.get("/{name}", handle)])


if __name__ == '__main__':
    # Create users
    User = get_user_model()
    for i in range(100):
        User.objects.get_or_create(username=f"user{i}", email=f"user{i}@aaa.com")

    web.run_app(app)