feat: add support for setting cursor result set position for multi-result queries

diff --git a/docs/api/cursors.rst b/docs/api/cursors.rst
index ec7c8249f586665cb5cf09b8be28f6a607d2d7e5..6abb2b86c7315c4a5a7cc711ab3087129df4c2ed 100644 (file)
--- a/docs/api/cursors.rst
+++ b/docs/api/cursors.rst
@@ -264,6 +264,15 @@ The `!Cursor` class
     .. automethod:: fetchmany
     .. automethod:: fetchall
     .. automethod:: nextset
+    .. automethod:: set_result
+
+        Move to a specific result set if `execute()` returned multiple result
+        sets. The parameter `result_no` specifies the zero-based index of the
+        desired result set. Negative values are supported and refer to result
+        sets counted from the end.
+
+        .. versionadded:: 3.3
+
     .. automethod:: results
 
         The iterator yields the cursor itself upon iteration, but the cursor

diff --git a/psycopg/psycopg/_cursor_base.py b/psycopg/psycopg/_cursor_base.py
index 31bb751220d93979c9881dab8d4312d1d71c1c39..5b76744dd15240dbf771e59a1db714cfa49dd338 100644 (file)
--- a/psycopg/psycopg/_cursor_base.py
+++ b/psycopg/psycopg/_cursor_base.py
@@ -165,6 +165,21 @@ class BaseCursor(Generic[ConnectionType, Row]):
         else:
             return None
 
+    def set_result(self, result_no: int) -> None:
+        """
+        Move to a specific result set if `execute()` returned multiple result sets.
+
+        Args:
+            result_no (int): Zero based index, supports negative values
+        """
+        total_results = len(self._results)
+        if result_no < 0:
+            result_no = total_results + result_no
+        if 0 <= result_no < total_results:
+            self._select_current_result(result_no)
+        else:
+            raise ValueError("result_no value not in available results range")
+
     @property
     def statusmessage(self) -> str | None:
         """

diff --git a/tests/test_cursor_common.py b/tests/test_cursor_common.py
index 56db3b0814a658a7ed640986d6f024ee88856f5d..845e8e4e0a6a4c6a01f6565ddf7305036999c78f 100644 (file)
--- a/tests/test_cursor_common.py
+++ b/tests/test_cursor_common.py
@@ -220,6 +220,20 @@ def test_execute_many_results(conn):
     assert cur.rowcount == 3
     assert cur.nextset() is None
 
+    cur.set_result(0)
+    assert cur.fetchall() == [("foo",)]
+    assert cur.rowcount == 1
+
+    cur.set_result(-1)
+    assert cur.fetchall() == [(1,), (2,), (3,)]
+    assert cur.rowcount == 3
+
+    with pytest.raises(ValueError):
+        cur.set_result(2)
+
+    with pytest.raises(ValueError):
+        cur.set_result(-3)
+
     cur.close()
     assert cur.nextset() is None
 

diff --git a/tests/test_cursor_common_async.py b/tests/test_cursor_common_async.py
index 3b7b50194d56849186899a9e381e3ca9eb711cb1..87765e623e2ca20bb1b648cfdcebd02c0dcba4ce 100644 (file)
--- a/tests/test_cursor_common_async.py
+++ b/tests/test_cursor_common_async.py
@@ -218,6 +218,20 @@ async def test_execute_many_results(aconn):
     assert cur.rowcount == 3
     assert cur.nextset() is None
 
+    cur.set_result(0)
+    assert cur.fetchall() == [("foo",)]
+    assert cur.rowcount == 1
+
+    cur.set_result(-1)
+    assert cur.fetchall() == [(1,), (2,), (3,)]
+    assert cur.rowcount == 3
+
+    with pytest.raises(ValueError):
+        cur.set_result(2)
+
+    with pytest.raises(ValueError):
+        cur.set_result(-3)
+
     await cur.close()
     assert cur.nextset() is None