diff --git a/pynuodb/cursor.py b/pynuodb/cursor.py index 4e515f5..ed28d7c 100644 --- a/pynuodb/cursor.py +++ b/pynuodb/cursor.py @@ -184,7 +184,11 @@ def executemany(self, operation, seq_of_parameters): def fetchone(self): # type: () -> Optional[result_set.Row] """Return the next row of results from the previous SQL operation.""" - self._check_closed() + # Inline _check_closed to avoid per-row function-call overhead. + if self.closed: + raise Error("cursor is closed") + if self.session.closed: + raise Error("connection is closed") if self._result_set is None: raise Error("Previous execute did not produce any results or no call was issued yet") self.rownumber += 1 @@ -192,41 +196,54 @@ def fetchone(self): self.session.fetch_result_set_next(self._result_set) return self._result_set.fetchone() + def _drain(self, limit): + # type: (Optional[int]) -> List[result_set.Row] + """Drain rows from the current result set. + + :param limit: maximum rows to return; None means "no cap". + """ + self._check_closed() + if self._result_set is None: + raise Error("Previous execute did not produce any results or no call was issued yet") + + fetched_rows = [] # type: List[result_set.Row] + remaining = limit # None == unlimited + while remaining is None or remaining > 0: + # Drain the current in-memory batch in one shot instead of + # calling fetchone() per row. + idx = self._result_set.results_idx + batch = self._result_set.results + avail = len(batch) - idx + if avail > 0: + if remaining is None: + take = avail + else: + take = avail if avail < remaining else remaining + remaining -= take + fetched_rows.extend(batch[idx:idx + take]) + self._result_set.results_idx = idx + take + if remaining == 0: + break + if self._result_set.complete: + break + self.session.fetch_result_set_next(self._result_set) + self.rownumber += len(fetched_rows) + return fetched_rows + def fetchmany(self, size=None): # type: (Optional[int]) -> List[result_set.Row] """Return SIZE rows from the previous SQL operation. If size is None, uses the default size for this Cursor. """ - self._check_closed() - if size is None: size = self.arraysize - - fetched_rows = [] - num_fetched_rows = 0 - while num_fetched_rows < size: - row = self.fetchone() - if row is None: - break - else: - fetched_rows.append(row) - num_fetched_rows += 1 - return fetched_rows + return self._drain(size) def fetchall(self): # type: () -> List[result_set.Row] """Return all rows generated by the previous SQL operation.""" - self._check_closed() - - fetched_rows = [] - while True: - row = self.fetchone() - if row is None: - break - else: - fetched_rows.append(row) - return fetched_rows + return self._drain(None) def nextset(self): # pylint: disable=no-self-use # type: () -> None