• Search by Module
• Search by Word
• Project Search
• Top Python APIs
• Popular Projects

• Java
• C++
• Python
• Scala
• Blog

• hiku-master
  • requirements-docs.in
  • hiku
    • endpoint
      • __init__.py
      • graphql.py
    • expr
      • nodes.py
      • refs.py
      • __init__.py
      • compiler.py
      • checker.py
      • core.py
    • engine.py
    • typedef
      • kinko.py
      • __init__.py
      • types.py
    • graph.py
    • executors
      • threads.py
      • queue.py
      • sync.py
      • asyncio.py
      • __init__.py
    • sources
      • sqlalchemy.py
      • graph.py
      • aiopg.py
      • __init__.py
    • protobuf
      • query.proto
      • query_pb2.py
      • __init__.py
    • validate
      • graph.py
      • query.py
      • __init__.py
      • errors.py
    • edn.py
    • query.py
    • export
      • simple.py
      • protobuf.py
      • __init__.py
      • graphql.py
    • console
      • ui.py
      • __init__.py
      • assets
        • console.jsx
        • console.css
        • console.html
    • introspection
      • __init__.py
      • graphql.py
    • result.py
    • __init__.py
    • utils.py
    • readers
      • simple.py
      • protobuf.py
      • __init__.py
      • graphql.py
    • compat.py
    • builder.py
    • types.py
    • denormalize
      • __init__.py
      • graphql.py
      • base.py
    • telemetry
      • prometheus.py
      • __init__.py
  • requirements-tests.txt
  • Makefile
  • pi.yaml
  • requirements-rtd.txt
  • examples
    • console.py
    • graphql_flask.py
    • graphql_aiohttp.py
  • readthedocs.yml
  • requirements-tests.in
  • webpack.config.js
  • README.rst
  • setup.py
  • tests_pg
    • test_source_aiopg.py
  • requirements.in
  • .travis.yml
  • package.json
  • tests
    • test_executor_asyncio.py
    • test_export_simple.py
    • test_builder.py
    • test_source_graph.py
    • test_export_graphql.py
    • test_edn.py
    • test_denormalize_base.py
    • test_source_sqlalchemy.py
    • test_console.py
    • protobuf
      • result_pb2.py
      • result.proto
      • __init__.py
    • test_query.py
    • test_introspection_graphql.py
    • test_endpoint_graphql.py
    • test_compiler.py
    • test_telemetry_prometheus.py
    • test_read_protobuf.py
    • test_executor_queue.py
    • test_refs.py
    • test_typedef_kinko.py
    • test_engine.py
    • __init__.py
    • test_checker.py
    • test_validate_graph.py
    • test_read_simple.py
    • test_export_protobuf.py
    • test_denormalize_graphql.py
    • test_validate_query.py
    • test_read_graphql.py
    • test_result.py
    • base.py
  • scripts
    • release_check.sh
  • requirements.txt
  • .gitignore
  • docs
    • test_protobuf.py
    • graphql.rst
    • basics.rst
    • example_pb2.py
    • subgraph.rst
    • protobuf.rst
    • test_database.py
    • changelog
      • changes_04.rst
      • changes_03.rst
      • changes_06.rst
      • changes_05.rst
      • changes_02.rst
      • index.rst
    • database.rst
    • reference
      • expr.rst
      • protobuf
        • query.rst
      • graph.rst
      • readers.rst
      • result.rst
      • index.rst
      • query.rst
    • test_subgraph.py
    • index.rst
    • conf.py
    • example.proto
    • asyncio.rst
    • test_asyncio.py
    • _static
      • style.css
      • fixes.css
    • basics
      • test_stage1.py
      • test_stage3.py
      • test_stage2.py
      • __init__.py
  • requirements-docs.txt
  • LICENSE.txt
  • MANIFEST.in
  • tox.ini
  • requirements-tox.txt

import time
from abc import abstractmethod
from functools import partial

from prometheus_client import Summary

from ..graph import GraphTransformer
from ..engine import pass_context, _do_pass_context
from ..sources.graph import CheckedExpr


_METRIC = None


def _get_default_metric():
    global _METRIC
    if _METRIC is None:
        _METRIC = Summary(
            'graph_field_time',
            'Graph field time (seconds)',
            ['graph', 'node', 'field'],
        )
    return _METRIC


def _func_field_names(func):
    fields_pos = 1 if _do_pass_context(func) else 0

    def wrapper(*args):
        return func([f.name for f in args[fields_pos]], *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'

    def __init__(self, name, *, metric=None):
        self._name = name
        self._metric = metric or _get_default_metric()
        self._node = None
        self._wrappers = {}

    @abstractmethod
    def field_wrapper(self, observe, func):
        raise NotImplementedError

    @abstractmethod
    def link_wrapper(self, observe, func):
        raise NotImplementedError

    @abstractmethod
    def subquery_wrapper(self, observe, subquery):
        raise NotImplementedError

    def _observe_fields(self, node_name):
        by_field = {}

        def observe(start_time, field_names):
            duration = time.perf_counter() - start_time
            for name in field_names:
                try:
                    field_metric = by_field[name]
                except KeyError:
                    field_metric = by_field[name] = \
                        self._metric.labels(self._name, node_name, name)
                field_metric.observe(duration)
        return observe

    def _wrap_field(self, node_name, func):
        observe = self._observe_fields(node_name)
        wrapper = self.field_wrapper(observe, func)
        if _do_pass_context(func):
            wrapper = pass_context(wrapper)
        wrapper = _func_field_names(wrapper)
        if _do_pass_context(func):
            wrapper = pass_context(wrapper)
        return wrapper

    def _wrap_link(self, node_name, link_name, func):
        observe = self._observe_fields(node_name)
        wrapper = self.link_wrapper(observe, func)
        if _do_pass_context(func):
            wrapper = pass_context(wrapper)
        wrapper = partial(wrapper, link_name)
        if _do_pass_context(func):
            wrapper = pass_context(wrapper)
        return wrapper

    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

    def visit_node(self, obj):
        self._node = obj
        try:
            return super().visit_node(obj)
        finally:
            self._node = None

    def visit_field(self, obj):
        obj = super().visit_field(obj)
        node_name = self.root_name if self._node is None else self._node.name
        if isinstance(obj.func, CheckedExpr):
            func = obj.func.__subquery__
        else:
            func = obj.func

        wrapper = self._wrappers.get(func)
        if wrapper is None:
            if isinstance(obj.func, CheckedExpr):
                wrapper = self._wrappers[func] = self._wrap_subquery(
                    node_name, func,
                )
            else:
                wrapper = self._wrappers[func] = self._wrap_field(
                    node_name, func,
                )

        if isinstance(obj.func, CheckedExpr):
            obj.func = CheckedExpr(
                wrapper,
                obj.func.expr,
                obj.func.reqs,
                obj.func.proc,
            )
        else:
            obj.func = wrapper
        return obj

    def visit_link(self, obj):
        obj = super().visit_link(obj)
        node_name = self.root_name if self._node is None else self._node.name
        obj.func = self._wrap_link(node_name, obj.name, obj.func)
        return obj


class _SubqueryMixin:

    def subquery_wrapper(self, observe, subquery):
        def wrapper(field_names, *args):
            start_time = time.perf_counter()
            result_proc = subquery(*args)

            def proc_wrapper():
                result = result_proc()
                observe(start_time, field_names)
                return result
            return proc_wrapper
        return wrapper


class GraphMetrics(_SubqueryMixin, GraphMetricsBase):

    def field_wrapper(self, observe, func):
        def wrapper(field_names, *args):
            start_time = time.perf_counter()
            result = func(*args)
            observe(start_time, field_names)
            return result
        return wrapper

    def link_wrapper(self, observe, func):
        def wrapper(link_name, *args):
            start_time = time.perf_counter()
            result = func(*args)
            observe(start_time, [link_name])
            return result
        return wrapper


class AsyncGraphMetrics(_SubqueryMixin, GraphMetricsBase):

    def field_wrapper(self, observe, func):
        async def wrapper(field_names, *args):
            start_time = time.perf_counter()
            result = await func(*args)
            observe(start_time, field_names)
            return result
        return wrapper

    def link_wrapper(self, observe, func):
        async def wrapper(link_name, *args):
            start_time = time.perf_counter()
            result = await func(*args)
            observe(start_time, [link_name])
            return result
        return wrapper