cqle: guard against formatting an invalid keyspace in a query

aholmberg · aholmberg · commit 85f32ec53284 · 2016-02-09T09:54:37.000-06:00
PYTHON-395

Resolves AttributeError when trying to query without a keyspace name,
instead raising a CQLEngineException.
diff --git a/cassandra/cqlengine/models.py b/cassandra/cqlengine/models.py
@@ -480,7 +480,10 @@ def column_family_name(cls, include_keyspace=True):
         """
         cf_name = protect_name(cls._raw_column_family_name())
         if include_keyspace:
-            return '{0}.{1}'.format(protect_name(cls._get_keyspace()), cf_name)
+            keyspace = cls._get_keyspace()
+            if not keyspace:
+                raise CQLEngineException("Model keyspace is not set and no default is available. Set model keyspace or setup connection before attempting to generate a query.")
+            return '{0}.{1}'.format(protect_name(keyspace), cf_name)
 
         return cf_name
 
diff --git a/tests/integration/cqlengine/model/test_model.py b/tests/integration/cqlengine/model/test_model.py
@@ -16,7 +16,9 @@
 except ImportError:
     import unittest  # noqa
 
-from cassandra.cqlengine import columns
+from mock import patch
+
+from cassandra.cqlengine import columns, CQLEngineException
 from cassandra.cqlengine.management import sync_table, drop_table, create_keyspace_simple, drop_keyspace
 from cassandra.cqlengine.models import Model, ModelDefinitionException
 
@@ -105,6 +107,25 @@ class table(Model):
 
         drop_keyspace('keyspace')
 
+    def test_column_family(self):
+        class TestModel(Model):
+            k = columns.Integer(primary_key=True)
+
+        from cassandra.cqlengine import models
+        # no model keyspace uses default
+        self.assertEqual(TestModel.column_family_name(), "%s.test_model" % (models.DEFAULT_KEYSPACE,))
+
+        # model keyspace overrides
+        TestModel.__keyspace__ = "my_test_keyspace"
+        self.assertEqual(TestModel.column_family_name(), "%s.test_model" % (TestModel.__keyspace__,))
+
+        # neither set should raise CQLEngineException before failing or formatting an invalid name
+        del TestModel.__keyspace__
+        with patch('cassandra.cqlengine.models.DEFAULT_KEYSPACE', None):
+            self.assertRaises(CQLEngineException, TestModel.column_family_name)
+            # .. but we can still get the bare CF name
+            self.assertEqual(TestModel.column_family_name(include_keyspace=False), "test_model")
+
 
 class BuiltInAttributeConflictTest(unittest.TestCase):
     """tests Model definitions that conflict with built-in attributes/methods"""
