connection support in ContextQuery

aboudreault · aboudreault · commit da10c634611f · 2016-08-25T16:22:57.000-04:00
diff --git a/cassandra/cqlengine/models.py b/cassandra/cqlengine/models.py
@@ -810,7 +810,7 @@ def _inst_batch(self, batch):
 
     @classmethod
     def _class_get_connection(cls):
-        return cls.__connection__
+        return cls._connection or cls.__connection__
 
     def _inst_get_connection(self):
         return self._connection or self.__connection__
diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py
@@ -287,7 +287,7 @@ class ContextQuery(object):
 
     """
 
-    def __init__(self, model, keyspace=None):
+    def __init__(self, model, keyspace=None, connection=None):
         """
         :param model: A model. This should be a class type, not an instance.
         :param keyspace: (optional) A keyspace name
@@ -304,10 +304,14 @@ def __init__(self, model, keyspace=None):
             ks = keyspace
             self.model = _copy_model_class(model, {'__keyspace__': ks})
 
+        if connection:
+            self.model._connection = connection
+
     def __enter__(self):
         return self.model
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        self.model._connection = None
         return
 
 
@@ -991,7 +995,7 @@ def using(self, keyspace=None, connection=None):
         clone = copy.deepcopy(self)
         if keyspace:
             from cassandra.cqlengine.models import _copy_model_class
-            clone.model = type(self.model, {'__keyspace__': keyspace})
+            clone.model = _copy_model_class(self.model, {'__keyspace__': keyspace})
 
         if connection:
             clone._connection = connection
