Skip to content

Commit

Permalink
Fixup: use hardened bind
Browse files Browse the repository at this point in the history
  • Loading branch information
hmstepanek committed Aug 22, 2024
1 parent ab6ad0b commit d677524
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion newrelic/hooks/mlmodel_langchain.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,12 +125,17 @@
}


def bind_submit(func, *args, **kwargs):
return {"func": func, "args": args, "kwargs": kwargs}


def wrap_ContextThreadPoolExecutor_submit(wrapped, instance, args, kwargs):
trace = current_trace()
if not trace:
return wrapped(*args, **kwargs)

bound_args = bind_args(wrapped, args, kwargs)
# Use hardened function signature bind so we have safety net catchall of args and kwargs.
bound_args = bind_submit(*args, **kwargs)
bound_args["func"] = context_wrapper(bound_args["func"], trace=trace, strict=True)
return wrapped(bound_args["func"], *bound_args["args"], **bound_args["kwargs"])

Expand Down

0 comments on commit d677524

Please sign in to comment.