diff --git a/drift/instrumentation/psycopg/instrumentation.py b/drift/instrumentation/psycopg/instrumentation.py index b79fa41..c4423d8 100644 --- a/drift/instrumentation/psycopg/instrumentation.py +++ b/drift/instrumentation/psycopg/instrumentation.py @@ -559,10 +559,32 @@ def _traced_execute( span_kind=OTelSpanKind.CLIENT, ) - def _noop_execute(self, cursor: Any) -> Any: + def _noop_execute(self, cursor: Any, is_async: bool = False) -> Any: """Handle background requests in REPLAY mode - return cursor with empty mock data.""" cursor._mock_rows = [] # pyright: ignore cursor._mock_index = 0 # pyright: ignore + fetchone, fetchmany, fetchall = self._create_fetch_methods(cursor, "_mock_rows", "_mock_index") + if is_async: + sync_fetchone, sync_fetchmany, sync_fetchall = fetchone, fetchmany, fetchall + + async def async_fetchone(): + return sync_fetchone() + + async def async_fetchmany(size=None): + if size is None: + size = cursor.arraysize + return sync_fetchmany(size) + + async def async_fetchall(): + return sync_fetchall() + + cursor.fetchone = async_fetchone # pyright: ignore[reportAttributeAccessIssue] + cursor.fetchmany = async_fetchmany # pyright: ignore[reportAttributeAccessIssue] + cursor.fetchall = async_fetchall # pyright: ignore[reportAttributeAccessIssue] + else: + cursor.fetchone = fetchone # pyright: ignore[reportAttributeAccessIssue] + cursor.fetchmany = fetchmany # pyright: ignore[reportAttributeAccessIssue] + cursor.fetchall = fetchall # pyright: ignore[reportAttributeAccessIssue] return cursor def _replay_execute(self, cursor: Any, sdk: TuskDrift, query_str: str, params: Any, is_async: bool = False) -> Any: @@ -581,9 +603,17 @@ def _replay_execute(self, cursor: Any, sdk: TuskDrift, query_str: str, params: A if mock_result is None: is_pre_app_start = not sdk.app_ready + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start psycopg query, returning empty result. " + f"Query: {query_str[:100]}..." + ) + self._noop_execute(cursor, is_async=is_async) + span_info.span.end() + return cursor raise RuntimeError( f"[Tusk REPLAY] No mock found for psycopg execute query. " - f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. " + f"This query was not recorded during the trace capture. " f"Query: {query_str[:100]}..." ) @@ -742,12 +772,17 @@ def _replay_executemany( if mock_result is None: is_pre_app_start = not sdk.app_ready - logger.error( - f"No mock found for {'pre-app-start ' if is_pre_app_start else ''}psycopg executemany query in REPLAY mode: {query_str[:100]}" - ) + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start psycopg executemany query, returning empty result. " + f"Query: {query_str[:100]}..." + ) + self._noop_execute(cursor) + span_info.span.end() + return cursor raise RuntimeError( f"[Tusk REPLAY] No mock found for psycopg executemany query. " - f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. " + f"This query was not recorded during the trace capture. " f"Query: {query_str[:100]}..." ) @@ -1150,9 +1185,16 @@ def _replay_stream(self, cursor: Any, sdk: TuskDrift, query_str: str, params: An if mock_result is None: is_pre_app_start = not sdk.app_ready + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start psycopg stream query, returning empty result. " + f"Query: {query_str[:100]}..." + ) + span_info.span.end() + return raise RuntimeError( f"[Tusk REPLAY] No mock found for psycopg stream query. " - f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded. " + f"This query was not recorded. " f"Query: {query_str[:100]}..." ) @@ -1295,9 +1337,17 @@ def _replay_copy(self, cursor: Any, sdk: TuskDrift, query_str: str) -> Iterator[ if mock_result is None: is_pre_app_start = not sdk.app_ready + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start psycopg copy operation, returning empty result. " + f"Query: {query_str[:100]}..." + ) + span_info.span.end() + yield MockCopy(data=[]) + return raise RuntimeError( f"[Tusk REPLAY] No mock found for psycopg copy operation. " - f"This {'pre-app-start ' if is_pre_app_start else ''}copy was not recorded. " + f"This copy was not recorded. " f"Query: {query_str[:100]}..." ) diff --git a/drift/instrumentation/psycopg2/instrumentation.py b/drift/instrumentation/psycopg2/instrumentation.py index cc722db..51847ff 100644 --- a/drift/instrumentation/psycopg2/instrumentation.py +++ b/drift/instrumentation/psycopg2/instrumentation.py @@ -604,9 +604,17 @@ def _replay_execute(self, cursor: Any, sdk: TuskDrift, query_str: str, params: A if mock_result is None: is_pre_app_start = not sdk.app_ready + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start psycopg2 query, returning empty result. " + f"Query: {query_str[:100]}..." + ) + self._noop_execute(cursor) + span_info.span.end() + return None raise RuntimeError( f"[Tusk REPLAY] No mock found for psycopg2 execute query. " - f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. " + f"This query was not recorded during the trace capture. " f"Query: {query_str[:100]}..." ) @@ -754,9 +762,17 @@ def _replay_executemany(self, cursor: Any, sdk: TuskDrift, query_str: str, param if mock_result is None: is_pre_app_start = not sdk.app_ready + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start psycopg2 executemany query, returning empty result. " + f"Query: {query_str[:100]}..." + ) + self._noop_execute(cursor) + span_info.span.end() + return None raise RuntimeError( f"[Tusk REPLAY] No mock found for psycopg2 executemany query. " - f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. " + f"This query was not recorded during the trace capture. " f"Query: {query_str[:100]}..." ) diff --git a/drift/instrumentation/redis/instrumentation.py b/drift/instrumentation/redis/instrumentation.py index 5b30362..7bbe4e5 100644 --- a/drift/instrumentation/redis/instrumentation.py +++ b/drift/instrumentation/redis/instrumentation.py @@ -431,9 +431,16 @@ def _replay_execute_command(self, sdk: TuskDrift, command_name: str, command_str if mock_result is None: is_pre_app_start = not sdk.app_ready + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start Redis command, returning default response. " + f"Command: {command_str}" + ) + span_info.span.end() + return self._get_default_response(command_name) raise RuntimeError( f"[Tusk REPLAY] No mock found for Redis command. " - f"This {'pre-app-start ' if is_pre_app_start else ''}command was not recorded during the trace capture. " + f"This command was not recorded during the trace capture. " f"Command: {command_str}" ) @@ -769,9 +776,16 @@ def _replay_pipeline_execute(self, sdk: TuskDrift, command_str: str, command_sta if mock_result is None: is_pre_app_start = not sdk.app_ready + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start Redis pipeline, returning empty result. " + f"Commands: {command_str}" + ) + span_info.span.end() + return [] raise RuntimeError( f"[Tusk REPLAY] No mock found for Redis pipeline. " - f"This {'pre-app-start ' if is_pre_app_start else ''}pipeline was not recorded during the trace capture. " + f"This pipeline was not recorded during the trace capture. " f"Commands: {command_str}" ) diff --git a/drift/instrumentation/sqlalchemy/instrumentation.py b/drift/instrumentation/sqlalchemy/instrumentation.py index 05f184d..3ed93a4 100644 --- a/drift/instrumentation/sqlalchemy/instrumentation.py +++ b/drift/instrumentation/sqlalchemy/instrumentation.py @@ -89,19 +89,26 @@ def before_cursor_execute(conn, cursor, statement, parameters, context, executem if mock_result is None: is_pre_app_start = not sdk.app_ready - span_info.span.set_status( - Status( - OTelStatusCode.ERROR, - "No mock found for sqlalchemy query in replay", + if is_pre_app_start: + logger.warning( + f"[Tusk REPLAY] No mock found for pre-app-start sqlalchemy query, returning empty result. " + f"Query: {query_str[:120]}..." + ) + mock_result = {"rows": [], "rowcount": 0, "description": None} + else: + span_info.span.set_status( + Status( + OTelStatusCode.ERROR, + "No mock found for sqlalchemy query in replay", + ) + ) + span_info.span.end() + instrumentation._reset_context_state(context) + raise RuntimeError( + f"[Tusk REPLAY] No mock found for sqlalchemy query. " + f"This query was not recorded during the trace capture. " + f"Query: {query_str[:120]}..." ) - ) - span_info.span.end() - instrumentation._reset_context_state(context) - raise RuntimeError( - f"[Tusk REPLAY] No mock found for sqlalchemy query. " - f"This {'pre-app-start ' if is_pre_app_start else ''}query was not recorded during the trace capture. " - f"Query: {query_str[:120]}..." - ) # Preserve error-path behavior in replay for queries recorded as failures. if isinstance(mock_result, dict):