Skip to content
Open
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions modules/cratedb/README.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
.. autoclass:: testcontainers.cratedb.CrateDBContainer
.. title:: testcontainers.cratedb.CrateDBContainer
16 changes: 16 additions & 0 deletions modules/cratedb/example_basic.py
Original file line number Diff line number Diff line change
@@ -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()
158 changes: 158 additions & 0 deletions modules/cratedb/testcontainers/cratedb/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
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


# DockerSkippingContainer, KeepaliveContainer,
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",
"node.attr.storage": "hot",
"path.repo": "/tmp/snapshots",
}

def __init__(
self,
image: str = "crate/crate:nightly",
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
88 changes: 88 additions & 0 deletions modules/cratedb/tests/test_cratedb.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
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 "
"-Cnode.attr.storage=hot "
"-Cpath.repo=/tmp/snapshots "
"-Cindices.breaker.total.limit=90%"
),
id="add_cmd_option",
),
pytest.param(
{"discovery.type": "zen", "indices.breaker.total.limit": "90%"},
(
"-Cdiscovery.type=zen "
"-Cnode.attr.storage=hot "
"-Cpath.repo=/tmp/snapshots "
"-Cindices.breaker.total.limit=90%"
),
id="override_defaults",
),
],
)
def test_build_command(opts, expected):
db = CrateDBContainer(cmd_opts=opts)
assert db._command == expected
Loading