diff --git a/modules/cratedb/README.rst b/modules/cratedb/README.rst new file mode 100644 index 00000000..250f7ff0 --- /dev/null +++ b/modules/cratedb/README.rst @@ -0,0 +1,2 @@ +.. autoclass:: testcontainers.cratedb.CrateDBContainer +.. title:: testcontainers.cratedb.CrateDBContainer diff --git a/modules/cratedb/example_basic.py b/modules/cratedb/example_basic.py new file mode 100644 index 00000000..946ca507 --- /dev/null +++ b/modules/cratedb/example_basic.py @@ -0,0 +1,16 @@ +import sqlalchemy + +from testcontainers import cratedb + + +def main(): + with cratedb.CrateDBContainer("crate:latest", ports={4200: None, 5432: None}) as container: + engine = sqlalchemy.create_engine(container.get_connection_url()) + with engine.begin() as conn: + result = conn.execute(sqlalchemy.text("select version()")) + version = result.fetchone() + print(version) + + +if __name__ == "__main__": + main() diff --git a/modules/cratedb/testcontainers/cratedb/__init__.py b/modules/cratedb/testcontainers/cratedb/__init__.py new file mode 100644 index 00000000..10139d03 --- /dev/null +++ b/modules/cratedb/testcontainers/cratedb/__init__.py @@ -0,0 +1,155 @@ +import os +import typing as t +from urllib.parse import quote + +from testcontainers.core.container import DockerContainer +from testcontainers.core.exceptions import ContainerStartException +from testcontainers.core.utils import raise_for_deprecated_parameter +from testcontainers.core.wait_strategies import HttpWaitStrategy + + +class CrateDBContainer(DockerContainer): + """ + CrateDB database container. + + Example: + + The example spins up a CrateDB database and connects to it using + SQLAlchemy and its Python driver. + + .. doctest:: + + >>> from testcontainers import cratedb import CrateDBContainer + >>> import sqlalchemy + + >>> cratedb_container = + >>> with CrateDBContainer("crate:6.0") as cratedb: + ... engine = sqlalchemy.create_engine(cratedb.get_connection_url()) + ... with engine.begin() as connection: + ... result = connection.execute(sqlalchemy.text("select version()")) + ... version, = result.fetchone() + >>> version + 'CrateDB 6.0.2..' + """ + + CMD_OPTS: t.ClassVar[dict[str, str]] = { + "discovery.type": "single-node", + } + + def __init__( + self, + image: str = "crate/crate:latest", + ports: t.Optional[dict] = None, + user: t.Optional[str] = None, + password: t.Optional[str] = None, + cmd_opts: t.Optional[dict] = None, + **kwargs, + ) -> None: + """ + :param image: docker hub image path with optional tag + :param ports: optional dict that maps a port inside the container to a port on the host machine; + `None` as a map value generates a random port; + Dicts are ordered. By convention, the first key-val pair is designated to the HTTP interface. + Example: {4200: None, 5432: 15432} - port 4200 inside the container will be mapped + to a random port on the host, internal port 5432 for PSQL interface will be mapped + to the 15432 port on the host. + :param user: optional username to access the DB; if None, try `CRATEDB_USER` environment variable + :param password: optional password to access the DB; if None, try `CRATEDB_PASSWORD` environment variable + :param cmd_opts: an optional dict with CLI arguments to be passed to the DB entrypoint inside the container + :param kwargs: misc keyword arguments + """ + super().__init__(image=image, **kwargs) + cmd_opts = cmd_opts or {} + self._command = self._build_cmd({**self.CMD_OPTS, **cmd_opts}) + + self.CRATEDB_USER = user or os.environ.get("CRATEDB_USER", "crate") + self.CRATEDB_PASSWORD = password or os.environ.get("CRATEDB_PASSWORD", "crate") + + self.port_mapping = ports if ports else {4200: None} + self.port_to_expose = next(iter(self.port_mapping.items())) + + self.waiting_for(HttpWaitStrategy(4200).for_status_code(200).with_startup_timeout(5)) + + def exposed_ports(self) -> dict[int, int]: + """Returns a dictionary with the ports that are currently exposed in the container. + + Contrary to the '--port' parameter used in docker cli, this returns {internal_port: external_port} + + Examples: + {4200: 19382} + + :returns: The exposed ports. + """ + return {port: self.get_exposed_port(port) for port in self.ports} + + @staticmethod + def _build_cmd(opts: dict) -> str: + """ + Return a string with command options concatenated and optimised for ES5 use + """ + cmd = [] + for key, val in opts.items(): + if isinstance(val, bool): + val = str(val).lower() + cmd.append(f"-C{key}={val}") + return " ".join(cmd) + + def _configure_ports(self) -> None: + """ + Bind all the ports exposed inside the container to the same port on the host + """ + # If host_port is `None`, a random port to be generated + for container_port, host_port in self.port_mapping.items(): + self.with_bind_ports(container=container_port, host=host_port) + + def _configure_credentials(self) -> None: + self.with_env("CRATEDB_USER", self.CRATEDB_USER) + self.with_env("CRATEDB_PASSWORD", self.CRATEDB_PASSWORD) + + def _configure(self) -> None: + self._configure_ports() + self._configure_credentials() + + def get_connection_url(self, dialect: str = "crate", host: t.Optional[str] = None) -> str: + # We should remove this method once the new DBContainer generic gets added to the library. + """ + Return a connection URL to the DB + + :param host: optional string + :param dialect: a string with the dialect name to generate a DB URI + :return: string containing a connection URL to te DB + """ + return self._create_connection_url( + dialect=dialect, + username=self.CRATEDB_USER, + password=self.CRATEDB_PASSWORD, + host=host, + port=self.port_to_expose[0], + ) + + def _create_connection_url( + self, + dialect: str, + username: str, + password: str, + host: t.Optional[str] = None, + port: t.Optional[int] = None, + dbname: t.Optional[str] = None, + **kwargs: t.Any, + ) -> str: + if raise_for_deprecated_parameter(kwargs, "db_name", "dbname"): + raise ValueError(f"Unexpected arguments: {','.join(kwargs)}") + + if self._container is None: + raise ContainerStartException("container has not been started") + + host = host or self.get_container_host_ip() + assert port is not None + + port = self.get_exposed_port(port) + quoted_password = quote(password, safe=" +") + + url = f"{dialect}://{username}:{quoted_password}@{host}:{port}" + if dbname: + url = f"{url}/{dbname}" + return url diff --git a/modules/cratedb/tests/test_cratedb.py b/modules/cratedb/tests/test_cratedb.py new file mode 100644 index 00000000..74dd0b91 --- /dev/null +++ b/modules/cratedb/tests/test_cratedb.py @@ -0,0 +1,78 @@ +import urllib.parse +import os + +import sqlalchemy +import pytest + +from testcontainers.cratedb import CrateDBContainer + + +@pytest.mark.parametrize("version", ["5.9", "5.10", "6.0", "latest"]) +def test_docker_run_cratedb_versions(version: str): + with CrateDBContainer(f"crate:{version}") as container: + engine = sqlalchemy.create_engine(container.get_connection_url()) + with engine.begin() as conn: + result = conn.execute(sqlalchemy.text("select 1+2+3+4+5")) + sum_result = result.fetchone()[0] + assert sum_result == 15 + + +@pytest.mark.parametrize( + "ports, expected", + [ + ({5432: None, 4200: None}, False), + ({5432: 5432, 4200: 4200}, {5432: 5432, 4200: 4200}), + ], +) +def test_docker_run_cratedb_ports(ports, expected): + with CrateDBContainer("crate:latest", ports=ports) as container: + exposed_ports = container.exposed_ports() + assert len(exposed_ports) == 2 + assert all(map(lambda port: isinstance(port, int), exposed_ports)) + if expected: + assert exposed_ports == expected + + +def test_docker_run_cratedb_credentials(): + expected_user, expected_password, expected_port = "user1", "pass1", 4200 + expected_default_dialect, expected_default_host = "crate", "localhost" + expected_defined_dialect, expected_defined_host = "somedialect", "somehost" + os.environ["CRATEDB_USER"], os.environ["CRATEDB_PASSWORD"] = expected_user, expected_password + + with CrateDBContainer("crate:latest", ports={4200: expected_port}) as container: + url = urllib.parse.urlparse(container.get_connection_url()) + user, password = url.netloc.split("@")[0].split(":") + host, port = url.netloc.split("@")[1].split(":") + assert user == expected_user + assert password == expected_password + assert url.scheme == expected_default_dialect + assert host == expected_default_host + assert int(port) == expected_port + + url = urllib.parse.urlparse( + container.get_connection_url(dialect=expected_defined_dialect, host=expected_defined_host) + ) + host, _ = url.netloc.split("@")[1].split(":") + + assert url.scheme == expected_defined_dialect + assert host == expected_defined_host + + +@pytest.mark.parametrize( + "opts, expected", + [ + pytest.param( + {"indices.breaker.total.limit": "90%"}, + ("-Cdiscovery.type=single-node -Cindices.breaker.total.limit=90%"), + id="add_cmd_option", + ), + pytest.param( + {"discovery.type": "zen", "indices.breaker.total.limit": "90%"}, + ("-Cdiscovery.type=zen -Cindices.breaker.total.limit=90%"), + id="override_defaults", + ), + ], +) +def test_build_command(opts, expected): + db = CrateDBContainer(cmd_opts=opts) + assert db._command == expected diff --git a/poetry.lock b/poetry.lock index 7411ad74..790b8256 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -1193,6 +1193,33 @@ type = "legacy" url = "https://pypi.org/simple" reference = "PyPI-public" +[[package]] +name = "crate" +version = "2.0.0" +description = "CrateDB Python Client" +optional = true +python-versions = ">=3.6" +groups = ["main"] +markers = "extra == \"cratedb\"" +files = [ + {file = "crate-2.0.0-py3-none-any.whl", hash = "sha256:692a2c42a4c49ecdb6a2e395175d111c6ac7db051e6362243b69adc26bf00a18"}, + {file = "crate-2.0.0.tar.gz", hash = "sha256:b9f08842f2ee0dd2fa5c6b7d7f6d627e9c86d8cf8894cf56d50b817d679860c6"}, +] + +[package.dependencies] +orjson = "<4" +urllib3 = "*" +verlib2 = "*" + +[package.extras] +doc = ["crate-docs-theme (>=0.26.5)", "sphinx (>=3.5,<9)"] +test = ["backports.zoneinfo (<1) ; python_version < \"3.9\"", "certifi", "createcoverage (>=1,<2)", "mypy (<1.15)", "poethepoet (<0.33)", "pytz", "ruff (<0.10)", "stopit (>=1.1.2,<2)", "zc.customdoctests (>=1.0.1,<2)", "zope.testing (>=4,<6)", "zope.testrunner (>=5,<7)"] + +[package.source] +type = "legacy" +url = "https://pypi.org/simple" +reference = "PyPI-public" + [[package]] name = "cryptography" version = "45.0.7" @@ -1543,6 +1570,24 @@ type = "legacy" url = "https://pypi.org/simple" reference = "PyPI-public" +[[package]] +name = "geojson" +version = "3.2.0" +description = "Python bindings and utilities for GeoJSON" +optional = true +python-versions = ">=3.7" +groups = ["main"] +markers = "extra == \"cratedb\"" +files = [ + {file = "geojson-3.2.0-py3-none-any.whl", hash = "sha256:69d14156469e13c79479672eafae7b37e2dcd19bdfd77b53f74fa8fe29910b52"}, + {file = "geojson-3.2.0.tar.gz", hash = "sha256:b860baba1e8c6f71f8f5f6e3949a694daccf40820fa8f138b3f712bd85804903"}, +] + +[package.source] +type = "legacy" +url = "https://pypi.org/simple" +reference = "PyPI-public" + [[package]] name = "geomet" version = "0.2.1.post1" @@ -3559,7 +3604,7 @@ description = "Fast, correct Python JSON library supporting dataclasses, datetim optional = true python-versions = ">=3.9" groups = ["main"] -markers = "extra == \"chroma\"" +markers = "extra == \"chroma\" or extra == \"cratedb\"" files = [ {file = "orjson-3.11.3-cp310-cp310-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:29cb1f1b008d936803e2da3d7cba726fc47232c45df531b29edf0b232dd737e7"}, {file = "orjson-3.11.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:97dceed87ed9139884a55db8722428e27bd8452817fbf1869c58b49fecab1120"}, @@ -6056,6 +6101,38 @@ type = "legacy" url = "https://pypi.org/simple" reference = "PyPI-public" +[[package]] +name = "sqlalchemy-cratedb" +version = "0.41.0" +description = "SQLAlchemy dialect for CrateDB." +optional = true +python-versions = ">=3.6" +groups = ["main"] +markers = "extra == \"cratedb\"" +files = [ + {file = "sqlalchemy_cratedb-0.41.0-py3-none-any.whl", hash = "sha256:542ffe29ffa8733ed50c658034c95a3b1379a440d423f297d97ecd4c0e3ff3aa"}, + {file = "sqlalchemy_cratedb-0.41.0.tar.gz", hash = "sha256:f24620f13baf6645c8500a01f29a16a3d13f2de2ef86e6cc766261592647ed6f"}, +] + +[package.dependencies] +crate = ">=2,<3" +geojson = ">=2.5,<4" +sqlalchemy = ">=1,<2.1" +verlib2 = "0.2" + +[package.extras] +all = ["sqlalchemy-cratedb[vector]"] +develop = ["mypy (<1.15)", "poethepoet (<0.33)", "pyproject-fmt (<2.6)", "ruff (<0.10)", "validate-pyproject (<0.24)"] +doc = ["crate-docs-theme (>=0.26.5)", "sphinx (>=3.5,<9)"] +release = ["build (<2)", "twine (<7)"] +test = ["cratedb-toolkit[testing]", "dask[dataframe] ; python_version < \"3.13\"", "pandas (<2.3)", "pueblo (>=0.0.7)", "pytest (<9)", "pytest-cov (<7)", "pytest-mock (<4)"] +vector = ["numpy"] + +[package.source] +type = "legacy" +url = "https://pypi.org/simple" +reference = "PyPI-public" + [[package]] name = "tenacity" version = "8.2.3" @@ -6476,6 +6553,29 @@ type = "legacy" url = "https://pypi.org/simple" reference = "PyPI-public" +[[package]] +name = "verlib2" +version = "0.2.0" +description = "A standalone variant of packaging.version, without anything else." +optional = true +python-versions = ">=3.6" +groups = ["main"] +markers = "extra == \"cratedb\"" +files = [ + {file = "verlib2-0.2.0-py3-none-any.whl", hash = "sha256:c0fb8dbae76e228e55bb6178207cda931d44c843d8485c873b7676835edf2e8f"}, + {file = "verlib2-0.2.0.tar.gz", hash = "sha256:0ab94087c094e27163948dfa817c98d62b6bdd01226cce5188f8cc3cd68345b9"}, +] + +[package.extras] +develop = ["black (<23)", "mypy (==1.6.0)", "poethepoet (<1)", "pyproject-fmt (<1.3)", "ruff (==0.0.292)", "validate-pyproject (<0.16)"] +release = ["build (<2)", "twine (<5)"] +test = ["pretend (<2)", "pytest (<8)", "pytest-cov (<5)"] + +[package.source] +type = "legacy" +url = "https://pypi.org/simple" +reference = "PyPI-public" + [[package]] name = "virtualenv" version = "20.25.1" @@ -6795,6 +6895,7 @@ chroma = ["chromadb-client"] clickhouse = ["clickhouse-driver"] cockroachdb = [] cosmosdb = ["azure-cosmos"] +cratedb = ["crate", "sqlalchemy-cratedb"] db2 = ["ibm_db_sa", "sqlalchemy"] elasticsearch = [] generic = ["httpx", "redis"] @@ -6836,4 +6937,4 @@ weaviate = ["weaviate-client"] [metadata] lock-version = "2.1" python-versions = ">=3.9.2,<4.0" -content-hash = "241e8b6ba610907adea4496fdeaef4c3fdc3315d222ab87004692aa9371698fa" +content-hash = "aec835483ff266ad4261ef82169300a639a8f04ba4055aee312a781b48fed376" diff --git a/pyproject.toml b/pyproject.toml index acdaa0d0..2a3ea077 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,7 @@ packages = [ { include = "testcontainers", from = "modules/trino" }, { include = "testcontainers", from = "modules/vault" }, { include = "testcontainers", from = "modules/weaviate" }, + { include = "testcontainers", from = "modules/cratedb" } ] [tool.poetry.urls] @@ -122,6 +123,8 @@ azure-cosmos = { version = "*", optional = true } cryptography = { version = "*", optional = true } trino = { version = "*", optional = true } ibm_db_sa = { version = "*", optional = true, markers = "platform_machine != 'aarch64' and platform_machine != 'arm64'" } +crate = { version = "*", optional = true } +sqlalchemy-cratedb = { version = "*", optional = true } [tool.poetry.extras] arangodb = ["python-arango"] @@ -172,6 +175,7 @@ vault = [] weaviate = ["weaviate-client"] chroma = ["chromadb-client"] trino = ["trino"] +cratedb = ["crate", "sqlalchemy-cratedb"] [tool.poetry.group.dev.dependencies] mypy = "1.11.2" @@ -338,6 +342,7 @@ mypy_path = [ "modules/sftp", # "modules/vault" # "modules/weaviate" + # "modules/cratedb" ] enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]