Merge pull request apache#596 from datastax/583

aboudreault · web-flow · commit f2df34076a70 · 2016-06-09T18:17:49.000-04:00
PYTHON-583: Add API to get the host metadata associated with the control connection node
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -1552,6 +1552,14 @@ def _target_type_from_refresh_args(keyspace, table, usertype, function, aggregat
             return SchemaTargetType.KEYSPACE
         return None
 
+    def get_control_connection_host(self):
+        """
+        Returns the control connection host metadata.
+        """
+        connection = self.control_connection._connection
+        host = connection.host if connection else None
+        return self.metadata.get_host(host) if host else None
+
     def refresh_schema_metadata(self, max_schema_agreement_wait=None):
         """
         Synchronously refresh all schema metadata.
@@ -3185,7 +3193,7 @@ def _on_timeout(self):
             if self.is_schema_agreed:
                 errors = {self._current_host.address: "Client request timeout. See Session.execute[_async](timeout)"}
             else:
-                connection = getattr(self.session.cluster.control_connection, '_connection')
+                connection = self.session.cluster.control_connection._connection
                 host = connection.host if connection else 'unknown'
                 errors = {host: "Request timed out while waiting for schema agreement. See Session.execute[_async](timeout) and Cluster.max_schema_agreement_wait."}
 
diff --git a/docs/api/cassandra/cluster.rst b/docs/api/cassandra/cluster.rst
@@ -92,6 +92,8 @@
 
    .. automethod:: set_max_connections_per_host
 
+   .. automethod:: get_control_connection_host
+
    .. automethod:: refresh_schema_metadata
 
    .. automethod:: refresh_keyspace_metadata
diff --git a/tests/integration/standard/test_control_connection.py b/tests/integration/standard/test_control_connection.py
@@ -39,7 +39,6 @@ def setUp(self):
                 "Native protocol 3,0+ is required for UDTs using %r"
                 % (PROTOCOL_VERSION,))
         self.cluster = Cluster(protocol_version=PROTOCOL_VERSION)
-        self.session = self.cluster.connect()
 
     def tearDown(self):
         try:
@@ -65,6 +64,7 @@ def test_drop_keyspace(self):
         @test_category connection
         """
 
+        self.session = self.cluster.connect()
         self.session.execute("""
             CREATE KEYSPACE keyspacetodrop
             WITH replication = { 'class' : 'SimpleStrategy', 'replication_factor': '1' }
@@ -76,3 +76,18 @@ def test_drop_keyspace(self):
         self.session.execute("DROP KEYSPACE keyspacetodrop")
         cc_id_post_drop = id(self.cluster.control_connection._connection)
         self.assertEqual(cc_id_post_drop, cc_id_pre_drop)
+
+    def test_get_control_connection_host(self):
+        """
+        Test to validate Cluster.get_control_connection_host() metadata
+        """
+
+        host = self.cluster.get_control_connection_host()
+        self.assertEqual(host, None)
+
+        self.session = self.cluster.connect()
+        cc_host = self.cluster.control_connection._connection.host
+
+        host = self.cluster.get_control_connection_host()
+        self.assertEqual(host.address, cc_host)
+        self.assertEqual(host.is_up, True)
