SNOW-870230: robust error stream handling (#1710)

Author: sfc-gh-aling

src/snowflake/connector/cpp/ArrowIterator/BinaryConverter.cpp

@@ -20,7 +20,7 @@ PyObject* BinaryConverter::toPyObject(int64_t rowIndex) const
     Py_RETURN_NONE;
   }
   ArrowStringView stringView = ArrowArrayViewGetStringUnsafe(m_array, rowIndex);
-    return PyByteArray_FromStringAndSize(stringView.data, stringView.size_bytes);
+  return PyByteArray_FromStringAndSize(stringView.data, stringView.size_bytes);
 }
 
 }  // namespace sf

src/snowflake/connector/cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -17,28 +17,15 @@
 #include <string>
 #include <vector>
 
-
-#define SF_CHECK_PYTHON_ERR() \
-  if (py::checkPyError())\
-  {\
-    PyObject *type, * val, *traceback;\
-    PyErr_Fetch(&type, &val, &traceback);\
-    PyErr_Clear();\
-    m_currentPyException.reset(val);\
-\
-    Py_XDECREF(type);\
-    Py_XDECREF(traceback);\
-\
-    return std::make_shared<ReturnVal>(nullptr, m_currentPyException.get());\
-  }
-
-
 namespace sf
 {
 
 CArrowChunkIterator::CArrowChunkIterator(PyObject* context, char* arrow_bytes, int64_t arrow_bytes_size, PyObject *use_numpy)
 : CArrowIterator(arrow_bytes, arrow_bytes_size), m_latestReturnedRow(nullptr), m_context(context)
 {
+  if (py::checkPyError()) {
+    return;
+  }
   m_currentBatchIndex = -1;
   m_rowIndexInBatch = -1;
   m_rowCountInBatch = 0;

src/snowflake/connector/cpp/ArrowIterator/CArrowChunkIterator.hpp

@@ -66,9 +66,6 @@ class CArrowChunkIterator : public CArrowIterator
   /** total number of rows inside current record batch */
   int64_t m_rowCountInBatch;
 
-  /** pointer to the current python exception object */
-  py::UniqueRef m_currentPyException;
-
   /** arrow format convert context for the current session */
   PyObject* m_context;
 

src/snowflake/connector/cpp/ArrowIterator/CArrowIterator.cpp

@@ -45,11 +45,18 @@ CArrowIterator::CArrowIterator(char* arrow_bytes, int64_t arrow_bytes_size)
       SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(returnCode, stream, "[Snowflake Exception] error setting ArrowArrayView from array : %s, error code: %d", ArrowErrorMessage(&error), returnCode);
       m_ipcArrowArrayViewVec.push_back(std::move(newUniqueArrayView));
     } else {
+      SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(retcode, stream, "[Snowflake Exception] error getting schema from stream, error code: %d", returnCode);
       break;
     }
   }
   stream.release(&stream);
   logger->debug(__FILE__, __func__, __LINE__, "Arrow BatchSize: %d", m_ipcArrowArrayVec.size());
 }
 
+std::shared_ptr<ReturnVal> CArrowIterator::checkInitializationStatus()
+{
+    SF_CHECK_PYTHON_ERR()
+    return std::make_shared<ReturnVal>(nullptr, nullptr);
+}
+
 }

src/snowflake/connector/cpp/ArrowIterator/CArrowIterator.hpp

@@ -77,13 +77,28 @@ static const char* NANOARROW_TYPE_ENUM_STRING[] = {
 #define SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(arrow_status, stream, format_string, ...) \
   if (arrow_status != NANOARROW_OK) \
   { \
-    std::string errorInfo = Logger::formatString(format_string, ##__VA_ARGS__); \
-    logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str()); \
-    PyErr_SetString(PyExc_Exception, errorInfo.c_str()); \
+    std::string errorInfo = std::string(format_string) + std::string(", error info: ") + std::string(stream.get_last_error(&stream)); \
+    std::string fullErrorInfo = Logger::formatString(errorInfo.c_str(), ##__VA_ARGS__); \
+    logger->error(__FILE__, __func__, __LINE__, fullErrorInfo.c_str()); \
+    PyErr_SetString(PyExc_Exception, fullErrorInfo.c_str()); \
     stream.release(&stream); \
     return; \
   }
 
+#define SF_CHECK_PYTHON_ERR() \
+  if (py::checkPyError())\
+  {\
+    PyObject *type, * val, *traceback;\
+    PyErr_Fetch(&type, &val, &traceback);\
+    PyErr_Clear();\
+    m_currentPyException.reset(val);\
+\
+    Py_XDECREF(type);\
+    Py_XDECREF(traceback);\
+\
+    return std::make_shared<ReturnVal>(nullptr, m_currentPyException.get());\
+  }
+
 namespace sf
 {
 
@@ -122,13 +137,19 @@ class CArrowIterator
   virtual std::vector<uintptr_t> getArrowArrayPtrs() { return {}; };
   virtual std::vector<uintptr_t> getArrowSchemaPtrs() { return {}; };
 
+  /** check whether initialization succeeded or encountered error */
+  std::shared_ptr<ReturnVal> checkInitializationStatus();
+
 protected:
   static Logger* logger;
 
   /** nanoarrow data */
   std::vector<nanoarrow::UniqueArray> m_ipcArrowArrayVec;
   std::vector<nanoarrow::UniqueArrayView> m_ipcArrowArrayViewVec;
   nanoarrow::UniqueSchema m_ipcArrowSchema;
+
+  /** pointer to the current python exception object */
+  py::UniqueRef m_currentPyException;
 };
 }
 

src/snowflake/connector/cpp/ArrowIterator/CArrowTableIterator.cpp

@@ -190,6 +190,9 @@ const bool number_to_decimal
 m_context(context),
 m_convert_number_to_decimal(number_to_decimal)
 {
+  if (py::checkPyError()) {
+    return;
+  }
   py::UniqueRef tz(PyObject_GetAttrString(m_context, "_timezone"));
   PyArg_Parse(tz.get(), "s", &m_timezone);
 }

src/snowflake/connector/cpp/ArrowIterator/arrow_iterator.pyx

@@ -39,6 +39,7 @@ cdef extern from "CArrowIterator.hpp" namespace "sf":
 
     cdef cppclass CArrowIterator:
         shared_ptr[ReturnVal] next() except +;
+        shared_ptr[ReturnVal] checkInitializationStatus() except +;
         vector[uintptr_t] getArrowArrayPtrs();
         vector[uintptr_t] getArrowSchemaPtrs();
 
@@ -114,11 +115,11 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
             object use_dict_result,
             object numpy,
             object number_to_decimal,
-
+            object is_table_unit = False,
     ):
         self.context = arrow_context
         self.cIterator = NULL
-        self.unit = ''
+        self.unit = IterUnit.TABLE_UNIT.value if is_table_unit else IterUnit.ROW_UNIT.value
         self.use_dict_result = use_dict_result
         self.cursor = cursor
         self.use_numpy = numpy
@@ -129,6 +130,11 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
         self.arrow_bytes_size = len(arrow_bytes)
         self.python_bytes = arrow_bytes
 
+        if self.unit == IterUnit.TABLE_UNIT.value:
+            self.init_table_unit()
+        else:
+            self.init_row_unit()
+
     def __dealloc__(self):
         del self.cIterator
 
@@ -174,6 +180,9 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
         self.unit = iter_unit
 
     def init_row_unit(self) -> None:
+        if self.cIterator is not NULL:
+            return
+
         self.cIterator = new CArrowChunkIterator(
             <PyObject *> self.context,
             self.arrow_bytes,
@@ -187,6 +196,16 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
             self.arrow_bytes_size,
             <PyObject *> self.use_numpy
             )
+        self.cret = self.cIterator.checkInitializationStatus()
+        if self.cret.get().exception:
+            Error.errorhandler_wrapper(
+                self.cursor.connection if self.cursor is not None else None,
+                self.cursor,
+                OperationalError,
+                {
+                    'msg': f'Failed to open arrow stream: {str(<object>self.cret.get().exception)}',
+                    'errno': ER_FAILED_TO_READ_ARROW_STREAM
+                })
         snow_logger.debug(msg=f"Batches read: {self.cIterator.getArrowArrayPtrs().size()}", path_name=__file__, func_name="init_row_unit")
 
     def init_table_unit(self) -> None:
@@ -202,13 +221,25 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
                 },
             )
 
+        if self.cIterator is not NULL:
+            return
+
         self.cIterator = new CArrowTableIterator(
             <PyObject *> self.context,
             self.arrow_bytes,
             self.arrow_bytes_size,
             self.number_to_decimal,
         )
-
+        self.cret = self.cIterator.checkInitializationStatus()
+        if self.cret.get().exception:
+            Error.errorhandler_wrapper(
+                self.cursor.connection if self.cursor is not None else None,
+                self.cursor,
+                OperationalError,
+                {
+                    'msg': f'Failed to open arrow stream: {str(<object>self.cret.get().exception)}',
+                    'errno': ER_FAILED_TO_READ_ARROW_STREAM
+                })
         self.cret = self.cIterator.next()
         self.unit = 'table'
         self.nanoarrow_Table = self.cIterator.getArrowArrayPtrs()

src/snowflake/connector/cpp/ArrowIterator/flatcc/flatcc_assert.h

@@ -25,7 +25,7 @@ extern "C" {
 * `FLATCC_NO_ASSERT`.
 */
 
-#ifdef FLATCC_NO_ASSERT 
+#ifdef FLATCC_NO_ASSERT
 /* NOTE: This will not affect inclusion of <assert.h> for static assertions. */
 #undef FLATCC_ASSERT
 #define FLATCC_ASSERT(x) ((void)0)

src/snowflake/connector/cpp/ArrowIterator/flatcc/flatcc_epilogue.h

@@ -5,4 +5,3 @@
 #endif
 
 #include "flatcc/portable/pdiagnostic_pop.h"
-

src/snowflake/connector/cpp/ArrowIterator/flatcc/portable/paligned_alloc.h

@@ -53,7 +53,7 @@ extern "C" {
 #if defined (__MINGW32__)
 /* MingW does not provide aligned_alloc despite defining _ISOC11_SOURCE */
 #define PORTABLE_C11_ALIGNED_ALLOC 0
-#elif defined (_ISOC11_SOURCE) 
+#elif defined (_ISOC11_SOURCE)
 /* glibc aligned_alloc detection, but MingW is not truthful */
 #define PORTABLE_C11_ALIGNED_ALLOC 1
 #elif defined (__GLIBC__)
@@ -178,7 +178,7 @@ static inline void *__portable_aligned_alloc(size_t alignment, size_t size)
 static inline void __portable_aligned_free(void *p)
 {
     char *raw;
-    
+
     if (p) {
         raw = (char*)((void **)p)[-1];
         free(raw);