Add a _use_sys_pgcon context manager
Makes the code nicer, and I'm planning to add some more uses.
@@ -24,6 +24,7 @@ from typing import *
import asyncio
import collections
+import contextlib
import ipaddress
import json
import logging
@@ -580,12 +581,9 @@ class Server(ha_base.ClusterProtocol):
raise
async def load_sys_config(self):
- syscon = await self._acquire_sys_pgcon()
- try:
+ async with self._use_sys_pgcon() as syscon:
query = self.get_sys_query('sysconfig')
sys_config_json = await syscon.sql_fetch_val(query)
- finally:
- self._release_sys_pgcon()
return config.from_json(config.get_settings(), sys_config_json)
@@ -627,11 +625,8 @@ class Server(ha_base.ClusterProtocol):
if conn is not None:
json_data = await conn.sql_fetch_val(intro_query)
else:
- syscon = await self._acquire_sys_pgcon()
- try:
+ async with self._use_sys_pgcon() as syscon:
json_data = await syscon.sql_fetch_val(intro_query)
- finally:
- self._release_sys_pgcon()
return s_refl.parse_into(
base_schema=self._std_schema,
@@ -795,13 +790,10 @@ class Server(ha_base.ClusterProtocol):
self.release_pgcon(dbname, conn)
async def _introspect_dbs(self):
- syscon = await self._acquire_sys_pgcon()
- try:
+ async with self._use_sys_pgcon() as syscon:
dbs_query = self.get_sys_query('listdbs')
json_data = await syscon.sql_fetch_val(dbs_query)
dbnames = json.loads(json_data)
- finally:
- self._release_sys_pgcon()
async with taskgroup.TaskGroup(name='introspect DB extensions') as g:
for dbname in dbnames:
@@ -825,8 +817,7 @@ class Server(ha_base.ClusterProtocol):
self._roles = immutables.Map(roles)
async def _load_instance_data(self):
- syscon = await self._acquire_sys_pgcon()
- try:
+ async with self._use_sys_pgcon() as syscon:
result = await syscon.sql_fetch_val(b'''\
SELECT json::json FROM edgedbinstdata.instdata
WHERE key = 'instancedata';
@@ -886,9 +877,6 @@ class Server(ha_base.ClusterProtocol):
WHERE key = 'report_configs_typedesc';
''')
- finally:
- self._release_sys_pgcon()
-
def get_roles(self):
return self._roles
@@ -1127,9 +1115,16 @@ class Server(ha_base.ClusterProtocol):
def _release_sys_pgcon(self):
self._sys_pgcon_waiter.release()
- async def _cancel_pgcon_operation(self, pgcon) -> bool:
- syscon = await self._acquire_sys_pgcon()
+ @contextlib.asynccontextmanager
+ async def _use_sys_pgcon(self):
+ conn = await self._acquire_sys_pgcon()
try:
+ yield conn
+ finally:
+ self._release_sys_pgcon()
+
+ async def _cancel_pgcon_operation(self, pgcon) -> bool:
+ async with self._use_sys_pgcon() as syscon:
if pgcon.idle:
# pgcon could have received the query results while we
# were acquiring a system connection to cancel it.
@@ -1152,8 +1147,6 @@ class Server(ha_base.ClusterProtocol):
pgcon.finish_pg_cancellation()
return result == b'\x01'
- finally:
- self._release_sys_pgcon()
async def _cancel_and_discard_pgcon(self, pgcon, dbname) -> None:
try:
