From 4aff6793f28fbeb8358adcc728283ea9a7b94e5f Mon Sep 17 00:00:00 2001 From: Joel Hess Date: Tue, 18 Jun 2024 04:36:31 -0500 Subject: [PATCH] fix: Add Cockroach DB Module to Testcontainers (#608) Adds [Cockroach DB] (https://www.cockroachlabs.com/) module to use with Test containers I had done this previously under https://github.com/testcontainers/testcontainers-python/pull/281, but opted to just redo it rather than try to rebase all the things. - [x] Create a new feature directory and populate it with the package structure [described in the documentation](https://testcontainers-python.readthedocs.io/en/latest/#package-structure). Copying one of the existing features is likely the best way to get started. - [x] Implement the new feature (typically in `__init__.py`) and corresponding tests. - [x] Update the feature `README.rst` and add it to the table of contents (`toctree` directive) in the top-level `README.rst`. - [] Add a line `[feature name]` to the list of components in the GitHub Action workflow in `.github/workflows/main.yml` to run tests, build, and publish your package when pushed to the `main` branch. - [x] Rebase your development branch on `main` (or merge `main` into your development branch). - [x] Add Package to pyproject.toml - [ ] Add a line `-e file:[feature name]` to `requirements.in` and open a pull request. Opening a pull request will automatically generate lock files to ensure reproducible builds (see the [pip-tools documentation](https://pip-tools.readthedocs.io/en/latest/) for details). Finally, run `python get_requirements.py --pr=[your PR number]` to fetch the updated requirement files (the build needs to have succeeded). --------- Co-authored-by: joelhess Co-authored-by: David Ankin --- modules/cockroachdb/README.rst | 2 + .../testcontainers/cockroachdb/__init__.py | 100 ++++++++++++++++++ modules/cockroachdb/tests/test_cockroachdb.py | 14 +++ poetry.lock | 17 ++- pyproject.toml | 3 + 5 files changed, 135 insertions(+), 1 deletion(-) create mode 100644 modules/cockroachdb/README.rst create mode 100644 modules/cockroachdb/testcontainers/cockroachdb/__init__.py create mode 100644 modules/cockroachdb/tests/test_cockroachdb.py diff --git a/modules/cockroachdb/README.rst b/modules/cockroachdb/README.rst new file mode 100644 index 00000000..7b53fc33 --- /dev/null +++ b/modules/cockroachdb/README.rst @@ -0,0 +1,2 @@ +.. autoclass:: testcontainers.cockroachdb.CockroachDBContainer +.. title:: testcontainers.cockroachdb.CockroachDBContainer diff --git a/modules/cockroachdb/testcontainers/cockroachdb/__init__.py b/modules/cockroachdb/testcontainers/cockroachdb/__init__.py new file mode 100644 index 00000000..13a17ed5 --- /dev/null +++ b/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 diff --git a/modules/cockroachdb/tests/test_cockroachdb.py b/modules/cockroachdb/tests/test_cockroachdb.py new file mode 100644 index 00000000..af20fd58 --- /dev/null +++ b/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] diff --git a/poetry.lock b/poetry.lock index 891c7bd7..d70ba7e3 100644 --- a/poetry.lock +++ b/poetry.lock @@ -4006,6 +4006,20 @@ postgresql-psycopgbinary = ["psycopg[binary] (>=3.0.7)"] pymysql = ["pymysql"] sqlcipher = ["sqlcipher3_binary"] +[[package]] +name = "sqlalchemy-cockroachdb" +version = "2.0.2" +description = "CockroachDB dialect for SQLAlchemy" +optional = false +python-versions = "*" +files = [ + {file = "sqlalchemy-cockroachdb-2.0.2.tar.gz", hash = "sha256:119756eb905855d6a11345b99cfe853031a3fe598a9c4bf35a8ddac9f89fe8cc"}, + {file = "sqlalchemy_cockroachdb-2.0.2-py3-none-any.whl", hash = "sha256:0d5d50e805b024cb2ccd85423a5c1a367d1a56a5cd0ea47765233fd47665070d"}, +] + +[package.dependencies] +SQLAlchemy = "*" + [[package]] name = "tenacity" version = "8.2.3" @@ -4447,6 +4461,7 @@ azurite = ["azure-storage-blob"] cassandra = [] chroma = ["chromadb-client"] clickhouse = ["clickhouse-driver"] +cockroachdb = [] elasticsearch = [] google = ["google-cloud-datastore", "google-cloud-pubsub"] influxdb = ["influxdb", "influxdb-client"] @@ -4479,4 +4494,4 @@ weaviate = ["weaviate-client"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0" -content-hash = "043c7eea4ca72646a19a705891b26577a27149673ba38c8a6dd4732d30ce081c" +content-hash = "040fa3576807a8bd7b129b889c934d5c4bca9e95376c50e15fa870f4d8336fdb" diff --git a/pyproject.toml b/pyproject.toml index afe841c8..4967a28e 100644 --- a/pyproject.toml +++ b/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"