shard_aware: adding shard_aware_options to Cluster options

In some cases users don't want the automatic opening of so many
connections (num of shard * num of nodes), this is adding
a new Cluster parameter that can disable shard awareness

```python
cluster = Cluster(contact_points=["127.0.0.1"], shard_aware_options=dict(disable=True),
                  load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()))
```

Author: fruch

cassandra/cluster.py

@@ -553,6 +553,20 @@ def default(self):
 """
 
 
+class ShardAwareOptions:
+    disable = None
+    disable_shardaware_port = False
+
+    def __init__(self, opts=None, disable=None, disable_shardaware_port=None):
+        self.disable = disable
+        self.disable_shardaware_port = disable_shardaware_port
+        if opts:
+            if isinstance(opts, ShardAwareOptions):
+                self.__dict__.update(opts.__dict__)
+            elif isinstance(opts, dict):
+                self.__dict__.update(opts)
+
+
 class _ConfigMode(object):
     UNCOMMITTED = 0
     LEGACY = 1
@@ -1003,6 +1017,12 @@ def default_retry_policy(self, policy):
     load the configuration and certificates.
     """
 
+    shard_aware_options = None
+    """
+    Can be set with :class:`ShardAwareOptions` or with a dict, to disable the automatic shardaware,
+    or to disable the shardaware port (advanced shardaware)
+    """
+
     @property
     def schema_metadata_enabled(self):
         """
@@ -1104,7 +1124,8 @@ def __init__(self,
                  monitor_reporting_enabled=True,
                  monitor_reporting_interval=30,
                  client_id=None,
-                 cloud=None):
+                 cloud=None,
+                 shard_aware_options=None):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
         extablishing connection pools or refreshing metadata.
@@ -1304,6 +1325,7 @@ def __init__(self,
         self.reprepare_on_up = reprepare_on_up
         self.monitor_reporting_enabled = monitor_reporting_enabled
         self.monitor_reporting_interval = monitor_reporting_interval
+        self.shard_aware_options = ShardAwareOptions(opts=shard_aware_options)
 
         self._listeners = set()
         self._listener_lock = Lock()

cassandra/pool.py

@@ -430,8 +430,7 @@ def __init__(self, host, host_distance, session):
 
         if self._keyspace:
             first_connection.set_keyspace_blocking(self._keyspace)
-
-        if first_connection.sharding_info:
+        if first_connection.sharding_info and not self._session.cluster.shard_aware_options.disable:
             self.host.sharding_info = first_connection.sharding_info
             self._open_connections_for_all_shards(first_connection.shard_id)
 
@@ -446,7 +445,7 @@ def _get_connection_for_routing_key(self, routing_key=None):
             raise NoConnectionsAvailable()
 
         shard_id = None
-        if self.host.sharding_info and routing_key:
+        if not self._session.cluster.shard_aware_options.disable and self.host.sharding_info and routing_key:
             t = self._session.cluster.metadata.token_map.token_class.from_key(routing_key)
             shard_id = self.host.sharding_info.shard_id_from_token(t.value)
 
@@ -585,7 +584,7 @@ def _replace(self, connection):
             try:
                 if connection.shard_id in self._connections.keys():
                     del self._connections[connection.shard_id]
-                if self.host.sharding_info:
+                if self.host.sharding_info and not self._session.cluster.shard_aware_options.disable:
                     self._connecting.add(connection.shard_id)
                     self._session.submit(self._open_connection_to_missing_shard, connection.shard_id)
                 else:
@@ -652,7 +651,8 @@ def disable_advanced_shard_aware(self, secs):
         self.advanced_shardaware_block_until = max(time.time() + secs, self.advanced_shardaware_block_until)
 
     def _get_shard_aware_endpoint(self):
-        if self.advanced_shardaware_block_until and self.advanced_shardaware_block_until < time.time():
+        if (self.advanced_shardaware_block_until and self.advanced_shardaware_block_until < time.time()) or \
+           self._session.cluster.shard_aware_options.disable_shardaware_port:
             return None
 
         endpoint = None
@@ -820,7 +820,7 @@ def _open_connections_for_all_shards(self, skip_shard_id=None):
                 return
 
             for shard_id in range(self.host.sharding_info.shards_count):
-                if skip_shard_id and skip_shard_id == shard_id:
+                if skip_shard_id is not None and skip_shard_id == shard_id:
                     continue
                 future = self._session.submit(self._open_connection_to_missing_shard, shard_id)
                 if isinstance(future, Future):

docs/scylla_specific.rst

@@ -26,6 +26,25 @@ https://github.com/scylladb/scylla/blob/master/docs/design-notes/protocols.md#cq
 New Cluster Helpers
 -------------------
 
+* ``shard_aware_options``
+
+  Setting it to ``dict(disable=True)`` would disable the shard aware functionally, for cases favoring once connection per host (example, lots of processes connecting from one client host, generating a big load of connections
+
+  Other option is to configure scylla by setting ``enable_shard_aware_drivers: false`` on scylla.yaml.
+
+.. code:: python
+
+    from cassandra.cluster import Cluster
+
+    cluster = Cluster(shard_aware_options=dict(disable=True))
+    session = cluster.connect()
+
+    assert not cluster.is_shard_aware(), "Shard aware should be disabled"
+
+    # or just disable the shard aware port logic
+    cluster = Cluster(shard_aware_options=dict(disable_shardaware_port=True))
+    session = cluster.connect()
+
 * ``cluster.is_shard_aware()``
 
   New method available on ``Cluster`` allowing to check whether the remote cluster supports shard awareness (bool)

tests/unit/test_control_connection.py

@@ -100,9 +100,6 @@ def on_up(self, host):
     def on_down(self, host, is_host_addition):
         self.down_host = host
 
-    def get_control_connection_host(self):
-        return self.added_hosts[0] if  self.added_hosts else None
-
 
 def _node_meta_results(local_results, peer_results):
     """

tests/unit/test_host_connection_pool.py

@@ -26,7 +26,7 @@
 from mock import Mock, NonCallableMagicMock, MagicMock
 from threading import Thread, Event, Lock
 
-from cassandra.cluster import Session
+from cassandra.cluster import Session, ShardAwareOptions
 from cassandra.connection import Connection
 from cassandra.pool import HostConnection, HostConnectionPool
 from cassandra.pool import Host, NoConnectionsAvailable
@@ -160,6 +160,7 @@ def test_return_defunct_connection_on_down_host(self):
         conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False,
                                     max_request_id=100, signaled_error=False)
         session.cluster.connection_factory.return_value = conn
+        session.cluster.shard_aware_options = ShardAwareOptions()
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
         session.cluster.connection_factory.assert_called_once_with(host.endpoint, owning_pool=pool)

tests/unit/test_shard_aware.py

@@ -21,6 +21,7 @@
 from unittest.mock import MagicMock
 from futures.thread import ThreadPoolExecutor
 
+from cassandra.cluster import ShardAwareOptions
 from cassandra.pool import HostConnection, HostDistance
 from cassandra.connection import ShardingInfo, DefaultEndPoint
 from cassandra.metadata import Murmur3Token
@@ -67,10 +68,11 @@ def __init__(self, is_ssl=False, *args, **kwargs):
                     self.cluster.ssl_options = {'some_ssl_options': True}
                 else:
                     self.cluster.ssl_options = None
+                self.cluster.shard_aware_options = ShardAwareOptions()
                 self.cluster.executor = ThreadPoolExecutor(max_workers=2)
                 self.cluster.signal_connection_failure = lambda *args, **kwargs: False
                 self.cluster.connection_factory = self.mock_connection_factory
-                self.connection_counter = -1
+                self.connection_counter = 0
                 self.futures = []
 
             def submit(self, fn, *args, **kwargs):
@@ -87,7 +89,8 @@ def mock_connection_factory(self, *args, **kwargs):
                 connection.is_closed = False
                 connection.orphaned_threshold_reached = False
                 connection.endpoint = args[0]
-                connection.shard_id = kwargs.get('shard_id', 0)
+                connection.shard_id = kwargs.get('shard_id', self.connection_counter)
+                self.connection_counter += 1
                 connection.sharding_info = ShardingInfo(shard_id=1, shards_count=4,
                                                          partitioner="", sharding_algorithm="", sharding_ignore_msb=0,
                                                          shard_aware_port=19042, shard_aware_port_ssl=19045)