rewrite ParadeDbExtension to use ProxiedDockerContainerExtension as base

Author: whummer

paradedb/localstack_paradedb/extension.py

@@ -1,7 +1,13 @@
 import os
+import socket
 import logging
 
-from localstack_paradedb.utils.docker import DatabaseDockerContainerExtension
+from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
+from localstack.extensions.api import http
+from localstack.utils.docker_utils import DOCKER_CLIENT
+from localstack.utils.container_utils.container_client import PortMappings
+from localstack.utils.sync import retry
+from werkzeug.datastructures import Headers
 
 LOG = logging.getLogger(__name__)
 
@@ -18,7 +24,7 @@
 DEFAULT_POSTGRES_PORT = 5432
 
 
-class ParadeDbExtension(DatabaseDockerContainerExtension):
+class ParadeDbExtension(ProxiedDockerContainerExtension):
     name = "paradedb"
 
     # Name of the Docker image to spin up
@@ -34,37 +40,97 @@ def __init__(self):
         postgres_port = int(os.environ.get(ENV_POSTGRES_PORT, DEFAULT_POSTGRES_PORT))
 
         # Environment variables to pass to the container
-        env_vars = {
+        self.env_vars = {
             "POSTGRES_USER": postgres_user,
             "POSTGRES_PASSWORD": postgres_password,
             "POSTGRES_DB": postgres_db,
         }
 
-        super().__init__(
-            image_name=self.DOCKER_IMAGE,
-            container_ports=[postgres_port],
-            env_vars=env_vars,
-        )
-
         # Store configuration for connection info
         self.postgres_user = postgres_user
         self.postgres_password = postgres_password
         self.postgres_db = postgres_db
         self.postgres_port = postgres_port
 
+        super().__init__(
+            image_name=self.DOCKER_IMAGE,
+            container_ports=[postgres_port],
+        )
+
+    def should_proxy_request(self, headers: Headers) -> bool:
+        """
+        Define whether a request should be proxied based on request headers.
+        For database extensions, this is not used as connections are direct TCP.
+        """
+        return False
+
+    def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
+        """
+        Override to start container without setting up HTTP gateway routes.
+        Database extensions don't need HTTP routing - clients connect directly via TCP.
+        """
+        self.start_container()
+
+    def start_container(self) -> None:
+        """Override to add env_vars support and database-specific health checking."""
+        LOG.debug("Starting extension container %s", self.container_name)
+
+        port_mapping = PortMappings()
+        for port in self.container_ports:
+            port_mapping.add(port)
+
+        try:
+            DOCKER_CLIENT.run_container(
+                self.image_name,
+                detach=True,
+                remove=True,
+                name=self.container_name,
+                ports=port_mapping,
+                env_vars=self.env_vars,
+            )
+        except Exception as e:
+            LOG.debug("Failed to start container %s: %s", self.container_name, e)
+            raise
+
+        def _check_health():
+            """Check if PostgreSQL port is accepting connections."""
+            self._check_tcp_port(self.container_host, self.postgres_port)
+
+        try:
+            retry(_check_health, retries=60, sleep=1)
+        except Exception as e:
+            LOG.info("Failed to connect to container %s: %s", self.container_name, e)
+            self._remove_container()
+            raise
+
+        LOG.info(
+            "Successfully started extension container %s on %s:%s",
+            self.container_name,
+            self.container_host,
+            self.postgres_port,
+        )
+
+    def _check_tcp_port(self, host: str, port: int, timeout: float = 2.0) -> None:
+        """Check if a TCP port is accepting connections."""
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(timeout)
+        try:
+            sock.connect((host, port))
+            sock.close()
+        except (socket.timeout, socket.error) as e:
+            raise AssertionError(f"Port {port} not ready: {e}")
+
     def get_connection_info(self) -> dict:
         """Return connection information for ParadeDB."""
-        info = super().get_connection_info()
-        info.update(
-            {
-                "database": self.postgres_db,
-                "user": self.postgres_user,
-                "password": self.postgres_password,
-                "port": self.postgres_port,
-                "connection_string": (
-                    f"postgresql://{self.postgres_user}:{self.postgres_password}"
-                    f"@{self.container_host}:{self.postgres_port}/{self.postgres_db}"
-                ),
-            }
-        )
-        return info
+        return {
+            "host": self.container_host,
+            "database": self.postgres_db,
+            "user": self.postgres_user,
+            "password": self.postgres_password,
+            "port": self.postgres_port,
+            "ports": {self.postgres_port: self.postgres_port},
+            "connection_string": (
+                f"postgresql://{self.postgres_user}:{self.postgres_password}"
+                f"@{self.container_host}:{self.postgres_port}/{self.postgres_db}"
+            ),
+        }

paradedb/localstack_paradedb/utils/__init__.py

@@ -13,7 +13,11 @@ authors = [
 ]
 keywords = ["LocalStack", "ParadeDB", "PostgreSQL", "Search", "Analytics"]
 classifiers = []
-dependencies = []
+dependencies = [
+    # TODO remove / replace prior to merge!
+#    "localstack-extensions-utils",
+    "localstack-extensions-utils @ git+https://github.com/localstack/localstack-extensions.git@extract-utils-package#subdirectory=utils"
+]
 
 [project.urls]
 Homepage = "https://github.com/localstack/localstack-extensions"