Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions paradedb/localstack_paradedb/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

from localstack_extensions.utils.docker import ProxiedDockerContainerExtension
from localstack import config
from localstack.extensions.api import http

# Environment variables for configuration
ENV_POSTGRES_USER = "PARADEDB_POSTGRES_USER"
Expand Down Expand Up @@ -56,6 +57,23 @@ def _tcp_health_check():
tcp_ports=[postgres_port], # Enable TCP proxying through gateway
)

# TODO: this should be migrated into the base class directly ..!
def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
"""
Override to set up only TCP routing without HTTP proxy.

ParadeDB uses the native PostgreSQL wire protocol (not HTTP), so we
only need TCP protocol routing - not HTTP proxying. Adding an HTTP
proxy without a host restriction would cause all HTTP requests to be
forwarded to the PostgreSQL container, breaking other services.
"""
# Start the container
self.start_container()

# Set up only TCP protocol routing (skip HTTP proxy from base class)
if self.tcp_ports:
self._setup_tcp_protocol_routing()

def tcp_connection_matcher(self, data: bytes) -> bool:
"""
Identify PostgreSQL/ParadeDB connections by protocol handshake.
Expand Down
78 changes: 78 additions & 0 deletions paradedb/tests/test_extension.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import boto3
import psycopg2
from localstack.utils.strings import short_uid

Expand Down Expand Up @@ -104,3 +105,80 @@ def test_paradedb_quickstart():
conn.commit()
cursor.close()
conn.close()


def test_mixed_tcp_and_http_traffic():
"""
Test that mixed TCP (ParadeDB) and HTTP (AWS) traffic works correctly.

This verifies that the ParadeDB extension only intercepts PostgreSQL wire
protocol connections and doesn't interfere with regular HTTP-based AWS
API requests to LocalStack.
"""
# First, verify ParadeDB TCP connection works
conn = get_connection()
cursor = conn.cursor()
cursor.execute("SELECT 1 as test_value;")
result = cursor.fetchone()
assert result[0] == 1, "ParadeDB TCP connection should work"
cursor.close()
conn.close()

# Now verify AWS HTTP requests still work (S3 and STS)
# These should NOT be intercepted by the ParadeDB extension
endpoint_url = f"http://localhost:{PORT}"

# Test S3 HTTP requests
s3_client = boto3.client(
"s3",
endpoint_url=endpoint_url,
aws_access_key_id="test",
aws_secret_access_key="test",
region_name="us-east-1",
)

bucket_name = f"test-bucket-{short_uid()}"
s3_client.create_bucket(Bucket=bucket_name)

# List buckets to verify HTTP API is working
buckets = s3_client.list_buckets()
bucket_names = [b["Name"] for b in buckets["Buckets"]]
assert bucket_name in bucket_names, "S3 HTTP API should work alongside ParadeDB TCP"

# Put and get an object
test_key = "test-object.txt"
test_content = b"Hello from mixed TCP/HTTP test!"
s3_client.put_object(Bucket=bucket_name, Key=test_key, Body=test_content)
response = s3_client.get_object(Bucket=bucket_name, Key=test_key)
retrieved_content = response["Body"].read()
assert retrieved_content == test_content, "S3 object operations should work"

# Clean up S3
s3_client.delete_object(Bucket=bucket_name, Key=test_key)
s3_client.delete_bucket(Bucket=bucket_name)

# Test STS HTTP requests
sts_client = boto3.client(
"sts",
endpoint_url=endpoint_url,
aws_access_key_id="test",
aws_secret_access_key="test",
region_name="us-east-1",
)

caller_identity = sts_client.get_caller_identity()
assert "Account" in caller_identity, (
"STS HTTP API should work alongside ParadeDB TCP"
)
assert "Arn" in caller_identity, "STS should return valid caller identity"

# Finally, verify ParadeDB still works after HTTP requests
conn = get_connection()
cursor = conn.cursor()
cursor.execute("SELECT 'tcp_works_after_http' as verification;")
result = cursor.fetchone()
assert result[0] == "tcp_works_after_http", (
"ParadeDB should still work after HTTP requests"
)
cursor.close()
conn.close()