fix cursor.scroll() position tracking in RECORD mode

Author: sohankshirsagar

drift/instrumentation/psycopg/e2e-tests/src/app.py

@@ -414,58 +414,11 @@ def test_statusmessage():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
-
-# ============================================================================
-# CONFIRMED BUG TEST ENDPOINTS
-# These endpoints expose confirmed bugs in the psycopg instrumentation.
-# See BUG_TRACKING.md for detailed documentation of each bug.
-#
-# Bug Summary:
-# 1. /test/cursor-scroll - scroll() broken during RECORD mode
-# 4. /test/nextset - nextset() iteration broken during RECORD mode
-# 5. /test/server-cursor-scroll - scroll() broken during RECORD mode
-# ============================================================================
-
-@app.route("/test/cursor-scroll")
-def test_cursor_scroll():
-    """Test cursor.scroll() method.
-
-    BUG: During RECORD mode, the instrumentation's _finalize_query_span calls
-    fetchall() which breaks the cursor position tracking. After fetchall(),
-    scroll(0, absolute) doesn't properly reset the cursor position because
-    the patched fetch methods use _tusk_index instead of _pos.
-    """
-    try:
-        with psycopg.connect(get_conn_string()) as conn, conn.cursor() as cur:
-            cur.execute("SELECT id, name FROM users ORDER BY id")
-
-            # Fetch first row
-            first = cur.fetchone()
-
-            # Scroll back to start
-            cur.scroll(0, mode='absolute')
-
-            # Fetch first row again
-            first_again = cur.fetchone()
-
-        return jsonify({
-            "first": {"id": first[0], "name": first[1]} if first else None,
-            "first_again": {"id": first_again[0], "name": first_again[1]} if first_again else None,
-            "match": first == first_again
-        })
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
 @app.route("/test/nextset")
 def test_nextset():
     """Test cursor.nextset() for multiple result sets.
 
-    BUG: During RECORD mode, the interaction between executemany with
-    returning=True and the fetch method patching breaks nextset() iteration.
-    The instrumentation's fetch patching may consume results before nextset()
-    can iterate through them, causing 0 results in RECORD but correct
-    results in REPLAY.
+    Tests whether the instrumentation correctly handles nextset() for multiple result sets.
     """
     try:
         with psycopg.connect(get_conn_string()) as conn, conn.cursor() as cur:
@@ -497,16 +450,38 @@ def test_nextset():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
+@app.route("/test/cursor-scroll")
+def test_cursor_scroll():
+    """Test cursor.scroll() method.
+
+    Tests whether the instrumentation correctly handles scroll() for cursor position tracking.
+    """
+    try:
+        with psycopg.connect(get_conn_string()) as conn, conn.cursor() as cur:
+            cur.execute("SELECT id, name FROM users ORDER BY id")
+
+            # Fetch first row
+            first = cur.fetchone()
+
+            # Scroll back to start
+            cur.scroll(0, mode='absolute')
+
+            # Fetch first row again
+            first_again = cur.fetchone()
+
+        return jsonify({
+            "first": {"id": first[0], "name": first[1]} if first else None,
+            "first_again": {"id": first_again[0], "name": first_again[1]} if first_again else None,
+            "match": first == first_again
+        })
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
 @app.route("/test/server-cursor-scroll")
 def test_server_cursor_scroll():
     """Test ServerCursor.scroll() method.
 
-    BUG: Same root cause as /test/cursor-scroll. During RECORD mode,
-    the instrumentation breaks scroll() functionality by consuming all
-    rows via fetchall() in _finalize_query_span. ServerCursor.scroll()
-    sends MOVE commands to the server, but the position tracking is
-    inconsistent after the instrumentation patches the fetch methods.
+    Tests whether the instrumentation correctly handles scroll() for server-side cursors.
     """
     try:
         with psycopg.connect(get_conn_string()) as conn:

drift/instrumentation/psycopg/e2e-tests/src/test_requests.py

@@ -86,9 +86,10 @@ def make_request(method, endpoint, **kwargs):
 
     make_request("GET", "/test/statusmessage")
 
-    # Bug-exposing test endpoints
     make_request("GET", "/test/nextset")
+
     make_request("GET", "/test/server-cursor-scroll")
+
     make_request("GET", "/test/cursor-scroll")
 
     print("\nAll requests completed successfully")

drift/instrumentation/psycopg/instrumentation.py

@@ -857,6 +857,8 @@ def _record_execute(
             cursor.fetchall = cursor._tusk_original_fetchall
             cursor._tusk_rows = None
             cursor._tusk_index = 0
+        if hasattr(cursor, '_tusk_original_scroll'):
+            cursor.scroll = cursor._tusk_original_scroll
 
         span_info = SpanUtils.create_span(
             CreateSpanOptions(
@@ -2090,16 +2092,36 @@ def patched_fetchall():
                                 cursor._tusk_index = len(cursor._tusk_rows)  # pyright: ignore[reportAttributeAccessIssue]
                                 return result
 
+                            def patched_scroll(value: int, mode: str = "relative") -> None:
+                                """Scroll the cursor to a new position in the captured result set."""
+                                if mode == "relative":
+                                    newpos = cursor._tusk_index + value  # pyright: ignore[reportAttributeAccessIssue]
+                                elif mode == "absolute":
+                                    newpos = value
+                                else:
+                                    raise ValueError(f"bad mode: {mode}. It should be 'relative' or 'absolute'")
+
+                                num_rows = len(cursor._tusk_rows)  # pyright: ignore[reportAttributeAccessIssue]
+                                if num_rows > 0:
+                                    if not (0 <= newpos < num_rows):
+                                        raise IndexError("cursor position out of range")
+                                elif newpos != 0:
+                                    raise IndexError("cursor position out of range")
+
+                                cursor._tusk_index = newpos  # pyright: ignore[reportAttributeAccessIssue]
+
                             # Save original fetch methods before patching (only if not already saved)
                             # These will be restored at the start of the next execute() call
                             if not hasattr(cursor, '_tusk_original_fetchone'):
                                 cursor._tusk_original_fetchone = cursor.fetchone  # pyright: ignore[reportAttributeAccessIssue]
                                 cursor._tusk_original_fetchmany = cursor.fetchmany  # pyright: ignore[reportAttributeAccessIssue]
                                 cursor._tusk_original_fetchall = cursor.fetchall  # pyright: ignore[reportAttributeAccessIssue]
+                                cursor._tusk_original_scroll = cursor.scroll  # pyright: ignore[reportAttributeAccessIssue]
 
                             cursor.fetchone = patched_fetchone  # pyright: ignore[reportAttributeAccessIssue]
                             cursor.fetchmany = patched_fetchmany  # pyright: ignore[reportAttributeAccessIssue]
                             cursor.fetchall = patched_fetchall  # pyright: ignore[reportAttributeAccessIssue]
+                            cursor.scroll = patched_scroll  # pyright: ignore[reportAttributeAccessIssue]
 
                         except Exception as fetch_error:
                             logger.debug(f"Could not fetch rows (query might not return rows): {fetch_error}")