Fix Connection.reset() on read-only connections

Fixes: #48

Author: elprans

asyncpg/_testbase.py

@@ -70,17 +70,19 @@ def wrapper(self, *args, __meth__=meth, **kwargs):
 
 class TestCase(unittest.TestCase, metaclass=TestCaseMeta):
 
- def setUp(self):
+ @classmethod
+ def setUpClass(cls):
  if os.environ.get('USE_UVLOOP'):
  import uvloop
  asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
  loop = asyncio.new_event_loop()
  asyncio.set_event_loop(None)
- self.loop = loop
+ cls.loop = loop
 
- def tearDown(self):
- self.loop.close()
+ @classmethod
+ def tearDownClass(cls):
+ cls.loop.close()
  asyncio.set_event_loop(None)
 
  @contextlib.contextmanager
@@ -97,7 +99,16 @@ def assertRunUnder(self, delta):
 _default_cluster = None
 
 
-def _start_cluster(server_settings={}):
+def _start_cluster(ClusterCls, cluster_kwargs, server_settings):
+ cluster = ClusterCls(**cluster_kwargs)
+ cluster.init()
+ cluster.trust_local_connections()
+ cluster.start(port='dynamic', server_settings=server_settings)
+ atexit.register(_shutdown_cluster, cluster)
+ return cluster
+
+
+def _start_default_cluster(server_settings={}):
  global _default_cluster
 
  if _default_cluster is None:
@@ -106,12 +117,8 @@ def _start_cluster(server_settings={}):
  # Using existing cluster, assuming it is initialized and running
  _default_cluster = pg_cluster.RunningCluster()
  else:
- _default_cluster = pg_cluster.TempCluster()
- _default_cluster.init()
- _default_cluster.trust_local_connections()
- _default_cluster.start(port='dynamic',
- server_settings=server_settings)
- atexit.register(_shutdown_cluster, _default_cluster)
+ _default_cluster = _start_cluster(
+ pg_cluster.TempCluster, {}, server_settings)
 
  return _default_cluster
 
@@ -122,9 +129,10 @@ def _shutdown_cluster(cluster):
 
 
 class ClusterTestCase(TestCase):
- def setUp(self):
- super().setUp()
- self.cluster = _start_cluster({
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+ cls.cluster = _start_default_cluster({
  'log_connections': 'on'
  })
 
@@ -133,6 +141,11 @@ def create_pool(self, **kwargs):
  conn_spec.update(kwargs)
  return pg_pool.create_pool(loop=self.loop, **conn_spec)
 
+ @classmethod
+ def start_cluster(cls, ClusterCls, *,
+ cluster_kwargs={}, server_settings={}):
+ return _start_cluster(ClusterCls, cluster_kwargs, server_settings)
+
 
 class ConnectedTestCase(ClusterTestCase):
 

asyncpg/cluster.py

@@ -16,6 +16,7 @@
 import socket
 import subprocess
 import tempfile
+import textwrap
 import time
 
 import asyncpg
@@ -332,6 +333,20 @@ def trust_local_connections(self):
  if status == 'running':
  self.reload()
 
+ def trust_local_replication_by(self, user):
+ if _system != 'Windows':
+ self.add_hba_entry(type='local', database='replication',
+ user=user, auth_method='trust')
+ self.add_hba_entry(type='host', address='127.0.0.1/32',
+ database='replication', user=user,
+ auth_method='trust')
+ self.add_hba_entry(type='host', address='::1/128',
+ database='replication', user=user,
+ auth_method='trust')
+ status = self.get_status()
+ if status == 'running':
+ self.reload()
+
  def _init_env(self):
  self._pg_config = self._find_pg_config(self._pg_config_path)
  self._pg_config_data = self._run_pg_config(self._pg_config)
@@ -489,6 +504,55 @@ def __init__(self, *,
  super().__init__(self._data_dir, pg_config_path=pg_config_path)
 
 
+class HotStandbyCluster(TempCluster):
+ def __init__(self, *,
+ master, replication_user,
+ data_dir_suffix=None, data_dir_prefix=None,
+ data_dir_parent=None, pg_config_path=None):
+ self._master = master
+ self._repl_user = replication_user
+ super().__init__(
+ data_dir_suffix=data_dir_suffix,
+ data_dir_prefix=data_dir_prefix,
+ data_dir_parent=data_dir_parent,
+ pg_config_path=pg_config_path)
+
+ def _init_env(self):
+ super()._init_env()
+ self._pg_basebackup = self._find_pg_binary('pg_basebackup')
+
+ def init(self, **settings):
+ """Initialize cluster."""
+ if self.get_status() != 'not-initialized':
+ raise ClusterError(
+ 'cluster in {!r} has already been initialized'.format(
+ self._data_dir))
+
+ process = subprocess.run(
+ [self._pg_basebackup, '-h', self._master['host'],
+ '-p', self._master['port'], '-D', self._data_dir,
+ '-U', self._repl_user],
+ stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+
+ output = process.stdout
+
+ if process.returncode != 0:
+ raise ClusterError(
+ 'pg_basebackup init exited with status {:d}:\n{}'.format(
+ process.returncode, output.decode()))
+
+ with open(os.path.join(self._data_dir, 'recovery.conf'), 'w') as f:
+ f.write(textwrap.dedent("""\
+ standby_mode = 'on'
+ primary_conninfo = 'host={host} port={port} user={user}'
+ """.format(
+ host=self._master['host'],
+ port=self._master['port'],
+ user=self._repl_user)))
+
+ return output.decode()
+
+
 class RunningCluster(Cluster):
  def __init__(self, **kwargs):
  self.conn_spec = kwargs

asyncpg/connection.py

@@ -369,11 +369,19 @@ def terminate(self):
 
  async def reset(self):
  self._listeners = {}
+
  await self.execute('''
+ DO $$
+ BEGIN
+ PERFORM * FROM pg_listening_channels() LIMIT 1;
+ IF FOUND THEN
+ UNLISTEN *;
+ END IF;
+ END;
+ $$;
  SET SESSION AUTHORIZATION DEFAULT;
  RESET ALL;
  CLOSE ALL;
- UNLISTEN *;
  SELECT pg_advisory_unlock_all();
  ''')
 

tests/test_pool.py

@@ -9,7 +9,8 @@
 import platform
 
 from asyncpg import _testbase as tb
-
+from asyncpg import cluster as pg_cluster
+from asyncpg import pool as pg_pool
 
 _system = platform.uname().system
 
@@ -148,3 +149,72 @@ async def worker():
  # Reset cluster's pg_hba.conf since we've meddled with it
  self.cluster.trust_local_connections()
  self.cluster.reload()
+
+
+class TestHostStandby(tb.ConnectedTestCase):
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+
+ cls.master_cluster = cls.start_cluster(
+ pg_cluster.TempCluster,
+ server_settings={
+ 'max_wal_senders': 10,
+ 'wal_level': 'hot_standby'
+ })
+
+ con = None
+
+ try:
+ con = cls.loop.run_until_complete(
+ cls.master_cluster.connect(database='postgres', loop=cls.loop))
+
+ cls.loop.run_until_complete(
+ con.execute('''
+ CREATE ROLE replication WITH LOGIN REPLICATION
+ '''))
+
+ cls.master_cluster.trust_local_replication_by('replication')
+
+ conn_spec = cls.master_cluster.get_connection_spec()
+
+ cls.standby_cluster = cls.start_cluster(
+ pg_cluster.HotStandbyCluster,
+ cluster_kwargs={
+ 'master': conn_spec,
+ 'replication_user': 'replication'
+ },
+ server_settings={
+ 'hot_standby': True
+ })
+
+ finally:
+ if con is not None:
+ cls.loop.run_until_complete(con.close())
+
+ @classmethod
+ def tearDownMethod(cls):
+ cls.standby_cluster.stop()
+ cls.standby_cluster.destroy()
+ cls.master_cluster.stop()
+ cls.master_cluster.destroy()
+
+ def create_pool(self, **kwargs):
+ conn_spec = self.standby_cluster.get_connection_spec()
+ conn_spec.update(kwargs)
+ return pg_pool.create_pool(loop=self.loop, **conn_spec)
+
+ async def test_standby_pool_01(self):
+ for n in {1, 3, 5, 10, 20, 100}:
+ with self.subTest(tasksnum=n):
+ pool = await self.create_pool(database='postgres',
+ min_size=5, max_size=10)
+
+ async def worker():
+ con = await pool.acquire()
+ self.assertEqual(await con.fetchval('SELECT 1'), 1)
+ await pool.release(con)
+
+ tasks = [worker() for _ in range(n)]
+ await asyncio.gather(*tasks, loop=self.loop)
+ await pool.close()