From: Daniele Varrazzo Date: Wed, 22 Sep 2021 20:10:09 +0000 (+0200) Subject: Clear the prepared statements if the catalog might have changed X-Git-Tag: 3.0~65 X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=d770a0fd19d991b1c11c993f9ff9ceabddf249ff;p=thirdparty%2Fpsycopg.git 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 32e2c1d90..08960e8a2 100644 --- 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 b834bcc4a..f1c879cc1 100644 --- 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 d70fcd748..ed54be95a 100644 --- 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 1395b14bb..2baa2b088 100644 --- 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 68bcf7c32..af4c51fc8 100644 --- 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()