diff --git a/docs/community/mongodb.rst b/docs/community/mongodb.rst index 40afcc0db..91c84ef80 100644 --- a/docs/community/mongodb.rst +++ b/docs/community/mongodb.rst @@ -1,3 +1,4 @@ .. autoclass:: testcontainers.community.mongodb.MongoDbContainer +.. autoclass:: testcontainers.community.mongodb.MongoDbReplicaSetContainer .. autoclass:: testcontainers.community.mongodb.MongoDBAtlasLocalContainer .. title:: testcontainers.community.mongodb.MongoDbContainer diff --git a/docs/modules/mongodb_example.py b/docs/modules/mongodb_example.py index 0ecfe89c0..ba6216807 100644 --- a/docs/modules/mongodb_example.py +++ b/docs/modules/mongodb_example.py @@ -3,7 +3,7 @@ from pymongo import MongoClient -from testcontainers.community.mongodb import MongoDbContainer +from testcontainers.community.mongodb import MongoDbContainer, MongoDbReplicaSetContainer def basic_example(): @@ -81,5 +81,45 @@ def basic_example(): print(json.dumps(stats, default=str, indent=2)) +def replica_set_example(): + with MongoDbReplicaSetContainer() as mongodb: + client = mongodb.get_connection_client() + db = client.test_db + + print("\nConnected to MongoDB replica set") + print(json.dumps(client.admin.command("hello"), default=str, indent=2)) + + # Run a multi-document transaction + accounts = db.accounts + with client.start_session() as session, session.start_transaction(): + accounts.insert_one({"name": "checking", "balance": 100}, session=session) + accounts.update_one( + {"name": "checking"}, + {"$inc": {"balance": -25}}, + session=session, + ) + + print("\nTransaction result:") + print(json.dumps(accounts.find_one({"name": "checking"}), default=str, indent=2)) + + # Observe an insert through a change stream + events = db.events + with events.watch([{"$match": {"operationType": "insert"}}]) as changes: + events.insert_one({"kind": "created", "value": 42}) + change = next(changes) + + print("\nChange stream event:") + print(json.dumps(change, default=str, indent=2)) + + +def unauthenticated_replica_set_example(): + with MongoDbReplicaSetContainer(auth_enabled=False) as mongodb: + client = mongodb.get_connection_client() + print("\nConnected to MongoDB replica set without authentication") + print(json.dumps(client.admin.command("hello"), default=str, indent=2)) + + if __name__ == "__main__": basic_example() + replica_set_example() + unauthenticated_replica_set_example() diff --git a/src/testcontainers/community/mongodb/__init__.py b/src/testcontainers/community/mongodb/__init__.py index 53eda1971..2ede8e8f8 100644 --- a/src/testcontainers/community/mongodb/__init__.py +++ b/src/testcontainers/community/mongodb/__init__.py @@ -10,16 +10,39 @@ # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the # License for the specific language governing permissions and limitations # under the License. +import base64 import os import re -from typing import Optional +import shlex +import time +from typing import Any, Optional +from urllib.parse import urlencode from pymongo import MongoClient +from pymongo.errors import OperationFailure, PyMongoError +from typing_extensions import Self +from testcontainers.core.config import testcontainers_config +from testcontainers.core.exceptions import ContainerStartException from testcontainers.core.generic import DbContainer from testcontainers.core.utils import raise_for_deprecated_parameter from testcontainers.core.wait_strategies import HealthcheckWaitStrategy, LogMessageWaitStrategy +_REPLICA_SET_KEYFILE_PATH = "/tmp/testcontainers-mongodb-keyfile" +_REPLICA_SET_ENTRYPOINT_PATH = "/tmp/testcontainers-mongodb-entrypoint.sh" +_REPLICA_SET_ENTRYPOINT = f"""#!/bin/bash +set -Eeuo pipefail +chown mongodb:mongodb {_REPLICA_SET_KEYFILE_PATH} +chmod 400 {_REPLICA_SET_KEYFILE_PATH} +exec /usr/local/bin/docker-entrypoint.sh "$@" +""".encode() + + +def _is_root_container_user(user: Any) -> bool: + if user in (None, "", 0, "0", "root"): + return True + return isinstance(user, str) and user.partition(":")[0] in ("0", "root") + class MongoDbContainer(DbContainer): """ @@ -89,6 +112,203 @@ def get_connection_client(self) -> MongoClient: return MongoClient(self.get_connection_url()) +class MongoDbReplicaSetContainer(MongoDbContainer): + """MongoDB container configured as a single-node replica set. + + Authentication is enabled by default and uses an ephemeral keyfile for internal + replica-set authentication. Set ``auth_enabled=False`` to run without + authentication or a keyfile. + + Automatic keyfile setup targets the official ``mongo`` image and requires the + container to start as root before the image drops privileges to ``mongodb``. + + Example: + + .. code-block:: python + + from testcontainers.community.mongodb import MongoDbReplicaSetContainer + + with MongoDbReplicaSetContainer("mongo:7.0.7") as mongo: + client = mongo.get_connection_client() + with client.start_session() as session, session.start_transaction(): + client.test.items.insert_one({"name": "example"}, session=session) + + with MongoDbReplicaSetContainer( + "mongo:7.0.7", + auth_enabled=False, + ) as mongo: + client = mongo.get_connection_client() + """ + + def __init__( + self, + image: str = "mongo:latest", + port: int = 27017, + username: Optional[str] = None, + password: Optional[str] = None, + dbname: Optional[str] = None, + replica_set: str = "docker-rs", + auth_enabled: bool = True, + **kwargs: Any, + ) -> None: + if not replica_set: + raise ValueError("replica_set must not be empty") + if not auth_enabled and (username is not None or password is not None): + raise ValueError("username and password cannot be set when authentication is disabled") + if auth_enabled: + if "entrypoint" in kwargs: + raise ValueError("entrypoint cannot be overridden for an authenticated replica set") + if "user" in kwargs and not _is_root_container_user(kwargs["user"]): + raise ValueError("authenticated replica sets must start as the root container user") + kwargs["entrypoint"] = _REPLICA_SET_ENTRYPOINT_PATH + + super().__init__( + image=image, + port=port, + username=username, + password=password, + dbname=dbname, + **kwargs, + ) + self.auth_enabled = auth_enabled + self.replica_set = replica_set + if not auth_enabled: + self.username = "" + self.password = "" + else: + super().with_copy_into_container(_REPLICA_SET_ENTRYPOINT, _REPLICA_SET_ENTRYPOINT_PATH, mode=0o755) + keyfile = base64.b64encode(os.urandom(756)) + super().with_copy_into_container(keyfile, _REPLICA_SET_KEYFILE_PATH, mode=0o400) + command = MongoDbReplicaSetContainer._replica_set_command(self, self._command) + super().with_command(command) + + def _replica_set_command(self, command: Optional[str | list[str]]) -> list[str]: + command_parts = shlex.split(command) if isinstance(command, str) else list(command or []) + if command_parts and not (command_parts[0].startswith("-") or os.path.basename(command_parts[0]) == "mongod"): + raise ValueError("replica set commands must contain mongod options or start with mongod") + if any( + part in ("--replSet", "--keyFile") or part.startswith(("--replSet=", "--keyFile=")) + for part in command_parts + ): + raise ValueError("replica set and keyfile options are managed by MongoDbReplicaSetContainer") + + command_parts.extend(["--replSet", self.replica_set]) + if self.auth_enabled: + command_parts.extend(["--keyFile", _REPLICA_SET_KEYFILE_PATH]) + return command_parts + + def with_command(self, command: str | list[str]) -> Self: + return super().with_command(self._replica_set_command(command)) + + def with_kwargs(self, **kwargs: Any) -> Self: + if self.auth_enabled: + if "entrypoint" in kwargs: + raise ValueError("entrypoint cannot be overridden for an authenticated replica set") + if "user" in kwargs and not _is_root_container_user(kwargs["user"]): + raise ValueError("authenticated replica sets must start as the root container user") + kwargs["entrypoint"] = _REPLICA_SET_ENTRYPOINT_PATH + return super().with_kwargs(**kwargs) + + def _configure(self) -> None: + if self.auth_enabled: + assert self.username is not None + assert self.password is not None + self.with_env("MONGO_INITDB_ROOT_USERNAME", self.username) + self.with_env("MONGO_INITDB_ROOT_PASSWORD", self.password) + else: + self.env.pop("MONGO_INITDB_ROOT_USERNAME", None) + self.env.pop("MONGO_INITDB_ROOT_PASSWORD", None) + self.with_env("MONGO_DB", self.dbname) + + def get_connection_url(self) -> str: + if self.auth_enabled: + url = super().get_connection_url() + else: + host = self.get_container_host_ip() + port = self.get_exposed_port(self.port) + url = f"mongodb://{host}:{port}" + + return f"{url}/?{urlencode({'replicaSet': self.replica_set, 'directConnection': 'true'})}" + + def _connect(self) -> None: + direct_url = self.get_connection_url().replace( + urlencode({"replicaSet": self.replica_set, "directConnection": "true"}), + urlencode({"directConnection": "true"}), + ) + client: MongoClient[dict[str, Any]] = MongoClient( + direct_url, + serverSelectionTimeoutMS=1000, + connectTimeoutMS=1000, + socketTimeoutMS=1000, + ) + deadline = time.monotonic() + testcontainers_config.timeout + + try: + self._wait_for_mongodb(client, deadline) + self._wait_for_replica_set_primary(client, deadline) + finally: + client.close() + + def _wait_for_mongodb(self, client: MongoClient[dict[str, Any]], deadline: float) -> None: + last_error: Optional[Exception] = None + while time.monotonic() < deadline: + self._raise_if_replica_set_container_stopped() + try: + client.admin.command("ping") + return + except PyMongoError as error: + last_error = error + time.sleep(testcontainers_config.sleep_time) + + raise ContainerStartException("MongoDB did not become ready") from last_error + + def _wait_for_replica_set_primary(self, client: MongoClient[dict[str, Any]], deadline: float) -> None: + last_error: Optional[Exception] = None + while time.monotonic() < deadline: + self._raise_if_replica_set_container_stopped() + try: + self._initiate_replica_set_if_needed(client) + if client.admin.command("hello").get("isWritablePrimary"): + return + except OperationFailure: + raise + except PyMongoError as error: + last_error = error + time.sleep(testcontainers_config.sleep_time) + + raise ContainerStartException("MongoDB replica set did not elect a primary") from last_error + + def _initiate_replica_set_if_needed(self, client: MongoClient[dict[str, Any]]) -> None: + try: + client.admin.command("replSetGetStatus") + return + except OperationFailure as error: + if error.code != 94: # NotYetInitialized + raise + + try: + client.admin.command( + { + "replSetInitiate": { + "_id": self.replica_set, + "members": [{"_id": 0, "host": f"localhost:{self.port}"}], + } + } + ) + except OperationFailure as error: + if error.code != 23: # AlreadyInitialized + raise + + def _raise_if_replica_set_container_stopped(self) -> None: + self.reload() + if self.status not in ("exited", "dead"): + return + + stdout, stderr = self.get_logs() + logs = (stdout + stderr).decode(errors="replace") + raise ContainerStartException(f"MongoDB stopped while initializing its replica set:\n{logs}") + + class MongoDBAtlasLocalContainer(DbContainer): """ MongoDB Atlas Local document-based database container. diff --git a/tests/community/mongodb/test_mongodb.py b/tests/community/mongodb/test_mongodb.py index e4cc9af9e..62f2342c3 100644 --- a/tests/community/mongodb/test_mongodb.py +++ b/tests/community/mongodb/test_mongodb.py @@ -1,8 +1,9 @@ import time +from urllib.parse import parse_qs, urlsplit import pytest -from testcontainers.community.mongodb import MongoDBAtlasLocalContainer, MongoDbContainer +from testcontainers.community.mongodb import MongoDBAtlasLocalContainer, MongoDbContainer, MongoDbReplicaSetContainer @pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"]) @@ -28,6 +29,80 @@ def test_docker_run_mongodb(version: str): assert cursor.next()["restaurant_id"] == doc["restaurant_id"] +@pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"]) +def test_docker_run_mongodb_replica_set_without_authentication(version: str): + with MongoDbReplicaSetContainer(f"mongo:{version}", replica_set="docker-rs", auth_enabled=False) as mongo: + connection_url = urlsplit(mongo.get_connection_url()) + assert connection_url.username is None + assert parse_qs(connection_url.query) == { + "directConnection": ["true"], + "replicaSet": ["docker-rs"], + } + + client = mongo.get_connection_client() + + hello = client.admin.command("hello") + assert hello["setName"] == "docker-rs" + assert hello["isWritablePrimary"] is True + + collection = client.test.transactions + with client.start_session() as session, session.start_transaction(): + result = collection.insert_one({"message": "replica sets support transactions"}, session=session) + + assert collection.find_one({"_id": result.inserted_id}) is not None + + +@pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"]) +def test_docker_run_authenticated_mongodb_replica_set(version: str): + container = MongoDbReplicaSetContainer( + f"mongo:{version}", + username="replica-user", + password="replica-password", + replica_set="secure-rs", + command=["--profile", "1"], + ).with_kwargs(labels={"testcontainers.mongodb.replica-set": "true"}) + + with container as mongo: + connection_url = urlsplit(mongo.get_connection_url()) + assert connection_url.username == "replica-user" + assert connection_url.password == "replica-password" + assert parse_qs(connection_url.query) == { + "directConnection": ["true"], + "replicaSet": ["secure-rs"], + } + + client = mongo.get_connection_client() + + hello = client.admin.command("hello") + assert hello["setName"] == "secure-rs" + assert hello["isWritablePrimary"] is True + + result = client.test.documents.insert_one({"authenticated": True}) + assert client.test.documents.find_one({"_id": result.inserted_id}) == { + "_id": result.inserted_id, + "authenticated": True, + } + + +@pytest.mark.parametrize("username,password", [("user", None), (None, "password")]) +def test_mongodb_rejects_credentials_when_authentication_is_disabled( + username: str | None, + password: str | None, +): + with pytest.raises(ValueError, match="authentication is disabled"): + MongoDbReplicaSetContainer(auth_enabled=False, username=username, password=password) + + +def test_authenticated_mongodb_replica_set_rejects_non_root_container_user(): + with pytest.raises(ValueError, match="root container user"): + MongoDbReplicaSetContainer(replica_set="secure-rs", user="mongodb") + + +def test_mongodb_replica_set_rejects_non_mongod_command(): + with pytest.raises(ValueError, match="mongod options"): + MongoDbReplicaSetContainer(replica_set="secure-rs", command=["bash", "-lc", "exec mongod"]) + + @pytest.mark.parametrize("version", ["8.0.13", "7.0.23"]) def test_docker_run_mongodb_atlas_local(version: str): with MongoDBAtlasLocalContainer(f"mongodb/mongodb-atlas-local:{version}") as mongo_atlas: