From: Aditya Toshniwal <aditya.toshniwal@enterprisedb.com>
Date: Mon, 6 Oct 2025 08:58:56 +0000 (+0530)
Subject: feat: add support for setting cursor result set position for multi-result queries
X-Git-Tag: 3.3.0~13^2~1
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=f3c32209d2d8b5da51b5f414ab4b1b6a9ad8782a;p=thirdparty%2Fpsycopg.git

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 ec7c8249f..6abb2b86c 100644
--- 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 31bb75122..5b76744dd 100644
--- 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 56db3b081..845e8e4e0 100644
--- 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 3b7b50194..87765e623 100644
--- 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