Invalidate statement cache on schema changes affecting statement result.

PostgreSQL will raise an exception when it detects that the result type of the
query has changed from when the statement was prepared.  This may happen, for
example, after an ALTER TABLE or SET search_path.

When this happens, and there is no transaction running, we can simply
re-prepare the statement and try again.

If the transaction _is_ running, this error will put it into an error state,
and we have no choice but to raise an exception.  The original error is
somewhat cryptic, so we raise a custom InvalidCachedStatementError with the
original server exception as context.

In either case we clear the statement cache for this connection and all other
connections of the pool this connection belongs to (if any).

See #72 and #76 for discussion.

Fixes: #72.

Author: elprans

asyncpg/connection.py

@@ -187,9 +187,7 @@ async def execute(self, query: str, *args, timeout: float=None) -> str:
         if not args:
             return await self._protocol.query(query, timeout)
 
-        stmt = await self._get_statement(query, timeout)
-        _, status, _ = await self._protocol.bind_execute(stmt, args, '', 0,
-                                                         True, timeout)
+        _, status, _ = await self._do_execute(query, args, 0, timeout, True)
         return status.decode()
 
     async def executemany(self, command: str, args, timeout: float=None):
@@ -283,10 +281,7 @@ async def fetch(self, query, *args, timeout=None) -> list:
 
         :return list: A list of :class:`Record` instances.
         """
-        stmt = await self._get_statement(query, timeout)
-        data = await self._protocol.bind_execute(stmt, args, '', 0,
-                                                 False, timeout)
-        return data
+        return await self._do_execute(query, args, 0, timeout)
 
     async def fetchval(self, query, *args, column=0, timeout=None):
         """Run a query and return a value in the first row.
@@ -302,9 +297,7 @@ async def fetchval(self, query, *args, column=0, timeout=None):
 
         :return: The value of the specified column of the first record.
         """
-        stmt = await self._get_statement(query, timeout)
-        data = await self._protocol.bind_execute(stmt, args, '', 1,
-                                                 False, timeout)
+        data = await self._do_execute(query, args, 1, timeout)
         if not data:
             return None
         return data[0][column]
@@ -318,9 +311,7 @@ async def fetchrow(self, query, *args, timeout=None):
 
         :return: The first row as a :class:`Record` instance.
         """
-        stmt = await self._get_statement(query, timeout)
-        data = await self._protocol.bind_execute(stmt, args, '', 1,
-                                                 False, timeout)
+        data = await self._do_execute(query, args, 1, timeout)
         if not data:
             return None
         return data[0]
@@ -551,6 +542,61 @@ def _set_proxy(self, proxy):
 
         self._proxy = proxy
 
+    def _drop_local_statement_cache(self):
+        self._stmt_cache.clear()
+
+    def _drop_global_statement_cache(self):
+        if self._proxy is not None:
+            # This connection is a member of a pool, so we delegate
+            # the cache drop to the pool.
+            pool = self._proxy._holder._pool
+            pool._drop_statement_cache()
+        else:
+            self._drop_local_statement_cache()
+
+    async def _do_execute(self, query, args, limit, timeout,
+                          return_status=False):
+        stmt = await self._get_statement(query, timeout)
+
+        try:
+            result = await self._protocol.bind_execute(
+                stmt, args, '', limit, return_status, timeout)
+
+        except exceptions.InvalidCachedStatementError as e:
+            # PostgreSQL will raise an exception when it detects
+            # that the result type of the query has changed from
+            # when the statement was prepared.  This may happen,
+            # for example, after an ALTER TABLE or SET search_path.
+            #
+            # When this happens, and there is no transaction running,
+            # we can simply re-prepare the statement and try again.
+            #
+            # If the transaction _is_ running, this error will put it
+            # into an error state, and we have no choice but to
+            # re-raise the exception.
+            #
+            # In either case we clear the statement cache for this
+            # connection and all other connections of the pool this
+            # connection belongs to (if any).
+            #
+            # Note that we specifically do not rely on the error
+            # message, as it is localizable.
+            #
+            # See https://github.com/MagicStack/asyncpg/issues/72
+            # and https://github.com/MagicStack/asyncpg/issues/76
+            # for discussion.
+            #
+            self._drop_global_statement_cache()
+
+            if self._protocol.is_in_transaction():
+                raise
+            else:
+                stmt = await self._get_statement(query, timeout)
+                result = await self._protocol.bind_execute(
+                    stmt, args, '', limit, return_status, timeout)
+
+        return result
+
 
 async def connect(dsn=None, *,
                   host=None, port=None,

asyncpg/exceptions/__init__.py

@@ -1,10 +1,3 @@
-# Copyright (C) 2016-present the ayncpg authors and contributors
-# <see AUTHORS file>
-#
-# This module is part of asyncpg and is released under
-# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
-
-
 # GENERATED FROM postgresql/src/backend/utils/errcodes.txt
 # DO NOT MODIFY, use tools/generate_exceptions.py to update
 
@@ -92,6 +85,10 @@ class FeatureNotSupportedError(_base.PostgresError):
     sqlstate = '0A000'
 
 
+class InvalidCachedStatementError(FeatureNotSupportedError):
+    pass
+
+
 class InvalidTransactionInitiationError(_base.PostgresError):
     sqlstate = '0B000'
 
@@ -1025,15 +1022,16 @@ class IndexCorruptedError(InternalServerError):
     'InvalidArgumentForPowerFunctionError',
     'InvalidArgumentForWidthBucketFunctionError',
     'InvalidAuthorizationSpecificationError',
-    'InvalidBinaryRepresentationError', 'InvalidCatalogNameError',
-    'InvalidCharacterValueForCastError', 'InvalidColumnDefinitionError',
-    'InvalidColumnReferenceError', 'InvalidCursorDefinitionError',
-    'InvalidCursorNameError', 'InvalidCursorStateError',
-    'InvalidDatabaseDefinitionError', 'InvalidDatetimeFormatError',
-    'InvalidEscapeCharacterError', 'InvalidEscapeOctetError',
-    'InvalidEscapeSequenceError', 'InvalidForeignKeyError',
-    'InvalidFunctionDefinitionError', 'InvalidGrantOperationError',
-    'InvalidGrantorError', 'InvalidIndicatorParameterValueError',
+    'InvalidBinaryRepresentationError', 'InvalidCachedStatementError',
+    'InvalidCatalogNameError', 'InvalidCharacterValueForCastError',
+    'InvalidColumnDefinitionError', 'InvalidColumnReferenceError',
+    'InvalidCursorDefinitionError', 'InvalidCursorNameError',
+    'InvalidCursorStateError', 'InvalidDatabaseDefinitionError',
+    'InvalidDatetimeFormatError', 'InvalidEscapeCharacterError',
+    'InvalidEscapeOctetError', 'InvalidEscapeSequenceError',
+    'InvalidForeignKeyError', 'InvalidFunctionDefinitionError',
+    'InvalidGrantOperationError', 'InvalidGrantorError',
+    'InvalidIndicatorParameterValueError',
     'InvalidLocatorSpecificationError', 'InvalidNameError',
     'InvalidObjectDefinitionError', 'InvalidParameterValueError',
     'InvalidPasswordError', 'InvalidPreparedStatementDefinitionError',

asyncpg/exceptions/_base.py

@@ -87,7 +87,26 @@ def new(cls, fields, query=None):
             if field:
                 mapped[field] = v
 
-        e = exccls(mapped.get('message', ''))
+        message = mapped.get('message', '')
+
+        # PostgreSQL will raise an exception when it detects
+        # that the result type of the query has changed from
+        # when the statement was prepared.
+        #
+        # The original error is somewhat cryptic and unspecific,
+        # so we raise a custom subclass that is easier to handle
+        # and identify.
+        if (exccls.__name__ == 'FeatureNotSupportedError' and
+                (exccls.__module__ == 'asyncpg' or
+                    exccls.__module__.startswith('asyncpg.')) and
+                mapped.get('server_source_function') ==
+                'RevalidateCachedQuery'):
+            exceptions = sys.modules[exccls.__module__]
+            exccls = exceptions.InvalidCachedStatementError
+            message = ('cached statement plan is invalid due to a database '
+                       'schema or configuration change')
+
+        e = exccls(message)
         e.__dict__.update(mapped)
 
         return e

asyncpg/pool.py

@@ -402,6 +402,12 @@ def _check_init(self):
         if self._closed:
             raise exceptions.InterfaceError('pool is closed')
 
+    def _drop_statement_cache(self):
+        # Drop statement cache for all connections in the pool.
+        for ch in self._holders:
+            if ch._con is not None:
+                ch._con._drop_local_statement_cache()
+
     def __await__(self):
         return self._async__init__().__await__()
 

asyncpg/prepared_stmt.py

@@ -154,11 +154,7 @@ async def fetch(self, *args, timeout=None):
 
         :return: A list of :class:`Record` instances.
         """
-        self.__check_open()
-        protocol = self._connection._protocol
-        data, status, _ = await protocol.bind_execute(
-            self._state, args, '', 0, True, timeout)
-        self._last_status = status
+        data = await self.__bind_execute(args, 0, timeout)
         return data
 
     async def fetchval(self, *args, column=0, timeout=None):
@@ -174,11 +170,7 @@ async def fetchval(self, *args, column=0, timeout=None):
 
         :return: The value of the specified column of the first record.
         """
-        self.__check_open()
-        protocol = self._connection._protocol
-        data, status, _ = await protocol.bind_execute(
-            self._state, args, '', 1, True, timeout)
-        self._last_status = status
+        data = await self.__bind_execute(args, 1, timeout)
         if not data:
             return None
         return data[0][column]
@@ -192,14 +184,18 @@ async def fetchrow(self, *args, timeout=None):
 
         :return: The first row as a :class:`Record` instance.
         """
+        data = await self.__bind_execute(args, 1, timeout)
+        if not data:
+            return None
+        return data[0]
+
+    async def __bind_execute(self, args, limit, timeout):
         self.__check_open()
         protocol = self._connection._protocol
         data, status, _ = await protocol.bind_execute(
-            self._state, args, '', 1, True, timeout)
+            self._state, args, '', limit, True, timeout)
         self._last_status = status
-        if not data:
-            return None
-        return data[0]
+        return data
 
     def __check_open(self):
         if self._state.closed:

asyncpg/protocol/protocol.pyx

@@ -121,7 +121,9 @@ cdef class BaseProtocol(CoreProtocol):
         return self.settings
 
     def is_in_transaction(self):
-        return self.xact_status == PQTRANS_INTRANS
+        # PQTRANS_INTRANS = idle, within transaction block
+        # PQTRANS_INERROR = idle, within failed transaction
+        return self.xact_status in (PQTRANS_INTRANS, PQTRANS_INERROR)
 
     async def prepare(self, stmt_name, query, timeout):
         if self.cancel_waiter is not None:

tests/test_cache_invalidation.py

@@ -0,0 +1,87 @@
+# Copyright (C) 2016-present the ayncpg authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of asyncpg and is released under
+# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
+
+
+import asyncpg
+from asyncpg import _testbase as tb
+
+
+class TestCacheInvalidation(tb.ConnectedTestCase):
+    async def test_prepare_cache_invalidation_silent(self):
+        await self.con.execute('CREATE TABLE tab1(a int, b int)')
+
+        try:
+            await self.con.execute('INSERT INTO tab1 VALUES (1, 2)')
+            result = await self.con.fetchrow('SELECT * FROM tab1')
+            self.assertEqual(result, (1, 2))
+
+            await self.con.execute(
+                'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
+
+            result = await self.con.fetchrow('SELECT * FROM tab1')
+            self.assertEqual(result, (1, '2'))
+        finally:
+            await self.con.execute('DROP TABLE tab1')
+
+    async def test_prepare_cache_invalidation_in_transaction(self):
+        await self.con.execute('CREATE TABLE tab1(a int, b int)')
+
+        try:
+            await self.con.execute('INSERT INTO tab1 VALUES (1, 2)')
+            result = await self.con.fetchrow('SELECT * FROM tab1')
+            self.assertEqual(result, (1, 2))
+
+            await self.con.execute(
+                'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
+
+            with self.assertRaisesRegex(asyncpg.InvalidCachedStatementError,
+                                        'cached statement plan is invalid'):
+                async with self.con.transaction():
+                    result = await self.con.fetchrow('SELECT * FROM tab1')
+
+            # This is now OK,
+            result = await self.con.fetchrow('SELECT * FROM tab1')
+            self.assertEqual(result, (1, '2'))
+        finally:
+            await self.con.execute('DROP TABLE tab1')
+
+    async def test_prepare_cache_invalidation_in_pool(self):
+        pool = await self.create_pool(database='postgres',
+                                      min_size=2, max_size=2)
+
+        con1 = await pool.acquire()
+        con2 = await pool.acquire()
+
+        await self.con.execute('CREATE TABLE tab1(a int, b int)')
+
+        try:
+            await self.con.execute('INSERT INTO tab1 VALUES (1, 2)')
+
+            result = await con1.fetchrow('SELECT * FROM tab1')
+            self.assertEqual(result, (1, 2))
+
+            result = await con2.fetchrow('SELECT * FROM tab1')
+            self.assertEqual(result, (1, 2))
+
+            await self.con.execute(
+                'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
+
+            # con1 tries the same plan, will invalidate the cache
+            # for the entire pool.
+            result = await con1.fetchrow('SELECT * FROM tab1')
+            self.assertEqual(result, (1, '2'))
+
+            async with con2.transaction():
+                # This should work, as con1 should have invalidated
+                # the plan cache.
+                result = await con2.fetchrow('SELECT * FROM tab1')
+                self.assertEqual(result, (1, '2'))
+
+        finally:
+            await self.con.execute('DROP TABLE tab1')
+            await pool.release(con2)
+            await pool.release(con1)
+            await pool.close()

tests/test_prepare.py

@@ -406,3 +406,21 @@ async def test_prepare_22_empty(self):
         result = await self.con.fetchrow('SELECT')
         self.assertEqual(result, ())
         self.assertEqual(repr(result), '<Record>')
+
+    async def test_prepare_statement_invalid(self):
+        await self.con.execute('CREATE TABLE tab1(a int, b int)')
+
+        try:
+            await self.con.execute('INSERT INTO tab1 VALUES (1, 2)')
+
+            stmt = await self.con.prepare('SELECT * FROM tab1')
+
+            await self.con.execute(
+                'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
+
+            with self.assertRaisesRegex(asyncpg.InvalidCachedStatementError,
+                                        'cached statement plan is invalid'):
+                await stmt.fetchrow()
+
+        finally:
+            await self.con.execute('DROP TABLE tab1')