diff --git a/localstack-typedb/localstack_typedb/extension.py b/localstack-typedb/localstack_typedb/extension.py index 487746a..98aba14 100644 --- a/localstack-typedb/localstack_typedb/extension.py +++ b/localstack-typedb/localstack_typedb/extension.py @@ -1,21 +1,43 @@ +import os +import shlex + +from localstack.utils.docker_utils import DOCKER_CLIENT from localstack_typedb.utils.docker import ProxiedDockerContainerExtension from rolo import Request +# environment variable for user-defined command args to pass to TypeDB +ENV_CMD_FLAGS = "TYPEDB_FLAGS" + class TypeDbExtension(ProxiedDockerContainerExtension): name = "localstack-typedb" HOST = "typedb." + # name of the Docker image to spin up DOCKER_IMAGE = "typedb/typedb" + # default command args to pass to TypeDB + DEFAULT_CMD_FLAGS = ["--diagnostics.reporting.metrics=false"] + # default port for TypeDB HTTP2/gRPC endpoint + TYPEDB_PORT = 1729 def __init__(self): + command_flags = (os.environ.get(ENV_CMD_FLAGS) or "").strip() + command_flags = self.DEFAULT_CMD_FLAGS + shlex.split(command_flags) + command = self._get_image_command() + command_flags super().__init__( image_name=self.DOCKER_IMAGE, container_ports=[8000, 1729], host=self.HOST, request_to_port_router=self.request_to_port_router, + command=command, + http2_ports=[self.TYPEDB_PORT], ) + def _get_image_command(self) -> list[str]: + result = DOCKER_CLIENT.inspect_image(self.DOCKER_IMAGE) + image_command = result["Config"]["Cmd"] + return image_command + def request_to_port_router(self, request: Request): # TODO add REST API / gRPC routing based on request return 1729 diff --git a/localstack-typedb/localstack_typedb/utils/docker.py b/localstack-typedb/localstack_typedb/utils/docker.py index 08e1586..f79e553 100644 --- a/localstack-typedb/localstack_typedb/utils/docker.py +++ b/localstack-typedb/localstack_typedb/utils/docker.py @@ -22,8 +22,6 @@ ) logging.basicConfig() -TYPEDB_PORT = 1729 - class ProxiedDockerContainerExtension(Extension): name: str @@ -41,9 +39,13 @@ class ProxiedDockerContainerExtension(Extension): """ path: str | None """Optional path on which to expose the container endpoints.""" + command: list[str] | None + """Optional command (and flags) to execute in the container.""" request_to_port_router: Callable[[Request], int] | None """Callable that returns the target port for a given request, for routing purposes""" + http2_ports: list[int] | None + """List of ports for which HTTP2 proxy forwarding into the container should be enabled.""" def __init__( self, @@ -52,14 +54,18 @@ def __init__( host: str | None = None, path: str | None = None, container_name: str | None = None, + command: list[str] | None = None, request_to_port_router: Callable[[Request], int] | None = None, + http2_ports: list[int] | None = None, ): self.image_name = image_name self.container_ports = container_ports self.host = host self.path = path self.container_name = container_name + self.command = command self.request_to_port_router = request_to_port_router + self.http2_ports = http2_ports def update_gateway_routes(self, router: http.Router[http.RouteHandler]): if self.path: @@ -72,10 +78,10 @@ def update_gateway_routes(self, router: http.Router[http.RouteHandler]): if self.host: resource = WithHost(self.host, [resource]) router.add(resource) + # apply patches to serve HTTP/2 requests - apply_http2_patches_for_grpc_support( - get_addressable_container_host(), TYPEDB_PORT - ) + for port in self.http2_ports or []: + apply_http2_patches_for_grpc_support(get_addressable_container_host(), port) def on_platform_shutdown(self): self._remove_container() @@ -95,12 +101,18 @@ def start_container(self) -> None: ports = PortMappings() for port in self.container_ports: ports.add(port) + + kwargs = {} + if self.command: + kwargs["command"] = self.command + DOCKER_CLIENT.run_container( self.image_name, detach=True, remove=True, name=container_name, ports=ports, + **kwargs, ) main_port = self.container_ports[0] @@ -118,15 +130,6 @@ def _ping_endpoint(): self._remove_container() raise - # TODO: enable support for TCP port proxying! - # for port in self.container_ports: - # proxy = TCPProxy( - # target_address="localhost", - # target_port=port, - # port=..., - # host="...", - # ) - LOG.debug("Successfully started extension container %s", container_name) def _remove_container(self):