diff --git a/drift/instrumentation/psycopg/instrumentation.py b/drift/instrumentation/psycopg/instrumentation.py index bbdcfad..4f7a448 100644 --- a/drift/instrumentation/psycopg/instrumentation.py +++ b/drift/instrumentation/psycopg/instrumentation.py @@ -79,24 +79,13 @@ def patched_connect(*args, **kwargs): # Create server cursor factory for named cursors (conn.cursor(name="...")) server_cursor_factory = instrumentation._create_server_cursor_factory(sdk) - # In REPLAY mode, try to connect but fall back to mock connection if DB is unavailable + # In REPLAY mode, skip real DB connection entirely - use MockConnection directly + # This enables true dependency-free replay without requiring a database server if sdk.mode == TuskDriftMode.REPLAY: - try: - kwargs["cursor_factory"] = cursor_factory - if user_row_factory is not None: - kwargs["row_factory"] = user_row_factory - connection = original_connect(*args, **kwargs) - # Set server cursor factory on the connection for named cursors - if server_cursor_factory: - connection.server_cursor_factory = server_cursor_factory - logger.info("[PATCHED_CONNECT] REPLAY mode: Successfully connected to database (psycopg3)") - return connection - except Exception as e: - logger.info( - f"[PATCHED_CONNECT] REPLAY mode: Database connection failed ({e}), using mock connection (psycopg3)" - ) - # Return mock connection that doesn't require a real database - return MockConnection(sdk, instrumentation, cursor_factory, row_factory=user_row_factory) + logger.info( + "[PATCHED_CONNECT] REPLAY mode: Using MockConnection (no real DB connection needed, psycopg3)" + ) + return MockConnection(sdk, instrumentation, cursor_factory, row_factory=user_row_factory) # In RECORD mode, always require real connection kwargs["cursor_factory"] = cursor_factory diff --git a/drift/instrumentation/psycopg2/instrumentation.py b/drift/instrumentation/psycopg2/instrumentation.py index 6145e48..0973d27 100644 --- a/drift/instrumentation/psycopg2/instrumentation.py +++ b/drift/instrumentation/psycopg2/instrumentation.py @@ -49,7 +49,15 @@ class MockConnection: def __init__(self, sdk: TuskDrift, instrumentation: Psycopg2Instrumentation, cursor_factory): self.sdk = sdk self.instrumentation = instrumentation + + # Set default cursor factory if None (needed for django_prometheus compatibility) + # django_prometheus does: conn.cursor_factory or Cursor() which fails if cursor_factory is None + if cursor_factory is None: + from psycopg2.extensions import cursor as DefaultCursor + + cursor_factory = DefaultCursor self.cursor_factory = cursor_factory + self.closed = False self.autocommit = False @@ -440,21 +448,11 @@ def patched_connect(*args, **kwargs): logger.debug("[PATCHED_CONNECT] SDK disabled, passing through") return original_connect(*args, **kwargs) - # In REPLAY mode, try to connect but fall back to mock connection if DB is unavailable + # In REPLAY mode, skip real DB connection entirely - use MockConnection directly + # This enables true dependency-free replay without requiring a database server if sdk.mode == TuskDriftMode.REPLAY: - try: - logger.debug("[PATCHED_CONNECT] REPLAY mode: Attempting real DB connection...") - connection = original_connect(*args, **kwargs) - logger.info("[PATCHED_CONNECT] REPLAY mode: Successfully connected to real database") - # Wrap connection to intercept cursor() calls - return InstrumentedConnection(connection, instrumentation, sdk) - except Exception as e: - logger.info( - f"[PATCHED_CONNECT] REPLAY mode: Database connection failed ({e}), using mock connection" - ) - # Return mock connection that doesn't require a real database - # MockConnection already handles cursor_factory correctly in its cursor() method - return MockConnection(sdk, instrumentation, None) + logger.info("[PATCHED_CONNECT] REPLAY mode: Using MockConnection (no real DB connection needed)") + return MockConnection(sdk, instrumentation, None) # In RECORD mode, always require real connection logger.debug("[PATCHED_CONNECT] RECORD mode: Connecting to database...")