coreapp.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. import asyncio
  2. import asyncpg
  3. import os
  4. import jinja2
  5. import orjson
  6. from logging import getLogger
  7. from random import randint
  8. from operator import itemgetter
  9. from apidaora.asgi.app import asgi_app
  10. from apidaora.asgi.responses import (
  11. JSON_RESPONSE, HTML_RESPONSE, PLAINTEXT_RESPONSE
  12. )
  13. from apidaora.asgi.router import Route, make_router
  14. from apidaora.method import MethodType
  15. logger = getLogger(__name__)
  16. READ_ROW_SQL = 'SELECT "randomnumber", "id" FROM "world" WHERE id = $1'
  17. READ_ROW_SQL_TO_UPDATE = 'SELECT "id", "randomnumber" FROM "world" WHERE id = $1'
  18. WRITE_ROW_SQL = 'UPDATE "world" SET "randomnumber"=$1 WHERE id=$2'
  19. ADDITIONAL_ROW = [0, 'Additional fortune added at request time.']
  20. async def setup_database():
  21. global connection_pool
  22. connection_pool = await asyncpg.create_pool(
  23. user=os.getenv('PGUSER', 'benchmarkdbuser'),
  24. password=os.getenv('PGPASS', 'benchmarkdbpass'),
  25. database='hello_world',
  26. host='tfb-database',
  27. port=5432
  28. )
  29. def load_fortunes_template():
  30. path = os.path.join('templates', 'fortune.html')
  31. with open(path, 'r') as template_file:
  32. template_text = template_file.read()
  33. return jinja2.Template(template_text)
  34. def get_num_queries(queries):
  35. try:
  36. query_count = int(queries[0])
  37. except (ValueError, TypeError):
  38. return 1
  39. if query_count < 1:
  40. return 1
  41. if query_count > 500:
  42. return 500
  43. return query_count
  44. connection_pool = None
  45. sort_fortunes_key = itemgetter(1)
  46. template = load_fortunes_template()
  47. loop = asyncio.get_event_loop()
  48. loop.run_until_complete(setup_database())
  49. def json_serialization(request):
  50. return JSON_RESPONSE, orjson.dumps({'message': 'Hello, world!'})
  51. async def single_database_query(request):
  52. row_id = randint(1, 10000)
  53. async with connection_pool.acquire() as connection:
  54. number = await connection.fetchval(READ_ROW_SQL, row_id)
  55. return JSON_RESPONSE, orjson.dumps(
  56. {'id': row_id, 'randomNumber': number}
  57. )
  58. async def multiple_database_queries(request):
  59. num_queries = get_num_queries(request.query_dict.get('queries', 1))
  60. row_ids = [randint(1, 10000) for _ in range(num_queries)]
  61. worlds = []
  62. async with connection_pool.acquire() as connection:
  63. statement = await connection.prepare(READ_ROW_SQL)
  64. for row_id in row_ids:
  65. number = await statement.fetchval(row_id)
  66. worlds.append(
  67. dict(
  68. id=row_id,
  69. randomNumber=number
  70. )
  71. )
  72. return JSON_RESPONSE, orjson.dumps(worlds)
  73. async def fortunes(request):
  74. async with connection_pool.acquire() as connection:
  75. fortunes = await connection.fetch('SELECT * FROM Fortune')
  76. fortunes.append(ADDITIONAL_ROW)
  77. fortunes.sort(key=sort_fortunes_key)
  78. content = template.render(fortunes=fortunes).encode('utf-8')
  79. return HTML_RESPONSE, content
  80. async def database_updates(request):
  81. worlds = []
  82. updates = set()
  83. queries = request.query_dict.get('queries', 1)
  84. async with connection_pool.acquire() as connection:
  85. statement = await connection.prepare(READ_ROW_SQL_TO_UPDATE)
  86. for _ in range(get_num_queries(queries)):
  87. record = await statement.fetchrow(randint(1, 10000))
  88. world = dict(
  89. id=record['id'], randomNumber=record['randomnumber']
  90. )
  91. world['randomNumber'] = randint(1, 10000)
  92. worlds.append(world)
  93. updates.add((world['id'], world['randomNumber']))
  94. await connection.executemany(WRITE_ROW_SQL, updates)
  95. return JSON_RESPONSE, orjson.dumps(worlds)
  96. def plaintext(request):
  97. return PLAINTEXT_RESPONSE, b'Hello, world!'
  98. routes = (
  99. Route('/json', MethodType.GET, json_serialization),
  100. Route('/db', MethodType.GET, single_database_query),
  101. Route('/queries', MethodType.GET, multiple_database_queries, has_query=True),
  102. Route('/fortunes', MethodType.GET, fortunes),
  103. Route('/updates', MethodType.GET, database_updates, has_query=True),
  104. Route('/plaintext', MethodType.GET, plaintext),
  105. )
  106. router = make_router(routes)
  107. app = asgi_app(router)