Evict cached UDT class when new type is registered.

aholmberg · aholmberg · commit e4c4c675f6b1 · 2015-03-23T10:52:40.000-05:00
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -47,6 +47,7 @@
                        UnsupportedOperation, Unauthorized)
 from cassandra.connection import (ConnectionException, ConnectionShutdown,
                                   ConnectionHeartbeat)
+from cassandra.cqltypes import UserType
 from cassandra.encoder import Encoder
 from cassandra.protocol import (QueryMessage, ResultMessage,
                                 ErrorMessage, ReadTimeoutErrorMessage,
@@ -615,6 +616,7 @@ def __init__(self, street, zipcode):
         self._user_types[keyspace][user_type] = klass
         for session in self.sessions:
             session.user_type_registered(keyspace, user_type, klass)
+        UserType.evict_udt_class(keyspace, user_type)
 
     def get_min_requests_per_connection(self, host_distance):
         return self._min_requests_per_connection[host_distance]
diff --git a/cassandra/cqltypes.py b/cassandra/cqltypes.py
@@ -899,6 +899,15 @@ def make_udt_class(cls, keyspace, udt_name, names_and_types, mapped_class):
             cls._cache[(keyspace, udt_name)] = instance
             return instance
 
+    @classmethod
+    def evict_udt_class(cls, keyspace, udt_name):
+        if six.PY2 and isinstance(udt_name, unicode):
+            udt_name = udt_name.encode('utf-8')
+        try:
+            del cls._cache[(keyspace, udt_name)]
+        except KeyError:
+            pass
+
     @classmethod
     def apply_parameters(cls, subtypes, names):
         keyspace = subtypes[0]
