-
Notifications
You must be signed in to change notification settings - Fork 292
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix(cosmosdb): Add support for the CosmosDB Emulator (#579)
Adds support for the [CosmosDB Emulator container](https://learn.microsoft.com/en-us/azure/cosmos-db/emulator) --------- Co-authored-by: Mehdi BEN ABDALLAH <@mbenabda> Co-authored-by: David Ankin <daveankin@gmail.com>
- Loading branch information
1 parent
e575b28
commit 8045a80
Showing
12 changed files
with
333 additions
and
62 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,5 @@ | ||
.. autoclass:: testcontainers.cosmosdb.CosmosDBMongoEndpointContainer | ||
.. title:: testcontainers.cosmosdb.CosmosDBMongoEndpointContainer | ||
|
||
.. autoclass:: testcontainers.cosmosdb.CosmosDBNoSQLEndpointContainer | ||
.. title:: testcontainers.cosmosdb.CosmosDBNoSQLEndpointContainer |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,4 @@ | ||
from .mongodb import CosmosDBMongoEndpointContainer | ||
from .nosql import CosmosDBNoSQLEndpointContainer | ||
|
||
__all__ = ["CosmosDBMongoEndpointContainer", "CosmosDBNoSQLEndpointContainer"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,110 @@ | ||
import os | ||
import socket | ||
import ssl | ||
from collections.abc import Iterable | ||
from distutils.util import strtobool | ||
from urllib.error import HTTPError, URLError | ||
from urllib.request import urlopen | ||
|
||
from typing_extensions import Self | ||
|
||
from testcontainers.core.container import DockerContainer | ||
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs | ||
|
||
from . import _grab as grab | ||
|
||
__all__ = ["CosmosDBEmulatorContainer"] | ||
|
||
EMULATOR_PORT = 8081 | ||
|
||
|
||
class CosmosDBEmulatorContainer(DockerContainer): | ||
""" | ||
Abstract class for CosmosDB Emulator endpoints. | ||
Concrete implementations for each endpoint is provided by a separate class: | ||
NoSQLEmulatorContainer and MongoDBEmulatorContainer. | ||
""" | ||
|
||
def __init__( | ||
self, | ||
image: str = os.getenv( | ||
"AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:latest" | ||
), | ||
partition_count: int = os.getenv("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", None), | ||
enable_data_persistence: bool = strtobool(os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")), | ||
key: str = os.getenv( | ||
"AZURE_COSMOS_EMULATOR_KEY", | ||
"C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==", | ||
), | ||
bind_ports: bool = strtobool(os.getenv("AZURE_COSMOS_EMULATOR_BIND_PORTS", "true")), | ||
endpoint_ports: Iterable[int] = [], | ||
**other_kwargs, | ||
): | ||
super().__init__(image=image, **other_kwargs) | ||
self.endpoint_ports = endpoint_ports | ||
self.partition_count = partition_count | ||
self.key = key | ||
self.enable_data_persistence = enable_data_persistence | ||
self.bind_ports = bind_ports | ||
|
||
@property | ||
def host(self) -> str: | ||
""" | ||
Emulator host | ||
""" | ||
return self.get_container_host_ip() | ||
|
||
@property | ||
def server_certificate_pem(self) -> bytes: | ||
""" | ||
PEM-encoded server certificate | ||
""" | ||
return self._cert_pem_bytes | ||
|
||
def start(self) -> Self: | ||
self._configure() | ||
super().start() | ||
self._wait_until_ready() | ||
self._cert_pem_bytes = self._download_cert() | ||
return self | ||
|
||
def _configure(self) -> None: | ||
all_ports = {EMULATOR_PORT, *self.endpoint_ports} | ||
if self.bind_ports: | ||
for port in all_ports: | ||
self.with_bind_ports(port, port) | ||
else: | ||
self.with_exposed_ports(*all_ports) | ||
|
||
( | ||
self.with_env("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", str(self.partition_count)) | ||
.with_env("AZURE_COSMOS_EMULATOR_IP_ADDRESS_OVERRIDE", socket.gethostbyname(socket.gethostname())) | ||
.with_env("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", str(self.enable_data_persistence)) | ||
.with_env("AZURE_COSMOS_EMULATOR_KEY", str(self.key)) | ||
) | ||
|
||
def _wait_until_ready(self) -> Self: | ||
wait_for_logs(container=self, predicate="Started\\s*$") | ||
|
||
if self.bind_ports: | ||
self._wait_for_url(f"https://{self.host}:{EMULATOR_PORT}/_explorer/index.html") | ||
self._wait_for_query_success() | ||
|
||
return self | ||
|
||
def _download_cert(self) -> bytes: | ||
with grab.file( | ||
self.get_wrapped_container(), | ||
"/tmp/cosmos/appdata/.system/profiles/Client/AppData/Local/CosmosDBEmulator/emulator.pem", | ||
) as cert: | ||
return cert.read() | ||
|
||
@wait_container_is_ready(HTTPError, URLError) | ||
def _wait_for_url(self, url: str) -> Self: | ||
with urlopen(url, context=ssl._create_unverified_context()) as response: | ||
response.read() | ||
return self | ||
|
||
def _wait_for_query_success(self) -> None: | ||
pass |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,26 @@ | ||
import tarfile | ||
import tempfile | ||
from contextlib import contextmanager | ||
from os import path | ||
from pathlib import Path | ||
|
||
from docker.models.containers import Container | ||
|
||
|
||
@contextmanager | ||
def file(container: Container, target: str): | ||
target_path = Path(target) | ||
assert target_path.is_absolute(), "target must be an absolute path" | ||
|
||
with tempfile.TemporaryDirectory() as tmp: | ||
archive = Path(tmp) / "grabbed.tar" | ||
|
||
# download from container as tar archive | ||
with open(archive, "wb") as f: | ||
tar_bits, _ = container.get_archive(target) | ||
for chunk in tar_bits: | ||
f.write(chunk) | ||
|
||
# extract target file from tar archive | ||
with tarfile.TarFile(archive) as tar: | ||
yield tar.extractfile(path.basename(target)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
import os | ||
|
||
from ._emulator import CosmosDBEmulatorContainer | ||
|
||
__all__ = ["CosmosDBMongoEndpointContainer"] | ||
|
||
ENDPOINT_PORT = 10255 | ||
|
||
|
||
class CosmosDBMongoEndpointContainer(CosmosDBEmulatorContainer): | ||
""" | ||
CosmosDB MongoDB enpoint Emulator. | ||
Example: | ||
.. code-block:: python | ||
>>> from testcontainers.cosmosdb import CosmosDBMongoEndpointContainer | ||
>>> with CosmosDBMongoEndpointContainer(mongodb_version="4.0") as emulator: | ||
... print(f"Point your MongoDB client at {emulator.host}:{emulator.port} using key {emulator.key}") | ||
... print(f"and eiher disable TLS server auth or trust the server's self signed cert (emulator.server_certificate_pem)") | ||
""" | ||
|
||
def __init__( | ||
self, | ||
mongodb_version: str, | ||
image: str = os.getenv( | ||
"AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:mongodb" | ||
), | ||
**other_kwargs, | ||
): | ||
super().__init__(image=image, endpoint_ports=[ENDPOINT_PORT], **other_kwargs) | ||
assert mongodb_version is not None, "A MongoDB version is required to use the MongoDB Endpoint" | ||
self.mongodb_version = mongodb_version | ||
|
||
@property | ||
def port(self) -> str: | ||
""" | ||
The exposed port to the MongoDB endpoint | ||
""" | ||
return self.get_exposed_port(ENDPOINT_PORT) | ||
|
||
def _configure(self) -> None: | ||
super()._configure() | ||
self.with_env("AZURE_COSMOS_EMULATOR_ENABLE_MONGODB_ENDPOINT", self.mongodb_version) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,69 @@ | ||
from azure.core.exceptions import ServiceRequestError | ||
from azure.cosmos import CosmosClient as SyncCosmosClient | ||
from azure.cosmos.aio import CosmosClient as AsyncCosmosClient | ||
|
||
from testcontainers.core.waiting_utils import wait_container_is_ready | ||
|
||
from ._emulator import CosmosDBEmulatorContainer | ||
|
||
__all__ = ["CosmosDBNoSQLEndpointContainer"] | ||
|
||
NOSQL_PORT = 8081 | ||
|
||
|
||
class CosmosDBNoSQLEndpointContainer(CosmosDBEmulatorContainer): | ||
""" | ||
CosmosDB NoSQL enpoint Emulator. | ||
Example: | ||
.. code-block:: python | ||
>>> from testcontainers.cosmosdb import CosmosDBNoSQLEndpointContainer | ||
>>> with CosmosDBNoSQLEndpointContainer() as emulator: | ||
... db = emulator.insecure_sync_client().create_database_if_not_exists("test") | ||
.. code-block:: python | ||
>>> from testcontainers.cosmosdb import CosmosDBNoSQLEndpointContainer | ||
>>> from azure.cosmos import CosmosClient | ||
>>> with CosmosDBNoSQLEndpointContainer() as emulator: | ||
... client = CosmosClient(url=emulator.url, credential=emulator.key, connection_verify=False) | ||
... db = client.create_database_if_not_exists("test") | ||
""" | ||
|
||
def __init__(self, **kwargs): | ||
super().__init__(endpoint_ports=[NOSQL_PORT], **kwargs) | ||
|
||
@property | ||
def port(self) -> str: | ||
""" | ||
The exposed port to the NoSQL endpoint | ||
""" | ||
return self.get_exposed_port(NOSQL_PORT) | ||
|
||
@property | ||
def url(self) -> str: | ||
""" | ||
The url to the NoSQL endpoint | ||
""" | ||
return f"https://{self.host}:{self.port}" | ||
|
||
def insecure_async_client(self): | ||
""" | ||
Returns an asynchronous CosmosClient instance | ||
""" | ||
return AsyncCosmosClient(url=self.url, credential=self.key, connection_verify=False) | ||
|
||
def insecure_sync_client(self): | ||
""" | ||
Returns a synchronous CosmosClient instance | ||
""" | ||
return SyncCosmosClient(url=self.url, credential=self.key, connection_verify=False) | ||
|
||
@wait_container_is_ready(ServiceRequestError) | ||
def _wait_for_query_success(self) -> None: | ||
with self.insecure_sync_client() as c: | ||
list(c.list_databases()) |
Oops, something went wrong.