Skip to content

Commit

Permalink
feat: Adding support for Cassandra and Scylla (#167)
Browse files Browse the repository at this point in the history
This add the support for those cassandra based dbs
and their drivers, cassandra-driver, scylla-driver

Ref: https://cassandra.apache.org/
Ref: https://www.scylladb.com/
Ref: https://pypi.org/project/cassandra-driver/
Ref: https://pypi.org/project/scylla-driver/

---------

Co-authored-by: David Ankin <daveankin@gmail.com>
  • Loading branch information
fruch and alexanderankin authored Aug 14, 2024
1 parent d5de0aa commit 2d8bc11
Show file tree
Hide file tree
Showing 5 changed files with 75 additions and 6 deletions.
2 changes: 2 additions & 0 deletions modules/scylla/README.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
.. autoclass:: testcontainers.scylla.ScyllaContainer
.. title:: testcontainers.scylla.ScyllaContainer
47 changes: 47 additions & 0 deletions modules/scylla/testcontainers/scylla/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
from testcontainers.core.config import MAX_TRIES
from testcontainers.core.generic import DockerContainer
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs


class ScyllaContainer(DockerContainer):
"""
Scylla database container.
Example
-------
.. doctest::
>>> from testcontainers.scylla import ScyllaContainer
>>> with ScyllaContainer() as scylla:
... cluster = scylla.get_cluster()
... with cluster.connect() as session:
... result = session.execute(
... "CREATE KEYSPACE keyspace1 WITH replication "
... "= {'class': 'SimpleStrategy', 'replication_factor': '1'};")
"""

def __init__(self, image="scylladb/scylla:latest", ports_to_expose=(9042,)):
super().__init__(image)
self.ports_to_expose = ports_to_expose
self.with_exposed_ports(*self.ports_to_expose)
self.with_command("--skip-wait-for-gossip-to-settle=0")

@wait_container_is_ready()
def _connect(self):
wait_for_logs(self, predicate="Starting listening for CQL clients", timeout=MAX_TRIES)
cluster = self.get_cluster()
cluster.connect()

def start(self):
super().start()
self._connect()
return self

def get_cluster(self, **kwargs):
from cassandra.cluster import Cluster

container = self.get_wrapped_container()
container.reload()
hostname = container.attrs["NetworkSettings"]["IPAddress"]
return Cluster(contact_points=[hostname], **kwargs)
18 changes: 18 additions & 0 deletions modules/scylla/tests/test_scylla.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
from testcontainers.scylla import ScyllaContainer


def test_docker_run_scylla():
with ScyllaContainer() as scylla:
cluster = scylla.get_cluster()
with cluster.connect() as session:
session.execute(
"CREATE KEYSPACE keyspace1 WITH replication = "
"{'class': 'SimpleStrategy', 'replication_factor': '1'};"
)
session.execute("CREATE TABLE keyspace1.table1 (key1 int, key2 int, PRIMARY KEY (key1));")
session.execute("INSERT INTO keyspace1.table1 (key1,key2) values (1,2);")

response = session.execute("SELECT * FROM keyspace1.table1")

assert response.one().key1 == 1
assert response.one().key2 == 2
10 changes: 5 additions & 5 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 3 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ packages = [
{ include = "testcontainers", from = "modules/registry" },
{ include = "testcontainers", from = "modules/sftp" },
{ include = "testcontainers", from = "modules/selenium" },
{ include = "testcontainers", from = "modules/scylla" },
{ include = "testcontainers", from = "modules/trino" },
{ include = "testcontainers", from = "modules/vault" },
{ include = "testcontainers", from = "modules/weaviate" },
Expand All @@ -86,6 +87,7 @@ typing-extensions = "*"
# community modules
python-arango = { version = "^7.8", optional = true }
azure-storage-blob = { version = "^12.19", optional = true }
cassandra-driver = { version = "3.29.1", optional = true }
clickhouse-driver = { version = "*", optional = true }
google-cloud-pubsub = { version = ">=2", optional = true }
google-cloud-datastore = { version = ">=2", optional = true }
Expand Down Expand Up @@ -156,6 +158,7 @@ rabbitmq = ["pika"]
redis = ["redis"]
registry = ["bcrypt"]
selenium = ["selenium"]
scylla = ["cassandra-driver"]
sftp = ["cryptography"]
vault = []
weaviate = ["weaviate-client"]
Expand All @@ -175,7 +178,6 @@ psycopg2-binary = "2.9.9"
pg8000 = "1.30.5"
sqlalchemy = "2.0.28"
psycopg = "3.1.18"
cassandra-driver = "3.29.1"
pytest-asyncio = "0.23.5"
kafka-python-ng = "^2.2.0"
hvac = "2.1.0"
Expand Down

0 comments on commit 2d8bc11

Please sign in to comment.