Don't set keyspace in HostConnection if it doesn't exit

beltran · beltran · commit f0daa87702c9 · 2017-10-19T14:31:04.000-04:00
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -42,7 +42,8 @@
 
 from cassandra import (ConsistencyLevel, AuthenticationFailed,
                        OperationTimedOut, UnsupportedOperation,
-                       SchemaTargetType, DriverException, ProtocolVersion)
+                       SchemaTargetType, DriverException, ProtocolVersion,
+                       InvalidRequest)
 from cassandra.connection import (ConnectionException, ConnectionShutdown,
                                   ConnectionHeartbeat, ProtocolVersionUnsupported)
 from cassandra.cqltypes import UserType
@@ -1275,13 +1276,16 @@ def _session_register_user_types(self, session):
             for udt_name, klass in six.iteritems(type_map):
                 session.user_type_registered(keyspace, udt_name, klass)
 
-    def _cleanup_failed_on_up_handling(self, host):
+    def _cleanup_failed_on_up_handling(self, host, set_keyspace_failed=False):
         self.profile_manager.on_down(host)
         self.control_connection.on_down(host)
+        keyspaces = []
         for session in tuple(self.sessions):
             session.remove_pool(host)
+            if set_keyspace_failed and session.keyspace:
+                keyspaces.append(session.keyspace)
 
-        self._start_reconnector(host, is_host_addition=False)
+        self._start_reconnector(host, is_host_addition=False, keyspaces=keyspaces)
 
     def _on_up_future_completed(self, host, futures, results, lock, finished_future):
         with lock:
@@ -1299,7 +1303,10 @@ def _on_up_future_completed(self, host, futures, results, lock, finished_future)
             # all futures have completed at this point
             for exc in [f for f in results if isinstance(f, Exception)]:
                 log.error("Unexpected failure while marking node %s up:", host, exc_info=exc)
-                self._cleanup_failed_on_up_handling(host)
+                set_keyspace_failed = False
+                if isinstance(exc, InvalidRequest):
+                    set_keyspace_failed = True
+                self._cleanup_failed_on_up_handling(host, set_keyspace_failed=set_keyspace_failed)
                 return
 
             if not all(results):
@@ -1392,7 +1399,7 @@ def on_up(self, host):
         # for testing purposes
         return futures
 
-    def _start_reconnector(self, host, is_host_addition):
+    def _start_reconnector(self, host, is_host_addition, keyspaces):
         if self.profile_manager.distance(host) == HostDistance.IGNORED:
             return
 
@@ -1405,7 +1412,7 @@ def _start_reconnector(self, host, is_host_addition):
 
         reconnector = _HostReconnectionHandler(
             host, conn_factory, is_host_addition, self.on_add, self.on_up,
-            self.scheduler, schedule, host.get_and_set_reconnection_handler,
+            keyspaces, self.scheduler, schedule, host.get_and_set_reconnection_handler,
             new_handler=None)
 
         old_reconnector = host.get_and_set_reconnection_handler(reconnector)
@@ -1453,7 +1460,7 @@ def on_down(self, host, is_host_addition, expect_host_to_be_down=False):
         for listener in self.listeners:
             listener.on_down(host)
 
-        self._start_reconnector(host, is_host_addition)
+        self._start_reconnector(host, is_host_addition, None)
 
     def on_add(self, host, refresh_nodes=True):
         if self.is_shutdown:
@@ -1765,7 +1772,12 @@ def _prepare_all_queries(self, host):
             else:
                 for keyspace, ks_statements in groupby(statements, lambda s: s.keyspace):
                     if keyspace is not None:
-                        connection.set_keyspace_blocking(keyspace)
+                        try:
+                            connection.set_keyspace_blocking(keyspace)
+                        except InvalidRequest:
+                            log.warning("Error trying to prepare statements on "
+                                            "host %s, keyspace % s doesn't exist", host, keyspace)
+                            continue
 
                     # prepare 10 statements at a time
                     ks_statements = list(ks_statements)
@@ -2012,10 +2024,10 @@ def __init__(self, cluster, hosts, keyspace=None):
                 self._initial_connect_futures.add(future)
 
         futures = wait_futures(self._initial_connect_futures, return_when=FIRST_COMPLETED)
-        while futures.not_done and not any(f.result() for f in futures.done):
+        while futures.not_done and not any(f.result() for f in futures.done if not f.exception()):
             futures = wait_futures(futures.not_done, return_when=FIRST_COMPLETED)
 
-        if not any(f.result() for f in self._initial_connect_futures):
+        if not any(f.result() for f in self._initial_connect_futures if not f.exception()):
             msg = "Unable to connect to any servers"
             if self.keyspace:
                 msg += " using keyspace '%s'" % self.keyspace
@@ -2391,15 +2403,15 @@ def run_add_or_renew_pool():
             except AuthenticationFailed as auth_exc:
                 conn_exc = ConnectionException(str(auth_exc), host=host)
                 self.cluster.signal_connection_failure(host, conn_exc, is_host_addition)
-                return False
+                raise
             except Exception as conn_exc:
                 log.warning("Failed to create connection pool for new host %s:",
                             host, exc_info=conn_exc)
                 # the host itself will still be marked down, so we need to pass
                 # a special flag to make sure the reconnector is created
                 self.cluster.signal_connection_failure(
                     host, conn_exc, is_host_addition, expect_host_to_be_down=True)
-                return False
+                raise
 
             previous = self._pools.get(host)
             with self._lock:
@@ -2411,15 +2423,18 @@ def run_add_or_renew_pool():
                     def callback(pool, errors):
                         errors_returned.extend(errors)
                         set_keyspace_event.set()
-
-                    new_pool._set_keyspace_for_all_conns(self.keyspace, callback)
+                    try:
+                        new_pool._set_keyspace_for_all_conns(self.keyspace, callback)
+                    except ConnectionShutdown:
+                        self._lock.acquire()
+                        raise
                     set_keyspace_event.wait(self.cluster.connect_timeout)
                     if not set_keyspace_event.is_set() or errors_returned:
                         log.warning("Failed setting keyspace for pool after keyspace changed during connect: %s", errors_returned)
                         self.cluster.on_down(host, is_host_addition)
                         new_pool.shutdown()
                         self._lock.acquire()
-                        return False
+                        raise errors_returned
                     self._lock.acquire()
                 self._pools[host] = new_pool
 
diff --git a/cassandra/connection.py b/cassandra/connection.py
@@ -37,7 +37,7 @@
 else:
     from six.moves.queue import Queue, Empty  # noqa
 
-from cassandra import ConsistencyLevel, AuthenticationFailed, OperationTimedOut, ProtocolVersion
+from cassandra import ConsistencyLevel, AuthenticationFailed, OperationTimedOut, InvalidRequest, ProtocolVersion
 from cassandra.marshal import int32_pack
 from cassandra.protocol import (ReadyMessage, AuthenticateMessage, OptionsMessage,
                                 StartupMessage, ErrorMessage, CredentialsMessage,
@@ -515,6 +515,8 @@ def wait_for_responses(self, *msgs, **kwargs):
             return waiter.deliver(timeout)
         except OperationTimedOut:
             raise
+        except InvalidRequest:
+            raise
         except Exception as exc:
             self.defunct(exc)
             raise
@@ -799,9 +801,9 @@ def set_keyspace_blocking(self, keyspace):
                              consistency_level=ConsistencyLevel.ONE)
         try:
             result = self.wait_for_response(query)
-        except InvalidRequestException as ire:
+        except InvalidRequest as ire:
             # the keyspace probably doesn't exist
-            raise ire.to_exception()
+            raise ire
         except Exception as exc:
             conn_exc = ConnectionException(
                 "Problem while setting keyspace: %r" % (exc,), self.host)
diff --git a/cassandra/pool.py b/cassandra/pool.py
@@ -27,7 +27,7 @@
 except ImportError:
     from cassandra.util import WeakSet  # NOQA
 
-from cassandra import AuthenticationFailed
+from cassandra import AuthenticationFailed, InvalidRequest
 from cassandra.connection import ConnectionException
 from cassandra.policies import HostDistance
 
@@ -267,16 +267,21 @@ def on_exception(self, exc, next_delay):
 
 class _HostReconnectionHandler(_ReconnectionHandler):
 
-    def __init__(self, host, connection_factory, is_host_addition, on_add, on_up, *args, **kwargs):
+    def __init__(self, host, connection_factory, is_host_addition, on_add, on_up, keyspaces, *args, **kwargs):
         _ReconnectionHandler.__init__(self, *args, **kwargs)
         self.is_host_addition = is_host_addition
         self.on_add = on_add
         self.on_up = on_up
         self.host = host
         self.connection_factory = connection_factory
+        self.keyspaces = keyspaces
 
     def try_reconnect(self):
-        return self.connection_factory()
+        conn = self.connection_factory()
+        if self.keyspaces:
+            for keyspace in self.keyspaces:
+                conn.set_keyspace_blocking(keyspace)
+        return conn
 
     def on_reconnection(self, connection):
         log.info("Successful reconnection to %s, marking node up if it isn't already", self.host)
