| # Copyright DataStax, Inc. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| from cassandra import InvalidRequest |
| from cassandra.cluster import NoHostAvailable |
| from cassandra.cqlengine import columns, CQLEngineException |
| from cassandra.cqlengine import connection as conn |
| from cassandra.cqlengine.management import drop_keyspace, sync_table, drop_table, create_keyspace_simple |
| from cassandra.cqlengine.models import Model, QuerySetDescriptor |
| from cassandra.cqlengine.query import ContextQuery, BatchQuery, ModelQuerySet |
| from tests.integration.cqlengine import setup_connection, DEFAULT_KEYSPACE |
| from tests.integration.cqlengine.base import BaseCassEngTestCase |
| from tests.integration.cqlengine.query import test_queryset |
| from tests.integration import local, CASSANDRA_IP, TestCluster |
| |
| |
| class TestModel(Model): |
| |
| __keyspace__ = 'ks1' |
| |
| partition = columns.Integer(primary_key=True) |
| cluster = columns.Integer(primary_key=True) |
| count = columns.Integer() |
| text = columns.Text() |
| |
| |
| class AnotherTestModel(Model): |
| |
| __keyspace__ = 'ks1' |
| |
| partition = columns.Integer(primary_key=True) |
| cluster = columns.Integer(primary_key=True) |
| count = columns.Integer() |
| text = columns.Text() |
| |
| class ContextQueryConnectionTests(BaseCassEngTestCase): |
| |
| @classmethod |
| def setUpClass(cls): |
| super(ContextQueryConnectionTests, cls).setUpClass() |
| create_keyspace_simple('ks1', 1) |
| |
| conn.unregister_connection('default') |
| conn.register_connection('fake_cluster', ['1.2.3.4'], lazy_connect=True, retry_connect=True, default=True) |
| conn.register_connection('cluster', [CASSANDRA_IP]) |
| |
| with ContextQuery(TestModel, connection='cluster') as tm: |
| sync_table(tm) |
| |
| @classmethod |
| def tearDownClass(cls): |
| super(ContextQueryConnectionTests, cls).tearDownClass() |
| |
| with ContextQuery(TestModel, connection='cluster') as tm: |
| drop_table(tm) |
| drop_keyspace('ks1', connections=['cluster']) |
| |
| |
| # reset the default connection |
| conn.unregister_connection('fake_cluster') |
| conn.unregister_connection('cluster') |
| setup_connection(DEFAULT_KEYSPACE) |
| |
| def setUp(self): |
| super(BaseCassEngTestCase, self).setUp() |
| |
| def test_context_connection_priority(self): |
| """ |
| Tests to ensure the proper connection priority is honored. |
| |
| Explicit connection should have higest priority, |
| Followed by context query connection |
| Default connection should be honored last. |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result priorities should be honored |
| |
| @test_category object_mapper |
| """ |
| # model keyspace write/read |
| |
| # Set the default connection on the Model |
| TestModel.__connection__ = 'cluster' |
| with ContextQuery(TestModel) as tm: |
| tm.objects.create(partition=1, cluster=1) |
| |
| # ContextQuery connection should have priority over default one |
| with ContextQuery(TestModel, connection='fake_cluster') as tm: |
| with self.assertRaises(NoHostAvailable): |
| tm.objects.create(partition=1, cluster=1) |
| |
| # Explicit connection should have priority over ContextQuery one |
| with ContextQuery(TestModel, connection='fake_cluster') as tm: |
| tm.objects.using(connection='cluster').create(partition=1, cluster=1) |
| |
| # Reset the default conn of the model |
| TestModel.__connection__ = None |
| |
| # No model connection and an invalid default connection |
| with ContextQuery(TestModel) as tm: |
| with self.assertRaises(NoHostAvailable): |
| tm.objects.create(partition=1, cluster=1) |
| |
| def test_context_connection_with_keyspace(self): |
| """ |
| Tests to ensure keyspace param is honored |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result Invalid request is thrown |
| |
| @test_category object_mapper |
| """ |
| |
| # ks2 doesn't exist |
| with ContextQuery(TestModel, connection='cluster', keyspace='ks2') as tm: |
| with self.assertRaises(InvalidRequest): |
| tm.objects.create(partition=1, cluster=1) |
| |
| |
| class ManagementConnectionTests(BaseCassEngTestCase): |
| |
| keyspaces = ['ks1', 'ks2'] |
| conns = ['cluster'] |
| |
| @classmethod |
| def setUpClass(cls): |
| super(ManagementConnectionTests, cls).setUpClass() |
| conn.unregister_connection('default') |
| conn.register_connection('fake_cluster', ['127.0.0.100'], lazy_connect=True, retry_connect=True, default=True) |
| conn.register_connection('cluster', [CASSANDRA_IP]) |
| |
| @classmethod |
| def tearDownClass(cls): |
| super(ManagementConnectionTests, cls).tearDownClass() |
| |
| # reset the default connection |
| conn.unregister_connection('fake_cluster') |
| conn.unregister_connection('cluster') |
| setup_connection(DEFAULT_KEYSPACE) |
| |
| def setUp(self): |
| super(BaseCassEngTestCase, self).setUp() |
| |
| def test_create_drop_keyspace(self): |
| """ |
| Tests drop and create keyspace with connections explicitly set |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result keyspaces should be created and dropped |
| |
| @test_category object_mapper |
| """ |
| |
| # No connection (default is fake) |
| with self.assertRaises(NoHostAvailable): |
| create_keyspace_simple(self.keyspaces[0], 1) |
| |
| # Explicit connections |
| for ks in self.keyspaces: |
| create_keyspace_simple(ks, 1, connections=self.conns) |
| |
| for ks in self.keyspaces: |
| drop_keyspace(ks, connections=self.conns) |
| |
| def test_create_drop_table(self): |
| """ |
| Tests drop and create Table with connections explicitly set |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result Tables should be created and dropped |
| |
| @test_category object_mapper |
| """ |
| for ks in self.keyspaces: |
| create_keyspace_simple(ks, 1, connections=self.conns) |
| |
| # No connection (default is fake) |
| with self.assertRaises(NoHostAvailable): |
| sync_table(TestModel) |
| |
| # Explicit connections |
| sync_table(TestModel, connections=self.conns) |
| |
| # Explicit drop |
| drop_table(TestModel, connections=self.conns) |
| |
| # Model connection |
| TestModel.__connection__ = 'cluster' |
| sync_table(TestModel) |
| TestModel.__connection__ = None |
| |
| # No connection (default is fake) |
| with self.assertRaises(NoHostAvailable): |
| drop_table(TestModel) |
| |
| # Model connection |
| TestModel.__connection__ = 'cluster' |
| drop_table(TestModel) |
| TestModel.__connection__ = None |
| |
| # Model connection |
| for ks in self.keyspaces: |
| drop_keyspace(ks, connections=self.conns) |
| |
| def test_connection_creation_from_session(self): |
| """ |
| Test to ensure that you can register a connection from a session |
| @since 3.8 |
| @jira_ticket PYTHON-649 |
| @expected_result queries should execute appropriately |
| |
| @test_category object_mapper |
| """ |
| cluster = TestCluster() |
| session = cluster.connect() |
| connection_name = 'from_session' |
| conn.register_connection(connection_name, session=session) |
| self.assertIsNotNone(conn.get_connection(connection_name).cluster.metadata.get_host(CASSANDRA_IP)) |
| self.addCleanup(conn.unregister_connection, connection_name) |
| cluster.shutdown() |
| |
| def test_connection_from_hosts(self): |
| """ |
| Test to ensure that you can register a connection from a list of hosts |
| @since 3.8 |
| @jira_ticket PYTHON-692 |
| @expected_result queries should execute appropriately |
| |
| @test_category object_mapper |
| """ |
| connection_name = 'from_hosts' |
| conn.register_connection(connection_name, hosts=[CASSANDRA_IP]) |
| self.assertIsNotNone(conn.get_connection(connection_name).cluster.metadata.get_host(CASSANDRA_IP)) |
| self.addCleanup(conn.unregister_connection, connection_name) |
| |
| def test_connection_param_validation(self): |
| """ |
| Test to validate that invalid parameter combinations for registering connections via session are not tolerated |
| @since 3.8 |
| @jira_ticket PYTHON-649 |
| @expected_result queries should execute appropriately |
| |
| @test_category object_mapper |
| """ |
| cluster = TestCluster() |
| session = cluster.connect() |
| with self.assertRaises(CQLEngineException): |
| conn.register_connection("bad_coonection1", session=session, consistency="not_null") |
| with self.assertRaises(CQLEngineException): |
| conn.register_connection("bad_coonection2", session=session, lazy_connect="not_null") |
| with self.assertRaises(CQLEngineException): |
| conn.register_connection("bad_coonection3", session=session, retry_connect="not_null") |
| with self.assertRaises(CQLEngineException): |
| conn.register_connection("bad_coonection4", session=session, cluster_options="not_null") |
| with self.assertRaises(CQLEngineException): |
| conn.register_connection("bad_coonection5", hosts="not_null", session=session) |
| cluster.shutdown() |
| |
| cluster.shutdown() |
| |
| |
| cluster.shutdown() |
| |
| class BatchQueryConnectionTests(BaseCassEngTestCase): |
| |
| conns = ['cluster'] |
| |
| @classmethod |
| def setUpClass(cls): |
| super(BatchQueryConnectionTests, cls).setUpClass() |
| |
| create_keyspace_simple('ks1', 1) |
| sync_table(TestModel) |
| sync_table(AnotherTestModel) |
| |
| conn.unregister_connection('default') |
| conn.register_connection('fake_cluster', ['127.0.0.100'], lazy_connect=True, retry_connect=True, default=True) |
| conn.register_connection('cluster', [CASSANDRA_IP]) |
| |
| @classmethod |
| def tearDownClass(cls): |
| super(BatchQueryConnectionTests, cls).tearDownClass() |
| |
| # reset the default connection |
| conn.unregister_connection('fake_cluster') |
| conn.unregister_connection('cluster') |
| setup_connection(DEFAULT_KEYSPACE) |
| |
| drop_keyspace('ks1') |
| |
| def setUp(self): |
| super(BaseCassEngTestCase, self).setUp() |
| |
| def test_basic_batch_query(self): |
| """ |
| Test Batch queries with connections explicitly set |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result queries should execute appropriately |
| |
| @test_category object_mapper |
| """ |
| |
| # No connection with a QuerySet (default is a fake one) |
| with self.assertRaises(NoHostAvailable): |
| with BatchQuery() as b: |
| TestModel.objects.batch(b).create(partition=1, cluster=1) |
| |
| # Explicit connection with a QuerySet |
| with BatchQuery(connection='cluster') as b: |
| TestModel.objects.batch(b).create(partition=1, cluster=1) |
| |
| # Get an object from the BD |
| with ContextQuery(TestModel, connection='cluster') as tm: |
| obj = tm.objects.get(partition=1, cluster=1) |
| obj.__connection__ = None |
| |
| # No connection with a model (default is a fake one) |
| with self.assertRaises(NoHostAvailable): |
| with BatchQuery() as b: |
| obj.count = 2 |
| obj.batch(b).save() |
| |
| # Explicit connection with a model |
| with BatchQuery(connection='cluster') as b: |
| obj.count = 2 |
| obj.batch(b).save() |
| |
| def test_batch_query_different_connection(self): |
| """ |
| Test BatchQuery with Models that have a different connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result queries should execute appropriately |
| |
| @test_category object_mapper |
| """ |
| |
| # Testing on a model class |
| TestModel.__connection__ = 'cluster' |
| AnotherTestModel.__connection__ = 'cluster2' |
| |
| with self.assertRaises(CQLEngineException): |
| with BatchQuery() as b: |
| TestModel.objects.batch(b).create(partition=1, cluster=1) |
| AnotherTestModel.objects.batch(b).create(partition=1, cluster=1) |
| |
| TestModel.__connection__ = None |
| AnotherTestModel.__connection__ = None |
| |
| with BatchQuery(connection='cluster') as b: |
| TestModel.objects.batch(b).create(partition=1, cluster=1) |
| AnotherTestModel.objects.batch(b).create(partition=1, cluster=1) |
| |
| # Testing on a model instance |
| with ContextQuery(TestModel, AnotherTestModel, connection='cluster') as (tm, atm): |
| obj1 = tm.objects.get(partition=1, cluster=1) |
| obj2 = atm.objects.get(partition=1, cluster=1) |
| |
| obj1.__connection__ = 'cluster' |
| obj2.__connection__ = 'cluster2' |
| |
| obj1.count = 4 |
| obj2.count = 4 |
| |
| with self.assertRaises(CQLEngineException): |
| with BatchQuery() as b: |
| obj1.batch(b).save() |
| obj2.batch(b).save() |
| |
| def test_batch_query_connection_override(self): |
| """ |
| Test that we cannot override a BatchQuery connection per model |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result Proper exceptions should be raised |
| |
| @test_category object_mapper |
| """ |
| |
| with self.assertRaises(CQLEngineException): |
| with BatchQuery(connection='cluster') as b: |
| TestModel.batch(b).using(connection='test').save() |
| |
| with self.assertRaises(CQLEngineException): |
| with BatchQuery(connection='cluster') as b: |
| TestModel.using(connection='test').batch(b).save() |
| |
| with ContextQuery(TestModel, AnotherTestModel, connection='cluster') as (tm, atm): |
| obj1 = tm.objects.get(partition=1, cluster=1) |
| obj1.__connection__ = None |
| |
| with self.assertRaises(CQLEngineException): |
| with BatchQuery(connection='cluster') as b: |
| obj1.using(connection='test').batch(b).save() |
| |
| with self.assertRaises(CQLEngineException): |
| with BatchQuery(connection='cluster') as b: |
| obj1.batch(b).using(connection='test').save() |
| |
| class UsingDescriptorTests(BaseCassEngTestCase): |
| |
| conns = ['cluster'] |
| keyspaces = ['ks1', 'ks2'] |
| |
| @classmethod |
| def setUpClass(cls): |
| super(UsingDescriptorTests, cls).setUpClass() |
| |
| conn.unregister_connection('default') |
| conn.register_connection('fake_cluster', ['127.0.0.100'], lazy_connect=True, retry_connect=True, default=True) |
| conn.register_connection('cluster', [CASSANDRA_IP]) |
| |
| @classmethod |
| def tearDownClass(cls): |
| super(UsingDescriptorTests, cls).tearDownClass() |
| |
| # reset the default connection |
| conn.unregister_connection('fake_cluster') |
| conn.unregister_connection('cluster') |
| setup_connection(DEFAULT_KEYSPACE) |
| |
| for ks in cls.keyspaces: |
| drop_keyspace(ks) |
| |
| def setUp(self): |
| super(BaseCassEngTestCase, self).setUp() |
| |
| def _reset_data(self): |
| |
| for ks in self.keyspaces: |
| drop_keyspace(ks, connections=self.conns) |
| for ks in self.keyspaces: |
| create_keyspace_simple(ks, 1, connections=self.conns) |
| sync_table(TestModel, keyspaces=self.keyspaces, connections=self.conns) |
| |
| def test_keyspace(self): |
| """ |
| Test keyspace segregation when same connection is used |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result Keyspace segration is honored |
| |
| @test_category object_mapper |
| """ |
| self._reset_data() |
| |
| with ContextQuery(TestModel, connection='cluster') as tm: |
| |
| # keyspace Model class |
| tm.objects.using(keyspace='ks2').create(partition=1, cluster=1) |
| tm.objects.using(keyspace='ks2').create(partition=2, cluster=2) |
| |
| with self.assertRaises(TestModel.DoesNotExist): |
| tm.objects.get(partition=1, cluster=1) # default keyspace ks1 |
| obj1 = tm.objects.using(keyspace='ks2').get(partition=1, cluster=1) |
| |
| obj1.count = 2 |
| obj1.save() |
| |
| with self.assertRaises(NoHostAvailable): |
| TestModel.objects.using(keyspace='ks2').get(partition=1, cluster=1) |
| |
| obj2 = TestModel.objects.using(connection='cluster', keyspace='ks2').get(partition=1, cluster=1) |
| self.assertEqual(obj2.count, 2) |
| |
| # Update test |
| TestModel.objects(partition=2, cluster=2).using(connection='cluster', keyspace='ks2').update(count=5) |
| obj3 = TestModel.objects.using(connection='cluster', keyspace='ks2').get(partition=2, cluster=2) |
| self.assertEqual(obj3.count, 5) |
| |
| TestModel.objects(partition=2, cluster=2).using(connection='cluster', keyspace='ks2').delete() |
| with self.assertRaises(TestModel.DoesNotExist): |
| TestModel.objects.using(connection='cluster', keyspace='ks2').get(partition=2, cluster=2) |
| |
| def test_connection(self): |
| """ |
| Test basic connection functionality |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| self._reset_data() |
| |
| # Model class |
| with self.assertRaises(NoHostAvailable): |
| TestModel.objects.create(partition=1, cluster=1) |
| |
| TestModel.objects.using(connection='cluster').create(partition=1, cluster=1) |
| TestModel.objects(partition=1, cluster=1).using(connection='cluster').update(count=2) |
| obj1 = TestModel.objects.using(connection='cluster').get(partition=1, cluster=1) |
| self.assertEqual(obj1.count, 2) |
| |
| obj1.using(connection='cluster').update(count=5) |
| obj1 = TestModel.objects.using(connection='cluster').get(partition=1, cluster=1) |
| self.assertEqual(obj1.count, 5) |
| |
| obj1.using(connection='cluster').delete() |
| with self.assertRaises(TestModel.DoesNotExist): |
| TestModel.objects.using(connection='cluster').get(partition=1, cluster=1) |
| |
| |
| class ModelQuerySetNew(ModelQuerySet): |
| def __init__(self, *args, **kwargs): |
| super(ModelQuerySetNew, self).__init__(*args, **kwargs) |
| self._connection = "cluster" |
| |
| class BaseConnectionTestNoDefault(object): |
| conns = ['cluster'] |
| |
| @classmethod |
| def setUpClass(cls): |
| conn.register_connection('cluster', [CASSANDRA_IP]) |
| test_queryset.TestModel.__queryset__ = ModelQuerySetNew |
| test_queryset.IndexedTestModel.__queryset__ = ModelQuerySetNew |
| test_queryset.CustomIndexedTestModel.__queryset__ = ModelQuerySetNew |
| test_queryset.IndexedCollectionsTestModel.__queryset__ = ModelQuerySetNew |
| test_queryset.TestMultiClusteringModel.__queryset__ = ModelQuerySetNew |
| |
| super(BaseConnectionTestNoDefault, cls).setUpClass() |
| conn.unregister_connection('default') |
| |
| @classmethod |
| def tearDownClass(cls): |
| conn.unregister_connection('cluster') |
| setup_connection(DEFAULT_KEYSPACE) |
| super(BaseConnectionTestNoDefault, cls).tearDownClass() |
| # reset the default connection |
| |
| def setUp(self): |
| super(BaseCassEngTestCase, self).setUp() |
| |
| |
| class TestQuerySetOperationConnection(BaseConnectionTestNoDefault, test_queryset.TestQuerySetOperation): |
| """ |
| Execute test_queryset.TestQuerySetOperation using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestQuerySetDistinctNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetDistinct): |
| """ |
| Execute test_queryset.TestQuerySetDistinct using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestQuerySetOrderingNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetOrdering): |
| """ |
| Execute test_queryset.TestQuerySetOrdering using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestQuerySetCountSelectionAndIterationNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetCountSelectionAndIteration): |
| """ |
| Execute test_queryset.TestQuerySetOrdering using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestQuerySetSlicingNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetSlicing): |
| """ |
| Execute test_queryset.TestQuerySetOrdering using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestQuerySetValidationNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetValidation): |
| """ |
| Execute test_queryset.TestQuerySetOrdering using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestQuerySetDeleteNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetDelete): |
| """ |
| Execute test_queryset.TestQuerySetDelete using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestValuesListNoDefault(BaseConnectionTestNoDefault, test_queryset.TestValuesList): |
| """ |
| Execute test_queryset.TestValuesList using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
| |
| |
| class TestObjectsPropertyNoDefault(BaseConnectionTestNoDefault, test_queryset.TestObjectsProperty): |
| """ |
| Execute test_queryset.TestObjectsProperty using non default connection |
| |
| @since 3.7 |
| @jira_ticket PYTHON-613 |
| @expected_result proper connection should be used |
| |
| @test_category object_mapper |
| """ |
| pass |
