app.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  1. import asyncio
  2. import asyncpg
  3. import os
  4. import jinja2
  5. from logging import getLogger
  6. from apidaora import (
  7. HTMLResponse,
  8. JSONRequestBody,
  9. JSONResponse,
  10. MethodType,
  11. PlainResponse,
  12. appdaora,
  13. header_param,
  14. path,
  15. )
  16. from jsondaora import jsondaora
  17. from random import randint
  18. from operator import itemgetter
  19. from http import HTTPStatus
  20. from typing import List, TypedDict, Optional
  21. logger = getLogger(__name__)
  22. READ_ROW_SQL = 'SELECT "randomnumber" FROM "world" WHERE id = $1'
  23. READ_ROW_SQL_TO_UPDATE = 'SELECT "id", "randomnumber" FROM "world" WHERE id = $1'
  24. WRITE_ROW_SQL = 'UPDATE "world" SET "randomnumber"=$1 WHERE id=$2'
  25. ADDITIONAL_ROW = [0, 'Additional fortune added at request time.']
  26. async def setup_database():
  27. global connection_pool
  28. connection_pool = await asyncpg.create_pool(
  29. user=os.getenv('PGUSER', 'benchmarkdbuser'),
  30. password=os.getenv('PGPASS', 'benchmarkdbpass'),
  31. database='hello_world',
  32. host='tfb-database',
  33. port=5432
  34. )
  35. def load_fortunes_template():
  36. path = os.path.join('templates', 'fortune.html')
  37. with open(path, 'r') as template_file:
  38. template_text = template_file.read()
  39. return jinja2.Template(template_text)
  40. def get_num_queries(queries):
  41. try:
  42. query_count = int(queries)
  43. except (ValueError, TypeError):
  44. return 1
  45. if query_count < 1:
  46. return 1
  47. if query_count > 500:
  48. return 500
  49. return query_count
  50. connection_pool = None
  51. sort_fortunes_key = itemgetter(1)
  52. template = load_fortunes_template()
  53. loop = asyncio.get_event_loop()
  54. loop.run_until_complete(setup_database())
  55. @jsondaora
  56. class JsonResponse(JSONResponse):
  57. class Body(TypedDict):
  58. message: str
  59. body: Body
  60. @path('/json', MethodType.GET)
  61. async def json_serialization() -> JsonResponse:
  62. return JsonResponse(
  63. body=JsonResponse.Body(message='Hello, world!')
  64. )
  65. @jsondaora
  66. class SingleDatabaseResponse(JSONResponse):
  67. class Body(TypedDict):
  68. id: int
  69. randomNumber: float
  70. body: Body
  71. @path('/db', MethodType.GET)
  72. async def single_database_query() -> SingleDatabaseResponse:
  73. row_id = randint(1, 10000)
  74. async with connection_pool.acquire() as connection:
  75. number = await connection.fetchval(READ_ROW_SQL, row_id)
  76. return SingleDatabaseResponse(
  77. body=SingleDatabaseResponse.Body(id=row_id, randomNumber=number)
  78. )
  79. @jsondaora
  80. class MultipleDatabaseObject(TypedDict):
  81. id: int
  82. randomNumber: float
  83. @jsondaora
  84. class MultipleDatabaseResponse(JSONResponse):
  85. body: List[MultipleDatabaseObject]
  86. @path('/queries', MethodType.GET)
  87. async def multiple_database_queries(
  88. queries: Optional[str] = None
  89. ) -> MultipleDatabaseResponse:
  90. num_queries = get_num_queries(queries)
  91. row_ids = [randint(1, 10000) for _ in range(num_queries)]
  92. worlds = []
  93. async with connection_pool.acquire() as connection:
  94. statement = await connection.prepare(READ_ROW_SQL)
  95. for row_id in row_ids:
  96. number = await statement.fetchval(row_id)
  97. worlds.append(
  98. MultipleDatabaseObject(
  99. id=row_id,
  100. randomNumber=number
  101. )
  102. )
  103. return MultipleDatabaseResponse(body=worlds)
  104. @jsondaora
  105. class FortunesResponse(HTMLResponse):
  106. body: str
  107. @path('/fortunes', MethodType.GET)
  108. async def fortunes() -> FortunesResponse:
  109. async with connection_pool.acquire() as connection:
  110. fortunes = await connection.fetch('SELECT * FROM Fortune')
  111. fortunes.append(ADDITIONAL_ROW)
  112. fortunes.sort(key=sort_fortunes_key)
  113. content = template.render(fortunes=fortunes)
  114. return FortunesResponse(
  115. body=content
  116. )
  117. @path('/updates', MethodType.GET)
  118. async def database_updates(
  119. queries: Optional[str] = None
  120. ) -> MultipleDatabaseResponse:
  121. worlds = []
  122. updates = set()
  123. async with connection_pool.acquire() as connection:
  124. statement = await connection.prepare(READ_ROW_SQL_TO_UPDATE)
  125. for _ in range(get_num_queries(queries)):
  126. record = await statement.fetchrow(randint(1, 10000))
  127. world = MultipleDatabaseObject(
  128. id=record['id'], randomNumber=record['randomnumber']
  129. )
  130. world['randomNumber'] = randint(1, 10000)
  131. worlds.append(world)
  132. updates.add((world['id'], world['randomNumber']))
  133. await connection.executemany(WRITE_ROW_SQL, updates)
  134. return MultipleDatabaseResponse(
  135. body=worlds
  136. )
  137. @jsondaora
  138. class TextResponse(PlainResponse):
  139. body: str
  140. @path('/plaintext', MethodType.GET)
  141. async def plaintext() -> TextResponse:
  142. return TextResponse(
  143. body='Hello, world!'
  144. )
  145. app = appdaora(operations=[
  146. json_serialization,
  147. single_database_query,
  148. multiple_database_queries,
  149. fortunes,
  150. database_updates,
  151. plaintext
  152. ])