main.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. import os
  2. from pathlib import Path
  3. import aiohttp_jinja2
  4. import aiopg.sa
  5. import asyncpg
  6. import jinja2
  7. from aiohttp import web
  8. from sqlalchemy.engine.url import URL
  9. from .views import (
  10. json,
  11. single_database_query_orm,
  12. multiple_database_queries_orm,
  13. fortunes,
  14. updates,
  15. plaintext,
  16. single_database_query_raw,
  17. multiple_database_queries_raw,
  18. fortunes_raw,
  19. updates_raw,
  20. )
  21. CONNECTION_ORM = os.getenv('CONNECTION', 'ORM').upper() == 'ORM'
  22. THIS_DIR = Path(__file__).parent
  23. def pg_dsn() -> str:
  24. """
  25. :return: DSN url suitable for sqlalchemy and aiopg.
  26. """
  27. return str(URL(
  28. database='hello_world',
  29. password=os.getenv('PGPASS', 'benchmarkdbpass'),
  30. host=os.getenv('DBHOST', 'localhost'),
  31. port='5432',
  32. username=os.getenv('PGUSER', 'benchmarkdbuser'),
  33. drivername='postgres',
  34. ))
  35. async def startup(app: web.Application):
  36. dsn = pg_dsn()
  37. min_size, max_size = 50, 100
  38. if CONNECTION_ORM:
  39. app['pg'] = await aiopg.sa.create_engine(dsn=dsn, minsize=min_size, maxsize=max_size, loop=app.loop)
  40. else:
  41. app['pg'] = await asyncpg.create_pool(dsn=dsn, min_size=min_size, max_size=max_size, loop=app.loop)
  42. async def cleanup(app: web.Application):
  43. if CONNECTION_ORM:
  44. app['pg'].close()
  45. await app['pg'].wait_closed()
  46. else:
  47. await app['pg'].close()
  48. def setup_routes(app):
  49. if CONNECTION_ORM:
  50. app.router.add_get('/json', json)
  51. app.router.add_get('/db', single_database_query_orm)
  52. app.router.add_get('/queries/{queries:.*}', multiple_database_queries_orm)
  53. app.router.add_get('/fortunes', fortunes)
  54. app.router.add_get('/updates/{queries:.*}', updates)
  55. app.router.add_get('/plaintext', plaintext)
  56. else:
  57. app.router.add_get('/db', single_database_query_raw)
  58. app.router.add_get('/queries/{queries:.*}', multiple_database_queries_raw)
  59. app.router.add_get('/fortunes', fortunes_raw)
  60. app.router.add_get('/updates/{queries:.*}', updates_raw)
  61. def create_app(loop):
  62. app = web.Application(loop=loop)
  63. jinja2_loader = jinja2.FileSystemLoader(str(THIS_DIR / 'templates'))
  64. aiohttp_jinja2.setup(app, loader=jinja2_loader)
  65. app.on_startup.append(startup)
  66. app.on_cleanup.append(cleanup)
  67. setup_routes(app)
  68. return app