refactor: optimize server packet parsing with pre-compiled structs and fast paths

Author: rusher

mariadb/impl/message/server/column_definition_packet.py

@@ -9,21 +9,29 @@
 
 import struct
 from typing import TYPE_CHECKING, Optional, Tuple
-# No longer need PacketBuffer import
+
 if TYPE_CHECKING:
     from ...client.context import Context
 
+# Pre-compile struct formats for faster unpacking
+_STRUCT_UINT16 = struct.Struct('<H')
+_STRUCT_FIXED_FIELDS = struct.Struct('<HIBHB')  # charset(H), column_length(I), type(B), flags(H), decimals(B)
+
 
 def read_small_length_encoded_bytes(data: memoryview, pos: int) -> Tuple[bytes, int]:
     """Read length-encoded bytes and advance position"""
     length = data[pos]
     pos += 1
 
-    if length >= 251:
-        length = struct.unpack('<H', data[pos:pos+2])[0]
-        pos += 2
-
-    result = data[pos:pos+length].tobytes()
+    # Fast path: most lengths are < 251
+    if length < 251:
+        result = bytes(data[pos:pos+length])
+        return result, pos + length
+    
+    # Slow path: 2-byte length
+    length = _STRUCT_UINT16.unpack_from(data, pos)[0]
+    pos += 2
+    result = bytes(data[pos:pos+length])
     return result, pos + length
 
 
@@ -128,7 +136,7 @@ def org_name(self) -> str:
 
     @staticmethod
     def decode(data: memoryview, context: 'Context') -> 'ColumnDefinitionPacket':
-        """Decode column definition packet from bytearray with context"""
+        """Decode column definition packet"""
 
         pos = 0
 
@@ -139,35 +147,35 @@ def decode(data: memoryview, context: 'Context') -> 'ColumnDefinitionPacket':
         name_bytes, pos = read_small_length_encoded_bytes(data, pos)
         org_name_bytes, pos = read_small_length_encoded_bytes(data, pos)
 
-        # Handle extended info only if EXTENDED_METADATA capability is enabled
+        # Fast path: no extended metadata (most common case)
         ext_type_name = None
         ext_type_format = None
         special_format = False
-
-        # Check if we have the length field (0x0C) or extended metadata
+        
         if context.hasExtendedMetadata():
-            # Has extended info - read length-encoded buffer
             ext_length = data[pos]
             pos += 1
-            ext_end = pos + ext_length
-            while pos < ext_end and pos < len(data):
-                ext_type = data[pos]
-                pos += 1
+            
+            if ext_length > 0:
                 special_format = True
-                if ext_type == 0:
-                    # Extended type name
-                    ext_type_name, pos = read_small_length_encoded_bytes(data, pos)
-                elif ext_type == 1:
-                    # Extended type format
-                    ext_type_format, pos = read_small_length_encoded_bytes(data, pos)
-                else:
-                    # Skip unknown extended data
-                    _, pos = read_small_length_encoded_bytes(data, pos)
+                ext_end = pos + ext_length
+                
+                while pos < ext_end:
+                    ext_type = data[pos]
+                    pos += 1
+                    
+                    if ext_type == 0:
+                        ext_type_name, pos = read_small_length_encoded_bytes(data, pos)
+                    elif ext_type == 1:
+                        ext_type_format, pos = read_small_length_encoded_bytes(data, pos)
+                    else:
+                        # Skip unknown extended data
+                        _, pos = read_small_length_encoded_bytes(data, pos)
 
-        # Skip length field (always 0x0C = 12 for fixed fields)
+        # Skip length field (0x0C) and unpack fixed fields using pre-compiled struct
         pos += 1
-        charset, column_length, type, flags, decimals = struct.unpack('<HIBHB', data[pos:pos+10])
-
+        charset, column_length, type, flags, decimals = _STRUCT_FIXED_FIELDS.unpack_from(data, pos)
+        
         return ColumnDefinitionPacket(
             catalog_bytes,
             schema_bytes,

mariadb/impl/message/server/eof_packet.py

@@ -9,12 +9,16 @@
 
 from typing import TYPE_CHECKING
 from ...completion import Completion
-# No longer need PacketBuffer import
+
 if TYPE_CHECKING:
     from ...client.context import Context
 
 from mariadb_shared import constants
 
+# Pre-compute constant
+_PS_OUT_PARAMS_MASK = constants.STATUS.PS_OUT_PARAMS
+
+
 class EofPacket(Completion):
     """
     EOF Packet from MariaDB server
@@ -29,27 +33,26 @@ class EofPacket(Completion):
     In that case, use OkPacket.decode() instead.
     """
     __slots__ = (
-        'warning_count',
         'server_status',
     )
+    
     def __init__(
         self,
-        warning_count: int = 0,
-        server_status: int = 0,
-        is_output_parameters: bool = False
+        warning_count: int,
+        server_status: int,
     ):
         """Initialize EOF packet with warning count and server status"""
         self.affected_rows = 0
         self.insert_id = 0
         self.warning_count = warning_count
         self.result_set = None
 
-        # EofPacket-specific fields
+        # EofPacket-specific field
         self.server_status = server_status
 
     def is_output_parameters(self) -> bool:
         """Check if completion has output parameters"""
-        return (self.server_status & constants.STATUS.PS_OUT_PARAMS) != 0
+        return (self.server_status & _PS_OUT_PARAMS_MASK) != 0
 
     @staticmethod
     def decode(data: memoryview, context: 'Context') -> 'EofPacket':
@@ -61,12 +64,4 @@ def decode(data: memoryview, context: 'Context') -> 'EofPacket':
         context.server_status = server_status
         context.warning_count = warning_count
 
-        # Check if this marks output parameters (PS_OUT_PARAMS flag)
-        is_output_parameters = (server_status & constants.STATUS.PS_OUT_PARAMS) != 0
-        
-        return EofPacket(
-            warning_count,
-            server_status,
-            is_output_parameters
-        )
-    
+        return EofPacket(warning_count, server_status)

mariadb/impl/message/server/error_packet.py

@@ -7,14 +7,16 @@
 Based on MySQL/MariaDB protocol error packet structure.
 """
 
-import struct
 from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
     from ...client.context import Context
     from ...client.exception_factory import ExceptionFactory
 from ...client.socket.payload_parser import PayloadParser
-# No longer need PacketBuffer import
+
+# Pre-compute constants
+_HASH_MARKER = 0x23  # '#'
+_DEFAULT_SQL_STATE = "HY000"
 
 
 class ErrorPacket:
@@ -35,11 +37,12 @@ class ErrorPacket:
         'sql_state',
         'error_message',
     )
+    
     def __init__(
         self,
         error_code: int,
-        sql_state: str = "HY000",
-        error_message: str = "",
+        sql_state: str,
+        error_message: str,
     ):
         """Initialize error packet with error code, SQL state, and message"""
         self.error_code = error_code
@@ -54,34 +57,41 @@ def is_output_parameters(self) -> bool:
     def decode(data: memoryview, context: Optional['Context'] = None) -> 'ErrorPacket':
         """Decode error packet from bytearray with optional context"""
         parser = PayloadParser(data)
-        parser.read_byte()
+        parser.skip(1)  # Skip error marker (0xFF) - skip is faster than read_byte if we don't use value
         error_code = parser.read_uint16()
-        sql_state = "HY000"  # Default SQL state
 
-        # Check for SQL state marker '#' (0x23)
-        if parser.has_remaining() and parser.get_byte() == 0x23:  # '#' symbol
-            parser.read_byte()  # Skip '#' marker
+        # Fast path: check for SQL state marker
+        if parser.has_remaining() and parser.get_byte() == _HASH_MARKER:
+            parser.skip(1)  # Skip '#' marker
+            
             # SQL state (5 bytes)
             if parser.remaining_bytes() >= 5:
-                sql_state = bytes(parser.read_bytes(5)).decode('ascii')
+                sql_state = parser.read_bytes(5).decode('ascii')
             else:
                 raise IOError("Invalid error packet: SQL state truncated")
+        else:
+            sql_state = _DEFAULT_SQL_STATE
 
-        error_message = bytes(parser.read_remaining()).decode('utf-8', errors='replace')
-        return ErrorPacket(
-            error_code,
-            sql_state,
-            error_message
-        )
-
+        # Decode error message - remove unnecessary bytes() wrapper if read_remaining returns bytes
+        error_message = parser.read_remaining().decode('utf-8', errors='replace')
+        
+        return ErrorPacket(error_code, sql_state, error_message)
 
     def toError(self, exception_factory: 'ExceptionFactory', sql: Optional[str] = None):
-        return exception_factory.create_exception(self.error_message, self.sql_state, self.error_code, sql)
+        """Convert to exception"""
+        return exception_factory.create_exception(
+            self.error_message, 
+            self.sql_state, 
+            self.error_code, 
+            sql
+        )
 
     def __repr__(self) -> str:
+        """String representation for debugging"""
         return (f"ErrorPacket(error_code={self.error_code}, "
                 f"sql_state='{self.sql_state}', "
                 f"error_message='{self.error_message}')")
 
     def __str__(self) -> str:
-        return f"[{self.error_code}] ({self.sql_state}): {self.error_message}"
+        """Human-readable string representation"""
+        return f"[{self.error_code}] ({self.sql_state}): {self.error_message}"

mariadb/impl/message/server/prepare_stmt_packet.py

@@ -7,14 +7,18 @@
 Based on MySQL/MariaDB protocol COM_STMT_PREPARE response structure.
 """
 
-from typing import TYPE_CHECKING, Optional
+import struct
+from typing import TYPE_CHECKING, List, Optional
 
 from .column_definition_packet import ColumnDefinitionPacket
-from ...client.socket.payload_parser import PayloadParser
-# No longer need PacketBuffer import
+
 if TYPE_CHECKING:
     from ...client.context import Context
 
+# Pre-compile struct format for faster unpacking
+# Format: skip(B), statement_id(I), column_count(H), parameter_count(H), reserved(B), warning_count(H)
+_STRUCT_PREPARE_RESPONSE = struct.Struct('<BIHHBH')
+
 
 class PrepareStmtPacket:
     """
@@ -49,10 +53,10 @@ def __init__(
         statement_id: int,
         column_count: int,
         parameter_count: int,
-        warning_count: int = 0,
-        sql: str = None
+        warning_count: int,
+        sql: Optional[str]
     ):
-        """Initialize prepare statement packet with statement ID, column count, and parameter count"""
+        """Initialize prepare statement packet"""
         self.statement_id = statement_id
         self.column_count = column_count
         self.parameter_count = parameter_count
@@ -63,16 +67,10 @@ def __init__(
         self.closed = False
 
     @staticmethod
-    def decode(data: memoryview, context: Optional['Context'] = None, sql: str = None) -> 'PrepareStmtPacket':
-        """Decode COM_STMT_PREPARE response packet from bytearray with optional context"""
-        parser = PayloadParser(data)
-        
-        parser.read_byte()  # Skip OK marker (0x00)
-        statement_id = parser.read_uint32()
-        column_count = parser.read_uint16()
-        parameter_count = parser.read_uint16()
-        parser.read_byte()  # Skip reserved byte (0x00)
-        warning_count = parser.read_uint16()
+    def decode(data: memoryview, context: Optional['Context'] = None, sql: Optional[str] = None) -> 'PrepareStmtPacket':
+        """Decode COM_STMT_PREPARE response packet (optimized)"""
+        # Unpack all fields in one operation using pre-compiled struct
+        _, statement_id, column_count, parameter_count, _, warning_count = _STRUCT_PREPARE_RESPONSE.unpack_from(data, 0)
 
         # Update context if provided
         if context: