register UDT with default keyspace if none specified

aholmberg · aholmberg · commit 84c3e08cf2bd · 2016-02-09T14:37:53.000-06:00
PYTHON-347
diff --git a/cassandra/cqlengine/connection.py b/cassandra/cqlengine/connection.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import namedtuple
+from collections import namedtuple, defaultdict
 import logging
 import six
 
@@ -37,7 +37,7 @@
 # Because type models may be registered before a connection is present,
 # and because sessions may be replaced, we must register UDTs here, in order
 # to have them registered when a new session is established.
-udt_by_keyspace = {}
+udt_by_keyspace = defaultdict(dict)
 
 
 class UndefinedKeyspaceException(CQLEngineException):
@@ -191,10 +191,7 @@ def handle_lazy_connect():
 
 
 def register_udt(keyspace, type_name, klass):
-    try:
-        udt_by_keyspace[keyspace][type_name] = klass
-    except KeyError:
-        udt_by_keyspace[keyspace] = {type_name: klass}
+    udt_by_keyspace[keyspace][type_name] = klass
 
     global cluster
     if cluster:
@@ -205,9 +202,10 @@ def register_udt(keyspace, type_name, klass):
 
 
 def _register_known_types(cluster):
+    from cassandra.cqlengine import models
     for ks_name, name_type_map in udt_by_keyspace.items():
         for type_name, klass in name_type_map.items():
             try:
-                cluster.register_user_type(ks_name, type_name, klass)
+                cluster.register_user_type(ks_name or models.DEFAULT_KEYSPACE, type_name, klass)
             except UserTypeDoesNotExist:
                 pass  # new types are covered in management sync functions
diff --git a/tests/integration/cqlengine/model/test_udts.py b/tests/integration/cqlengine/model/test_udts.py
@@ -18,6 +18,7 @@
 
 from datetime import datetime, date, time
 from decimal import Decimal
+from mock import Mock
 from uuid import UUID, uuid4
 
 from cassandra.cqlengine.models import Model
@@ -426,4 +427,21 @@ class UserModelText(Model):
         UserModelText.create(id=unicode_name, info=user_template_ascii)
         UserModelText.create(id=unicode_name, info=user_template_unicode)
 
+    def test_register_default_keyspace(self):
+        class User(UserType):
+            age = columns.Integer()
+            name = columns.Text()
+
+        from cassandra.cqlengine import models
+        from cassandra.cqlengine import connection
+
+        # None emulating no model and no default keyspace before connecting
+        connection.udt_by_keyspace.clear()
+        User.register_for_keyspace(None)
+        self.assertEqual(len(connection.udt_by_keyspace), 1)
+        self.assertIn(None, connection.udt_by_keyspace)
 
+        # register should be with default keyspace, not None
+        cluster = Mock()
+        connection._register_known_types(cluster)
+        cluster.register_user_type.assert_called_with(models.DEFAULT_KEYSPACE, User.type_name(), User)
