bgpwtf/cccampix/frontend/channel.py

import logging

from flask import current_app
from flask import _app_ctx_stack as stack
import grpc


logger = logging.getLogger(__name__)


class Channel:
    def __init__(self, app, address):
        self.app = app
        self.address = address
        self.stubs = {}
        app.teardown_appcontext(self.teardown)

    def _connect(self):
        logger.info("Connecting to {}...".format(self.address))
        return grpc.insecure_channel(self.address)

    @property
    def conn(self):
        ctx = stack.top
        if ctx is not None:
            if not hasattr(ctx, 'conn'):
                ctx.conn = self._connect()
            return ctx.conn

    def stub(self, stub):
        if stub not in self.stubs:
            self.stubs[stub] = stub(self.conn)
        return self.stubs[stub]

    def teardown(self, exception):
        ctx = stack.top
        if hasattr(ctx, 'conn'):
            del ctx.conn