extract common extension utils into a separate package (#123)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Steve Purcell <steve@sanityinc.com>

Author: 3 people

.github/workflows/miniflare.yml

@@ -33,9 +33,6 @@ jobs:
           docker pull localstack/localstack-pro &
           pip install localstack localstack-ext
 
-          # TODO remove
-          mkdir ~/.localstack; echo '{"token":"test"}' > ~/.localstack/auth.json
-
           branchName=${GITHUB_HEAD_REF##*/}
           if [ "$branchName" = "" ]; then branchName=main; fi
           echo "Installing from branch name $branchName"

.github/workflows/utils.yml

@@ -0,0 +1,66 @@
+name: LocalStack Extensions Utils Tests
+
+on:
+  push:
+    paths:
+      - utils/**
+    branches:
+      - main
+  pull_request:
+    paths:
+      - .github/workflows/utils.yml
+      - utils/**
+  workflow_dispatch:
+
+jobs:
+  unit-tests:
+    name: Run Unit Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install dependencies
+        run: |
+          cd utils
+          pip install -e .[dev,test]
+
+      - name: Lint
+        run: |
+          cd utils
+          make lint
+
+      - name: Run unit tests
+        run: |
+          cd utils
+          make test-unit
+
+  integration-tests:
+    name: Run Integration Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install dependencies
+        run: |
+          cd utils
+          pip install -e .[dev,test]
+
+      - name: Run integration tests
+        run: |
+          docker pull moul/grpcbin &
+          cd utils
+          make test-integration

.github/workflows/wiremock.yml

@@ -39,6 +39,7 @@ jobs:
           pip install localstack terraform-local awscli-local[ver1]
 
           make install
+          make lint
           make dist
           localstack extensions -v install file://$(ls ./dist/localstack_wiremock-*.tar.gz)
 

flake.nix

@@ -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/Makefile

@@ -1,15 +1,13 @@
 import os
-import logging
+import socket
 
-from localstack_paradedb.utils.docker import DatabaseDockerContainerExtension
-
-LOG = logging.getLogger(__name__)
+from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
+from localstack import config
 
 # Environment variables for configuration
 ENV_POSTGRES_USER = "PARADEDB_POSTGRES_USER"
 ENV_POSTGRES_PASSWORD = "PARADEDB_POSTGRES_PASSWORD"
 ENV_POSTGRES_DB = "PARADEDB_POSTGRES_DB"
-ENV_POSTGRES_PORT = "PARADEDB_POSTGRES_PORT"
 
 # Default values
 DEFAULT_POSTGRES_USER = "myuser"
@@ -18,7 +16,7 @@
 DEFAULT_POSTGRES_PORT = 5432
 
 
-class ParadeDbExtension(DatabaseDockerContainerExtension):
+class ParadeDbExtension(ProxiedDockerContainerExtension):
     name = "paradedb"
 
     # Name of the Docker image to spin up
@@ -31,7 +29,13 @@ def __init__(self):
             ENV_POSTGRES_PASSWORD, DEFAULT_POSTGRES_PASSWORD
         )
         postgres_db = os.environ.get(ENV_POSTGRES_DB, DEFAULT_POSTGRES_DB)
-        postgres_port = int(os.environ.get(ENV_POSTGRES_PORT, DEFAULT_POSTGRES_PORT))
+        postgres_port = DEFAULT_POSTGRES_PORT
+
+        # 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
 
         # Environment variables to pass to the container
         env_vars = {
@@ -40,31 +44,70 @@ def __init__(self):
             "POSTGRES_DB": postgres_db,
         }
 
+        def _tcp_health_check():
+            """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
         )
 
-        # 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
+    def tcp_connection_matcher(self, data: bytes) -> bool:
+        """
+        Identify PostgreSQL/ParadeDB connections by protocol handshake.
+
+        PostgreSQL can start with either:
+        1. SSL request: protocol code 80877103 (0x04D2162F)
+        2. Startup message: protocol version 3.0 (0x00030000)
+
+        Both use the same format:
+        - 4 bytes: message length
+        - 4 bytes: protocol version/code
+        """
+        if len(data) < 8:
+            return False
+
+        # Check for SSL request (80877103 = 0x04D2162F)
+        if data[4:8] == b"\x04\xd2\x16\x2f":
+            return True
+
+        # Check for protocol version 3.0 (0x00030000)
+        if data[4:8] == b"\x00\x03\x00\x00":
+            return True
+
+        return False
+
+    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
+        # Clients should connect through the LocalStack gateway
+        gateway_host = "paradedb.localhost.localstack.cloud"
+        gateway_port = config.LOCALSTACK_HOST.port
+
+        return {
+            "host": gateway_host,
+            "database": self.postgres_db,
+            "user": self.postgres_user,
+            "password": self.postgres_password,
+            "port": gateway_port,
+            "connection_string": (
+                f"postgresql://{self.postgres_user}:{self.postgres_password}"
+                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,
+        }