fix: Add Cockroach DB Module to Testcontainers

modules/cockroachdb/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.cockroachdb.CockroachDBContainer
+.. title:: testcontainers.cockroachdb.CockroachDBContainer

modules/cockroachdb/testcontainers/cockroachdb/__init__.py

@@ -0,0 +1,100 @@
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+from os import environ
+from typing import Optional
+from urllib.error import HTTPError, URLError
+from urllib.request import urlopen
+
+from testcontainers.core.generic import DbContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+
+
+class CockroachDBContainer(DbContainer):
+    """
+    CockroachDB database container.
+
+    Example:
+
+        The example will spin up a CockroachDB database to which you can connect with the credentials
+        passed in the constructor. Alternatively, you may use the :code:`get_connection_url()`
+        method which returns a sqlalchemy-compatible url in format
+        :code:`dialect+driver://username:password@host:port/database`.
+
+        .. doctest::
+
+            >>> import sqlalchemy
+            >>> from testcontainers.cockroachdb import CockroachDBContainer
+
+            >>> with CockroachDBContainer('cockroachdb/cockroach:v24.1.1') as crdb:
+            ...     engine = sqlalchemy.create_engine(crdb.get_connection_url())
+            ...     with engine.begin() as connection:
+            ...         result = connection.execute(sqlalchemy.text("select version()"))
+            ...         version, = result.fetchone()
+
+    """
+
+    COCKROACH_DB_PORT: int = 26257
+    COCKROACH_API_PORT: int = 8080
+
+    def __init__(
+        self,
+        image: str = "cockroachdb/cockroach:v24.1.1",
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        dbname: Optional[str] = None,
+        dialect="cockroachdb+psycopg2",
+        **kwargs,
+    ) -> None:
+        super().__init__(image, **kwargs)
+
+        self.with_exposed_ports(self.COCKROACH_DB_PORT, self.COCKROACH_API_PORT)
+        self.username = username or environ.get("COCKROACH_USER", "cockroach")
+        self.password = password or environ.get("COCKROACH_PASSWORD", "arthropod")
+        self.dbname = dbname or environ.get("COCKROACH_DATABASE", "roach")
+        self.dialect = dialect
+
+    def _configure(self) -> None:
+        self.with_env("COCKROACH_DATABASE", self.dbname)
+        self.with_env("COCKROACH_USER", self.username)
+        self.with_env("COCKROACH_PASSWORD", self.password)
+
+        cmd = "start-single-node"
+        if not self.password:
+            cmd += " --insecure"
+        self.with_command(cmd)
+
+    @wait_container_is_ready(HTTPError, URLError)
+    def _connect(self) -> None:
+        host = self.get_container_host_ip()
+        url = f"http://{host}:{self.get_exposed_port(self.COCKROACH_API_PORT)}/health"
+        self._wait_for_health(url)
+        wait_for_logs(self, "finished creating default user*")
+
+    @staticmethod
+    def _wait_for_health(url):
+        with urlopen(url) as response:
+            response.read()
+
+    def get_connection_url(self) -> str:
+        conn_str = super()._create_connection_url(
+            dialect=self.dialect,
+            username=self.username,
+            password=self.password,
+            dbname=self.dbname,
+            port=self.COCKROACH_DB_PORT,
+        )
+
+        if self.password:
+            conn_str += "?sslmode=require"
+
+        return conn_str

modules/cockroachdb/tests/test_cockroachdb.py

@@ -0,0 +1,14 @@
+import sqlalchemy
+
+from testcontainers.cockroachdb import CockroachDBContainer
+
+
+def test_docker_run_mysql():
+    config = CockroachDBContainer("cockroachdb/cockroach:v24.1.1")
+    with config as crdb:
+        engine = sqlalchemy.create_engine(crdb.get_connection_url())
+        with engine.begin() as connection:
+            result = connection.execute(sqlalchemy.text("select version()"))
+            for row in result:
+                assert "CockroachDB" in row[0]
+                assert "v24.1.1" in row[0]

pyproject.toml

@@ -33,6 +33,7 @@ packages = [
     { include = "testcontainers", from = "modules/cassandra" },
     { include = "testcontainers", from = "modules/chroma" },
     { include = "testcontainers", from = "modules/clickhouse" },
+    { include = "testcontainers", from = "modules/cockroachdb" },
     { include = "testcontainers", from = "modules/elasticsearch" },
     { include = "testcontainers", from = "modules/google" },
     { include = "testcontainers", from = "modules/influxdb" },
@@ -107,6 +108,7 @@ arangodb = ["python-arango"]
 azurite = ["azure-storage-blob"]
 cassandra = []
 clickhouse = ["clickhouse-driver"]
+cockroachdb = []
 elasticsearch = []
 google = ["google-cloud-pubsub", "google-cloud-datastore"]
 influxdb = ["influxdb", "influxdb-client"]
@@ -157,6 +159,7 @@ hvac = "2.1.0"
 pymilvus = "2.4.3"
 httpx = "0.27.0"
 paho-mqtt = "2.1.0"
+sqlalchemy-cockroachdb = "2.0.2"
 
 [[tool.poetry.source]]
 name = "PyPI"