Merge pull request datastax#1046 from datastax/python-1009-take2

[PYTHON-1009] Support execution profiles in cqlengine

Author: TheRealFalcon

CHANGELOG.rst

@@ -1,3 +1,11 @@
+3.20.0
+======
+Unreleased
+
+Bug Fixes
+---------
+* Connection setup methods prevent using ExecutionProfile in cqlengine (PYTHON-1009)
+
 3.19.0
 ======
 August 26, 2019

cassandra/cqlengine/connection.py

@@ -17,7 +17,7 @@
 import six
 import threading
 
-from cassandra.cluster import Cluster, _NOT_SET, NoHostAvailable, UserTypeDoesNotExist, ConsistencyLevel
+from cassandra.cluster import Cluster, _ConfigMode, _NOT_SET, NoHostAvailable, UserTypeDoesNotExist, ConsistencyLevel
 from cassandra.query import SimpleStatement, dict_factory
 
 from cassandra.cqlengine import CQLEngineException
@@ -108,17 +108,21 @@ def setup(self):
                 self.lazy_connect = True
             raise
 
-        if self.consistency is not None:
-            self.session.default_consistency_level = self.consistency
-
         if DEFAULT_CONNECTION in _connections and _connections[DEFAULT_CONNECTION] == self:
             cluster = _connections[DEFAULT_CONNECTION].cluster
             session = _connections[DEFAULT_CONNECTION].session
 
         self.setup_session()
 
     def setup_session(self):
-        self.session.row_factory = dict_factory
+        if self.cluster._config_mode == _ConfigMode.PROFILES:
+            self.cluster.profile_manager.default.row_factory = dict_factory
+            if self.consistency is not None:
+                self.cluster.profile_manager.default.consistency_level = self.consistency
+        else:
+            self.session.row_factory = dict_factory
+            if self.consistency is not None:
+                self.session.default_consistency_level = self.consistency
         enc = self.session.encoder
         enc.mapping[tuple] = enc.cql_encode_tuple
         _register_known_types(self.session.cluster)
@@ -182,10 +186,7 @@ def register_connection(name, hosts=None, consistency=None, lazy_connect=False,
                 "Session configuration arguments and 'session' argument are mutually exclusive"
             )
         conn = Connection.from_session(name, session=session)
-        conn.setup_session()
     else:  # use hosts argument
-        if consistency is None:
-            consistency = ConsistencyLevel.LOCAL_ONE
         conn = Connection(
             name, hosts=hosts,
             consistency=consistency, lazy_connect=lazy_connect,
@@ -281,8 +282,12 @@ def set_session(s):
         if conn.session:
             log.warning("configuring new default session for cqlengine when one was already set")
 
-    if s.row_factory is not dict_factory:
-        raise CQLEngineException("Failed to initialize: 'Session.row_factory' must be 'dict_factory'.")
+    if not any([
+        s.cluster.profile_manager.default.row_factory is dict_factory and s.cluster._config_mode in [_ConfigMode.PROFILES, _ConfigMode.UNCOMMITTED],
+        s.row_factory is dict_factory and s.cluster._config_mode in [_ConfigMode.LEGACY, _ConfigMode.UNCOMMITTED],
+    ]):
+        raise CQLEngineException("Failed to initialize: row_factory must be 'dict_factory'")
+
     conn.session = s
     conn.cluster = s.cluster
 

tests/integration/cqlengine/connections/test_connection.py

@@ -18,18 +18,17 @@
     import unittest  # noqa
 
 
+from cassandra import ConsistencyLevel
 from cassandra.cqlengine.models import Model
-from cassandra.cqlengine import columns, connection
+from cassandra.cqlengine import columns, connection, models
 from cassandra.cqlengine.management import sync_table
-from cassandra.cluster import Cluster, _clusters_for_shutdown
+from cassandra.cluster import Cluster, ExecutionProfile, _clusters_for_shutdown, _ConfigMode, EXEC_PROFILE_DEFAULT
+from cassandra.policies import RoundRobinPolicy
 from cassandra.query import dict_factory
 
-from tests.integration import PROTOCOL_VERSION, execute_with_long_wait_retry, local
+from tests.integration import CASSANDRA_IP, PROTOCOL_VERSION, execute_with_long_wait_retry, local
 from tests.integration.cqlengine.base import BaseCassEngTestCase
 from tests.integration.cqlengine import DEFAULT_KEYSPACE, setup_connection
-from cassandra.cqlengine import models
-
-from mock import patch
 
 
 class TestConnectModel(Model):
@@ -127,3 +126,73 @@ def test_connection_session_switch(self):
         connection.set_session(self.session2)
         self.assertEqual(1, TestConnectModel.objects.count())
         self.assertEqual(TestConnectModel.objects.first(), TCM2)
+
+
+class ConnectionModel(Model):
+    key = columns.Integer(primary_key=True)
+    some_data = columns.Text()
+
+
+class ConnectionInitTest(unittest.TestCase):
+    def test_default_connection_uses_legacy(self):
+        connection.default()
+        conn = connection.get_connection()
+        self.assertEqual(conn.cluster._config_mode, _ConfigMode.LEGACY)
+
+    def test_connection_with_legacy_settings(self):
+        connection.setup(
+            hosts=[CASSANDRA_IP],
+            default_keyspace=DEFAULT_KEYSPACE,
+            consistency=ConsistencyLevel.LOCAL_ONE
+        )
+        conn = connection.get_connection()
+        self.assertEqual(conn.cluster._config_mode, _ConfigMode.LEGACY)
+
+    def test_connection_from_session_with_execution_profile(self):
+        cluster = Cluster(execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=dict_factory)})
+        session = cluster.connect()
+        connection.default()
+        connection.set_session(session)
+        conn = connection.get_connection()
+        self.assertEqual(conn.cluster._config_mode, _ConfigMode.PROFILES)
+
+    def test_connection_from_session_with_legacy_settings(self):
+        cluster = Cluster(load_balancing_policy=RoundRobinPolicy())
+        session = cluster.connect()
+        session.row_factory = dict_factory
+        connection.set_session(session)
+        conn = connection.get_connection()
+        self.assertEqual(conn.cluster._config_mode, _ConfigMode.LEGACY)
+
+    def test_uncommitted_session_uses_legacy(self):
+        cluster = Cluster()
+        session = cluster.connect()
+        session.row_factory = dict_factory
+        connection.set_session(session)
+        conn = connection.get_connection()
+        self.assertEqual(conn.cluster._config_mode, _ConfigMode.LEGACY)
+
+    def test_legacy_insert_query(self):
+        connection.setup(
+            hosts=[CASSANDRA_IP],
+            default_keyspace=DEFAULT_KEYSPACE,
+            consistency=ConsistencyLevel.LOCAL_ONE
+        )
+        self.assertEqual(connection.get_connection().cluster._config_mode, _ConfigMode.LEGACY)
+
+        sync_table(ConnectionModel)
+        ConnectionModel.objects.create(key=0, some_data='text0')
+        ConnectionModel.objects.create(key=1, some_data='text1')
+        self.assertEqual(ConnectionModel.objects(key=0)[0].some_data, 'text0')
+
+    def test_execution_profile_insert_query(self):
+        cluster = Cluster(execution_profiles={EXEC_PROFILE_DEFAULT: ExecutionProfile(row_factory=dict_factory)})
+        session = cluster.connect()
+        connection.default()
+        connection.set_session(session)
+        self.assertEqual(connection.get_connection().cluster._config_mode, _ConfigMode.PROFILES)
+
+        sync_table(ConnectionModel)
+        ConnectionModel.objects.create(key=0, some_data='text0')
+        ConnectionModel.objects.create(key=1, some_data='text1')
+        self.assertEqual(ConnectionModel.objects(key=0)[0].some_data, 'text0')

tests/unit/cqlengine/test_connection.py

@@ -17,6 +17,7 @@
 except ImportError:
     import unittest  # noqa
 
+from cassandra.cluster import _ConfigMode
 from cassandra.cqlengine import connection
 from cassandra.query import dict_factory
 
@@ -38,9 +39,13 @@ def test_set_session_without_existing_connection(self):
         """
         Users can set the default session without having a default connection set.
         """
+        mock_cluster = Mock(
+            _config_mode=_ConfigMode.LEGACY,
+        )
         mock_session = Mock(
             row_factory=dict_factory,
-            encoder=Mock(mapping={})
+            encoder=Mock(mapping={}),
+            cluster=mock_cluster,
         )
         connection.set_session(mock_session)