|
| 1 | +import os |
| 2 | +import shlex |
| 3 | + |
| 4 | +from localstack.config import is_env_not_false |
| 5 | +from localstack.utils.docker_utils import DOCKER_CLIENT |
| 6 | +from localstack_typedb.utils.docker import ProxiedDockerContainerExtension |
| 7 | +from rolo import Request |
| 8 | +from werkzeug.datastructures import Headers |
| 9 | + |
| 10 | +# environment variable for user-defined command args to pass to TypeDB |
| 11 | +ENV_CMD_FLAGS = "TYPEDB_FLAGS" |
| 12 | +# environment variable for flag to enable/disable HTTP2 proxy for gRPC traffic |
| 13 | +ENV_HTTP2_PROXY = "TYPEDB_HTTP2_PROXY" |
| 14 | + |
| 15 | + |
| 16 | +class TypeDbExtension(ProxiedDockerContainerExtension): |
| 17 | + name = "typedb" |
| 18 | + |
| 19 | + # pattern of the hostname under which the extension is accessible |
| 20 | + HOST = "typedb.<domain>" |
| 21 | + # name of the Docker image to spin up |
| 22 | + DOCKER_IMAGE = "typedb/typedb" |
| 23 | + # default command args to pass to TypeDB |
| 24 | + DEFAULT_CMD_FLAGS = ["--diagnostics.reporting.metrics=false"] |
| 25 | + # default port for TypeDB HTTP2/gRPC endpoint |
| 26 | + TYPEDB_PORT = 1729 |
| 27 | + |
| 28 | + def __init__(self): |
| 29 | + command_flags = (os.environ.get(ENV_CMD_FLAGS) or "").strip() |
| 30 | + command_flags = self.DEFAULT_CMD_FLAGS + shlex.split(command_flags) |
| 31 | + http2_ports = [self.TYPEDB_PORT] if is_env_not_false(ENV_HTTP2_PROXY) else [] |
| 32 | + super().__init__( |
| 33 | + image_name=self.DOCKER_IMAGE, |
| 34 | + container_ports=[8000, 1729], |
| 35 | + host=self.HOST, |
| 36 | + request_to_port_router=self.request_to_port_router, |
| 37 | + command=command_flags, |
| 38 | + http2_ports=http2_ports, |
| 39 | + ) |
| 40 | + |
| 41 | + def should_proxy_request(self, headers: Headers) -> bool: |
| 42 | + # determine if this is a gRPC request targeting TypeDB |
| 43 | + content_type = headers.get("content-type") or "" |
| 44 | + req_path = headers.get(":path") or "" |
| 45 | + is_typedb_grpc_request = ( |
| 46 | + "grpc" in content_type and "/typedb.protocol.TypeDB" in req_path |
| 47 | + ) |
| 48 | + return is_typedb_grpc_request |
| 49 | + |
| 50 | + def request_to_port_router(self, request: Request) -> int: |
| 51 | + # TODO add REST API / gRPC routing based on request |
| 52 | + return 1729 |
0 commit comments