SNOW-119349: handle year out of range in arrow format

Author: sfc-gh-stakeda

arrow_iterator.pyx

@@ -12,8 +12,8 @@ from libcpp cimport bool as c_bool
 from libcpp.memory cimport shared_ptr
 from libcpp.string cimport string as c_string
 from libcpp.vector cimport vector
-from .errors import (Error, OperationalError)
-from .errorcode import ER_FAILED_TO_READ_ARROW_STREAM
+from .errors import (Error, OperationalError, InterfaceError)
+from .errorcode import (ER_FAILED_TO_READ_ARROW_STREAM, ER_FAILED_TO_CONVERT_ROW_TO_PYTHON_TYPE)
 
 logger = getLogger(__name__)
 
@@ -27,8 +27,13 @@ ROW_UNIT, TABLE_UNIT, EMPTY_UNIT = 'row', 'table', ''
 
 
 cdef extern from "cpp/ArrowIterator/CArrowIterator.hpp" namespace "sf":
+    cdef cppclass ReturnVal:
+        PyObject * successObj;
+
+        PyObject * exception;
+
     cdef cppclass CArrowIterator:
-        PyObject* next();
+        shared_ptr[ReturnVal] next();
 
 
 cdef extern from "cpp/ArrowIterator/CArrowChunkIterator.hpp" namespace "sf":
@@ -132,20 +137,21 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
     cdef object context
     cdef CArrowIterator* cIterator
     cdef str unit
-    cdef PyObject* cret
+    cdef shared_ptr[ReturnVal] cret
     cdef vector[shared_ptr[CRecordBatch]] batches
     cdef object use_dict_result
+    cdef object cursor
 
-    def __cinit__(self, object py_inputstream, object arrow_context, object use_dict_result):
+    def __cinit__(self, object cursor, object py_inputstream, object arrow_context, object use_dict_result):
         cdef shared_ptr[InputStream] input_stream
         cdef shared_ptr[CRecordBatchReader] reader
         cdef shared_ptr[CRecordBatch] record_batch
         input_stream.reset(new PyReadableFile(py_inputstream))
         cdef CStatus ret = CRecordBatchStreamReader.Open(input_stream.get(), &reader)
         if not ret.ok():
             Error.errorhandler_wrapper(
-                None,
-                None,
+                cursor.connection,
+                cursor,
                 OperationalError,
                 {
                     u'msg': u'Failed to open arrow stream: ' + ret.message(),
@@ -156,8 +162,8 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
             ret = reader.get().ReadNext(&record_batch)
             if not ret.ok():
                 Error.errorhandler_wrapper(
-                    None,
-                    None,
+                    cursor.connection,
+                    cursor,
                     OperationalError,
                     {
                         u'msg': u'Failed to read next arrow batch: ' + ret.message(),
@@ -175,18 +181,24 @@ cdef class PyArrowIterator(EmptyPyArrowIterator):
         self.cIterator = NULL
         self.unit = ''
         self.use_dict_result = use_dict_result
+        self.cursor = cursor
 
     def __dealloc__(self):
         del self.cIterator
 
     def __next__(self):
         self.cret = self.cIterator.next()
 
-        if not self.cret:
-            logger.error("Internal error from CArrowIterator\n")
+        if not self.cret.get().successObj:
+            msg = u'Failed to convert current row, cause: ' + str(<object>self.cret.get().exception)
+            Error.errorhandler_wrapper(self.cursor.connection, self.cursor, InterfaceError,
+                                       {
+                                           u'msg': msg,
+                                           u'errno': ER_FAILED_TO_CONVERT_ROW_TO_PYTHON_TYPE
+                                       })
             # it looks like this line can help us get into python and detect the global variable immediately
             # however, this log will not show up for unclear reason
-        ret = <object>self.cret
+        ret = <object>self.cret.get().successObj
 
         if ret is None:
             raise StopIteration

arrow_result.pyx

@@ -6,7 +6,6 @@
 # cython: language_level=3
 
 from base64 import b64decode
-from libcpp cimport bool
 import io
 from logging import getLogger
 from .telemetry import TelemetryField
@@ -57,7 +56,8 @@ cdef class ArrowResult:
         if rowset_b64:
             arrow_bytes = b64decode(rowset_b64)
             self._arrow_context = ArrowConverterContext(self._connection._session_parameters)
-            self._current_chunk_row = PyArrowIterator(io.BytesIO(arrow_bytes), self._arrow_context, self._use_dict_result)
+            self._current_chunk_row = PyArrowIterator(self._cursor, io.BytesIO(arrow_bytes),
+                                                      self._arrow_context, self._use_dict_result)
         else:
             logger.debug("Data from first gs response is empty")
             self._current_chunk_row = EmptyPyArrowIterator()

chunk_downloader.py

@@ -326,5 +326,5 @@ def __init__(self, cursor, connection):
     def to_iterator(self, raw_data_fd, download_time):
         from .arrow_iterator import PyArrowIterator
         gzip_decoder = GzipFile(fileobj=raw_data_fd, mode='r')
-        it = PyArrowIterator(gzip_decoder, self._arrow_context, self._cursor._use_dict_result)
+        it = PyArrowIterator(self._cursor, gzip_decoder, self._arrow_context, self._cursor._use_dict_result)
         return it

cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -14,6 +14,21 @@
 #include "TimeConverter.hpp"
 #include <string>
 
+#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
 {
 
@@ -31,18 +46,15 @@ CArrowChunkIterator::CArrowChunkIterator(PyObject* context, std::vector<std::sha
                m_columnCount);
 }
 
-PyObject* CArrowChunkIterator::next()
+std::shared_ptr<ReturnVal> CArrowChunkIterator::next()
 {
   m_rowIndexInBatch++;
 
   if (m_rowIndexInBatch < m_rowCountInBatch)
   {
     this->createRowPyObject();
-    if (py::checkPyError())
-    {
-      return nullptr;
-    }
-    return m_latestReturnedRow.get();
+    SF_CHECK_PYTHON_ERR()
+    return std::make_shared<ReturnVal>(m_latestReturnedRow.get(), nullptr);
   }
   else
   {
@@ -52,26 +64,21 @@ PyObject* CArrowChunkIterator::next()
       m_rowIndexInBatch = 0;
       m_rowCountInBatch = (*m_cRecordBatches)[m_currentBatchIndex]->num_rows();
       this->initColumnConverters();
-      if (py::checkPyError())
-      {
-        return nullptr;
-      }
+      SF_CHECK_PYTHON_ERR()
 
       logger.debug("Current batch index: %d, rows in current batch: %d",
                   m_currentBatchIndex, m_rowCountInBatch);
 
       this->createRowPyObject();
-      if (py::checkPyError())
-      {
-        return nullptr;
-      }
-      return m_latestReturnedRow.get();
+      SF_CHECK_PYTHON_ERR()
+
+      return std::make_shared<ReturnVal>(m_latestReturnedRow.get(), nullptr);
     }
   }
 
   /** It looks like no one will decrease the ref of this Py_None, so we don't
    * increament the ref count here */
-  return Py_None;
+  return std::make_shared<ReturnVal>(Py_None, nullptr);
 }
 
 void CArrowChunkIterator::createRowPyObject()

cpp/ArrowIterator/CArrowChunkIterator.hpp

@@ -32,7 +32,7 @@ class CArrowChunkIterator : public CArrowIterator
   /**
    * @return a python tuple object which contains all data in current row
    */
-  PyObject* next() override;
+  std::shared_ptr<ReturnVal> next() override;
 
 protected:
   /**
@@ -65,6 +65,9 @@ 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;
 

cpp/ArrowIterator/CArrowIterator.hpp

@@ -32,6 +32,23 @@
 namespace sf
 {
 
+/**
+ * A simple struct to contain return data back cython.
+ * PyObject would be nullptr if failed and cause string will be populated
+ */
+class ReturnVal
+{
+public:
+  ReturnVal(PyObject * obj, PyObject *except) :
+    successObj(obj), exception(except)
+  {
+  }
+
+  PyObject * successObj;
+
+  PyObject * exception;
+};
+
 /**
  * Arrow base iterator implementation in C++.
  */
@@ -46,7 +63,7 @@ class CArrowIterator
   /**
    * @return a python object which might be current row or an Arrow Table
    */
-  virtual PyObject* next() = 0;
+  virtual std::shared_ptr<ReturnVal> next() = 0;
 
 protected:
    /** list of all record batch in current chunk */

cpp/ArrowIterator/CArrowTableIterator.cpp

@@ -137,17 +137,17 @@ CArrowTableIterator::CArrowTableIterator(PyObject* context, std::vector<std::sha
   Py_XDECREF(tz);
 }
 
-PyObject* CArrowTableIterator::next()
+std::shared_ptr<ReturnVal> CArrowTableIterator::next()
 {
   bool firstDone = this->convertRecordBatchesToTable();
   if (firstDone && m_cTable)
   {
     m_pyTableObjRef.reset(arrow::py::wrap_table(m_cTable));
-    return m_pyTableObjRef.get();
+    return std::make_shared<ReturnVal>(m_pyTableObjRef.get(), nullptr);
   }
   else
   {
-    return Py_None;
+    return std::make_shared<ReturnVal>(Py_None, nullptr);
   }
 }
 

cpp/ArrowIterator/CArrowTableIterator.hpp

@@ -32,7 +32,7 @@ class CArrowTableIterator : public CArrowIterator
   /**
    * @return an arrow table containing all data in all record batches
    */
-  PyObject* next() override;
+  std::shared_ptr<ReturnVal> next() override;
 
 private:
   /* arrow table of all record batches in current chunk */

test/test_cursor.py

@@ -8,6 +8,7 @@
 import os
 import time
 from datetime import datetime
+from sys import platform
 
 import pytest
 import pytz
@@ -835,11 +836,17 @@ def test_close_twice(conn_testaccount):
     conn_testaccount.close()
 
 
-def test_fetch_out_of_range_timestamp_value(conn):
-    with conn() as cnx:
-        cur = cnx.cursor()
-        cur.execute("""
-select '12345-01-02'::timestamp_ntz
+@pytest.mark.skipif(not platform.startswith('linux'), reason="""
+Running tests only in linux env. Once 3.52.2 is deployed to preprod2,
+we can remove this skip and running this tests on all platform
 """)
-        with pytest.raises(errors.InterfaceError):
-            cur.fetchone()
+def test_fetch_out_of_range_timestamp_value(conn):
+    for result_format in ['arrow', 'json']:
+        with conn() as cnx:
+            cur = cnx.cursor()
+            cur.execute("alter session set python_connector_query_result_format='{}'".format(result_format))
+            cur.execute("""
+    select '12345-01-02'::timestamp_ntz
+    """)
+            with pytest.raises(errors.InterfaceError):
+                cur.fetchone()

test/test_unit_arrow_chunk_iterator.py

@@ -538,7 +538,7 @@ def iterate_over_test_chunk(pyarrow_type, column_meta, source_data_generator, ex
     # seek stream to begnning so that we can read from stream
     stream.seek(0)
     context = ArrowConverterContext()
-    it = PyArrowIterator(stream, context, False)
+    it = PyArrowIterator(None, stream, context, False)
     it.init(ROW_UNIT)
 
     count = 0