Use RateLimitReached error

Now, the connection negotiates protocol features and uses them later in
decoding. RateLimitReached is used instead of deafault.

Author: sylwiaszunejko

cassandra/connection.py

@@ -31,6 +31,8 @@
 import random
 import itertools
 
+from cassandra.protocol_features import ProtocolFeatures
+
 if 'gevent.monkey' in sys.modules:
     from gevent.queue import Queue, Empty
 else:
@@ -772,6 +774,8 @@ class Connection(object):
 
     _on_orphaned_stream_released = None
 
+    features = None
+
     @property
     def _iobuf(self):
         # backward compatibility, to avoid any change in the reactors
@@ -1263,7 +1267,7 @@ def process_msg(self, header, body):
                     return
 
         try:
-            response = decoder(header.version, self.user_type_map, stream_id,
+            response = decoder(header.version, self.features, self.user_type_map, stream_id,
                                header.flags, header.opcode, body, self.decompressor, result_metadata)
         except Exception as exc:
             log.exception("Error decoding response from Cassandra. "
@@ -1338,6 +1342,11 @@ def _handle_options_response(self, options_response):
         remote_supported_compressions = options_response.options['COMPRESSION']
         self._product_type = options_response.options.get('PRODUCT_TYPE', [None])[0]
 
+        protocol_features = ProtocolFeatures.parse_from_supported(options_response.options)
+        options = {}
+        protocol_features.add_startup_options(options)
+        self.features = protocol_features
+
         if self.cql_version:
             if self.cql_version not in supported_cql_versions:
                 raise ProtocolError(
@@ -1388,13 +1397,14 @@ def _handle_options_response(self, options_response):
                     self._compressor, self.decompressor = \
                         locally_supported_compressions[compression_type]
 
-        self._send_startup_message(compression_type, no_compact=self.no_compact)
+        self._send_startup_message(compression_type, no_compact=self.no_compact, extra_options=options)
 
     @defunct_on_error
-    def _send_startup_message(self, compression=None, no_compact=False):
+    def _send_startup_message(self, compression=None, no_compact=False, extra_options=None):
         log.debug("Sending StartupMessage on %s", self)
         opts = {'DRIVER_NAME': DRIVER_NAME,
-                'DRIVER_VERSION': DRIVER_VERSION}
+                'DRIVER_VERSION': DRIVER_VERSION,
+                **extra_options}
         if compression:
             opts['COMPRESSION'] = compression
         if no_compact:

cassandra/protocol.py

@@ -126,10 +126,13 @@ def __init__(self, code, message, info):
         self.info = info
 
     @classmethod
-    def recv_body(cls, f, protocol_version, *args):
+    def recv_body(cls, f, protocol_version, protocol_features, *args):
         code = read_int(f)
         msg = read_string(f)
-        subcls = error_classes.get(code, cls)
+        if code == protocol_features.rate_limit_error:
+            subcls = RateLimitReachedException
+        else:
+            subcls = error_classes.get(code, cls)
         extra_info = subcls.recv_error_info(f, protocol_version)
         return subcls(code=code, message=msg, info=extra_info)
 
@@ -751,7 +754,7 @@ def recv(self, f, protocol_version, user_type_map, result_metadata):
             raise DriverException("Unknown RESULT kind: %d" % self.kind)
 
     @classmethod
-    def recv_body(cls, f, protocol_version, user_type_map, result_metadata):
+    def recv_body(cls, f, protocol_version, protocol_features, user_type_map, result_metadata):
         kind = read_int(f)
         msg = cls(kind)
         msg.recv(f, protocol_version, user_type_map, result_metadata)
@@ -1160,7 +1163,7 @@ def _write_header(f, version, flags, stream_id, opcode, length):
         write_int(f, length)
 
     @classmethod
-    def decode_message(cls, protocol_version, user_type_map, stream_id, flags, opcode, body,
+    def decode_message(cls, protocol_version, protocol_features, user_type_map, stream_id, flags, opcode, body,
                        decompressor, result_metadata):
         """
         Decodes a native protocol message body
@@ -1206,7 +1209,7 @@ def decode_message(cls, protocol_version, user_type_map, stream_id, flags, opcod
             log.warning("Unknown protocol flags set: %02x. May cause problems.", flags)
 
         msg_class = cls.message_types_by_opcode[opcode]
-        msg = msg_class.recv_body(body, protocol_version, user_type_map, result_metadata)
+        msg = msg_class.recv_body(body, protocol_version, protocol_features, user_type_map, result_metadata)
         msg.stream_id = stream_id
         msg.trace_id = trace_id
         msg.custom_payload = custom_payload