|
17 | 17 |
|
18 | 18 | db_pool = None |
19 | 19 |
|
| 20 | +PG_POOL_SIZE = 4 |
| 21 | + |
| 22 | +class NoResetConnection(asyncpg.Connection): |
| 23 | + __slots__ = () |
| 24 | + |
| 25 | + def get_reset_query(self): |
| 26 | + return "" |
| 27 | + |
20 | 28 | async def db_setup(): |
21 | 29 | global db_pool |
22 | 30 | db_pool = await asyncpg.create_pool( |
23 | 31 | user=os.getenv('PGUSER', 'benchmarkdbuser'), |
24 | 32 | password=os.getenv('PGPASS', 'benchmarkdbpass'), |
25 | 33 | database='hello_world', |
26 | 34 | host='tfb-database', |
27 | | - port=5432 |
| 35 | + port=5432, |
| 36 | + min_size=PG_POOL_SIZE, |
| 37 | + max_size=PG_POOL_SIZE, |
| 38 | + connection_class=NoResetConnection, |
28 | 39 | ) |
29 | 40 |
|
30 | 41 | READ_ROW_SQL = 'SELECT "randomnumber", "id" FROM "world" WHERE id = $1' |
@@ -102,15 +113,10 @@ async def multiple_database_queries(scope, receive, send): |
102 | 113 | row_ids = random.sample(range(1, 10000), num_queries) |
103 | 114 | worlds = [ ] |
104 | 115 |
|
105 | | - db_conn = await db_pool.acquire() |
106 | | - try: |
107 | | - statement = await db_conn.prepare(READ_ROW_SQL) |
108 | | - for row_id in row_ids: |
109 | | - number = await statement.fetchval(row_id) |
110 | | - worlds.append( {'id': row_id, 'randomNumber': number} ) |
111 | | - finally: |
112 | | - await db_pool.release(db_conn) |
| 116 | + async with db_pool.acquire() as db_conn: |
| 117 | + rows = await db_conn.fetchmany(READ_ROW_SQL, [ (v, ) for v in row_ids ] ) |
113 | 118 |
|
| 119 | + worlds = [ { 'id': row_id, 'randomNumber': number[0] } for row_id, number in zip(row_ids, rows) ] |
114 | 120 | content = jsonify(worlds) |
115 | 121 | await send(JSON_RESPONSE) |
116 | 122 | await send({ |
@@ -154,14 +160,9 @@ async def database_updates(scope, receive, send): |
154 | 160 |
|
155 | 161 | worlds = [ {"id": row_id, "randomNumber": number} for row_id, number in updates ] |
156 | 162 |
|
157 | | - db_conn = await db_pool.acquire() |
158 | | - try: |
159 | | - statement = await db_conn.prepare(READ_ROW_SQL) |
160 | | - for row_id, _ in updates: |
161 | | - await statement.fetchval(row_id) |
| 163 | + async with db_pool.acquire() as db_conn: |
| 164 | + await db_conn.executemany(READ_ROW_SQL, [ (i[0], ) for i in updates ] ) |
162 | 165 | await db_conn.executemany(WRITE_ROW_SQL, updates) |
163 | | - finally: |
164 | | - await db_pool.release(db_conn) |
165 | 166 |
|
166 | 167 | content = jsonify(worlds) |
167 | 168 | await send(JSON_RESPONSE) |
|
0 commit comments