From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 6 Apr 2024 16:36:16 +0000 (+0000)
Subject: refactor(prepare): make maintenance operation a PQGen generator
X-Git-Tag: 3.2.0~55^2~5
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=f8cb9d9250755bbf1f3824502abbfba575b9a917;p=thirdparty%2Fpsycopg.git

refactor(prepare): make maintenance operation a PQGen generator

This will allow us to replace the SQL commands with protocol messages.
---

diff --git a/psycopg/psycopg/_connection_base.py b/psycopg/psycopg/_connection_base.py
index 12ae7395f..c5a6c2c1e 100644
--- a/psycopg/psycopg/_connection_base.py
+++ b/psycopg/psycopg/_connection_base.py
@@ -552,8 +552,7 @@ class BaseConnection(Generic[Row]):
 
         yield from self._exec_command(b"ROLLBACK")
         self._prepared.clear()
-        for cmd in self._prepared.get_maintenance_commands():
-            yield from self._exec_command(cmd)
+        yield from self._prepared.maintain_gen(self)
 
         if self._pipeline:
             yield from self._pipeline._sync_gen()
diff --git a/psycopg/psycopg/_cursor_base.py b/psycopg/psycopg/_cursor_base.py
index 1f53a8821..24a2cb1a2 100644
--- a/psycopg/psycopg/_cursor_base.py
+++ b/psycopg/psycopg/_cursor_base.py
@@ -195,9 +195,7 @@ class BaseCursor(Generic[ConnectionType, Row]):
             yield from self._conn._pipeline._communicate_gen()
 
         self._last_query = query
-
-        for cmd in self._conn._prepared.get_maintenance_commands():
-            yield from self._conn._exec_command(cmd)
+        yield from self._conn._prepared.maintain_gen(self._conn)
 
     def _executemany_gen_pipeline(
         self, query: Query, params_seq: Iterable[Params], returning: bool
@@ -232,8 +230,7 @@ class BaseCursor(Generic[ConnectionType, Row]):
         if returning:
             yield from pipeline._fetch_gen(flush=True)
 
-        for cmd in self._conn._prepared.get_maintenance_commands():
-            yield from self._conn._exec_command(cmd)
+        yield from self._conn._prepared.maintain_gen(self._conn)
 
     def _executemany_gen_no_pipeline(
         self, query: Query, params_seq: Iterable[Params], returning: bool
@@ -260,9 +257,7 @@ class BaseCursor(Generic[ConnectionType, Row]):
             yield from self._maybe_prepare_gen(pgq, prepare=True)
 
         self._last_query = query
-
-        for cmd in self._conn._prepared.get_maintenance_commands():
-            yield from self._conn._exec_command(cmd)
+        yield from self._conn._prepared.maintain_gen(self._conn)
 
     def _maybe_prepare_gen(
         self,
diff --git a/psycopg/psycopg/_preparing.py b/psycopg/psycopg/_preparing.py
index 465de53a4..dd969f88d 100644
--- a/psycopg/psycopg/_preparing.py
+++ b/psycopg/psycopg/_preparing.py
@@ -5,15 +5,18 @@ Support for prepared statements
 # Copyright (C) 2020 The Psycopg Team
 
 from enum import IntEnum, auto
-from typing import Iterator, Optional, Sequence, Tuple, TYPE_CHECKING
+from typing import Optional, Sequence, Tuple, TYPE_CHECKING
 from collections import OrderedDict
 
 from . import pq
+from .abc import PQGen
 from ._compat import Deque, TypeAlias
 from ._queries import PostgresQuery
 
 if TYPE_CHECKING:
+    from typing import Any
     from .pq.abc import PGresult
+    from ._connection_base import BaseConnection
 
 Key: TypeAlias = Tuple[bytes, Tuple[int, ...]]
 
@@ -185,9 +188,13 @@ class PrepareManager:
         else:
             return False
 
-    def get_maintenance_commands(self) -> Iterator[bytes]:
+    def maintain_gen(self, conn: "BaseConnection[Any]") -> PQGen[None]:
         """
-        Iterate over the commands needed to align the server state to our state
+        Generator to send the commands to perform periodic maintenance
+
+        Deallocate unneeded command in the server, or flush the prepared
+        statements server state entirely if necessary.
         """
         while self._maint_commands:
-            yield self._maint_commands.popleft()
+            cmd = self._maint_commands.popleft()
+            yield from conn._exec_command(cmd)
diff --git a/psycopg/psycopg/transaction.py b/psycopg/psycopg/transaction.py
index c6405aa43..56b710445 100644
--- a/psycopg/psycopg/transaction.py
+++ b/psycopg/psycopg/transaction.py
@@ -145,6 +145,10 @@ class BaseTransaction(Generic[ConnectionType]):
         for command in self._get_rollback_commands():
             yield from self._conn._exec_command(command)
 
+        # Also clear the prepared statements cache.
+        self._conn._prepared.clear()
+        yield from self._conn._prepared.maintain_gen(self._conn)
+
         if isinstance(exc_val, Rollback):
             if not exc_val.transaction or exc_val.transaction is self:
                 return True  # Swallow the exception
@@ -191,10 +195,6 @@ class BaseTransaction(Generic[ConnectionType]):
             assert not self._conn._num_transactions
             yield b"ROLLBACK"
 
-        # Also clear the prepared statements cache.
-        if self._conn._prepared.clear():
-            yield from self._conn._prepared.get_maintenance_commands()
-
     def _push_savepoint(self) -> None:
         """
         Push the transaction on the connection transactions stack.