WL12501: Connection compression

This worklog adds support for connection compression using the X Protocol.

The supported compression algorithms are lz4_message and deflate_stream via
third-party software python-lz4.

Author: nmariz

cpyint

@@ -65,7 +65,7 @@
 _SESS_OPTS = _SSL_OPTS + ["user", "password", "schema", "host", "port",
                           "routers", "socket", "ssl-mode", "auth", "use-pure",
                           "connect-timeout", "connection-attributes",
-                          "dns-srv"]
+                          "compression", "dns-srv"]
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())
 
@@ -82,7 +82,7 @@
 TLS_VERSIONS = ["TLSv1", "TLSv1.1", "TLSv1.2", "TLSv1.3"]
 
 TLS_V1_3_SUPPORTED = False
-if hasattr(ssl, "HAS_TLSv1_3") and ssl.HAS_TLSv1_3: 
+if hasattr(ssl, "HAS_TLSv1_3") and ssl.HAS_TLSv1_3:
     TLS_V1_3_SUPPORTED = True
 
 
@@ -256,6 +256,15 @@ def _validate_settings(settings):
         except (AttributeError, ValueError):
             raise InterfaceError("Invalid Auth '{0}'".format(settings["auth"]))
 
+    if "compression" in settings:
+        compression = settings["compression"].lower().strip()
+        if compression not in ("preferred", "required", "disabled"):
+            raise InterfaceError(
+                "The connection property 'compression' acceptable values are: "
+                "'preferred', 'required', or 'disabled'. The value '{0}' is "
+                "not acceptable".format(settings["compression"]))
+        settings["compression"] = compression
+
     if "connection-attributes" in settings:
         validate_connection_attributes(settings)
 

lib/mysqlx/__init__.py

@@ -78,7 +78,7 @@
 from .crud import Schema
 from .constants import SSLMode, Auth
 from .helpers import escape, get_item_or_attr, iani_to_openssl_cs_name
-from .protocol import Protocol, MessageReaderWriter
+from .protocol import Protocol, MessageReader, MessageWriter, HAVE_LZ4
 from .result import Result, RowResult, SqlResult, DocResult
 from .statement import SqlStatement, AddStatement, quote_identifier
 from .protobuf import Protobuf
@@ -558,7 +558,6 @@ class Connection(object):
     def __init__(self, settings):
         self.settings = settings
         self.stream = SocketStream()
-        self.reader_writer = None
         self.protocol = None
         self.keep_open = None
         self._user = settings.get("user")
@@ -617,10 +616,31 @@ def connect(self):
                 router = self.router_manager.get_next_router()
                 self.stream.connect(router.get_connection_params(),
                                     self._connect_timeout)
-                self.reader_writer = MessageReaderWriter(self.stream)
-                self.protocol = Protocol(self.reader_writer)
-                self._handle_capabilities()
+                reader = MessageReader(self.stream)
+                writer = MessageWriter(self.stream)
+                self.protocol = Protocol(reader, writer)
+
+                caps_data = self.protocol.get_capabilites().capabilities
+                caps = {
+                    get_item_or_attr(cap, "name").lower():
+                        cap for cap in caps_data
+                } if caps_data else {}
+
+                # Set TLS capabilities
+                self._set_tls_capabilities(caps)
+
+                # Set connection attributes capabilities
+                if "attributes" in self.settings:
+                    conn_attrs = self.settings["attributes"]
+                    self.protocol.set_capabilities(
+                        session_connect_attrs=conn_attrs)
+
+                # Set compression capabilities
+                compression = self.settings.get("compression", "preferred")
+                algorithm = None if compression == "disabled" \
+                    else self._set_compression_capabilities(caps, compression)
                 self._authenticate()
+                self.protocol.set_compression(algorithm)
                 return
             except (socket.error, RuntimeError) as err:
                 error = err
@@ -643,30 +663,31 @@ def connect(self):
             raise InterfaceError("Cannot connect to host: {0}".format(error))
         raise InterfaceError("Unable to connect to any of the target hosts", 4001)
 
-    def _handle_capabilities(self):
-        """Handle capabilities.
+    def _set_tls_capabilities(self, caps):
+        """Sets the TLS capabilities.
+
+        Args:
+            caps (dict): Dictionary with the server capabilities.
 
         Raises:
             :class:`mysqlx.OperationalError`: If SSL is not enabled at the
                                              server.
             :class:`mysqlx.RuntimeError`: If support for SSL is not available
                                           in Python.
+
+        .. versionadded:: 8.0.21
         """
         if self.settings.get("ssl-mode") == SSLMode.DISABLED:
             return
+
         if self.stream.is_socket():
             if self.settings.get("ssl-mode"):
                 _LOGGER.warning("SSL not required when using Unix socket.")
             return
 
-        try:
-            data = self.protocol.get_capabilites().capabilities
-            if not (get_item_or_attr(data[0], "name").lower() == "tls"
-                    if data else False):
-                self.close_connection()
-                raise OperationalError("SSL not enabled at server")
-        except (AttributeError, KeyError):
-            pass
+        if "tls" not in caps:
+            self.close_connection()
+            raise OperationalError("SSL not enabled at server")
 
         is_ol7 = False
         if platform.system() == "Linux":
@@ -694,6 +715,65 @@ def _handle_capabilities(self):
             conn_attrs = self.settings["attributes"]
             self.protocol.set_capabilities(session_connect_attrs=conn_attrs)
 
+    def _set_compression_capabilities(self, caps, compression):
+        """Sets the compression capabilities.
+
+        If compression is available, negociates client and server algorithms.
+        Using the following priority:
+
+        1) lz4_message
+        2) deflate_stream
+
+        Args:
+            caps (dict): Dictionary with the server capabilities.
+            compression (str): The compression connection setting.
+
+        Returns:
+            str: The compression algorithm.
+
+        .. versionadded:: 8.0.21
+        """
+        compression_data = caps.get("compression")
+        if compression_data is None:
+            msg = "Compression requested but the server does not support it"
+            if compression == "required":
+                raise NotSupportedError(msg)
+            else:
+                _LOGGER.warning(msg)
+            return None
+
+        compression_dict = {}
+        if isinstance(compression_data, dict):  # C extension is being used
+            for fld in compression_data["value"]["obj"]["fld"]:
+                compression_dict[fld["key"]] = [
+                    value["scalar"]["v_string"]["value"].decode("utf-8")
+                    for value in fld["value"]["array"]["value"]
+                ]
+        else:
+            for fld in compression_data.value.obj.fld:
+                compression_dict[fld.key] = [
+                    value.scalar.v_string.value.decode("utf-8")
+                    for value in fld.value.array.value
+                ]
+
+        server_algorithms = compression_dict.get("algorithm", [])
+        if HAVE_LZ4 and "lz4_message" in server_algorithms:
+            algorithm = "lz4_message"
+        else:
+            algorithm = "deflate_stream"
+
+        if algorithm not in server_algorithms:
+            msg = ("Compression requested but the compression algorithm "
+                   "negotiation failed")
+            if compression == "required":
+                raise InterfaceError(msg)
+            else:
+                _LOGGER.warning(msg)
+            return None
+
+        self.protocol.set_capabilities(compression={"algorithm": algorithm})
+        return algorithm
+
     def _authenticate(self):
         """Authenticate with the MySQL server."""
         auth = self.settings.get("auth")

lib/mysqlx/connection.py

@@ -56,6 +56,8 @@
      "Mysqlx.Sql.StmtExecuteOk"),
     ("Mysqlx.ServerMessages.Type.RESULTSET_FETCH_DONE_MORE_OUT_PARAMS",
      "Mysqlx.Resultset.FetchDoneMoreOutParams"),
+    ("Mysqlx.ServerMessages.Type.COMPRESSION",
+     "Mysqlx.Connection.Compression"),
 )
 
 PROTOBUF_REPEATED_TYPES = [list]
@@ -74,6 +76,7 @@
 try:
     from . import mysqlx_connection_pb2
     from . import mysqlx_crud_pb2
+    from . import mysqlx_cursor_pb2
     from . import mysqlx_datatypes_pb2
     from . import mysqlx_expect_pb2
     from . import mysqlx_expr_pb2
@@ -168,6 +171,8 @@
         mysqlx_connection_pb2.DESCRIPTOR.serialized_pb))
     _DESCRIPTOR_DB.Add(descriptor_pb2.FileDescriptorProto.FromString(
         mysqlx_crud_pb2.DESCRIPTOR.serialized_pb))
+    _DESCRIPTOR_DB.Add(descriptor_pb2.FileDescriptorProto.FromString(
+        mysqlx_cursor_pb2.DESCRIPTOR.serialized_pb))
     _DESCRIPTOR_DB.Add(descriptor_pb2.FileDescriptorProto.FromString(
         mysqlx_datatypes_pb2.DESCRIPTOR.serialized_pb))
     _DESCRIPTOR_DB.Add(descriptor_pb2.FileDescriptorProto.FromString(
@@ -213,6 +218,10 @@ def enum_value(key):
         def serialize_message(msg):
             return msg.SerializeToString()
 
+        @staticmethod
+        def serialize_partial_message(msg):
+            return msg.SerializePartialToString()
+
         @staticmethod
         def parse_message(msg_type_name, payload):
             msg = _mysqlxpb_pure.new_message(msg_type_name)
@@ -346,10 +355,21 @@ def serialize_to_string(self):
         """Serializes a message to a string.
 
         Returns:
-           string: A string representing a message containing parsed data.
+            str: A string representing a message containing parsed data.
         """
         return Protobuf.mysqlxpb.serialize_message(self._msg)
 
+    def serialize_partial_to_string(self):
+        """Serializes the protocol message to a binary string.
+
+        This method is similar to serialize_to_string but doesn't check if the
+        message is initialized.
+
+        Returns:
+            str: A string representation of the partial message.
+        """
+        return Protobuf.mysqlxpb.serialize_partial_message(self._msg)
+
     @property
     def type(self):
         """string: Message type name."""
@@ -366,9 +386,26 @@ def parse(msg_type_name, payload):
 
         Returns:
             dict: The dictionary representing a message containing parsed data.
+
+        .. versionadded:: 8.0.21
         """
         return Protobuf.mysqlxpb.parse_message(msg_type_name, payload)
 
+    @staticmethod
+    def byte_size(msg):
+        """Returns the size of the message in bytes.
+
+        Args:
+            msg (mysqlx.protobuf.Message): MySQL X Protobuf Message.
+
+        Returns:
+            int: Size of the message in bytes.
+
+        .. versionadded:: 8.0.21
+        """
+        return msg.ByteSize() if Protobuf.use_pure \
+            else len(encode_to_bytes(msg.serialize_to_string()))
+
     @staticmethod
     def parse_from_server(msg_type, payload):
         """Creates a new server-side message, initialized with parsed data.