From e38f8580d1b691ddd0c561d54c209abc221e58f3 Mon Sep 17 00:00:00 2001
From: Tyler Hobbs <tyler@datastax.com>
Date: Fri, 18 Jul 2014 12:36:54 -0500
Subject: [PATCH] Friendlier exception when user type does not exist

---
 cassandra/cluster.py                    | 19 ++++++++++++++++++-
 docs/api/cassandra/cluster.rst          |  2 ++
 tests/integration/standard/test_udts.py |  9 ++++++++-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/cassandra/cluster.py b/cassandra/cluster.py
index c8996c96..b0609061 100644
--- a/cassandra/cluster.py
+++ b/cassandra/cluster.py
@@ -1551,7 +1551,17 @@ class Session(object):
         mapping from a user-defined type to a class.  Intended for internal
         use only.
         """
-        type_meta = self.cluster.metadata.keyspaces[keyspace].user_types[user_type]
+        try:
+            ks_meta = self.cluster.metadata.keyspaces[keyspace]
+        except KeyError:
+            raise UserTypeDoesNotExist(
+                'Keyspace %s does not exist or has not been discovered by the driver' % (keyspace,))
+
+        try:
+            type_meta = ks_meta.user_types[user_type]
+        except KeyError:
+            raise UserTypeDoesNotExist(
+                'User type %s does not exist in keyspace %s' % (user_type, keyspace))
 
         def encode(val):
             return '{ %s }' % ' , '.join('%s : %s' % (
@@ -1570,6 +1580,13 @@ class Session(object):
         return dict((host, pool.get_state()) for host, pool in self._pools.items())
 
 
+class UserTypeDoesNotExist(Exception):
+    """
+    An attempt was made to use a user-defined type that does not exist.
+    """
+    pass
+
+
 class _ControlReconnectionHandler(_ReconnectionHandler):
     """
     Internal
diff --git a/docs/api/cassandra/cluster.rst b/docs/api/cassandra/cluster.rst
index 50fc36ea..da28eef0 100644
--- a/docs/api/cassandra/cluster.rst
+++ b/docs/api/cassandra/cluster.rst
@@ -104,3 +104,5 @@
 
 .. autoexception:: NoHostAvailable ()
    :members:
+
+.. autoexception:: UserTypeDoesNotExist ()
diff --git a/tests/integration/standard/test_udts.py b/tests/integration/standard/test_udts.py
index 4bf7f9ba..012a7fc7 100644
--- a/tests/integration/standard/test_udts.py
+++ b/tests/integration/standard/test_udts.py
@@ -22,7 +22,7 @@ log = logging.getLogger(__name__)
 
 from collections import namedtuple
 
-from cassandra.cluster import Cluster
+from cassandra.cluster import Cluster, UserTypeDoesNotExist
 
 from tests.integration import get_server_versions, PROTOCOL_VERSION
 
@@ -227,3 +227,10 @@ class TypeTests(unittest.TestCase):
         self.assertTrue(type(row.b) is User)
 
         c.shutdown()
+
+    def test_non_existing_types(self):
+        c = Cluster(protocol_version=PROTOCOL_VERSION)
+        c.connect()
+        User = namedtuple('user', ('age', 'name'))
+        self.assertRaises(UserTypeDoesNotExist, c.register_user_type, "some_bad_keyspace", "user", User)
+        self.assertRaises(UserTypeDoesNotExist, c.register_user_type, "system", "user", User)