Clear the prepared statements if the catalog might have changed

This could happen on DROP (and re-creation) of some object or on the
implicit drop caused by a rollback.

Issues happen if a prepared transaction is executed after an object has
been dropped: see
https://github.com/sqlalchemy/sqlalchemy/issues/6842#issuecomment-925131836

Postgres devs suggestion is to avoid to do that: see
https://www.postgresql.org/message-id/2220999.1632335261%40sss.pgh.pa.us

diff --git a/psycopg/psycopg/_preparing.py b/psycopg/psycopg/_preparing.py
index 32e2c1d90255280dbf8de128c416326ff170bfe3..08960e8a22b1e8fcb6dd989c681ab7e47d8fb1bb 100644 (file)
--- a/psycopg/psycopg/_preparing.py
+++ b/psycopg/psycopg/_preparing.py
@@ -79,6 +79,20 @@ class PrepareManager:
         if self.prepare_threshold is None:
             return None
 
+        # Check if we need to discard our entire state: it should happen on
+        # rollback or on dropping objects, because the same object may get
+        # recreated and postgres would fail internal lookups.
+        if self._prepared or prep == Prepare.SHOULD:
+            for result in results:
+                if result.status != ExecStatus.COMMAND_OK:
+                    continue
+                cmdstat = result.command_status
+                if cmdstat and (
+                    cmdstat.startswith(b"DROP ") or cmdstat == b"ROLLBACK"
+                ):
+                    self._prepared.clear()
+                    return b"DEALLOCATE ALL"
+
         key = (query.query, query.types)
 
         # If we know the query already the cache size won't change
@@ -97,11 +111,8 @@ class PrepareManager:
             # We cannot prepare a multiple statement
             return None
 
-        result = results[0]
-        if (
-            result.status != ExecStatus.TUPLES_OK
-            and result.status != ExecStatus.COMMAND_OK
-        ):
+        status = results[0].status
+        if ExecStatus.COMMAND_OK != status != ExecStatus.TUPLES_OK:
             # We don't prepare failed queries or other weird results
             return None
 
@@ -118,3 +129,11 @@ class PrepareManager:
             return b"DEALLOCATE " + old_val
         else:
             return None
+
+    def clear(self) -> Optional[bytes]:
+        if self._prepared_idx:
+            self._prepared.clear()
+            self._prepared_idx = 0
+            return b"DEALLOCATE ALL"
+        else:
+            return None

diff --git a/psycopg/psycopg/connection.py b/psycopg/psycopg/connection.py
index b834bcc4aefc98107e8024da8394dd3466d58428..f1c879cc198915dfac38eed5609f77d7f9682849 100644 (file)
--- a/psycopg/psycopg/connection.py
+++ b/psycopg/psycopg/connection.py
@@ -515,6 +515,9 @@ class BaseConnection(Generic[Row]):
             return
 
         yield from self._exec_command(b"ROLLBACK")
+        cmd = self._prepared.clear()
+        if cmd:
+            yield from self._exec_command(cmd)
 
 
 class Connection(BaseConnection[Row]):

diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py
index d70fcd748b4140e2d8de335ffc5adef737f2dd65..ed54be95ad7655d94202b7ad06be1ab34876953c 100644 (file)
--- a/psycopg/psycopg/cursor.py
+++ b/psycopg/psycopg/cursor.py
@@ -247,10 +247,11 @@ class BaseCursor(Generic[ConnectionType, Row]):
         results = yield from execute(self._conn.pgconn)
 
         # Update the prepare state of the query
-        if prepare is not False:
-            cmd = self._conn._prepared.maintain(pgq, results, prep, name)
-            if cmd:
-                yield from self._conn._exec_command(cmd)
+        # If an operation requires to flush our prepared statements cache,
+        # do it. Note that there is an off-by-one error because
+        cmd = self._conn._prepared.maintain(pgq, results, prep, name)
+        if cmd:
+            yield from self._conn._exec_command(cmd)
 
         return results
 

diff --git a/psycopg/psycopg/transaction.py b/psycopg/psycopg/transaction.py
index 1395b14bbe0c75162f8f5e1f12dc9cbe0d6be49f..2baa2b088daaa324304ec9c43d2868429e3c8d03 100644 (file)
--- a/psycopg/psycopg/transaction.py
+++ b/psycopg/psycopg/transaction.py
@@ -164,6 +164,11 @@ class BaseTransaction(Generic[ConnectionType]):
             assert not self._conn._savepoints
             commands.append(b"ROLLBACK")
 
+        # Also clear the prepared statements cache.
+        cmd = self._conn._prepared.clear()
+        if cmd:
+            commands.append(cmd)
+
         yield from self._conn._exec_command(b"; ".join(commands))
 
         if isinstance(exc_val, Rollback):

diff --git a/tests/test_prepared.py b/tests/test_prepared.py
index 68bcf7c3247f8d9d788d04fa8d2e0420f6e802d6..af4c51fc85b7604ce848cf980d9c4a6ebbef1f1b 100644 (file)
--- a/tests/test_prepared.py
+++ b/tests/test_prepared.py
@@ -191,3 +191,69 @@ def test_untyped_json(conn):
 
     cur = conn.execute("select parameter_types from pg_prepared_statements")
     assert cur.fetchall() == [(["jsonb"],)]
+
+
+def test_change_type_execute(conn):
+    conn.prepare_threshold = 0
+    for i in range(3):
+        conn.execute("CREATE TYPE prepenum AS ENUM ('foo', 'bar', 'baz')")
+        conn.execute("CREATE TABLE preptable(id integer, bar prepenum[])")
+        conn.cursor().execute(
+            "INSERT INTO preptable (bar) VALUES (%(enum_col)s::prepenum[])",
+            {"enum_col": ["foo"]},
+        )
+        conn.rollback()
+
+
+def test_change_type_executemany(conn):
+    for i in range(3):
+        conn.execute("CREATE TYPE prepenum AS ENUM ('foo', 'bar', 'baz')")
+        conn.execute("CREATE TABLE preptable(id integer, bar prepenum[])")
+        conn.cursor().executemany(
+            "INSERT INTO preptable (bar) VALUES (%(enum_col)s::prepenum[])",
+            [{"enum_col": ["foo"]}, {"enum_col": ["foo", "bar"]}],
+        )
+        conn.rollback()
+
+
+def test_change_type(conn):
+    conn.prepare_threshold = 0
+    conn.execute("CREATE TYPE prepenum AS ENUM ('foo', 'bar', 'baz')")
+    conn.execute("CREATE TABLE preptable(id integer, bar prepenum[])")
+    conn.cursor().execute(
+        "INSERT INTO preptable (bar) VALUES (%(enum_col)s::prepenum[])",
+        {"enum_col": ["foo"]},
+    )
+    conn.execute("DROP TABLE preptable")
+    conn.execute("DROP TYPE prepenum")
+    conn.execute("CREATE TYPE prepenum AS ENUM ('foo', 'bar', 'baz')")
+    conn.execute("CREATE TABLE preptable(id integer, bar prepenum[])")
+    conn.cursor().execute(
+        "INSERT INTO preptable (bar) VALUES (%(enum_col)s::prepenum[])",
+        {"enum_col": ["foo"]},
+    )
+
+    cur = conn.execute(
+        "select count(*) from pg_prepared_statements", prepare=False
+    )
+    assert cur.fetchone()[0] == 3
+
+
+def test_change_type_savepoint(conn):
+    conn.prepare_threshold = 0
+    with conn.transaction():
+        for i in range(3):
+            with pytest.raises(ZeroDivisionError):
+                with conn.transaction():
+                    conn.execute(
+                        "CREATE TYPE prepenum AS ENUM ('foo', 'bar', 'baz')"
+                    )
+                    conn.execute(
+                        "CREATE TABLE preptable(id integer, bar prepenum[])"
+                    )
+                    conn.cursor().execute(
+                        "INSERT INTO preptable (bar) "
+                        "VALUES (%(enum_col)s::prepenum[])",
+                        {"enum_col": ["foo"]},
+                    )
+                    raise ZeroDivisionError()