add initial version of cockroachdb extension

Author: HarshCasper

cockroachdb/.gitignore

@@ -0,0 +1,5 @@
+.venv
+dist
+build
+**/*.egg-info
+.eggs

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

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.

cockroachdb/localstack_cockroachdb/__init__.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"
+            ),
+        }

cockroachdb/localstack_cockroachdb/extension.py

@@ -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"