summaryrefslogtreecommitdiff
path: root/docker/api/buildkit/grpc.py
blob: 9ee640dc2bafe596c613a74fde5d7ea10fda3dbe (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import grpclib.server


class Handler(grpclib.server.Handler):
    def __init__(self, headers, *args, **kwargs):
        self.headers = headers
        super(Handler, self).__init__(*args, **kwargs)

    def accept(self, stream, headers, release_stream):
        for k, v in self.headers.items():
            headers.setdefault(k, v)


class Server(grpclib.server.Server):
    headers = {}

    def add_headers(self, h):
        self.headers.update(h)

    def _protocol_factory(self):
        self.__gc_step__()
        handler = Handler(
            self.headers, self._mapping, self._codec, loop=self._loop
        )
        self._handlers.add(handler)
        return grpclib.server.H2Protocol(
            handler, self._config, loop=self._loop
        )


class Attachable(object):
    def register(self, grpc_server):
        grpc_server._mapping.update(self.__mapping__())