Correct issue with fetching data frames that contain multiple packets

with asyncio.

Author: anthony-tuininga

doc/src/release_notes.rst

@@ -74,6 +74,8 @@ Common Changes
     - Fixed bug when fetching dates that are in the year 2038 or later.
     - Fixed bug when fetching numeric data with precision that exceeds 38 as
       decimal data.
+    - Fixed bug when fetching large amounts of data in one round-trip when
+      using asyncio with Oracle Database versions before 23ai.
 
     Note the data frame support in python-oracledb 3.3 is a pre-release, and
     may change in a future version.

src/oracledb/impl/thin/messages/base.pyx

@@ -734,6 +734,13 @@ cdef class Message:
         buf.write_ub8(state | TNS_SESSION_STATE_EXPLICIT_BOUNDARY)
         self.conn_impl._session_state_desired = 0
 
+    cdef int on_out_of_packets(self) except -1:
+        """
+        Called when an OufOfPackets exception is raised indicating that further
+        packets are required to continue processing of this message.
+        """
+        pass
+
     cdef int postprocess(self) except -1:
         pass
 
@@ -1307,6 +1314,7 @@ cdef class MessageWithData(Message):
             self.cursor_impl._last_row_index = self.row_index - 1
             self.cursor_impl._buffer_rowcount = self.row_index
             self.bit_vector = NULL
+        self.on_row_completed()
 
     cdef int _process_row_header(self, ReadBuffer buf) except -1:
         cdef uint32_t num_bytes
@@ -1511,6 +1519,39 @@ cdef class MessageWithData(Message):
                     continue
                 self._write_bind_params_column(buf, var_impl, pos + offset)
 
+    cdef int on_out_of_packets(self) except -1:
+        """
+        Called when an OufOfPackets exception is raised indicating that further
+        packets are required to continue processing of this message.
+        """
+        cdef ThinVarImpl var_impl
+
+        # when fetching Arrow data, if the column has already been processed
+        # and no saved array already exists, the array is saved so that
+        # subsequent processing will not append to the array further; once the
+        # complete row has been processed, the saved arrays are restored and
+        # processing continues
+        if self.cursor_impl.fetching_arrow:
+            for var_impl in self.cursor_impl.fetch_var_impls:
+                if var_impl._saved_arrow_array is not None:
+                    continue
+                elif var_impl._arrow_array.arrow_array.length > self.row_index:
+                    var_impl._saved_arrow_array = var_impl._arrow_array
+                    var_impl._arrow_array = None
+                    var_impl._create_arrow_array()
+
+    cdef int on_row_completed(self) except -1:
+        """
+        Called when a row has been successfully completed. This allows for any
+        saved Arrow arrays to be restored.
+        """
+        cdef ThinVarImpl var_impl
+        if self.cursor_impl.fetching_arrow:
+            for var_impl in self.cursor_impl.fetch_var_impls:
+                if var_impl._saved_arrow_array is not None:
+                    var_impl._arrow_array = var_impl._saved_arrow_array
+                    var_impl._saved_arrow_array = None
+
     cdef int postprocess(self) except -1:
         """
         Run any variable out converter functions on all non-null values that

src/oracledb/impl/thin/protocol.pyx

@@ -844,6 +844,7 @@ cdef class BaseAsyncProtocol(BaseProtocol):
                 break
             except OutOfPackets:
                 await self._receive_packet(message)
+                message.on_out_of_packets()
                 self._read_buf.restore_point()
 
     async def _process_single_message(self, Message message):

src/oracledb/impl/thin/var.pyx

@@ -33,6 +33,7 @@ cdef class ThinVarImpl(BaseVarImpl):
     cdef:
         object _last_raw_value
         ArrowArrayImpl _last_arrow_array
+        ArrowArrayImpl _saved_arrow_array
         list _coroutine_indexes
 
     cdef int _bind(self, object conn, BaseCursorImpl cursor_impl,