diff --git a/.github/workflows/cockroachdb.yml b/.github/workflows/cockroachdb.yml new file mode 100644 index 00000000..63a6d041 --- /dev/null +++ b/.github/workflows/cockroachdb.yml @@ -0,0 +1,53 @@ +name: LocalStack CockroachDB Extension Tests + +on: + push: + paths: + - cockroachdb/** + branches: + - main + pull_request: + paths: + - .github/workflows/cockroachdb.yml + - cockroachdb/** + workflow_dispatch: + +env: + LOCALSTACK_DISABLE_EVENTS: "1" + LOCALSTACK_AUTH_TOKEN: ${{ secrets.LOCALSTACK_AUTH_TOKEN }} + +jobs: + integration-tests: + name: Run Integration Tests + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup LocalStack and extension + run: | + cd cockroachdb + + docker pull localstack/localstack-pro & + docker pull cockroachdb/cockroach & + pip install localstack + + make install + make lint + make dist + localstack extensions -v install file://$(ls ./dist/localstack_extension_cockroachdb-*.tar.gz) + + DEBUG=1 localstack start -d + localstack wait + + - name: Run integration tests + run: | + cd cockroachdb + make test + + - name: Print logs + if: always() + run: | + localstack logs + localstack stop diff --git a/cockroachdb/.gitignore b/cockroachdb/.gitignore new file mode 100644 index 00000000..1808cca0 --- /dev/null +++ b/cockroachdb/.gitignore @@ -0,0 +1,5 @@ +.venv +dist +build +**/*.egg-info +.eggs diff --git a/cockroachdb/Makefile b/cockroachdb/Makefile new file mode 100644 index 00000000..9c4db1ec --- /dev/null +++ b/cockroachdb/Makefile @@ -0,0 +1,48 @@ +VENV_BIN = python3 -m venv +VENV_DIR ?= .venv +VENV_ACTIVATE = $(VENV_DIR)/bin/activate +VENV_RUN = . $(VENV_ACTIVATE) +TEST_PATH ?= tests + +usage: ## Shows usage for this Makefile + @cat Makefile | grep -E '^[a-zA-Z_-]+:.*?## .*$$' | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-15s\033[0m %s\n", $$1, $$2}' + +venv: $(VENV_ACTIVATE) + +$(VENV_ACTIVATE): pyproject.toml + test -d .venv || $(VENV_BIN) .venv + $(VENV_RUN); pip install --upgrade pip setuptools plux + $(VENV_RUN); pip install -e .[dev] + touch $(VENV_DIR)/bin/activate + +clean: + rm -rf .venv/ + rm -rf build/ + rm -rf .eggs/ + rm -rf *.egg-info/ + +install: venv ## Install dependencies + $(VENV_RUN); python -m plux entrypoints + +dist: venv ## Create distribution + $(VENV_RUN); python -m build + +publish: clean-dist venv dist ## Publish extension to pypi + $(VENV_RUN); pip install --upgrade twine; twine upload dist/* + +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 .; python -m ruff check --fix . + +lint: ## Run ruff to lint the codebase + $(VENV_RUN); python -m ruff check --output-format=full . + +test: ## Run integration tests (requires LocalStack running with the Extension installed) + $(VENV_RUN); pytest $(PYTEST_ARGS) $(TEST_PATH) + +clean-dist: clean + rm -rf dist/ + +.PHONY: clean clean-dist dist install publish usage venv format test diff --git a/cockroachdb/README.md b/cockroachdb/README.md new file mode 100644 index 00000000..1e62167f --- /dev/null +++ b/cockroachdb/README.md @@ -0,0 +1,86 @@ +# CockroachDB on LocalStack + +This repo contains a [LocalStack Extension](https://github.com/localstack/localstack-extensions) that facilitates developing [CockroachDB](https://www.cockroachlabs.com)-based applications locally. + +CockroachDB is a distributed SQL database built for cloud applications. It is PostgreSQL wire-protocol compatible, making it easy to use existing PostgreSQL drivers and tools. + +After installing the extension, a CockroachDB server instance will become available and can be accessed using standard PostgreSQL clients or CockroachDB-specific drivers. + +## Connection Details + +Once the extension is running, you can connect to CockroachDB using any PostgreSQL-compatible client: + +- **Host**: `cockroachdb.localhost.localstack.cloud` +- **Port**: `4566` (LocalStack gateway) +- **Database**: `defaultdb` +- **Username**: `root` +- **Password**: none (insecure mode) + +Example connection using `psql`: +```bash +psql "postgresql://root@cockroachdb.localhost.localstack.cloud:4566/defaultdb?sslmode=disable" +``` + +Example connection using Python with psycopg2: +```python +import psycopg2 + +conn = psycopg2.connect( + host="cockroachdb.localhost.localstack.cloud", + port=4566, + user="root", + database="defaultdb", + sslmode="disable", +) +cursor = conn.cursor() +cursor.execute("SELECT version()") +print(cursor.fetchone()[0]) +conn.close() +``` + +## Configuration + +The following environment variables can be passed to the LocalStack container to configure the extension: + +* `COCKROACHDB_IMAGE`: Docker image to use (default: `cockroachdb/cockroach:latest`) +* `COCKROACHDB_FLAGS`: Extra flags appended to the CockroachDB startup command (default: none) +* `COCKROACHDB_USER`: User for connection string reference (default: `root`) +* `COCKROACHDB_DB`: Database for connection string reference (default: `defaultdb`) + +Example: +```bash +COCKROACHDB_FLAGS="--cache=.25 --max-sql-memory=.25" localstack start +``` + +## Known Limitations + +* **Single-node only** — this extension runs CockroachDB in `start-single-node` mode. Multi-node clusters are not supported. +* **Insecure mode only** — TLS and authentication are disabled. This is intentional for local development. Do not use in production. +* **Ephemeral data** — data is lost when the CockroachDB container stops, matching LocalStack's stateless default behavior. + +## Prerequisites + +* Docker +* LocalStack Pro (free trial available) +* `localstack` CLI +* `make` + +## Install from GitHub repository + +This extension can be installed directly from this Github repo via: + +```bash +localstack extensions install "git+https://github.com/localstack/localstack-extensions.git#egg=localstack-extension-cockroachdb&subdirectory=cockroachdb" +``` + +## Install local development version + +Please refer to the docs [here](https://github.com/localstack/localstack-extensions?tab=readme-ov-file#start-localstack-with-the-extension) for instructions on how to start the extension in developer mode. + +## Change Log + +* `0.1.0`: Initial version of the extension + +## License + +The code in this repo is available under the Apache 2.0 license. diff --git a/cockroachdb/localstack_cockroachdb/__init__.py b/cockroachdb/localstack_cockroachdb/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cockroachdb/localstack_cockroachdb/extension.py b/cockroachdb/localstack_cockroachdb/extension.py new file mode 100644 index 00000000..6bc09012 --- /dev/null +++ b/cockroachdb/localstack_cockroachdb/extension.py @@ -0,0 +1,112 @@ +import os +import shlex +import socket + +from localstack import config +from localstack.extensions.api import http +from localstack_extensions.utils.docker import ProxiedDockerContainerExtension + +# Environment variables for configuration +ENV_IMAGE = "COCKROACHDB_IMAGE" +ENV_FLAGS = "COCKROACHDB_FLAGS" +ENV_USER = "COCKROACHDB_USER" +ENV_DB = "COCKROACHDB_DB" + +# Defaults +DEFAULT_IMAGE = "cockroachdb/cockroach:latest" +DEFAULT_USER = "root" +DEFAULT_DB = "defaultdb" +DEFAULT_PORT = 26257 + + +class CockroachDbExtension(ProxiedDockerContainerExtension): + name = "cockroachdb" + + # Base command args passed to the cockroachdb/cockroach Docker image entrypoint. + # --store=type=mem,size=1GiB: in-memory store — faster startup, truly ephemeral, + # avoids filesystem permission issues inside the container. + # Note: CockroachDB requires at least 640 MiB for in-memory store. + BASE_COMMAND = ["start-single-node", "--insecure", "--store=type=mem,size=1GiB"] + + def __init__(self): + image = os.environ.get(ENV_IMAGE, DEFAULT_IMAGE) + extra_flags = shlex.split((os.environ.get(ENV_FLAGS) or "").strip()) + + # Store for connection info (not passed to container — insecure mode + # auto-creates the root user and defaultdb database) + self.cockroach_user = os.environ.get(ENV_USER, DEFAULT_USER) + self.cockroach_db = os.environ.get(ENV_DB, DEFAULT_DB) + + def _health_check(): + self._check_tcp_port(self.container_host, DEFAULT_PORT) + + super().__init__( + image_name=image, + container_ports=[DEFAULT_PORT], + command=self.BASE_COMMAND + extra_flags, + health_check_fn=_health_check, + health_check_retries=120, # 2 minutes — CockroachDB can be slow on first start + tcp_ports=[DEFAULT_PORT], + ) + + def update_gateway_routes(self, router: http.Router[http.RouteHandler]): + """ + Override to set up only TCP routing without HTTP proxy. + + CockroachDB 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 CockroachDB container, breaking other services. + """ + self.start_container() + + if self.tcp_ports: + self._setup_tcp_protocol_routing() + + def tcp_connection_matcher(self, data: bytes) -> bool: + """ + Identify CockroachDB/PostgreSQL connections by protocol handshake. + + CockroachDB speaks the PostgreSQL wire protocol. Connections start with: + 1. SSL request: protocol code 80877103 (0x04D2162F) + 2. Startup message: protocol version 3.0 (0x00030000) + """ + if len(data) < 8: + return False + + # SSL request (80877103 = 0x04D2162F) + if data[4:8] == b"\x04\xd2\x16\x2f": + return True + + # 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 CockroachDB.""" + gateway_host = "cockroachdb.localhost.localstack.cloud" + gateway_port = config.LOCALSTACK_HOST.port + + return { + "host": gateway_host, + "port": gateway_port, + "user": self.cockroach_user, + "database": self.cockroach_db, + "connection_string": ( + f"cockroachdb+psycopg2://{self.cockroach_user}" + f"@{gateway_host}:{gateway_port}/{self.cockroach_db}" + f"?sslmode=disable" + ), + } diff --git a/cockroachdb/pyproject.toml b/cockroachdb/pyproject.toml new file mode 100644 index 00000000..e4c766a7 --- /dev/null +++ b/cockroachdb/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = ["setuptools", "wheel", "plux>=1.3.1"] +build-backend = "setuptools.build_meta" + +[project] +name = "localstack-extension-cockroachdb" +version = "0.1.0" +description = "LocalStack Extension: CockroachDB on LocalStack" +readme = {file = "README.md", content-type = "text/markdown; charset=UTF-8"} +requires-python = ">=3.10" +authors = [ + { name = "LocalStack team"} +] +keywords = ["LocalStack", "CockroachDB", "PostgreSQL", "SQL", "Distributed"] +classifiers = [] +dependencies = [ + "localstack-extensions-utils" +] + +[project.urls] +Homepage = "https://github.com/localstack/localstack-extensions" + +[project.optional-dependencies] +dev = [ + "boto3", + "build", + "jsonpatch", + "psycopg2-binary", + "pytest", + "rolo", + "ruff", +] + +[project.entry-points."localstack.extensions"] +localstack_cockroachdb = "localstack_cockroachdb.extension:CockroachDbExtension" diff --git a/cockroachdb/tests/__init__.py b/cockroachdb/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cockroachdb/tests/test_extension.py b/cockroachdb/tests/test_extension.py new file mode 100644 index 00000000..b545a919 --- /dev/null +++ b/cockroachdb/tests/test_extension.py @@ -0,0 +1,127 @@ +import uuid + +import boto3 +import psycopg2 +import psycopg2.extras + + +def short_uid() -> str: + return str(uuid.uuid4())[:8] + + +# Connection details for CockroachDB +# Connect through LocalStack gateway with TCP proxying +HOST = "cockroachdb.localhost.localstack.cloud" +PORT = 4566 +USER = "root" +DATABASE = "defaultdb" + + +def get_connection(): + """Create a psycopg2 connection to CockroachDB.""" + return psycopg2.connect( + host=HOST, + port=PORT, + user=USER, + database=DATABASE, + sslmode="disable", + ) + + +def test_connect_to_cockroachdb(): + """Test basic connection to CockroachDB and verify it's actually CockroachDB.""" + conn = get_connection() + try: + cursor = conn.cursor() + cursor.execute("SELECT version()") + version = cursor.fetchone()[0] + assert "CockroachDB" in version, f"Expected CockroachDB version, got: {version}" + cursor.close() + finally: + conn.close() + + +def test_cockroachdb_crud(): + """Test basic CRUD operations: CREATE TABLE, INSERT, SELECT, DROP TABLE.""" + conn = get_connection() + table = f"test_items_{short_uid()}" + try: + cursor = conn.cursor() + cursor.execute( + f"CREATE TABLE {table} (id INT PRIMARY KEY, name STRING NOT NULL)" + ) + cursor.execute( + f"INSERT INTO {table} (id, name) VALUES (1, 'hello'), (2, 'world')" + ) + cursor.execute(f"SELECT id, name FROM {table} ORDER BY id") + rows = cursor.fetchall() + + assert len(rows) == 2 + assert rows[0][0] == 1 + assert rows[0][1] == "hello" + assert rows[1][0] == 2 + assert rows[1][1] == "world" + + cursor.execute(f"DROP TABLE IF EXISTS {table}") + conn.commit() + cursor.close() + except Exception: + conn.rollback() + raise + finally: + conn.close() + + +def test_mixed_tcp_and_http_traffic(): + """ + Test that mixed TCP (CockroachDB) and HTTP (AWS) traffic works correctly. + + Verifies that the CockroachDB extension only intercepts PostgreSQL wire + protocol connections and doesn't interfere with regular HTTP-based AWS + API requests to LocalStack. + """ + # Verify CockroachDB TCP connection works + conn = get_connection() + try: + cursor = conn.cursor() + cursor.execute("SELECT 1 AS test_value") + assert cursor.fetchone()[0] == 1, "CockroachDB TCP connection should work" + cursor.close() + finally: + conn.close() + + # Verify AWS HTTP requests still work (S3) + endpoint_url = f"http://localhost:{PORT}" + + 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) + + bucket_names = [b["Name"] for b in s3_client.list_buckets()["Buckets"]] + assert bucket_name in bucket_names, "S3 HTTP API should work alongside CockroachDB TCP" + + 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) + assert response["Body"].read() == test_content, "S3 object operations should work" + + s3_client.delete_object(Bucket=bucket_name, Key=test_key) + s3_client.delete_bucket(Bucket=bucket_name) + + # Verify CockroachDB still works after HTTP requests + conn = get_connection() + try: + cursor = conn.cursor() + cursor.execute("SELECT 'tcp_works_after_http' AS verification") + assert cursor.fetchone()[0] == "tcp_works_after_http" + cursor.close() + finally: + conn.close()