From 7b710e3c99ad8cf4cf398756e9181b14f7d0cdab Mon Sep 17 00:00:00 2001
From: Denis Laxalde <denis.laxalde@dalibo.com>
Date: Mon, 18 Oct 2021 11:28:03 +0200
Subject: [PATCH] Add a setdefault() method to PrepareManager

---
 psycopg/psycopg/_preparing.py | 40 +++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/psycopg/psycopg/_preparing.py b/psycopg/psycopg/_preparing.py
index 3a58c1634..c90512e32 100644
--- a/psycopg/psycopg/_preparing.py
+++ b/psycopg/psycopg/_preparing.py
@@ -22,6 +22,7 @@ class Prepare(IntEnum):
 
 
 Key = Tuple[bytes, Tuple[int, ...]]
+Value = Union[int, bytes]
 
 
 class PrepareManager:
@@ -38,7 +39,7 @@ class PrepareManager:
         # Note: with this implementation we keep the tally of up to 100
         # queries, but most likely we will prepare way less than that. We might
         # change that if we think it would be better.
-        self._prepared: OrderedDict[Key, Union[int, bytes]] = OrderedDict()
+        self._prepared: OrderedDict[Key, Value] = OrderedDict()
 
         # Counter to generate prepared statements names
         self._prepared_idx = 0
@@ -91,6 +92,27 @@ class PrepareManager:
                     return b"DEALLOCATE ALL"
         return None
 
+    def setdefault(
+        self, query: PostgresQuery, prep: Prepare, name: bytes
+    ) -> Optional[Tuple[Key, Value]]:
+        """Check if the query is already in cache.
+
+        If not, return a new 'key' matching given query. Otherwise, just
+        update the count for that query and record as last used.
+        """
+        key = self.key(query)
+        if key in self._prepared:
+            if isinstance(self._prepared[key], int):
+                if prep is Prepare.SHOULD:
+                    self._prepared[key] = name
+                else:
+                    self._prepared[key] += 1  # type: ignore[operator]
+            self._prepared.move_to_end(key)
+            return None
+
+        value: Value = name if prep is Prepare.SHOULD else 1
+        return key, value
+
     def maintain(
         self,
         query: PostgresQuery,
@@ -107,17 +129,8 @@ class PrepareManager:
         if cmd:
             return cmd
 
-        key = self.key(query)
-
-        # If we know the query already the cache size won't change
-        # So just update the count and record as last used
-        if key in self._prepared:
-            if isinstance(self._prepared[key], int):
-                if prep is Prepare.SHOULD:
-                    self._prepared[key] = name
-                else:
-                    self._prepared[key] += 1  # type: ignore  # operator
-            self._prepared.move_to_end(key)
+        cached = self.setdefault(query, prep, name)
+        if cached is None:
             return None
 
         # The query is not in cache. Let's see if we must add it
@@ -131,7 +144,8 @@ class PrepareManager:
             return None
 
         # Ok, we got to the conclusion that this query is genuinely to prepare
-        self._prepared[key] = name if prep is Prepare.SHOULD else 1
+        key, value = cached
+        self._prepared[key] = value
 
         # Evict an old value from the cache; if it was prepared, deallocate it
         # Do it only once: if the cache was resized, deallocate gradually
-- 
2.47.2