Add new error for rate limit

Adds RateLimitReached error, which is a Scylla-specific error returned
when a per-partition rate limit is exceeded.

Author: sylwiaszunejko

cassandra/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from enum import Enum
 import logging
 
 
@@ -728,3 +729,21 @@ class UnresolvableContactPoints(DriverException):
     contact points, only when lookup fails for all hosts
     """
     pass
+
+
+class OperationType(Enum):
+    Read = 0
+    Write = 1
+
+class RateLimitReached(ConfigurationException):
+    '''
+    Rate limit was exceeded for a partition affected by the request.
+    '''
+    op_type = None
+    rejected_by_coordinator = False
+
+    def __init__(self, op_type=None, rejected_by_coordinator=False):
+        self.op_type = op_type
+        self.rejected_by_coordinator = rejected_by_coordinator
+        message = f"[request_error_rate_limit_reached OpType={op_type.name} RejectedByCoordinator={rejected_by_coordinator}]"
+        Exception.__init__(self, message)

cassandra/protocol.py

@@ -22,9 +22,9 @@
 from six.moves import range
 import io
 
-from cassandra import ProtocolVersion
+from cassandra import OperationType, ProtocolVersion
 from cassandra import type_codes, DriverException
-from cassandra import (Unavailable, WriteTimeout, ReadTimeout,
+from cassandra import (Unavailable, WriteTimeout, RateLimitReached, ReadTimeout,
                        WriteFailure, ReadFailure, FunctionFailure,
                        AlreadyExists, InvalidRequest, Unauthorized,
                        UnsupportedOperation, UserFunctionDescriptor,
@@ -390,6 +390,19 @@ def recv_error_info(f, protocol_version):
     def to_exception(self):
         return AlreadyExists(**self.info)
 
+class RateLimitReachedException(ConfigurationException):
+    summary= 'Rate limit was exceeded for a partition affected by the request'
+    error_code = 0x4321
+
+    @staticmethod
+    def recv_error_info(f, protocol_version):
+        return {
+            'op_type': OperationType(read_byte(f)),
+            'rejected_by_coordinator': read_byte(f) != 0
+        }
+
+    def to_exception(self):
+        return RateLimitReached(**self.info)
 
 class ClientWriteError(RequestExecutionException):
     summary = 'Client write failure.'

docs/scylla-specific.rst

@@ -104,3 +104,39 @@ New Table Attributes
 
     cluster.refresh_table_metadata("keyspace1", "standard1")
     assert cluster.metadata.keyspaces["keyspace1"].tables["standard1"].options["in_memory"] == True
+
+
+New Error Types
+--------------------
+
+* ``SCYLLA_RATE_LIMIT_ERROR`` Error
+
+  The ScyllaDB 5.1 introduced a feature called per-partition rate limiting. In case the (user defined) per-partition rate limit is exceeded, the database will start returning a Scylla-specific type of error: RateLimitReached.
+
+.. code:: python
+
+    from cassandra import RateLimitReached
+    from cassandra.cluster import Cluster
+
+    cluster = Cluster()
+    session = cluster.connect()
+    session.execute("""
+        CREATE KEYSPACE IF NOT EXISTS keyspace1 
+        WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'}
+    """)
+
+    session.execute("USE keyspace1")
+    session.execute("""
+        CREATE TABLE tbl (pk int PRIMARY KEY, v int) 
+        WITH per_partition_rate_limit = {'max_writes_per_second': 1}
+    """)
+
+    prepared = session.prepare("""
+        INSERT INTO tbl (pk, v) VALUES (?, ?)
+    """)
+    
+    try:
+        for _ in range(1000):
+            self.session.execute(prepared.bind((123, 456)))
+    except RateLimitReached:
+        raise