Make schema agreement waiting code renew connection on each iteration

When schema agreement is started it could happen that control connection
is getting disconnected/reconnected, when it happens schema agreement
code used to use disconnected connection to run all the queries.
As result, it could lead to schema agreement timeout, even if all nodes
got schema updated long time ago.

This commit updates connection on every iteration and makes it iterate
when underlying connection is clused

Author: dkropachev

cassandra/cluster.py

@@ -4220,15 +4220,14 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
             if self._is_shutdown:
                 return
 
-            if not connection:
-                connection = self._connection
+            current_connection = connection or self._connection
 
             if preloaded_results:
                 log.debug("[control connection] Attempting to use preloaded results for schema agreement")
 
                 peers_result = preloaded_results[0]
                 local_result = preloaded_results[1]
-                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, connection.endpoint)
+                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, current_connection.endpoint)
                 if schema_mismatches is None:
                     return True
 
@@ -4237,16 +4236,18 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
             elapsed = 0
             cl = ConsistencyLevel.ONE
             schema_mismatches = None
-            select_peers_query = self._get_peers_query(self.PeersQueryType.PEERS_SCHEMA, connection)
+            select_peers_query = self._get_peers_query(self.PeersQueryType.PEERS_SCHEMA, current_connection)
 
             while elapsed < total_timeout:
+                current_connection = connection or self._connection
+
                 peers_query = QueryMessage(query=maybe_add_timeout_to_query(select_peers_query, self._metadata_request_timeout),
                                            consistency_level=cl)
                 local_query = QueryMessage(query=maybe_add_timeout_to_query(self._SELECT_SCHEMA_LOCAL, self._metadata_request_timeout),
                                            consistency_level=cl)
                 try:
                     timeout = min(self._timeout, total_timeout - elapsed)
-                    peers_result, local_result = connection.wait_for_responses(
+                    peers_result, local_result = current_connection.wait_for_responses(
                         peers_query, local_query, timeout=timeout)
                 except OperationTimedOut as timeout:
                     log.debug("[control connection] Timed out waiting for "
@@ -4258,9 +4259,9 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
                         log.debug("[control connection] Aborting wait for schema match due to shutdown")
                         return None
                     else:
-                        raise
+                        continue
 
-                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, connection.endpoint)
+                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, current_connection.endpoint)
                 if schema_mismatches is None:
                     return True
 
@@ -4269,7 +4270,7 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
                 elapsed = self._time.time() - start
 
             log.warning("Node %s is reporting a schema disagreement: %s",
-                        connection.endpoint, schema_mismatches)
+                        current_connection.endpoint, schema_mismatches)
             return False
 
     def _get_schema_mismatches(self, peers_result, local_result, local_address):