app.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. import asyncio
  2. import asyncpg
  3. import jinja2
  4. import os
  5. import ujson
  6. from random import randint
  7. from operator import itemgetter
  8. from urllib.parse import parse_qs
  9. from routerling import Router
  10. router = Router()
  11. async def setup():
  12. global pool
  13. pool = await asyncpg.create_pool(
  14. user=os.getenv('PGUSER', 'benchmarkdbuser'),
  15. password=os.getenv('PGPASS', 'benchmarkdbpass'),
  16. database='hello_world',
  17. host='tfb-database',
  18. port=5432
  19. )
  20. CONTENT_TYPE = 'Content-Type'
  21. JSON = 'application/json'
  22. READ_ROW_SQL = 'SELECT "randomnumber", "id" FROM "world" WHERE id = $1'
  23. WRITE_ROW_SQL = 'UPDATE "world" SET "randomnumber"=$1 WHERE id=$2'
  24. ADDITIONAL_ROW = [0, 'Additional fortune added at request time.']
  25. pool = None
  26. key = itemgetter(1)
  27. json_dumps = ujson.dumps
  28. template = None
  29. path = os.path.join('templates', 'fortune.html')
  30. with open(path, 'r') as template_file:
  31. template_text = template_file.read()
  32. template = jinja2.Template(template_text)
  33. loop = asyncio.get_event_loop()
  34. loop.run_until_complete(setup())
  35. def get_num_queries(scope):
  36. try:
  37. query_string = scope['query_string']
  38. query_count = int(parse_qs(query_string)[b'queries'][0])
  39. except (KeyError, IndexError, ValueError):
  40. return 1
  41. if query_count < 1:
  42. return 1
  43. if query_count > 500:
  44. return 500
  45. return query_count
  46. async def json_serialization(r, w, c):
  47. """
  48. Test type 1: JSON Serialization
  49. """
  50. content = json_dumps({'message': 'Hello, world!'})
  51. w.headers = CONTENT_TYPE, JSON
  52. w.body = content
  53. async def single_database_query(r, w, c):
  54. """
  55. Test type 2: Single database object
  56. """
  57. row_id = randint(1, 10000)
  58. connection = await pool.acquire()
  59. try:
  60. number = await connection.fetchval(READ_ROW_SQL, row_id)
  61. world = {'id': row_id, 'randomNumber': number}
  62. finally:
  63. await pool.release(connection)
  64. content = json_dumps(world).encode('utf-8')
  65. w.headers = CONTENT_TYPE, JSON
  66. w.body = content
  67. async def multiple_database_queries(r, w, c):
  68. """
  69. Test type 3: Multiple database queries
  70. """
  71. num_queries = get_num_queries(r._scope)
  72. row_ids = [randint(1, 10000) for _ in range(num_queries)]
  73. worlds = []
  74. connection = await pool.acquire()
  75. try:
  76. statement = await connection.prepare(READ_ROW_SQL)
  77. for row_id in row_ids:
  78. number = await statement.fetchval(row_id)
  79. worlds.append({'id': row_id, 'randomNumber': number})
  80. finally:
  81. await pool.release(connection)
  82. content = json_dumps(worlds).encode('utf-8')
  83. w.headers = CONTENT_TYPE, JSON
  84. w.body = content
  85. async def fortunes(r, w, c):
  86. """
  87. Test type 4: Fortunes
  88. """
  89. connection = await pool.acquire()
  90. try:
  91. fortunes = await connection.fetch('SELECT * FROM Fortune')
  92. finally:
  93. await pool.release(connection)
  94. fortunes.append(ADDITIONAL_ROW)
  95. fortunes.sort(key=key)
  96. content = template.render(fortunes=fortunes).encode('utf-8')
  97. w.headers = CONTENT_TYPE, 'text/html; charset=utf-8'
  98. w.body = content
  99. async def database_updates(r, w, c):
  100. """
  101. Test type 5: Database updates
  102. """
  103. num_queries = get_num_queries(r._scope)
  104. updates = [(randint(1, 10000), randint(1, 10000)) for _ in range(num_queries)]
  105. worlds = [{'id': row_id, 'randomNumber': number} for row_id, number in updates]
  106. connection = await pool.acquire()
  107. try:
  108. statement = await connection.prepare(READ_ROW_SQL)
  109. for row_id, _ in updates:
  110. await statement.fetchval(row_id)
  111. await connection.executemany(WRITE_ROW_SQL, updates)
  112. finally:
  113. await pool.release(connection)
  114. content = json_dumps(worlds).encode('utf-8')
  115. w.headers = CONTENT_TYPE, JSON
  116. w.body = content
  117. async def plaintext(r, w, c):
  118. """
  119. Test type 6: Plaintext
  120. """
  121. content = 'Hello, world!'
  122. w.headers = CONTENT_TYPE, 'text/plain; charset=utf-8'
  123. w.body = content
  124. async def handle_404(r, w, c):
  125. content = b'Not found'
  126. w.headers = CONTENT_TYPE, b'text/plain; charset=utf-8'
  127. w.body = content
  128. router.HTTP('/json', json_serialization)
  129. router.HTTP('/db', single_database_query)
  130. router.HTTP('/queries', multiple_database_queries)
  131. router.HTTP('/fortunes', fortunes)
  132. router.HTTP('/updates', database_updates)
  133. router.HTTP('/plaintext', plaintext)