diff --git a/cassandra/policies.py b/cassandra/policies.py index 691287745..a1495f366 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -364,46 +364,39 @@ def distance(self, *args, **kwargs): return self._child_policy.distance(*args, **kwargs) def make_query_plan(self, working_keyspace=None, query=None): - if query and query.keyspace: - keyspace = query.keyspace - else: - keyspace = working_keyspace + keyspace = query.keyspace if query and query.keyspace else working_keyspace child = self._child_policy - if query is None: + if query is None or query.routing_key is None or keyspace is None: for host in child.make_query_plan(keyspace, query): yield host - else: - routing_key = query.routing_key - if routing_key is None or keyspace is None: - for host in child.make_query_plan(keyspace, query): - yield host - else: - replicas = [] - if self._tablets_routing_v1: - tablet = self._cluster_metadata._tablets.get_tablet_for_key(keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(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 replicas == []: - replicas = self._cluster_metadata.get_replicas(keyspace, routing_key) - - if self.shuffle_replicas: - shuffle(replicas) - for replica in replicas: - if replica.is_up and \ - child.distance(replica) == HostDistance.LOCAL: - yield replica - - for host in child.make_query_plan(keyspace, query): - # skip if we've already listed this host - if host not in replicas or \ - child.distance(host) == HostDistance.REMOTE: - yield host + return + + replicas = [] + if self._tablets_routing_v1: + tablet = self._cluster_metadata._tablets.get_tablet_for_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: + replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key) + + if self.shuffle_replicas: + shuffle(replicas) + + for replica in replicas: + if replica.is_up and child.distance(replica) == HostDistance.LOCAL: + yield replica + + for host in child.make_query_plan(keyspace, query): + # skip if we've already listed this host + if host not in replicas or child.distance(host) == HostDistance.REMOTE: + yield host def on_up(self, *args, **kwargs): return self._child_policy.on_up(*args, **kwargs) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index aeba7fa8a..5e638d78c 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -1,6 +1,7 @@ # Experimental, this interface and use may change from threading import Lock + class Tablet(object): """ Represents a single ScyllaDB tablet. @@ -11,7 +12,7 @@ class Tablet(object): last_token = 0 replicas = None - def __init__(self, first_token = 0, last_token = 0, replicas = None): + def __init__(self, first_token=0, last_token=0, replicas=None): self.first_token = first_token self.last_token = last_token self.replicas = replicas @@ -28,10 +29,11 @@ def _is_valid_tablet(replicas): @staticmethod def from_row(first_token, last_token, replicas): if Tablet._is_valid_tablet(replicas): - tablet = Tablet(first_token, last_token,replicas) + tablet = Tablet(first_token, last_token, replicas) return tablet return None + # Experimental, this interface and use may change class Tablets(object): _lock = None @@ -43,10 +45,10 @@ def __init__(self, tablets): def get_tablet_for_key(self, keyspace, table, t): tablet = self._tablets.get((keyspace, table), []) - if tablet == []: + if not tablet: return None - id = bisect_left(tablet, t.value, key = lambda tablet: tablet.last_token) + id = bisect_left(tablet, t.value, key=lambda tablet: tablet.last_token) if id < len(tablet) and t.value > tablet[id].first_token: return tablet[id] return None @@ -55,13 +57,13 @@ def add_tablet(self, keyspace, table, tablet): with self._lock: tablets_for_table = self._tablets.setdefault((keyspace, table), []) - # find first overlaping range - start = bisect_left(tablets_for_table, tablet.first_token, key = lambda t: t.first_token) + # find first overlapping range + start = bisect_left(tablets_for_table, tablet.first_token, key=lambda t: t.first_token) if start > 0 and tablets_for_table[start - 1].last_token > tablet.first_token: start = start - 1 - # find last overlaping range - end = bisect_left(tablets_for_table, tablet.last_token, key = lambda t: t.last_token) + # find last overlapping range + end = bisect_left(tablets_for_table, tablet.last_token, key=lambda t: t.last_token) if end < len(tablets_for_table) and tablets_for_table[end].first_token >= tablet.last_token: end = end - 1 @@ -70,6 +72,7 @@ def add_tablet(self, keyspace, table, tablet): tablets_for_table.insert(start, tablet) + # bisect.bisect_left implementation from Python 3.11, needed untill support for # Python < 3.10 is dropped, it is needed to use `key` to extract last_token from # Tablet list - better solution performance-wise than materialize list of last_tokens @@ -97,11 +100,11 @@ def bisect_left(a, x, lo=0, hi=None, *, key=None): lo = mid + 1 else: hi = mid - else: - while lo < hi: - mid = (lo + hi) // 2 - if key(a[mid]) < x: - lo = mid + 1 - else: - hi = mid + return + while lo < hi: + mid = (lo + hi) // 2 + if key(a[mid]) < x: + lo = mid + 1 + else: + hi = mid return lo