Merge pull request #249 from sylwiaszunejko/introduce_tablets

Introduce support for tablets

Author: constantfold

.github/workflows/integration-tests.yml

@@ -32,4 +32,11 @@ jobs:
  - name: Test with pytest
 run: |
  export EVENT_LOOP_MANAGER=${{ matrix.event_loop_manager }}
+ export SCYLLA_VERSION='release:5.1'
  ./ci/run_integration_test.sh tests/integration/standard/ tests/integration/cqlengine/
+
+ - name: Test tablets
+run: |
+ export EVENT_LOOP_MANAGER=${{ matrix.event_loop_manager }}
+ export SCYLLA_VERSION='unstable/master:2024-01-03T08:06:57Z'
+ ./ci/run_integration_test.sh tests/integration/experiments/

README.rst

@@ -26,6 +26,7 @@ Features
 * `Concurrent execution utilities <http://python-driver.docs.scylladb.com/stable/api/cassandra/concurrent.html>`_
 * `Object mapper <http://python-driver.docs.scylladb.com/stable/object-mapper.html>`_
 * `Shard awareness <http://python-driver.docs.scylladb.com/stable/scylla-specific.html#shard-awareness>`_
+* `Tablet awareness <http://python-driver.docs.scylladb.com/stable/scylla-specific.html#tablet-awareness>`_
 
 Installation
 ------------

cassandra/cluster.py

@@ -41,7 +41,7 @@
 importweakref
 fromweakrefimportWeakValueDictionary
 
-fromcassandraimport (ConsistencyLevel, AuthenticationFailed,
+fromcassandraimport (ConsistencyLevel, AuthenticationFailed,InvalidRequest,
 OperationTimedOut, UnsupportedOperation,
 SchemaTargetType, DriverException, ProtocolVersion,
 UnresolvableContactPoints)
@@ -51,6 +51,7 @@
 EndPoint, DefaultEndPoint, DefaultEndPointFactory,
 ContinuousPagingState, SniEndPointFactory, ConnectionBusy)
 fromcassandra.cqltypesimportUserType
+importcassandra.cqltypesastypes
 fromcassandra.encoderimportEncoder
 fromcassandra.protocolimport (QueryMessage, ResultMessage,
 ErrorMessage, ReadTimeoutErrorMessage,
@@ -79,6 +80,7 @@
 named_tuple_factory, dict_factory, tuple_factory, FETCH_SIZE_UNSET,
 HostTargetingStatement)
 fromcassandra.marshalimportint64_pack
+fromcassandra.tabletsimportTablet, Tablets
 fromcassandra.timestampsimportMonotonicTimestampGenerator
 fromcassandra.compatimportMapping
 fromcassandra.utilimport_resolve_contact_points_to_string_map, Version
@@ -1784,6 +1786,14 @@ 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_policyifhasattr(self.load_balancing_policy, 'child_policy') elseNone
+whilechild_policyisnotNone:
+ifhasattr(child_policy, '_tablet_routing_v1'):
+child_policy._tablet_routing_v1=self.control_connection._tablets_routing_v1
+child_policy=child_policy.child_policyifhasattr(child_policy, 'child_policy') elseNone
+
 self.profile_manager.check_supported() # todo: rename this method
 
 ifself.idle_heartbeat_interval:
@@ -2398,7 +2408,6 @@ def add_prepared(self, query_id, prepared_statement):
 withself._prepared_statement_lock:
 self._prepared_statements[query_id] =prepared_statement
 
-
 classSession(object):
 """
  A collection of connection pools for each host in the cluster.
@@ -3550,6 +3559,7 @@ class PeersQueryType(object):
 _schema_meta_page_size=1000
 
 _uses_peers_v2=True
+_tablets_routing_v1=False
 
 # for testing purposes
 _time=time
@@ -3683,6 +3693,8 @@ def _try_connect(self, host):
 # If sharding information is available, it's a ScyllaDB cluster, so do not use peers_v2 table.
 ifconnection.features.sharding_infoisnotNone:
 self._uses_peers_v2=False
+
+self._tablets_routing_v1=connection.features.tablets_routing_v1
 
 # use weak references in both directions
 # _clear_watcher will be called when this ControlConnection is about to be finalized
@@ -4609,7 +4621,10 @@ def _query(self, host, message=None, cb=None):
 connection=None
 try:
 # TODO get connectTimeout from cluster settings
-connection, request_id=pool.borrow_connection(timeout=2.0, routing_key=self.query.routing_keyifself.queryelseNone)
+ifself.query:
+connection, request_id=pool.borrow_connection(timeout=2.0, routing_key=self.query.routing_key, keyspace=self.query.keyspace, table=self.query.table)
+else:
+connection, request_id=pool.borrow_connection(timeout=2.0)
 self._connection=connection
 result_meta=self.prepared_statement.result_metadataifself.prepared_statementelse []
 
@@ -4728,6 +4743,19 @@ def _set_result(self, host, connection, pool, response):
 self._warnings=getattr(response, 'warnings', None)
 self._custom_payload=getattr(response, 'custom_payload', None)
 
+ifself._custom_payloadandself.session.cluster.control_connection._tablets_routing_v1and'tablets-routing-v1'inself._custom_payload:
+protocol=self.session.cluster.protocol_version
+info=self._custom_payload.get('tablets-routing-v1')
+ctype=types.lookup_casstype('TupleType(LongType, LongType, ListType(TupleType(UUIDType, Int32Type)))')
+tablet_routing_info=ctype.from_binary(info, protocol)
+first_token=tablet_routing_info[0]
+last_token=tablet_routing_info[1]
+tablet_replicas=tablet_routing_info[2]
+tablet=Tablet.from_row(first_token, last_token, tablet_replicas)
+keyspace=self.query.keyspace
+table=self.query.table
+self.session.cluster.metadata._tablets.add_tablet(keyspace, table, tablet)
+
 ifisinstance(response, ResultMessage):
 ifresponse.kind==RESULT_KIND_SET_KEYSPACE:
 session=getattr(self, 'session', None)

cassandra/metadata.py

@@ -44,6 +44,7 @@
 fromcassandra.poolimportHostDistance
 fromcassandra.connectionimportEndPoint
 fromcassandra.compatimportMapping
+fromcassandra.tabletsimportTablets
 
 log=logging.getLogger(__name__)
 
@@ -126,6 +127,7 @@ def __init__(self):
 self._hosts= {}
 self._host_id_by_endpoint= {}
 self._hosts_lock=RLock()
+self._tablets=Tablets({})
 
 defexport_schema_as_string(self):
 """

cassandra/policies.py

@@ -335,6 +335,7 @@ class TokenAwarePolicy(LoadBalancingPolicy):
 
 _child_policy=None
 _cluster_metadata=None
+_tablets_routing_v1=False
 shuffle_replicas=False
 """
  Yield local replicas in a random order.
@@ -346,6 +347,7 @@ def __init__(self, child_policy, shuffle_replicas=False):
 
 defpopulate(self, cluster, hosts):
 self._cluster_metadata=cluster.metadata
+self._tablets_routing_v1=cluster.control_connection._tablets_routing_v1
 self._child_policy.populate(cluster, hosts)
 
 defcheck_supported(self):
@@ -376,7 +378,19 @@ def make_query_plan(self, working_keyspace=None, query=None):
 forhostinchild.make_query_plan(keyspace, query):
 yieldhost
 else:
-replicas=self._cluster_metadata.get_replicas(keyspace, routing_key)
+replicas= []
+ifself._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))
+
+iftabletisnotNone:
+replicas_mapped=set(map(lambdar: r[0], tablet.replicas))
+child_plan=child.make_query_plan(keyspace, query)
+
+replicas= [hostforhostinchild_planifhost.host_idinreplicas_mapped]
+
+ifreplicas== []:
+replicas=self._cluster_metadata.get_replicas(keyspace, routing_key)
+
 ifself.shuffle_replicas:
 shuffle(replicas)
 forreplicainreplicas:

cassandra/pool.py

@@ -392,6 +392,8 @@ class HostConnection(object):
 # the number below, all excess connections will be closed.
 max_excess_connections_per_shard_multiplier=3
 
+tablets_routing_v1=False
+
 def__init__(self, host, host_distance, session):
 self.host=host
 self.host_distance=host_distance
@@ -436,10 +438,11 @@ def __init__(self, host, host_distance, session):
 iffirst_connection.features.sharding_infoandnotself._session.cluster.shard_aware_options.disable:
 self.host.sharding_info=first_connection.features.sharding_info
 self._open_connections_for_all_shards(first_connection.features.shard_id)
+self.tablets_routing_v1=first_connection.features.tablets_routing_v1
 
 log.debug("Finished initializing connection for host %s", self.host)
 
-def_get_connection_for_routing_key(self, routing_key=None):
+def_get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None):
 ifself.is_shutdown:
 raiseConnectionException(
 "Pool for %s is shutdown"% (self.host,), self.host)
@@ -450,7 +453,22 @@ def _get_connection_for_routing_key(self, routing_key=None):
 shard_id=None
 ifnotself._session.cluster.shard_aware_options.disableandself.host.sharding_infoandrouting_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)
+
+shard_id=None
+ifself.tablets_routing_v1andtableisnotNone:
+ifkeyspaceisNone:
+keyspace=self._keyspace
+
+tablet=self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t)
+
+iftabletisnotNone:
+forreplicaintablet.replicas:
+ifreplica[0] ==self.host.host_id:
+shard_id=replica[1]
+break
+
+ifshard_idisNone:
+shard_id=self.host.sharding_info.shard_id_from_token(t.value)
 
 conn=self._connections.get(shard_id)
 
@@ -496,15 +514,15 @@ def _get_connection_for_routing_key(self, routing_key=None):
 returnrandom.choice(active_connections)
 returnrandom.choice(list(self._connections.values()))
 
-defborrow_connection(self, timeout, routing_key=None):
-conn=self._get_connection_for_routing_key(routing_key)
+defborrow_connection(self, timeout, routing_key=None, keyspace=None, table=None):
+conn=self._get_connection_for_routing_key(routing_key, keyspace, table)
 start=time.time()
 remaining=timeout
 last_retry=False
 whileTrue:
 ifconn.is_closed:
 # The connection might have been closed in the meantime - if so, try again
-conn=self._get_connection_for_routing_key(routing_key)
+conn=self._get_connection_for_routing_key(routing_key, keyspace, table)
 withconn.lock:
 if (notconn.is_closedorlast_retry) andconn.in_flight<conn.max_request_id:
 # On last retry we ignore connection status, since it is better to return closed connection than

cassandra/protocol_features.py

@@ -6,22 +6,26 @@
 
 
 RATE_LIMIT_ERROR_EXTENSION="SCYLLA_RATE_LIMIT_ERROR"
+TABLETS_ROUTING_V1="TABLETS_ROUTING_V1"
 
 classProtocolFeatures(object):
 rate_limit_error=None
 shard_id=0
 sharding_info=None
+tablets_routing_v1=False
 
-def__init__(self, rate_limit_error=None, shard_id=0, sharding_info=None):
+def__init__(self, rate_limit_error=None, shard_id=0, sharding_info=None, tablets_routing_v1=False):
 self.rate_limit_error=rate_limit_error
 self.shard_id=shard_id
 self.sharding_info=sharding_info
+self.tablets_routing_v1=tablets_routing_v1
 
 @staticmethod
 defparse_from_supported(supported):
 rate_limit_error=ProtocolFeatures.maybe_parse_rate_limit_error(supported)
 shard_id, sharding_info=ProtocolFeatures.parse_sharding_info(supported)
-returnProtocolFeatures(rate_limit_error, shard_id, sharding_info)
+tablets_routing_v1=ProtocolFeatures.parse_tablets_info(supported)
+returnProtocolFeatures(rate_limit_error, shard_id, sharding_info, tablets_routing_v1)
 
 @staticmethod
 defmaybe_parse_rate_limit_error(supported):
@@ -43,6 +47,8 @@ def get_cql_extension_field(vals, key):
 defadd_startup_options(self, options):
 ifself.rate_limit_errorisnotNone:
 options[RATE_LIMIT_ERROR_EXTENSION] =""
+ifself.tablets_routing_v1:
+options[TABLETS_ROUTING_V1] =""
 
 @staticmethod
 defparse_sharding_info(options):
@@ -63,3 +69,6 @@ def parse_sharding_info(options):
 shard_aware_port, shard_aware_port_ssl)
 
 
+@staticmethod
+defparse_tablets_info(options):
+returnTABLETS_ROUTING_V1inoptions

cassandra/query.py

@@ -253,6 +253,13 @@ class Statement(object):
  .. versionadded:: 2.1.3
  """
 
+table=None
+"""
+ The string name of the table this query acts on. This is used when the tablet
+ experimental feature is enabled and in the same time :class`~.TokenAwarePolicy`
+ is configured in the profile load balancing policy.
+ """
+
 custom_payload=None
 """
  :ref:`custom_payload` to be passed to the server.
@@ -272,7 +279,7 @@ class Statement(object):
 
 def__init__(self, retry_policy=None, consistency_level=None, routing_key=None,
 serial_consistency_level=None, fetch_size=FETCH_SIZE_UNSET, keyspace=None, custom_payload=None,
-is_idempotent=False):
+is_idempotent=False, table=None):
 ifretry_policyandnothasattr(retry_policy, 'on_read_timeout'): # just checking one method to detect positional parameter errors
 raiseValueError('retry_policy should implement cassandra.policies.RetryPolicy')
 ifretry_policyisnotNone:
@@ -286,6 +293,8 @@ def __init__(self, retry_policy=None, consistency_level=None, routing_key=None,
 self.fetch_size=fetch_size
 ifkeyspaceisnotNone:
 self.keyspace=keyspace
+iftableisnotNone:
+self.table=table
 ifcustom_payloadisnotNone:
 self.custom_payload=custom_payload
 self.is_idempotent=is_idempotent
@@ -548,6 +557,7 @@ def __init__(self, prepared_statement, retry_policy=None, consistency_level=None
 meta=prepared_statement.column_metadata
 ifmeta:
 self.keyspace=meta[0].keyspace_name
+self.table=meta[0].table_name
 
 Statement.__init__(self, retry_policy, consistency_level, routing_key,
 serial_consistency_level, fetch_size, keyspace, custom_payload,