From db1bb3f4a3cb77de639e16c0e145768908c47697 Mon Sep 17 00:00:00 2001 From: joelhess Date: Mon, 17 Jun 2024 09:18:42 -0500 Subject: [PATCH 1/2] Add CRDB Module --- modules/cockroachdb/README.rst | 2 + .../cockroachdb/testcontainers/__init__.py | 80 +++++++++++++++++++ modules/cockroachdb/tests/test_cockroachdb.py | 18 +++++ pyproject.toml | 1 + 4 files changed, 101 insertions(+) create mode 100644 modules/cockroachdb/README.rst create mode 100644 modules/cockroachdb/testcontainers/__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/__init__.py b/modules/cockroachdb/testcontainers/__init__.py new file mode 100644 index 00000000..b86dfabf --- /dev/null +++ b/modules/cockroachdb/testcontainers/__init__.py @@ -0,0 +1,80 @@ +# +# 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 testcontainers.core.generic import DbContainer + + +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:latest') 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() + + """ + + def __init__( + self, + image: str = "cockroachdb/cockroach:latest", + username: Optional[str] = None, + password: Optional[str] = None, + dbname: Optional[str] = None, + port: int = 26257, + dialect="cockroachdb+psycopg2", + **kwargs, + ) -> None: + super().__init__(image, **kwargs) + + self.port = port + self.with_exposed_ports(self.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) + + 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.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..d46fd77f --- /dev/null +++ b/modules/cockroachdb/tests/test_cockroachdb.py @@ -0,0 +1,18 @@ +from pathlib import Path +import re +from unittest import mock + +import pytest +import sqlalchemy + +from testcontainers.core.utils import is_arm +from testcontainers.cockroachdb import CockroachDBContainer + +def test_docker_run_mysql(): + config = CockroachDBContainer("cockroachdb/cockroach:24.0.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 row[0].startswith("24.0.1") \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index afe841c8..99f20779 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" }, From 053e7066d2cbdb31b7a96360378f2d2db1f57315 Mon Sep 17 00:00:00 2001 From: David Ankin Date: Tue, 18 Jun 2024 04:51:38 -0400 Subject: [PATCH 2/2] adjustments --- .../{ => cockroachdb}/__init__.py | 32 +++++++++++++++---- modules/cockroachdb/tests/test_cockroachdb.py | 12 +++---- poetry.lock | 17 +++++++++- pyproject.toml | 2 ++ 4 files changed, 48 insertions(+), 15 deletions(-) rename modules/cockroachdb/testcontainers/{ => cockroachdb}/__init__.py (72%) diff --git a/modules/cockroachdb/testcontainers/__init__.py b/modules/cockroachdb/testcontainers/cockroachdb/__init__.py similarity index 72% rename from modules/cockroachdb/testcontainers/__init__.py rename to modules/cockroachdb/testcontainers/cockroachdb/__init__.py index b86dfabf..13a17ed5 100644 --- a/modules/cockroachdb/testcontainers/__init__.py +++ b/modules/cockroachdb/testcontainers/cockroachdb/__init__.py @@ -12,8 +12,11 @@ # 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): @@ -32,7 +35,7 @@ class CockroachDBContainer(DbContainer): >>> import sqlalchemy >>> from testcontainers.cockroachdb import CockroachDBContainer - >>> with CockroachDBContainer('cockroachdb/cockroach:latest') as crdb: + >>> 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()")) @@ -40,20 +43,21 @@ class CockroachDBContainer(DbContainer): """ + COCKROACH_DB_PORT: int = 26257 + COCKROACH_API_PORT: int = 8080 + def __init__( self, - image: str = "cockroachdb/cockroach:latest", + image: str = "cockroachdb/cockroach:v24.1.1", username: Optional[str] = None, password: Optional[str] = None, dbname: Optional[str] = None, - port: int = 26257, dialect="cockroachdb+psycopg2", **kwargs, ) -> None: super().__init__(image, **kwargs) - self.port = port - self.with_exposed_ports(self.port) + 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") @@ -69,9 +73,25 @@ def _configure(self) -> None: 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.port + dialect=self.dialect, + username=self.username, + password=self.password, + dbname=self.dbname, + port=self.COCKROACH_DB_PORT, ) if self.password: diff --git a/modules/cockroachdb/tests/test_cockroachdb.py b/modules/cockroachdb/tests/test_cockroachdb.py index d46fd77f..af20fd58 100644 --- a/modules/cockroachdb/tests/test_cockroachdb.py +++ b/modules/cockroachdb/tests/test_cockroachdb.py @@ -1,18 +1,14 @@ -from pathlib import Path -import re -from unittest import mock - -import pytest import sqlalchemy -from testcontainers.core.utils import is_arm from testcontainers.cockroachdb import CockroachDBContainer + def test_docker_run_mysql(): - config = CockroachDBContainer("cockroachdb/cockroach:24.0.1") + 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 row[0].startswith("24.0.1") \ No newline at end of file + 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 99f20779..4967a28e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,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"] @@ -158,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"