Use anonymous prepared statements when the cache is off.

`Connection.fetch()` and similar methods will now use anonymous
prepared statements when the statements cache is off.  This makes
it possible to use a limited subset of asyncpg functionality with
pgbouncer, and optimizes use of PostgreSQL server resources.

Author: 1st1

asyncpg/connection.py

@@ -41,7 +41,7 @@ class Connection(metaclass=ConnectionMeta):
                  '_stmt_cache_max_size', '_stmt_cache', '_stmts_to_close',
                  '_addr', '_opts', '_command_timeout', '_listeners',
                  '_server_version', '_server_caps', '_intro_query',
-                 '_reset_query', '_proxy')
+                 '_reset_query', '_proxy', '_stmt_exclusive_section')
 
     def __init__(self, protocol, transport, loop, addr, opts, *,
                  statement_cache_size, command_timeout):
@@ -97,6 +97,15 @@ def __init__(self, protocol, transport, loop, addr, opts, *,
         self._reset_query = None
         self._proxy = None
 
+        # Used to serialize operations that might involve anonymous
+        # statements.  Specifically, we want to make the following
+        # operation atomic:
+        #    ("prepare an anonymous statement", "use the statement")
+        #
+        # Used for `con.fetchval()`, `con.fetch()`, `con.fetchrow()`,
+        # `con.execute()`, and `con.executemany()`.
+        self._stmt_exclusive_section = _Atomic()
+
     async def add_listener(self, channel, callback):
         """Add a listener for Postgres notifications.
 
@@ -227,10 +236,9 @@ async def executemany(self, command: str, args, timeout: float=None):
         """
         return await self._executemany(command, args, timeout)
 
-    async def _get_statement(self, query, timeout):
-        cache = self._stmt_cache_max_size > 0
-
-        if cache:
+    async def _get_statement(self, query, timeout, *, named: bool=False):
+        use_cache = self._stmt_cache_max_size > 0
+        if use_cache:
             try:
                 state = self._stmt_cache[query]
             except KeyError:
@@ -241,7 +249,13 @@ async def _get_statement(self, query, timeout):
                     return state
 
         protocol = self._protocol
-        state = await protocol.prepare(None, query, timeout)
+
+        if use_cache or named:
+            stmt_name = self._get_unique_id('stmt')
+        else:
+            stmt_name = ''
+
+        state = await protocol.prepare(stmt_name, query, timeout)
 
         ready = state._init_types()
         if ready is not True:
@@ -251,7 +265,7 @@ async def _get_statement(self, query, timeout):
             types = await self._types_stmt.fetch(list(ready))
             protocol.get_settings().register_data_types(types)
 
-        if cache:
+        if use_cache:
             if len(self._stmt_cache) > self._stmt_cache_max_size - 1:
                 old_query, old_state = self._stmt_cache.popitem(last=False)
                 self._maybe_gc_stmt(old_state)
@@ -285,7 +299,7 @@ async def prepare(self, query, *, timeout=None):
 
         :return: A :class:`~prepared_stmt.PreparedStatement` instance.
         """
-        stmt = await self._get_statement(query, timeout)
+        stmt = await self._get_statement(query, timeout, named=True)
         return prepared_stmt.PreparedStatement(self, query, stmt)
 
     async def fetch(self, query, *args, timeout=None) -> list:
@@ -423,9 +437,9 @@ async def reset(self):
         if reset_query:
             await self.execute(reset_query)
 
-    def _get_unique_id(self):
+    def _get_unique_id(self, prefix):
         self._uid += 1
-        return 'id{}'.format(self._uid)
+        return '__asyncpg_{}_{}__'.format(prefix, self._uid)
 
     def _close_stmts(self):
         for stmt in self._stmt_cache.values():
@@ -450,9 +464,6 @@ async def _cleanup_stmts(self):
             # so we ignore the timeout.
             await self._protocol.close_statement(stmt, protocol.NO_TIMEOUT)
 
-    def _request_portal_name(self):
-        return self._get_unique_id()
-
     def _cancel_current_command(self, waiter):
         async def cancel():
             try:
@@ -572,17 +583,19 @@ def _drop_global_statement_cache(self):
         else:
             self._drop_local_statement_cache()
 
-    def _execute(self, query, args, limit, timeout, return_status=False):
+    async def _execute(self, query, args, limit, timeout, return_status=False):
         executor = lambda stmt, timeout: self._protocol.bind_execute(
             stmt, args, '', limit, return_status, timeout)
         timeout = self._protocol._get_timeout(timeout)
-        return self._do_execute(query, executor, timeout)
+        with self._stmt_exclusive_section:
+            return await self._do_execute(query, executor, timeout)
 
-    def _executemany(self, query, args, timeout):
+    async def _executemany(self, query, args, timeout):
         executor = lambda stmt, timeout: self._protocol.bind_execute_many(
             stmt, args, '', timeout)
         timeout = self._protocol._get_timeout(timeout)
-        return self._do_execute(query, executor, timeout)
+        with self._stmt_exclusive_section:
+            return await self._do_execute(query, executor, timeout)
 
     async def _do_execute(self, query, executor, timeout, retry=True):
         if timeout is None:
@@ -747,6 +760,22 @@ async def connect(dsn=None, *,
     return con
 
 
+class _Atomic:
+    __slots__ = ('_acquired',)
+
+    def __init__(self):
+        self._acquired = 0
+
+    def __enter__(self):
+        if self._acquired:
+            raise exceptions.InterfaceError(
+                'cannot perform operation: another operation is in progress')
+        self._acquired = 1
+
+    def __exit__(self, t, e, tb):
+        self._acquired = 0
+
+
 class _ConnectionProxy:
     # Base class to enable `isinstance(Connection)` check.
     __slots__ = ()

asyncpg/cursor.py

@@ -91,7 +91,7 @@ async def _bind_exec(self, n, timeout):
         con = self._connection
         protocol = con._protocol
 
-        self._portal_name = con._request_portal_name()
+        self._portal_name = con._get_unique_id('portal')
         buffer, _, self._exhausted = await protocol.bind_execute(
             self._state, self._args, self._portal_name, n, True, timeout)
         return buffer
@@ -106,7 +106,7 @@ async def _bind(self, timeout):
         con = self._connection
         protocol = con._protocol
 
-        self._portal_name = con._request_portal_name()
+        self._portal_name = con._get_unique_id('portal')
         buffer = await protocol.bind(self._state, self._args,
                                      self._portal_name,
                                      timeout)
@@ -168,7 +168,7 @@ def __aiter__(self):
     async def __anext__(self):
         if self._state is None:
             self._state = await self._connection._get_statement(
-                self._query, self._timeout)
+                self._query, self._timeout, named=True)
             self._state.attach()
 
         if not self._portal_name:
@@ -193,7 +193,7 @@ class Cursor(BaseCursor):
     async def _init(self, timeout):
         if self._state is None:
             self._state = await self._connection._get_statement(
-                self._query, timeout)
+                self._query, timeout, named=True)
             self._state.attach()
         self._check_ready()
         await self._bind(timeout)

asyncpg/protocol/protocol.pxd

@@ -40,7 +40,6 @@ cdef class BaseProtocol(CoreProtocol):
 
         str last_query
 
-        int uid_counter
         bint closing
 
         readonly uint64_t queries_count

asyncpg/protocol/protocol.pyx

@@ -95,7 +95,6 @@ cdef class BaseProtocol(CoreProtocol):
         self.settings = ConnectionSettings(
             (self.address, con_args.get('database')))
 
-        self.uid_counter = 0
         self.statement = None
         self.return_extra = False
 
@@ -138,10 +137,6 @@ cdef class BaseProtocol(CoreProtocol):
         self._check_state()
         timeout = self._get_timeout_impl(timeout)
 
-        if stmt_name is None:
-            self.uid_counter += 1
-            stmt_name = '__asyncpg_stmt_{}__'.format(self.uid_counter)
-
         self._prepare(stmt_name, query)
         self.last_query = query
         self.statement = PreparedStatementState(stmt_name, query, self)

asyncpg/transaction.py

@@ -100,7 +100,7 @@ async def start(self):
             self._nested = True
 
         if self._nested:
-            self._id = con._get_unique_id()
+            self._id = con._get_unique_id('savepoint')
             query = 'SAVEPOINT {};'.format(self._id)
         else:
             if self._isolation == 'read_committed':

tests/test_prepare.py

@@ -424,3 +424,35 @@ async def test_prepare_statement_invalid(self):
 
         finally:
             await self.con.execute('DROP TABLE tab1')
+
+    async def test_prepare_23_no_stmt_cache_seq(self):
+        # Disable cache, which will force connections to use
+        # anonymous prepared statements.
+        self.con._stmt_cache_max_size = 0
+
+        async def check_simple():
+            # Run a simple query a few times.
+            self.assertEqual(await self.con.fetchval('SELECT 1'), 1)
+            self.assertEqual(await self.con.fetchval('SELECT 2'), 2)
+            self.assertEqual(await self.con.fetchval('SELECT 1'), 1)
+
+        await check_simple()
+
+        # Run a query that timeouts.
+        with self.assertRaises(asyncio.TimeoutError):
+            await self.con.fetchrow('select pg_sleep(10)', timeout=0.02)
+
+        # Check that we can run new queries after a timeout.
+        await check_simple()
+
+        # Try a cursor/timeout combination. Cursors should always use
+        # named prepared statements.
+        async with self.con.transaction():
+            with self.assertRaises(asyncio.TimeoutError):
+                async for _ in self.con.cursor(   # NOQA
+                        'select pg_sleep(10)', timeout=0.1):
+                    pass
+
+        # Check that we can run queries after a failed cursor
+        # operation.
+        await check_simple()