Skip to content

Commit f778c6c

Browse files
Extract _ExecuteManyAccumulator for the RETURNING-aware loop
Both the sync and async cursor executemany implementations maintained a near-identical per-iteration accumulation (rowcount sum, inherit first-seen description, extend the row list). Factor the state into a shared _ExecuteManyAccumulator class that both cursors drive: push() after each execute/_execute_async and apply() at the end. Behaviour is preserved exactly — first non-None description wins, rowcount is total-affected, rows concatenate in iteration order. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 87b5fde commit f778c6c

File tree

2 files changed

+55
-28
lines changed

2 files changed

+55
-28
lines changed

src/dqlitedbapi/aio/cursor.py

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
_call_client,
99
_convert_params,
1010
_convert_row,
11+
_ExecuteManyAccumulator,
1112
_is_row_returning,
1213
)
1314
from dqlitedbapi.exceptions import InterfaceError, NotSupportedError, ProgrammingError
@@ -169,22 +170,11 @@ async def executemany(
169170
self._description = None
170171
self._rows = []
171172
self._row_index = 0
172-
total_affected = 0
173-
accumulated_rows: list[tuple[Any, ...]] = []
174-
accumulated_desc: list[tuple[str, int | None, None, None, None, None, None]] | None = None
173+
acc = _ExecuteManyAccumulator()
175174
for params in seq_of_parameters:
176175
await self.execute(operation, params)
177-
if self._rowcount >= 0:
178-
total_affected += self._rowcount
179-
if self._description is not None:
180-
if accumulated_desc is None:
181-
accumulated_desc = self._description
182-
accumulated_rows.extend(self._rows)
183-
self._rowcount = total_affected
184-
if accumulated_desc is not None:
185-
self._description = accumulated_desc
186-
self._rows = accumulated_rows
187-
self._row_index = 0
176+
acc.push(self)
177+
acc.apply(self)
188178
return self
189179

190180
def _check_result_set(self) -> None:

src/dqlitedbapi/cursor.py

Lines changed: 51 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,54 @@ def _strip_leading_comments(sql: str) -> str:
151151
_ROW_RETURNING_PREFIXES = ("SELECT", "VALUES", "PRAGMA", "EXPLAIN", "WITH")
152152

153153

154+
class _ExecuteManyAccumulator:
155+
"""Shared state for the RETURNING-aware ``executemany`` loop.
156+
157+
Both the sync and async cursor implementations iterate
158+
``seq_of_parameters`` calling their respective single-statement
159+
helper. For statements with a RETURNING clause, rows produced on
160+
each iteration must accumulate so a subsequent ``fetchall`` yields
161+
every returned row across parameter sets. The bodies differ only
162+
by the ``await`` on the inner call, so both flavours drive this
163+
accumulator and then apply it to the cursor.
164+
"""
165+
166+
__slots__ = ("total_affected", "rows", "description")
167+
168+
def __init__(self) -> None:
169+
self.total_affected = 0
170+
self.rows: list[tuple[Any, ...]] = []
171+
self.description: list[tuple[str, int | None, None, None, None, None, None]] | None = None
172+
173+
def push(self, cursor: Any) -> None:
174+
"""Record one iteration's output into the accumulator.
175+
176+
Accepts either :class:`Cursor` or :class:`AsyncCursor`; both
177+
expose the same ``_rowcount`` / ``_description`` / ``_rows``
178+
attributes.
179+
"""
180+
if cursor._rowcount >= 0:
181+
self.total_affected += cursor._rowcount
182+
if cursor._description is not None:
183+
if self.description is None:
184+
self.description = cursor._description
185+
self.rows.extend(cursor._rows)
186+
187+
def apply(self, cursor: Any) -> None:
188+
"""Materialise the accumulator's state onto the cursor.
189+
190+
``description is None`` means none of the iterations produced a
191+
result set (plain DML without RETURNING); leave ``_description``
192+
/ ``_rows`` as reset. Inherit the first-seen description
193+
otherwise.
194+
"""
195+
cursor._rowcount = self.total_affected
196+
if self.description is not None:
197+
cursor._description = self.description
198+
cursor._rows = self.rows
199+
cursor._row_index = 0
200+
201+
154202
def _is_row_returning(sql: str) -> bool:
155203
"""Heuristic for "does this statement return a result set?"
156204
@@ -331,22 +379,11 @@ async def _executemany_async(
331379
self._description = None
332380
self._rows = []
333381
self._row_index = 0
334-
total_affected = 0
335-
accumulated_rows: list[tuple[Any, ...]] = []
336-
accumulated_desc: list[tuple[str, int | None, None, None, None, None, None]] | None = None
382+
acc = _ExecuteManyAccumulator()
337383
for params in seq_of_parameters:
338384
await self._execute_async(operation, params)
339-
if self._rowcount >= 0:
340-
total_affected += self._rowcount
341-
if self._description is not None:
342-
if accumulated_desc is None:
343-
accumulated_desc = self._description
344-
accumulated_rows.extend(self._rows)
345-
self._rowcount = total_affected
346-
if accumulated_desc is not None:
347-
self._description = accumulated_desc
348-
self._rows = accumulated_rows
349-
self._row_index = 0
385+
acc.push(self)
386+
acc.apply(self)
350387

351388
def _check_result_set(self) -> None:
352389
if self._description is None:

0 commit comments

Comments
 (0)