TechEmpower.FrameworkBenchmarks/frameworks/Python/aiohttp/app/views.py

import platform
from operator import attrgetter, itemgetter
from pathlib import Path
from random import randint, sample

import jinja2
from aiohttp.web import Response
from sqlalchemy import bindparam, select
from sqlalchemy.orm.attributes import flag_modified

from .models import Fortune, World

if platform.python_implementation() == "PyPy":
    from aiohttp.web import json_response
else:
    from orjson import dumps

    def json_response(payload):
        return Response(
            body=dumps(payload),
            content_type="application/json",
        )

ADDITIONAL_FORTUNE_ORM = Fortune(id=0, message='Additional fortune added at request time.')
ADDITIONAL_FORTUNE_ROW = {'id': 0, 'message': 'Additional fortune added at request time.'}
READ_ROW_SQL = 'SELECT "randomnumber", "id" FROM "world" WHERE id = $1'
READ_SELECT_ORM = select(World.randomnumber).where(World.id == bindparam("id"))
READ_FORTUNES_ORM = select(Fortune.id, Fortune.message)
WRITE_ROW_SQL = 'UPDATE "world" SET "randomnumber"=$2 WHERE id=$1'
WRITE_ROW_BATCH_SQL = 'UPDATE "world" w SET "randomnumber"=u.new_val FROM (SELECT unnest($1::int[]) as id, unnest($2::int[]) as new_val) u WHERE w.id = u.id'

template_path = Path(__file__).parent / 'templates' / 'fortune.jinja'
template = jinja2.Template(template_path.read_text())
sort_fortunes_orm = attrgetter('message')
sort_fortunes_raw = itemgetter('message')


def get_num_queries(request):
    try:
        num_queries = int(request.match_info['queries'])
    except (KeyError, ValueError):
        return 1
    if num_queries < 1:
        return 1
    if num_queries > 500:
        return 500
    return num_queries


async def json(request):
    """
    Test 1
    """
    return json_response({'message': 'Hello, World!'})


async def single_database_query_orm(request):
    """
    Test 2 ORM
    """
    id_ = randint(1, 10000)
    async with request.app['db_session']() as sess:
        num = await sess.scalar(READ_SELECT_ORM, {"id": id_})
    return json_response({'id': id_, 'randomNumber': num})


async def single_database_query_raw(request):
    """
    Test 2 RAW
    """
    id_ = randint(1, 10000)

    async with request.app['pg'].acquire() as conn:
        r = await conn.fetchval(READ_ROW_SQL, id_)
    return json_response({'id': id_, 'randomNumber': r})


async def multiple_database_queries_orm(request):
    """
    Test 3 ORM
    """
    num_queries = get_num_queries(request)

    ids = [randint(1, 10000) for _ in range(num_queries)]

    result = []
    async with request.app['db_session']() as sess:
        for id_ in ids:
            num = await sess.scalar(READ_SELECT_ORM, {"id": id_})
            result.append({'id': id_, 'randomNumber': num})
    return json_response(result)


async def multiple_database_queries_raw(request):
    """
    Test 3 RAW
    """
    num_queries = get_num_queries(request)

    ids = [(randint(1, 10000), ) for _ in range(num_queries)]

    async with request.app['pg'].acquire() as conn:
        rows = await conn.fetchmany(READ_ROW_SQL, ids)
        result = [{'id': id_[0], 'randomNumber': row[0]} for id_, row in zip(ids, rows)]
    return json_response(result)


async def fortunes(request):
    """
    Test 4 ORM
    """
    async with request.app['db_session']() as sess:
        ret = await sess.execute(READ_FORTUNES_ORM)
        fortunes = ret.all()
    fortunes.append(ADDITIONAL_FORTUNE_ORM)
    fortunes.sort(key=sort_fortunes_orm)
    content = template.render(fortunes=fortunes)
    return Response(text=content, content_type='text/html')


async def fortunes_raw(request):
    """
    Test 4 RAW
    """
    async with request.app['pg'].acquire() as conn:
        fortunes = await conn.fetch('SELECT * FROM Fortune')
    fortunes.append(ADDITIONAL_FORTUNE_ROW)
    fortunes.sort(key=sort_fortunes_raw)
    content = template.render(fortunes=fortunes)
    return Response(text=content, content_type='text/html')


async def updates(request):
    """
    Test 5 ORM
    """
    num_queries = get_num_queries(request)
    update_ids = sample(range(1, 10001), num_queries)
    update_ids.sort()
    updates = tuple(zip(update_ids, sample(range(1, 10001), num_queries)))
    worlds = [{'id': row_id, 'randomNumber': number} for row_id, number in updates]

    async with request.app['db_session'].begin() as sess:
        for id_, number in updates:
            world = await sess.get(World, id_, populate_existing=True)
            world.randomnumber = number
            # Force sqlalchemy to UPDATE entry even if the value has not changed
            # doesn't make sense in a real application, added only to pass tests.
            flag_modified(world, "randomnumber")
    return json_response(worlds)

async def updates_raw(request):
    """
    Test 5 RAW
    """
    num_queries = get_num_queries(request)
    update_ids = sample(range(1, 10001), num_queries)
    update_ids.sort()

    numbers = sample(range(1, 10001), num_queries)
    fetch_params = [(i,) for i in update_ids]
    row_updates = [*zip(update_ids, numbers)]

    async with request.app['pg'].acquire() as conn:
        # the result of this is the int previous random number which we don't actually use
        await conn.executemany(READ_ROW_SQL, fetch_params)
        if num_queries <= 5:
            await conn.executemany(WRITE_ROW_SQL, row_updates)
        else:
            await conn.execute(WRITE_ROW_BATCH_SQL, update_ids, numbers)

    worlds = [{'id': row_id, 'randomNumber': number} for row_id, number in row_updates]
    return json_response(worlds)


async def plaintext(request):
    """
    Test 6
    """
    return Response(body=b'Hello, World!', content_type='text/plain')