From: Denis Laxalde <denis.laxalde@dalibo.com>
Date: Mon, 18 Oct 2021 09:18:12 +0000 (+0200)
Subject: Add a should_discard() method to PrepareManager
X-Git-Tag: pool-3.1~98^2~11
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=570aa164926867e4df1c733d5fdd2819b072034a;p=thirdparty%2Fpsycopg.git

Add a should_discard() method to PrepareManager
---

diff --git a/psycopg/psycopg/_preparing.py b/psycopg/psycopg/_preparing.py
index 39ef4c444..3a58c1634 100644
--- a/psycopg/psycopg/_preparing.py
+++ b/psycopg/psycopg/_preparing.py
@@ -72,6 +72,25 @@ class PrepareManager:
             # The query is not to be prepared yet
             return Prepare.NO, b""
 
+    def should_discard(
+        self, prep: Prepare, results: Sequence["PGresult"]
+    ) -> Optional[bytes]:
+        """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"
+        return None
+
     def maintain(
         self,
         query: PostgresQuery,
@@ -84,19 +103,9 @@ 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"
+        cmd = self.should_discard(prep, results)
+        if cmd:
+            return cmd
 
         key = self.key(query)