Fix class model copy

aboudreault · aboudreault · commit a0f86f1cdcd6 · 2016-08-25T15:03:30.000-04:00
diff --git a/cassandra/cqlengine/models.py b/cassandra/cqlengine/models.py
@@ -29,6 +29,14 @@
 log = logging.getLogger(__name__)
 
 
+def _copy_model_class(model, attrs):
+    new_type = type(model.__name__, (model,), attrs)
+    new_type.__abstract__ = model.__abstract__
+    new_type.__discriminator_value__ = model.__discriminator_value__
+    new_type.__default_ttl__ = model.__default_ttl__
+    return new_type
+
+
 class ModelException(CQLEngineException):
     pass
 
diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py
@@ -300,9 +300,9 @@ def __init__(self, model, keyspace=None):
         self.model = model
 
         if keyspace:
+            from cassandra.cqlengine.models import _copy_model_class
             ks = keyspace
-            new_type = type(model.__name__, (model,), {'__keyspace__': ks})
-            self.model = new_type
+            self.model = _copy_model_class(model, {'__keyspace__': ks})
 
     def __enter__(self):
         return self.model
@@ -990,8 +990,8 @@ def using(self, keyspace=None, connection=None):
 
         clone = copy.deepcopy(self)
         if keyspace:
-            new_type = type(self.model.__name__, (self.model,), {'__keyspace__': keyspace})
-            clone.model = new_type
+            from cassandra.cqlengine.models import _copy_model_class
+            clone.model = type(self.model, {'__keyspace__': keyspace})
 
         if connection:
             clone._connection = connection
