From 74ac1128aab729d00e963ac9fbdbb0ef02d52e00 Mon Sep 17 00:00:00 2001 From: Israel Fruchter Date: Thu, 30 Dec 2021 22:21:49 +0200 Subject: [PATCH] Adding support for Cassandra and Scylla 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/ --- .github/workflows/main.yml | 2 ++ docs/database.rst | 2 ++ requirements.in | 2 +- setup.py | 2 ++ testcontainers/cassandra.py | 40 ++++++++++++++++++++++++++++++++++++ testcontainers/scylla.py | 41 +++++++++++++++++++++++++++++++++++++ tests/test_cassandra.py | 18 ++++++++++++++++ tests/test_scylla.py | 18 ++++++++++++++++ 8 files changed, 124 insertions(+), 1 deletion(-) create mode 100644 testcontainers/cassandra.py create mode 100644 testcontainers/scylla.py create mode 100644 tests/test_cassandra.py create mode 100644 tests/test_scylla.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index a88eb31c8..6c86c2f03 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -31,6 +31,8 @@ jobs: - webdriver.py - keycloak.py - arangodb.py + - scylladb.py + - cassandra.py runs-on: ubuntu-latest steps: - uses: actions/checkout@v2 diff --git a/docs/database.rst b/docs/database.rst index 49d82eeb2..0dda5bb72 100644 --- a/docs/database.rst +++ b/docs/database.rst @@ -13,3 +13,5 @@ Allows to spin up database images such as MySQL, PostgreSQL, MariaDB, Oracle XE, .. autoclass:: testcontainers.clickhouse.ClickHouseContainer .. autoclass:: testcontainers.neo4j.Neo4jContainer .. autoclass:: testcontainers.arangodb.ArangoDbContainer +.. autoclass:: testcontainers.cassandra.CassandraContainer +.. autoclass:: testcontainers.scylla.ScyllaContainer diff --git a/requirements.in b/requirements.in index 4309c1d7f..c915fc604 100644 --- a/requirements.in +++ b/requirements.in @@ -1,4 +1,4 @@ --e file:.[docker-compose,mysql,oracle,postgresql,selenium,google-cloud-pubsub,mongo,redis,mssqlserver,neo4j,kafka,rabbitmq,clickhouse,keycloak,arangodb] +-e file:.[docker-compose,mysql,oracle,postgresql,selenium,google-cloud-pubsub,mongo,redis,mssqlserver,neo4j,kafka,rabbitmq,clickhouse,keycloak,arangodb,cassandra,scylla] codecov>=2.1.0 cryptography<37 flake8<3.8.0 # 3.8.0 adds a dependency on importlib-metadata which conflicts with other packages. diff --git a/setup.py b/setup.py index 5eb485455..86b40354a 100644 --- a/setup.py +++ b/setup.py @@ -68,6 +68,8 @@ 'clickhouse': ['clickhouse-driver'], 'keycloak': ['python-keycloak'], 'arangodb': ['python-arango'], + 'cassandra': ['cassandra-driver'], + 'scylla': ['scylla-driver'] }, long_description_content_type="text/x-rst", long_description=long_description, diff --git a/testcontainers/cassandra.py b/testcontainers/cassandra.py new file mode 100644 index 000000000..dabd1d74a --- /dev/null +++ b/testcontainers/cassandra.py @@ -0,0 +1,40 @@ +from testcontainers.core.generic import DockerContainer +from testcontainers.core.waiting_utils import wait_container_is_ready + + +class CassandraContainer(DockerContainer): + """ + Cassandra database container. + + Example + ------- + :: + + with CassandraContainer() as cassandra: + cluster = cassandra.get_cluster() + with cluster.connect() as session: + session.execute( + "CREATE KEYSPACE keyspace1 WITH replication = " + "{'class': 'SimpleStrategy', 'replication_factor': '1'};") + + """ + def __init__(self, image="rinscy/cassandra:latest", ports_to_expose=[9042]): + super(CassandraContainer, self).__init__(image) + self.ports_to_expose = ports_to_expose + self.with_exposed_ports(*self.ports_to_expose) + + @wait_container_is_ready() + def _connect(self): + cluster = self.get_cluster() + cluster.connect() + + def start(self): + super(CassandraContainer, self).start() + self._connect() + return self + + def get_cluster(self, **kwargs): + from cassandra.cluster import Cluster + hostname = self.get_container_host_ip() + port = self.get_exposed_port(9042) + return Cluster(contact_points=[hostname], port=port, **kwargs) diff --git a/testcontainers/scylla.py b/testcontainers/scylla.py new file mode 100644 index 000000000..6e1aa9ed7 --- /dev/null +++ b/testcontainers/scylla.py @@ -0,0 +1,41 @@ +from testcontainers.core.generic import DockerContainer +from testcontainers.core.waiting_utils import wait_container_is_ready + + +class ScyllaContainer(DockerContainer): + """ + Scylla database container. + + Example + ------- + :: + + 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'};") + + """ + def __init__(self, image="scylladb/scylla:4.6.rc1", ports_to_expose=[9042]): + super(ScyllaContainer, self).__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): + cluster = self.get_cluster() + cluster.connect() + + def start(self): + super(ScyllaContainer, self).start() + self._connect() + return self + + def get_cluster(self, **kwargs): + from cassandra.cluster import Cluster + hostname = self.get_container_host_ip() + port = self.get_exposed_port(9042) + return Cluster(contact_points=[hostname], port=port, **kwargs) diff --git a/tests/test_cassandra.py b/tests/test_cassandra.py new file mode 100644 index 000000000..bf2aca017 --- /dev/null +++ b/tests/test_cassandra.py @@ -0,0 +1,18 @@ +from testcontainers.cassandra import CassandraContainer + + +def test_docker_run_cassandra(): + with CassandraContainer() as cassandra: + cluster = cassandra.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 diff --git a/tests/test_scylla.py b/tests/test_scylla.py new file mode 100644 index 000000000..a8515e931 --- /dev/null +++ b/tests/test_scylla.py @@ -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