Make redirect_messages return count (1 Day Business Cards)

ReadBuffer.redirect_messages now returns the number of messages redirected. This is needed for counting RowDescription messages to figure out how many rows were modified by DML statements.

Author: LinkedIn Share Button

README.rst

@@ -1 +1 @@
-Low-level IO utilities for PosgtreSQL drivers.
+Low-level IO utilities for PostgreSQL drivers.

buffer.pxd

@@ -127,7 +127,7 @@ cdef class ReadBuffer:
  cdef inline const char* try_consume_message(self, ssize_t* len)
  cdef bytes consume_message(self)
  cdef discard_message(self)
- cdef redirect_messages(self, WriteBuffer buf, char mtype, int stop_at=?)
+ cdef int32_t redirect_messages(self, WriteBuffer buf, char mtype, int stop_at=?)
  cdef bytearray consume_messages(self, char mtype)
  cdef finish_message(self)
  cdef inline _finish_message(self)

buffer.pyx

@@ -667,8 +667,13 @@ cdef class ReadBuffer:
  self._finish_message()
  return mem
 
- cdef redirect_messages(self, WriteBuffer buf, char mtype,
+ cdef int32_t redirect_messages(self, WriteBuffer buf, char mtype,
  int stop_at=0):
+ # Redirects messages from self into buf until either
+ # a message with a type different than mtype is encountered, or
+ # buf contains stop_at bytes.
+ # Returns the number of messages redirected.
+
  if not self._current_message_ready:
  raise BufferError(
  'consume_full_messages called on a buffer without a '
@@ -687,8 +692,11 @@ cdef class ReadBuffer:
  ssize_t new_pos0
  ssize_t pos_delta
  int32_t done
+ int32_t count
 
+ count = 0
  while True:
+ count += 1
  buf.write_byte(mtype)
  buf.write_int32(self._current_message_len)
 
@@ -701,10 +709,10 @@ cdef class ReadBuffer:
  if self._length > 0:
  self._ensure_first_buf()
  else:
- return
+ return count
 
  if stop_at and buf._length >= stop_at:
- return
+ return count
 
  # Fast path: exhaust buf0 as efficiently as possible.
  if self._pos0 + 5 <= self._len0:
@@ -727,14 +735,16 @@ cdef class ReadBuffer:
  if new_pos0 + msg_len > cbuf_len:
  break
  new_pos0 += msg_len
+ count += 1
 
  if new_pos0 != self._pos0:
  assert self._pos0 < new_pos0 <= self._len0
 
  pos_delta = new_pos0 - self._pos0
  buf.write_cstr(
  cbuf + self._pos0,
- pos_delta)
+ pos_delta
+ )
 
  self._pos0 = new_pos0
  self._length -= pos_delta
@@ -743,11 +753,11 @@ cdef class ReadBuffer:
 
  if done:
  # The next message is of a different type.
- return
+ return count
 
  # Back to slow path.
  if not self.take_message_type(mtype):
- return
+ return count
 
  cdef bytearray consume_messages(self, char mtype):
  """Consume consecutive messages of the same type."""