refactor: optimize executemany to use prepared statements and batch execution

Author: rusher

mariadb/async_cursor.py

@@ -176,27 +176,6 @@ async def execute(self, sql: str, data: Optional[Union[Sequence[Any], dict]] = N
             )
 
     async def executemany(self, sql: str, data: Sequence[Union[Sequence[Any], dict]], buffered: Optional[bool] = None) -> None:
-        """
-        Execute a SQL statement multiple times with different parameter sets
-        
-        More efficient than calling execute() multiple times as it can
-        batch operations and reduce round trips.
-        
-        Args:
-            sql: SQL statement to execute (typically INSERT, UPDATE, DELETE)
-            data: Sequence of parameter sequences, one for each execution
-            buffered: Override cursor's buffered setting
-                
-        Raises:
-            ProgrammingError: If cursor is closed
-            DatabaseError: If execution fails
-            
-        Example:
-            >>> await cursor.executemany(
-            ...     "INSERT INTO users VALUES (?, ?)",
-            ...     [(1, 'Alice'), (2, 'Bob'), (3, 'Charlie')]
-            ... )
-        """
         """
         Execute a statement multiple times with different parameter sets
         
@@ -211,38 +190,34 @@ async def executemany(self, sql: str, data: Sequence[Union[Sequence[Any], dict]]
         if not isinstance(sql, str):
             raise TypeError("SQL statement must be a string")
 
-        # Check if data is None or not an array-like type
-        if data is None or not hasattr(data, '__iter__') or isinstance(data, (str, bytes)):
-            raise ProgrammingError("No data provided")
-        
         # Consume any pending streaming results before executing new query
         if self._result is not None and self._result.streaming():
             await self._result.fetch_remaining()
 
-        # If data is an empty list/tuple, return early with rowcount=0
-        if not data:
-            self._rowcount = 0
-            return
+        # Check if data is None or not an array-like type
+        if data is None or not hasattr(data, '__iter__') or isinstance(data, (str, bytes)):
+            raise ProgrammingError("No data provided")
 
         # Reset result state
         self._result = None
 
         try:
-            completions = list()
             if data and len(data) > 0 and not isinstance(data, (list, tuple)):
                 raise ProgrammingError(f"wrong parameter type")
 
-            # Pre-parse SQL once for optimization (avoid re-parsing for each row)
-            sql_bytes = None
-            param_positions = None
-            param_names = None
-            placeholder_count = 0
-            
-            # For positional parameters, parse SQL once
-            sql_bytes, param_positions = split_sql_parts(sql)
-            placeholder_count = len(param_positions) // 2
-            
+            if (self._stmt is not None):
+                if (self._stmt.sql != sql):
+                    await self.connection._client.close_prepared_statement(self._stmt)
+                    self._stmt = None
+
+            if (self._stmt is None):
+                self._stmt = await self.connection._client.prepare_statement(sql)
+
+            # Execute with parameters using ExecutePacket
+            from .impl.message.client.execute_packet import ExecutePacket
+
             # Execute the statement for each parameter set
+            commands = []
             for params in data:
                 # Execute with current parameter set
                 # Convert data to list format for parameter binding
@@ -252,18 +227,16 @@ async def executemany(self, sql: str, data: Sequence[Union[Sequence[Any], dict]]
 
                 # Validate parameter count matches placeholders
                 if parameters:
-                    if len(parameters) < placeholder_count:
+                    if len(parameters) < self._stmt.parameter_count:
                         raise ProgrammingError(
-                            f"Parameter count mismatch: SQL has {placeholder_count} placeholders, "
+                            f"Parameter count mismatch: SQL has {self._stmt.parameter_count} placeholders, "
                             f"but only {len(parameters)} parameters provided"
                         )
 
-                # Create query packet and execute with bytes
-                query_packet = QueryWithParamPacket(sql_bytes, param_positions, parameters)
-                # Use provided buffered parameter or fall back to cursor default
-                effective_buffered = buffered if buffered is not None else self._buffered
-                compl = await self.connection._client.execute(query_packet, self._config, effective_buffered)
-                completions.extend(compl)
+                execute_packet = ExecutePacket(self._stmt.statement_id, parameters, sql)
+                commands.append(execute_packet)
+                
+            completions = await self.connection._client.execute_many(commands, self._config, True, self._stmt)
 
             # Process the completions - aggregate result sets with compatible metadata
             self._process_executemany_completions(completions)
@@ -276,6 +249,7 @@ async def executemany(self, sql: str, data: Sequence[Union[Sequence[Any], dict]]
                 errno=2013,
                 sql_state='HY000'
             )
+        
 
     # =========================================================================
     # Result Fetching Methods

mariadb/base_cursor.py

@@ -67,9 +67,6 @@ def __init__(self, connection: TConnection, **kwargs):
         self._exception_factory = ExceptionFactory()
         self._buffered: bool = bool(kwargs.pop('buffered', True))
         self._result: Optional[TResult] = None
-        self._executemany_mode: bool = False
-        self._executemany_rowcount: int = 0
-        self._executemany_lastrowid: Optional[int] = None
         self._force_binary: bool = False
         self._stmt: Optional[PrepareStmtPacket] = None
         if kwargs:
@@ -91,9 +88,6 @@ def __init__(self, connection: TConnection, **kwargs):
     @property
     def rowcount(self) -> int:
         """Get the number of rows (read-only property)"""
-        # For executemany, return the aggregated rowcount
-        if self._executemany_mode:
-            return self._executemany_rowcount
         # Get the current completion
         if self._completion_index < len(self._completions):
             completion = self._completions[self._completion_index]
@@ -147,8 +141,6 @@ def sp_outparams(self) -> bool:
     def lastrowid(self) -> Optional[int]:
         """Get the last insert ID from the current completion"""
         # For executemany, return the last insert ID from all executions
-        if self._executemany_mode:
-            return self._executemany_lastrowid
         if self._completion_index < len(self._completions):
             completion = self._completions[self._completion_index]
             return completion.insert_id or None
@@ -287,11 +279,6 @@ def _process_completions(self, completions: List[Any]) -> None:
         Args:
             completions: List of completion objects
         """
-        # Reset executemany mode for regular execute
-        self._executemany_mode = False
-        self._executemany_rowcount = 0
-        self._executemany_lastrowid = None
-        
         # Store all completions for nextset() functionality
         self._completions = completions
         self._completion_index = 0
@@ -321,53 +308,33 @@ def _process_rows_set_completion(self, result_set: Result) -> None:
         """
         self._result = result_set
 
-    def _process_executemany_completions(self, completions: List[Any]) -> None:
+    def _process_executemany_completions(self, completions: List[List[Completion]]) -> None:
         """
         Process completions from executemany - aggregate all result sets.
         Since executemany runs the same query multiple times, all completions have identical metadata.
-        """
-        # Enable executemany mode and calculate aggregated values
-        self._executemany_mode = True
-        self._executemany_rowcount = 0
-        self._executemany_lastrowid = None
-        
-        # Calculate total affected rows and last insert ID from all completions
-        for c in completions:
-            if c.affected_rows >= 0:
-                self._executemany_rowcount += c.affected_rows
-            if c.insert_id is not None and c.insert_id > 0:
-                self._executemany_lastrowid = c.insert_id
         
+        Args:
+            completions: List[List[Completion]] - one list per executed message
+        """
         if not completions:
             self._result = None
             return
-        
-        # Find completions with result sets
-        result_set_completions = [c for c in completions if c.has_result_set()]
-        
-        if not result_set_completions:
-            # No result sets - just update counts (e.g., INSERT/UPDATE/DELETE)
-            # The aggregated values are already set above
+
+        firstCompletion = completions[0]
+        if not firstCompletion:
             self._result = None
             return
-        
-        # All completions have identical metadata - use the first one
-        first_rs = result_set_completions[0].get_result_set()
-        first_columns = first_rs.columns
-        
-        # Aggregate all rows from all result sets
-        aggregated_rows = []
-        for completion in result_set_completions:
-            rs = completion.get_result_set()
-            rows = rs.rows if hasattr(rs, 'rows') else rs.get('rows', [])
-            aggregated_rows.extend(rows)
-        
-        # Build result from aggregated data (description will be computed on-demand)
-        self._result = self._create_complete_result(
-            columns=first_columns,
-            column_count=len(first_columns),
-            rows=aggregated_rows
-        )
+
+        for u in range(1, len(completions)):
+            unit_completions = completions[u]
+            for i, c in enumerate(unit_completions):
+                if c.affected_rows >= 0:
+                    firstCompletion[i].affected_rows += c.affected_rows
+                if c.insert_id is not None and c.insert_id > 0:
+                    firstCompletion[i].insert_id = c.insert_id
+                if c.has_result_set():
+                    firstCompletion[i].result_set.rows.extend(c.result_set.rows)
+        self._process_completions(firstCompletion)
 
     def _build_description(self, columns: List[ColumnDefinitionPacket]) -> Optional[tuple]:
         """Build cursor description tuple from column definitions"""