From b3027786239ec1c397aa7b4e06a5b157f8f07f44 Mon Sep 17 00:00:00 2001
From: Mike Okner <mike@mikeokner.com>
Date: Fri, 12 Jun 2015 14:30:47 -0500
Subject: [PATCH] Set models.DEFAULT_KEYSPACE when calling set_session()

Sets the default keyspace for cqlengine models when passing in an
existing session to cqlengine.connection.set_session() if the session's
keyspace is set and models.DEFAULT_KEYSPACE isn't. Without this, the
user will see an error when trying to perform Cassandra interactions
with a Model instance like "Keyspace none does not exist" unless they
explicitly set models.DEFAULT_KEYSPACE themselves ahead of time or have
a __keyspace__ attribute set on the model.
Also no longer raising an exception in cqlengine.connection.setup() if
default_keyspace isn't set since individual models might still have
__keyspace__ set.
---
 cassandra/cqlengine/connection.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/cassandra/cqlengine/connection.py b/cassandra/cqlengine/connection.py
index ead241c6..c6540d49 100644
--- a/cassandra/cqlengine/connection.py
+++ b/cassandra/cqlengine/connection.py
@@ -82,6 +82,11 @@ def set_session(s):
     session = s
     cluster = s.cluster
 
+    # Set default keyspace from given session's keyspace if not already set
+    from cassandra.cqlengine import models
+    if not models.DEFAULT_KEYSPACE and session.keyspace:
+        models.DEFAULT_KEYSPACE = session.keyspace
+
     _register_known_types(cluster)
 
     log.debug("cqlengine connection initialized with %s", s)
@@ -109,9 +114,6 @@ def setup(
     if 'username' in kwargs or 'password' in kwargs:
         raise CQLEngineException("Username & Password are now handled by using the native driver's auth_provider")
 
-    if not default_keyspace:
-        raise UndefinedKeyspaceException()
-
     from cassandra.cqlengine import models
     models.DEFAULT_KEYSPACE = default_keyspace