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

@@ -226,6 +226,9 @@ async def _get_statement(self, query, timeout):
                 if not state.closed:
                     return state
 
+        return await self._prepare_statement(query, timeout, cache)
+
+    async def _prepare_statement(self, query, timeout, cache):
         protocol = self._protocol
         state = await protocol.prepare(None, query, timeout)
 
@@ -283,10 +286,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 +302,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 +316,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 +547,69 @@ 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):
+        stmt = await self._get_statement(query, timeout)
+
+        try:
+            data = await self._protocol.bind_execute(
+                stmt, args, '', limit, False, timeout)
+
+        except exceptions.FeatureNotSupportedError as e:
+            if e.server_source_function == 'RevalidateCachedQuery':
+                # 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 error with the 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).
+                #
+                # 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 exceptions.InvalidCachedStatementError(
+                        'cached statement plan is invalid due to a database '
+                        'schema or configuration change'
+                    ) from e
+                else:
+                    stmt = await self._prepare_statement(
+                        query, timeout, True)
+                    data = await self._protocol.bind_execute(
+                        stmt, args, '', limit, False, timeout)
+            else:
+                raise
+
+        return data
+
 
 async def connect(dsn=None, *,
                   host=None, port=None,

asyncpg/exceptions/_base.py

@@ -9,7 +9,7 @@
 
 
 __all__ = ('PostgresError', 'FatalPostgresError', 'UnknownPostgresError',
-           'InterfaceError')
+           'InterfaceError', 'InvalidCachedStatementError')
 
 
 class PostgresMessageMeta(type):
@@ -105,5 +105,13 @@ class UnknownPostgresError(FatalPostgresError):
     """An error with an unknown SQLSTATE code."""
 
 
-class InterfaceError(Exception):
+class AsyncpgError(Exception):
+    """An error raised by asyncpg, not PostgreSQL."""
+
+
+class InterfaceError(AsyncpgError):
     """An error caused by improper use of asyncpg API."""
+
+
+class InvalidCachedStatementError(AsyncpgError):
+    """Cached statement is not longer valid."""

asyncpg/pool.py

@@ -404,6 +404,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,15 +184,32 @@ async def fetchrow(self, *args, timeout=None):
 
         :return: The first row as a :class:`Record` instance.
         """
-        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]
 
+    async def __bind_execute(self, args, limit, timeout):
+        self.__check_open()
+        protocol = self._connection._protocol
+
+        try:
+            data, status, _ = await protocol.bind_execute(
+                self._state, args, '', limit, True, timeout)
+        except exceptions.FeatureNotSupportedError as e:
+            if e.server_source_function == 'RevalidateCachedQuery':
+                # Handle query plan rejection due to schema or config changes.
+                # See comment in Connection._do_execute.
+                raise exceptions.InvalidCachedStatementError(
+                    'prepared statement is invalid due to a database '
+                    'schema or configuration change'
+                ) from e
+            else:
+                raise
+
+        self._last_status = status
+        return data
+
     def __check_open(self):
         if self._state.closed:
             raise exceptions.InterfaceError('prepared statement is closed')

asyncpg/protocol/protocol.pyx

@@ -121,7 +121,7 @@ cdef class BaseProtocol(CoreProtocol):
         return self.settings
 
     def is_in_transaction(self):
-        return self.xact_status == PQTRANS_INTRANS
+        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,
+                                        'prepared statement is invalid'):
+                await stmt.fetchrow()
+
+        finally:
+            await self.con.execute('DROP TABLE tab1')