token-aware-policy: drop _tablets_routing_v1 flag

This flag was introduced to check if server supports tablets.
As result driver needs to sync it to policy when control connection is
established.
Which is an unwanted problem.
Let's relay on presence of tablets for given table instead, which will
not require any syncing.

Author: dkropachev

cassandra/cluster.py

@@ -1734,14 +1734,6 @@ def connect(self, keyspace=None, wait_for_all_pools=False):
                     self.shutdown()
                     raise
 
-                # Update the information about tablet support after connection handshake.
-                self.load_balancing_policy._tablets_routing_v1 = self.control_connection._tablets_routing_v1
-                child_policy = self.load_balancing_policy.child_policy if hasattr(self.load_balancing_policy, 'child_policy') else None
-                while child_policy is not None:
-                    if hasattr(child_policy, '_tablet_routing_v1'):
-                        child_policy._tablet_routing_v1 = self.control_connection._tablets_routing_v1
-                    child_policy = child_policy.child_policy if hasattr(child_policy, 'child_policy') else None
-
                 self.profile_manager.check_supported()  # todo: rename this method
 
                 if self.idle_heartbeat_interval:

cassandra/policies.py

@@ -476,7 +476,6 @@ class TokenAwarePolicy(LoadBalancingPolicy):
 
     _child_policy = None
     _cluster_metadata = None
-    _tablets_routing_v1 = False
     shuffle_replicas = False
     """
     Yield local replicas in a random order.
@@ -488,7 +487,6 @@ def __init__(self, child_policy, shuffle_replicas=False):
 
     def populate(self, cluster, hosts):
         self._cluster_metadata = cluster.metadata
-        self._tablets_routing_v1 = cluster.control_connection._tablets_routing_v1
         self._child_policy.populate(cluster, hosts)
 
     def check_supported(self):
@@ -513,17 +511,16 @@ def make_query_plan(self, working_keyspace=None, query=None):
             return
 
         replicas = []
-        if self._tablets_routing_v1:
+        if self._cluster_metadata._tablets.table_has_tablets(keyspace, query.table):
             tablet = self._cluster_metadata._tablets.get_tablet_for_key(
-                keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(query.routing_key))
+            keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(query.routing_key))
 
             if tablet is not None:
                 replicas_mapped = set(map(lambda r: r[0], tablet.replicas))
                 child_plan = child.make_query_plan(keyspace, query)
 
                 replicas = [host for host in child_plan if host.host_id in replicas_mapped]
-
-        if not replicas:
+        else:
             replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key)
 
         if self.shuffle_replicas:

cassandra/tablets.py

@@ -49,6 +49,9 @@ def __init__(self, tablets):
         self._tablets = tablets
         self._lock = Lock()
 
+    def table_has_tablets(self, keyspace, table) -> bool:
+        return bool(self._tablets.get((keyspace, table), []))
+
     def get_tablet_for_key(self, keyspace, table, t):
         tablet = self._tablets.get((keyspace, table), [])
         if not tablet:

tests/integration/standard/test_tablets.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from cassandra.cluster import Cluster
+from cassandra.cluster import Cluster, EXEC_PROFILE_DEFAULT, ExecutionProfile
 from cassandra.policies import ConstantReconnectionPolicy, RoundRobinPolicy, TokenAwarePolicy
 
 from tests.integration import PROTOCOL_VERSION, use_cluster, get_cluster
@@ -163,6 +163,36 @@ def test_tablets_shard_awareness(self):
         self.query_data_shard_select(self.session)
         self.query_data_shard_insert(self.session)
 
+    def test_tablets_lbp_in_profile(self):
+        cluster = Cluster(contact_points=["127.0.0.1", "127.0.0.2", "127.0.0.3"], protocol_version=PROTOCOL_VERSION,
+                          execution_profiles={
+                              EXEC_PROFILE_DEFAULT: ExecutionProfile(
+                                  load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
+                              )},
+                          reconnection_policy=ConstantReconnectionPolicy(1))
+        session = cluster.connect()
+        try:
+            self.query_data_host_select(self.session)
+            self.query_data_host_insert(self.session)
+        finally:
+            session.shutdown()
+            cluster.shutdown()
+
+    def test_tablets_shard_awareness_lbp_in_profile(self):
+        cluster = Cluster(contact_points=["127.0.0.1", "127.0.0.2", "127.0.0.3"], protocol_version=PROTOCOL_VERSION,
+                          execution_profiles={
+                              EXEC_PROFILE_DEFAULT: ExecutionProfile(
+                                  load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
+                              )},
+                          reconnection_policy=ConstantReconnectionPolicy(1))
+        session = cluster.connect()
+        try:
+            self.query_data_shard_select(self.session)
+            self.query_data_shard_insert(self.session)
+        finally:
+            session.shutdown()
+            cluster.shutdown()
+
     def test_tablets_invalidation_drop_ks_while_reconnecting(self):
         def recreate_while_reconnecting(_):
             # Kill control connection

tests/unit/test_policies.py

@@ -582,7 +582,6 @@ class TokenAwarePolicyTest(unittest.TestCase):
     def test_wrap_round_robin(self):
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
-        cluster.control_connection._tablets_routing_v1 = False
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)]
         for host in hosts:
             host.set_up()
@@ -614,7 +613,6 @@ def get_replicas(keyspace, packed_key):
     def test_wrap_dc_aware(self):
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
-        cluster.control_connection._tablets_routing_v1 = False
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)]
         for host in hosts:
             host.set_up()
@@ -744,7 +742,6 @@ def test_statement_keyspace(self):
 
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
-        cluster.control_connection._tablets_routing_v1 = False
         replicas = hosts[2:]
         cluster.metadata.get_replicas.return_value = replicas
 
@@ -835,7 +832,6 @@ def _assert_shuffle(self, patched_shuffle, keyspace, routing_key):
 
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
-        cluster.control_connection._tablets_routing_v1 = False
         replicas = hosts[2:]
         cluster.metadata.get_replicas.return_value = replicas
 
@@ -1538,7 +1534,6 @@ def test_query_plan_deferred_to_child(self):
 
     def test_wrap_token_aware(self):
         cluster = Mock(spec=Cluster)
-        cluster.control_connection._tablets_routing_v1 = False
         hosts = [Host(DefaultEndPoint("127.0.0.{}".format(i)), SimpleConvictionPolicy) for i in range(1, 6)]
         for host in hosts:
             host.set_up()