Skip to content

fix: gracefully handle missing mocks for pre-app-start queries in replay mode #92

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sohankshirsagar merged 3 commits into from
Apr 13, 2026
Merged

fix: gracefully handle missing mocks for pre-app-start queries in replay mode #92

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e6b9998
fix: gracefully handle missing mocks for pre-app-start queries in rep…
sohankshirsagar Apr 13, 2026
e2cfa88
Address pr comment
sohankshirsagar Apr 13, 2026
89b815b
add async support for noop
sohankshirsagar Apr 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 58 additions & 8 deletions drift/instrumentation/psycopg/instrumentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
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:
Expand All @@ -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]}..."
)

Expand Down Expand Up @@ -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)
Comment thread
sohankshirsagar marked this conversation as resolved.
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]}..."
)

Expand Down Expand Up @@ -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]}..."
)

Expand Down Expand Up @@ -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]}..."
)

Expand Down
20 changes: 18 additions & 2 deletions drift/instrumentation/psycopg2/instrumentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Comment thread
sohankshirsagar marked this conversation as resolved.
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]}..."
)

Expand Down Expand Up @@ -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]}..."
)

Expand Down
18 changes: 16 additions & 2 deletions drift/instrumentation/redis/instrumentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}"
)

Expand Down Expand Up @@ -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}"
)

Expand Down
31 changes: 19 additions & 12 deletions drift/instrumentation/sqlalchemy/instrumentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading