add TCP proxy logic for ParadeDB

Author: whummer

paradedb/Makefile

@@ -34,7 +34,7 @@ entrypoints: venv   ## Generate plugin entrypoints for Python package
 	$(VENV_RUN); python -m plux entrypoints
 
 format:		    ## Run ruff to format the codebase
-	$(VENV_RUN); python -m ruff format .; make lint
+	$(VENV_RUN); python -m ruff format .; python -m ruff check --fix .
 
 lint:		    ## Run ruff to lint the codebase
 	$(VENV_RUN); python -m ruff check --output-format=full .

paradedb/localstack_paradedb/extension.py

@@ -3,8 +3,8 @@
 import logging
 
 from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
-from localstack.extensions.api import http
 from werkzeug.datastructures import Headers
+from localstack import config
 
 LOG = logging.getLogger(__name__)
 
@@ -50,29 +50,33 @@ def __init__(self):
         }
 
         def _tcp_health_check():
-            """Check if PostgreSQL port is accepting connections."""
+            """Check if ParadeDB port is accepting connections."""
             self._check_tcp_port(self.container_host, self.postgres_port)
 
         super().__init__(
             image_name=self.DOCKER_IMAGE,
             container_ports=[postgres_port],
             env_vars=env_vars,
             health_check_fn=_tcp_health_check,
+            tcp_ports=[postgres_port],  # Enable TCP proxying through gateway
         )
 
-    def should_proxy_request(self, headers: Headers) -> bool:
+    def tcp_connection_matcher(self, data: bytes) -> bool:
         """
-        Define whether a request should be proxied based on request headers.
-        For database extensions, this is not used as connections are direct TCP.
+        Identify PostgreSQL/ParadeDB connections by protocol handshake.
+
+        PostgreSQL startup message format:
+        - 4 bytes: message length
+        - 4 bytes: protocol version (3.0 = 0x00030000)
         """
-        return False
+        return len(data) >= 8 and data[4:8] == b"\x00\x03\x00\x00"
 
-    def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
+    def should_proxy_request(self, headers: Headers) -> bool:
         """
-        Override to start container without setting up HTTP gateway routes.
-        Database extensions don't need HTTP routing - clients connect directly via TCP.
+        Define whether a request should be proxied based on request headers.
+        Not used for TCP connections - see tcp_connection_matcher instead.
         """
-        self.start_container()
+        return False
 
     def _check_tcp_port(self, host: str, port: int, timeout: float = 2.0) -> None:
         """Check if a TCP port is accepting connections."""
@@ -86,15 +90,21 @@ def _check_tcp_port(self, host: str, port: int, timeout: float = 2.0) -> None:
 
     def get_connection_info(self) -> dict:
         """Return connection information for ParadeDB."""
+        # Clients should connect through the LocalStack gateway
+        gateway_host = "paradedb.localhost.localstack.cloud"
+        gateway_port = config.LOCALSTACK_HOST.port
+
         return {
-            "host": self.container_host,
+            "host": gateway_host,
             "database": self.postgres_db,
             "user": self.postgres_user,
             "password": self.postgres_password,
-            "port": self.postgres_port,
-            "ports": {self.postgres_port: self.postgres_port},
+            "port": gateway_port,
             "connection_string": (
                 f"postgresql://{self.postgres_user}:{self.postgres_password}"
-                f"@{self.container_host}:{self.postgres_port}/{self.postgres_db}"
+                f"@{gateway_host}:{gateway_port}/{self.postgres_db}"
             ),
+            # Also include container connection details for debugging
+            "container_host": self.container_host,
+            "container_port": self.postgres_port,
         }

paradedb/tests/test_extension.py

@@ -3,8 +3,9 @@
 
 
 # Connection details for ParadeDB
-HOST = "localhost"
-PORT = 5432
+# Connect through LocalStack gateway with TCP proxying
+HOST = "paradedb.localhost.localstack.cloud"
+PORT = 4566
 USER = "myuser"
 PASSWORD = "mypassword"
 DATABASE = "mydatabase"

utils/localstack_extensions/utils/docker.py

@@ -19,8 +19,6 @@
 from rolo.proxy import Proxy
 from rolo.routing import RuleAdapter, WithHost
 from werkzeug.datastructures import Headers
-from twisted.internet import reactor
-from twisted.protocols.portforward import ProxyFactory
 
 LOG = logging.getLogger(__name__)
 

utils/localstack_extensions/utils/tcp_protocol_detector.py

@@ -35,9 +35,7 @@ def matcher(data: bytes) -> bool:
     return matcher
 
 
-def create_signature_matcher(
-    signature: bytes, offset: int = 0
-) -> ConnectionMatcher:
+def create_signature_matcher(signature: bytes, offset: int = 0) -> ConnectionMatcher:
     """
     Create a matcher that matches bytes at a specific offset.
 

utils/localstack_extensions/utils/tcp_protocol_router.py

@@ -7,7 +7,7 @@
 """
 
 import logging
-from twisted.internet import protocol, reactor
+from twisted.internet import reactor
 from twisted.protocols.portforward import ProxyClient, ProxyClientFactory
 from twisted.web.http import HTTPChannel
 
@@ -21,99 +21,6 @@
 _gateway_patched = False
 
 
-class ProtocolDetectingChannel(HTTPChannel):
-    """
-    HTTP channel wrapper that detects TCP protocols before HTTP processing.
-
-    This wraps the standard HTTPChannel and intercepts the first dataReceived call
-    to check if it's a known TCP protocol. If so, it morphs into a TCP proxy.
-    Otherwise, it passes through to normal HTTP handling.
-    """
-
-    def __init__(self):
-        super().__init__()
-        self._detection_buffer = []
-        self._detecting = True
-        self._tcp_peer = None
-        self._detection_buffer_size = 512
-
-    def dataReceived(self, data):
-        """Intercept data to detect protocol before HTTP processing."""
-        if not self._detecting:
-            # Already decided - either proxying TCP or processing HTTP
-            if self._tcp_peer:
-                # TCP proxying mode
-                self._tcp_peer.transport.write(data)
-            else:
-                # HTTP mode - pass to parent
-                super().dataReceived(data)
-            return
-
-        # Still detecting - buffer data
-        self._detection_buffer.append(data)
-        buffered_data = b"".join(self._detection_buffer)
-
-        # Try detection once we have enough bytes
-        if len(buffered_data) >= 8:
-            protocol_name = detect_protocol(buffered_data)
-
-            if protocol_name and protocol_name not in ("http", "http2"):
-                # Known TCP protocol (not HTTP) - check if we have a backend
-                backend_info = _protocol_backends.get(protocol_name)
-
-                if backend_info:
-                    LOG.info(
-                        f"Detected {protocol_name} on gateway port, routing to "
-                        f"{backend_info['host']}:{backend_info['port']}"
-                    )
-                    self._switch_to_tcp_proxy(
-                        backend_info["host"], backend_info["port"], buffered_data
-                    )
-                    self._detecting = False
-                    return
-
-            # Not a known TCP protocol, or no backend configured
-            # Check if we've buffered enough
-            if (
-                len(buffered_data) >= self._detection_buffer_size
-                or protocol_name in ("http", "http2")
-            ):
-                LOG.debug(
-                    f"Protocol detected as {protocol_name or 'unknown'}, using HTTP handler"
-                )
-                self._detecting = False
-                # Feed buffered data to HTTP handler
-                for chunk in self._detection_buffer:
-                    super().dataReceived(chunk)
-                self._detection_buffer = []
-
-    def _switch_to_tcp_proxy(self, host, port, initial_data):
-        """Switch this connection to TCP proxy mode."""
-        # Pause reading while we establish backend connection
-        self.transport.pauseProducing()
-
-        # Create backend connection
-        client_factory = ProxyClientFactory()
-        client_factory.server = self
-        client_factory.initial_data = initial_data
-
-        # Connect to backend
-        reactor.connectTCP(host, port, client_factory)
-
-    def set_tcp_peer(self, peer):
-        """Called when backend connection is established."""
-        self._tcp_peer = peer
-        # Resume reading from client
-        self.transport.resumeProducing()
-
-    def connectionLost(self, reason):
-        """Handle connection close."""
-        if self._tcp_peer:
-            self._tcp_peer.transport.loseConnection()
-            self._tcp_peer = None
-        super().connectionLost(reason)
-
-
 class TcpProxyClient(ProxyClient):
     """Backend TCP connection for protocol-detected connections."""
 
@@ -261,7 +168,9 @@ def register_tcp_extension(
         backend_port: Backend port to route to
     """
     _tcp_extensions.append((extension_name, matcher, backend_host, backend_port))
-    LOG.info(f"Registered TCP extension {extension_name} -> {backend_host}:{backend_port}")
+    LOG.info(
+        f"Registered TCP extension {extension_name} -> {backend_host}:{backend_port}"
+    )
 
 
 def unregister_tcp_extension(extension_name: str):

utils/tests/unit/test_tcp_protocol_detector.py

@@ -2,13 +2,14 @@
 Unit tests for TCP connection matcher helpers.
 """
 
-import pytest
 from localstack_extensions.utils.tcp_protocol_detector import (
     create_prefix_matcher,
     create_signature_matcher,
     create_custom_matcher,
     combine_matchers,
 )
+from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
+from werkzeug.datastructures import Headers
 
 
 class TestMatcherFactories:
@@ -26,13 +27,13 @@ def test_create_prefix_matcher(self):
     def test_create_signature_matcher(self):
         """Test creating a signature matcher with offset."""
         # Match signature at offset 4
-        matcher = create_signature_matcher(b"\xAA\xBB", offset=4)
+        matcher = create_signature_matcher(b"\xaa\xbb", offset=4)
 
-        assert matcher(b"\x00\x00\x00\x00\xAA\xBB\xCC")
-        assert matcher(b"\x00\x00\x00\x00\xAA\xBB")
-        assert not matcher(b"\xAA\xBB\xCC")  # Wrong offset
-        assert not matcher(b"\x00\x00\x00\x00\xCC\xDD")  # Wrong signature
-        assert not matcher(b"\x00\x00\x00\x00\xAA")  # Incomplete
+        assert matcher(b"\x00\x00\x00\x00\xaa\xbb\xcc")
+        assert matcher(b"\x00\x00\x00\x00\xaa\xbb")
+        assert not matcher(b"\xaa\xbb\xcc")  # Wrong offset
+        assert not matcher(b"\x00\x00\x00\x00\xcc\xdd")  # Wrong signature
+        assert not matcher(b"\x00\x00\x00\x00\xaa")  # Incomplete
 
     def test_create_custom_matcher(self):
         """Test creating a custom matcher."""
@@ -42,8 +43,8 @@ def my_check(data):
 
         matcher = create_custom_matcher(my_check)
 
-        assert matcher(b"\x00\x00\x00\x00\x00\xFF")
-        assert matcher(b"\x00\x00\x00\x00\x00\xFF\xFF")
+        assert matcher(b"\x00\x00\x00\x00\x00\xff")
+        assert matcher(b"\x00\x00\x00\x00\x00\xff\xff")
         assert not matcher(b"\x00\x00\x00\x00\x00\x00")
         assert not matcher(b"\x00\x00\x00\x00\x00")  # Too short
 
@@ -88,16 +89,14 @@ def test_matcher_with_extra_data(self):
         matcher = create_prefix_matcher(b"PREFIX")
 
         # Should match even with lots of extra data
-        assert matcher(b"PREFIX" + b"\xFF" * 1000)
+        assert matcher(b"PREFIX" + b"\xff" * 1000)
 
 
 class TestRealWorldUsage:
     """Tests for real-world usage patterns."""
 
     def test_extension_with_custom_protocol_matcher(self):
         """Test using custom matchers in an extension context."""
-        from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
-        from werkzeug.datastructures import Headers
 
         class CustomProtocolExtension(ProxiedDockerContainerExtension):
             name = "custom"
@@ -111,7 +110,7 @@ def __init__(self):
 
             def tcp_connection_matcher(self, data: bytes) -> bool:
                 # Match custom protocol with magic bytes at offset 4
-                matcher = create_signature_matcher(b"\xDE\xAD\xBE\xEF", offset=4)
+                matcher = create_signature_matcher(b"\xde\xad\xbe\xef", offset=4)
                 return matcher(data)
 
             def should_proxy_request(self, headers: Headers) -> bool:
@@ -121,16 +120,14 @@ def should_proxy_request(self, headers: Headers) -> bool:
         assert hasattr(extension, "tcp_connection_matcher")
 
         # Test the matcher
-        valid_data = b"\x00\x00\x00\x00\xDE\xAD\xBE\xEF\xFF"
+        valid_data = b"\x00\x00\x00\x00\xde\xad\xbe\xef\xff"
         assert extension.tcp_connection_matcher(valid_data)
 
-        invalid_data = b"\x00\x00\x00\x00\xFF\xFF\xFF\xFF"
+        invalid_data = b"\x00\x00\x00\x00\xff\xff\xff\xff"
         assert not extension.tcp_connection_matcher(invalid_data)
 
     def test_extension_with_combined_matchers(self):
         """Test using combined matchers in an extension."""
-        from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
-        from werkzeug.datastructures import Headers
 
         class MultiProtocolExtension(ProxiedDockerContainerExtension):
             name = "multi-protocol"
@@ -160,8 +157,6 @@ def should_proxy_request(self, headers: Headers) -> bool:
 
     def test_extension_with_inline_matcher(self):
         """Test using an inline matcher function."""
-        from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
-        from werkzeug.datastructures import Headers
 
         class InlineMatcherExtension(ProxiedDockerContainerExtension):
             name = "inline"
@@ -175,11 +170,7 @@ def __init__(self):
 
             def tcp_connection_matcher(self, data: bytes) -> bool:
                 # Inline custom logic without helper functions
-                return (
-                    len(data) >= 8
-                    and data.startswith(b"MAGIC")
-                    and data[7] == 0x42
-                )
+                return len(data) >= 8 and data.startswith(b"MAGIC") and data[7] == 0x42
 
             def should_proxy_request(self, headers: Headers) -> bool:
                 return False