fix: psycopg instrumentation fixes #29
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
1 issue found across 6 files
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="drift/instrumentation/psycopg/instrumentation.py">
<violation number="1" location="drift/instrumentation/psycopg/instrumentation.py:810">
P2: Span resource leak: `span_info.span.end()` is placed after the yield loop in this generator. If the generator is not fully consumed or is abandoned, the span will never be ended. Use try/finally to ensure the span is always ended, similar to the pattern used in `_record_stream`.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
1 issue found across 3 files (changes from recent commits).
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="drift/instrumentation/psycopg/instrumentation.py">
<violation number="1" location="drift/instrumentation/psycopg/instrumentation.py:1803">
P2: Initial state setup for direct fetch access is incomplete. The comment says 'in case user calls fetch without results()' but `fetchone`/`fetchmany`/`fetchall` methods are not patched here - they're only set inside `mock_results()`. Direct fetch calls will use default implementations that ignore `_mock_rows`.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
There was a problem hiding this comment.
2 issues found across 7 files (changes from recent commits).
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="drift/instrumentation/psycopg/mocks.py">
<violation number="1" location="drift/instrumentation/psycopg/mocks.py:113">
P2: `MockConnection.cursor()` ignores the `row_factory` kwarg, so cursor-level row factory overrides (e.g., `conn.cursor(row_factory=dict_row)`) never work in REPLAY mode.</violation>
</file>
<file name="drift/instrumentation/psycopg/wrappers.py">
<violation number="1" location="drift/instrumentation/psycopg/wrappers.py:66">
P2: `write()` stores the original buffer object, so a reused `memoryview` mutates previously recorded COPY FROM chunks, corrupting captured data. Convert memoryviews to `bytes` before appending (as already done in the read path).</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
There was a problem hiding this comment.
1 issue found across 5 files (changes from recent commits).
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="drift/instrumentation/utils/serialization.py">
<violation number="1" location="drift/instrumentation/utils/serialization.py:45">
P2: UUID values are now serialized to plain strings but the replay deserializer never converts those strings back into `uuid.UUID` objects, so recorded rows lose their original types during replay.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
There was a problem hiding this comment.
1 issue found across 15 files (changes from recent commits).
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="drift/instrumentation/psycopg/instrumentation.py">
<violation number="1" location="drift/instrumentation/psycopg/instrumentation.py:1768">
P1: Exception in `do_lazy_capture` is swallowed instead of being re-raised. When `original_fetchall(cursor)` fails, the original exception is lost and subsequent access to `cursor._tusk_rows` will raise an `AttributeError`, hiding the actual database error from the user.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
NULL values handling now passes thanks to the lazy capture mechanism that defers fetchall() until user code actually calls fetch.
Remove bug investigation comments from the null-values endpoint since the issue is now resolved and the test is part of the E2E suite.
The conn.transaction() context manager now works correctly thanks to the lazy capture mechanism that defers fetchall() until fetch is called.
Remove bug investigation comments from the endpoint since the issue is now resolved and the test is part of the E2E suite.
…h returning=True The instrumentation was patching results() and nextset() methods for navigating result sets, but set_result(index) was missing. This caused REPLAY mode to fail with "index out of range" errors when user code called set_result() to jump to a specific result set. Added patched_set_result() for both REPLAY mode (in _mock_executemany_returning_with_data) and RECORD mode (in _finalize_executemany_returning_span) that: - Validates index bounds and supports negative indices - Updates the current result set state and fetch methods - Returns the cursor as the real implementation does
- Created _CursorInstrumentationMixin class to share common cursor properties (description, rownumber, statusmessage) and iteration methods (__iter__, __next__) between InstrumentedCursor and InstrumentedServerCursor - Created _set_span_attributes() helper method to consolidate the repeated pattern of setting span attributes (input/output values, schemas, and hashes) that was duplicated in 5 different places This reduces code duplication by ~110 lines while maintaining the same functionality (all 37 E2E tests pass).
…hashing
Added proper serialization for Decimal and timedelta types in
serialize_value() to ensure consistent hashing between RECORD and
REPLAY modes:
- Decimal: Serialized as {"__decimal__": str(val)} to preserve precision
- timedelta: Serialized as {"__timedelta__": total_seconds} for consistency
Also added corresponding deserialization in deserialize_db_value() to
reconstruct the original Python types when reading from traces.
This fixes REPLAY mismatch errors ("No mock found") when queries use
Decimal or timedelta parameters.
Created helper methods to reduce code duplication: - _create_query_span(): Centralizes span creation (used 8 times) - _create_fetch_methods(): Creates fetch closures for mock cursors - _create_scroll_method(): Creates scroll method for cursor mocking - _get_row_factory_from_cursor(): Extracts row factory consistently - _set_cursor_description(): Sets cursor description with error handling - _create_row_transformer(): Creates row transform functions This reduces code duplication and improves maintainability while maintaining all existing functionality (39 E2E tests pass).
Added support for Python ipaddress module types in serialize_value(): - IPv4Address, IPv6Address - IPv4Interface, IPv6Interface - IPv4Network, IPv6Network These types are returned by psycopg when querying PostgreSQL inet and cidr columns. Without proper serialization, REPLAY mode failed to match recorded mocks.
Added serialization and deserialization support for psycopg Range
objects (INT4RANGE, TSRANGE, etc.):
Serialization (in serialize_value):
- Serializes to {"__range__": {"lower": ..., "upper": ..., "bounds": ...}}
- Handles empty ranges with {"__range__": {"empty": True}}
- Recursively serializes bounds for datetime/nested types
Deserialization (in deserialize_db_value):
- Reconstructs Range objects from tagged dict structure
- Recursively deserializes bounds
- Converts JSON floats back to ints when appropriate
This fixes REPLAY mismatch when Range types are used as query
parameters or returned in results.
There was a problem hiding this comment.
1 issue found across 9 files (changes from recent commits).
Prompt for AI agents (all issues)
Check if these issues are valid — if so, understand the root cause of each and fix them.
<file name="drift/instrumentation/psycopg/e2e-tests/src/app.py">
<violation number="1" location="drift/instrumentation/psycopg/e2e-tests/src/app.py:4">
P3: `UTC` is imported but never used; drop the unused import to keep the module clean and prevent lint failures.</violation>
</file>
Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.
jy-tan
approved these changes
Jan 16, 2026
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Note
Improves reliability of DB replay and makes e2e tests consistent and verifiable.
cursor.stream(),cursor.copy()(COPY TO/FROM), pipeline mode deferral/finalization, server cursors,executemany(..., returning=True), cursor iteration/scrolling/rownumber/statusmessage; introduces lazy result capture to avoid hangs; centralizes mock lookup; and robust type (bytes/UUID/Decimal/timedelta/inet/cidr/range) serialization/deserializatione2e_common.test_utilsprovidesmake_requestand printsTOTAL_REQUESTS_SENT;E2ETestRunnerBaseparses this count, setsPYTHONPATH, and validates passed test count vs requestsis_pre_app_startin HTTPX/requests; psycopg2 instrumentation now uses centralized mock matching and shared serializationWritten by Cursor Bugbot for commit 26f92f7. This will update automatically on new commits. Configure here.