diff --git a/hiku/telemetry/prometheus.py b/hiku/telemetry/prometheus.py index c080e5fe..9831e466 100644 --- a/hiku/telemetry/prometheus.py +++ b/hiku/telemetry/prometheus.py @@ -38,12 +38,6 @@ def wrapper(*args): return wrapper -def _subquery_field_names(func): - def wrapper(fields, *args): - return func([f.name for _, f in fields], fields, *args) - return wrapper - - class GraphMetricsBase(GraphTransformer): root_name = 'Root' @@ -102,7 +96,6 @@ def _wrap_link(self, node_name, link_name, func): def _wrap_subquery(self, node_name, subquery): observe = self._observe_fields(node_name) wrapper = self.subquery_wrapper(observe, subquery) - wrapper = _subquery_field_names(wrapper) wrapper.__subquery__ = lambda: wrapper return wrapper @@ -153,15 +146,23 @@ def visit_link(self, obj): class _SubqueryMixin: def subquery_wrapper(self, observe, subquery): - def wrapper(field_names, *args): + def wrap_proc(start_time, field_name, proc): + def _proc_wrapper(*args): + result = proc(*args) + observe(start_time, [field_name]) + return result + + return _proc_wrapper + + def wrapper(fields, *args): start_time = time.perf_counter() - result_proc = subquery(*args) + wrapped_fields = [] + for gf, qf in fields: + gf.func.proc = wrap_proc(start_time, gf.name, gf.func.proc) + wrapped_fields.append((gf, qf)) + + return subquery(wrapped_fields, *args) - def proc_wrapper(): - result = result_proc() - observe(start_time, field_names) - return result - return proc_wrapper return wrapper @@ -200,4 +201,4 @@ async def wrapper(link_name, *args): result = await func(*args) observe(start_time, [link_name]) return result - return wrapper + return wrapper \ No newline at end of file