SQLExtendedFetch Implementation

cpp/src/arrow/flight/sql/odbc/entry_points.cc

@@ -137,6 +137,13 @@ SQLRETURN SQL_API SQLExecDirect(SQLHSTMT stmt, SQLWCHAR* queryText,
 
 SQLRETURN SQL_API SQLFetch(SQLHSTMT stmt) { return arrow::SQLFetch(stmt); }
 
+SQLRETURN SQL_API SQLExtendedFetch(SQLHSTMT stmt, SQLUSMALLINT fetchOrientation,
+                                   SQLLEN fetchOffset, SQLULEN* rowCountPtr,
+                                   SQLUSMALLINT* rowStatusArray) {
+  return arrow::SQLExtendedFetch(stmt, fetchOrientation, fetchOffset, rowCountPtr,
+                                 rowStatusArray);
+}
+
 SQLRETURN SQL_API SQLFetchScroll(SQLHSTMT stmt, SQLSMALLINT fetchOrientation,
                                  SQLLEN fetchOffset) {
   return arrow::SQLFetchScroll(stmt, fetchOrientation, fetchOffset);

cpp/src/arrow/flight/sql/odbc/odbc.def

@@ -33,6 +33,7 @@ EXPORTS
 	SQLErrorW
 	SQLExecDirectW
 	SQLExecute
+	SQLExtendedFetch
 	SQLFetch
 	SQLFetchScroll
 	SQLForeignKeysW

cpp/src/arrow/flight/sql/odbc/odbc_api.cc

@@ -983,6 +983,39 @@ SQLRETURN SQLFetch(SQLHSTMT stmt) {
   });
 }
 
+SQLRETURN SQLExtendedFetch(SQLHSTMT stmt, SQLUSMALLINT fetchOrientation,
+                           SQLLEN fetchOffset, SQLULEN* rowCountPtr,
+                           SQLUSMALLINT* rowStatusArray) {
+  // GH-47110: SQLExtendedFetch should return SQL_SUCCESS_WITH_INFO for certain diag
+  // states
+  LOG_DEBUG(
+      "SQLExtendedFetch called with stmt: {}, fetchOrientation: {}, fetchOffset: {}, "
+      "rowCountPtr: {}, rowStatusArray: {}",
+      stmt, fetchOrientation, fetchOffset, fmt::ptr(rowCountPtr),
+      fmt::ptr(rowStatusArray));
+  using ODBC::ODBCDescriptor;
+  using ODBC::ODBCStatement;
+  return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
+    if (fetchOrientation != SQL_FETCH_NEXT) {
+      throw DriverException("Optional feature not supported.", "HYC00");
+    }
+    // fetchOffset is ignored as only SQL_FETCH_NEXT is supported
+
+    ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
+
+    // The SQL_ROWSET_SIZE statement attribute specifies the number of rows in the
+    // rowset.
+    SQLULEN rowSetSize = statement->GetRowsetSize();
+    LOG_DEBUG("SQL_ROWSET_SIZE value for SQLExtendedFetch: {}", rowSetSize);
+    if (statement->Fetch(static_cast<size_t>(rowSetSize), rowCountPtr, rowStatusArray)) {
+      return SQL_SUCCESS;
+    } else {
+      // Reached the end of rowset
+      return SQL_NO_DATA;
+    }
+  });
+}
+
 SQLRETURN SQLFetchScroll(SQLHSTMT stmt, SQLSMALLINT fetchOrientation,
                          SQLLEN fetchOffset) {
   LOG_DEBUG("SQLFetchScroll called with stmt: {}, fetchOrientation: {}, fetchOffset: {}",

cpp/src/arrow/flight/sql/odbc/odbc_api.h

@@ -68,6 +68,9 @@ SQLRETURN SQLExecDirect(SQLHSTMT stmt, SQLWCHAR* queryText, SQLINTEGER textLengt
 SQLRETURN SQLPrepare(SQLHSTMT stmt, SQLWCHAR* queryText, SQLINTEGER textLength);
 SQLRETURN SQLExecute(SQLHSTMT stmt);
 SQLRETURN SQLFetch(SQLHSTMT stmt);
+SQLRETURN SQLExtendedFetch(SQLHSTMT stmt, SQLUSMALLINT fetchOrientation,
+                           SQLLEN fetchOffset, SQLULEN* rowCountPtr,
+                           SQLUSMALLINT* rowStatusArray);
 SQLRETURN SQLFetchScroll(SQLHSTMT stmt, SQLSMALLINT fetchOrientation, SQLLEN fetchOffset);
 SQLRETURN SQLBindCol(SQLHSTMT stmt, SQLUSMALLINT recordNumber, SQLSMALLINT cType,
                      SQLPOINTER dataPtr, SQLLEN bufferLength, SQLLEN* indicatorPtr);

cpp/src/arrow/flight/sql/odbc/odbcabstraction/include/odbcabstraction/odbc_impl/odbc_statement.h

@@ -66,8 +66,10 @@ class ODBCStatement : public ODBCHandle<ODBCStatement> {
 
   /**
    * @brief Returns true if the number of rows fetch was greater than zero.
+   * rowCountPtr and rowStatusArray are optional arguments, they are only needed for
+   * SQLExtendedFetch
    */
-  bool Fetch(size_t rows);
+  bool Fetch(size_t rows, SQLULEN* rowCountPtr = 0, SQLUSMALLINT* rowStatusArray = 0);
   bool isPrepared() const;
 
   void GetStmtAttr(SQLINTEGER statementAttribute, SQLPOINTER output,

cpp/src/arrow/flight/sql/odbc/odbcabstraction/odbc_impl/odbc_statement.cc

@@ -317,7 +317,8 @@ void ODBCStatement::ExecuteDirect(const std::string& query) {
   m_isPrepared = false;
 }
 
-bool ODBCStatement::Fetch(size_t rows) {
+bool ODBCStatement::Fetch(size_t rows, SQLULEN* rowCountPtr,
+                          SQLUSMALLINT* rowStatusArray) {
   if (m_hasReachedEndOfResult) {
     m_ird->SetRowsProcessed(0);
     return false;
@@ -350,11 +351,24 @@ bool ODBCStatement::Fetch(size_t rows) {
     m_currentArd->NotifyBindingsHavePropagated();
   }
 
-  size_t rowsFetched = m_currenResult->Move(rows, m_currentArd->GetBindOffset(),
-                                            m_currentArd->GetBoundStructOffset(),
-                                            m_ird->GetArrayStatusPtr());
+  uint16_t* arrayStatusPtr;
+  if (rowStatusArray) {
+    // For SQLExtendedFetch only
+    arrayStatusPtr = rowStatusArray;
+  } else {
+    arrayStatusPtr = m_ird->GetArrayStatusPtr();
+  }
+
+  size_t rowsFetched =
+      m_currenResult->Move(rows, m_currentArd->GetBindOffset(),
+                           m_currentArd->GetBoundStructOffset(), arrayStatusPtr);
   m_ird->SetRowsProcessed(static_cast<SQLULEN>(rowsFetched));
 
+  if (rowCountPtr) {
+    // For SQLExtendedFetch only
+    *rowCountPtr = rowsFetched;
+  }
+
   m_rowNumber += rowsFetched;
   m_hasReachedEndOfResult = rowsFetched != rows;
   return rowsFetched != 0;

cpp/src/arrow/flight/sql/odbc/tests/statement_test.cc

@@ -1187,7 +1187,6 @@ TYPED_TEST(FlightSQLODBCTestBase, TestSQLExecDirectFloatTruncation) {
   EXPECT_EQ(ret, SQL_SUCCESS);
 
   int16_t ssmall_int_val;
-  SQLLEN buf_len = sizeof(ssmall_int_val);
 
   ret = SQLGetData(this->stmt, 1, SQL_C_SSHORT, &ssmall_int_val, 0, 0);
   EXPECT_EQ(ret, SQL_SUCCESS_WITH_INFO);
@@ -2151,4 +2150,90 @@ TYPED_TEST(FlightSQLODBCTestBase, TestSQLBindColIndicatorOnlySQLUnbind) {
 
   this->disconnect();
 }
+
+TYPED_TEST(FlightSQLODBCTestBase, TestSQLExtendedFetchRowFetching) {
+  // Set SQL_ROWSET_SIZE to fetch 3 rows at once
+  this->connect();
+
+  constexpr SQLULEN rows = 3;
+  SQLINTEGER val[rows];
+  SQLLEN buf_len = sizeof(val);
+  SQLLEN ind[rows];
+
+  // Same variable will be used for column 1, the value of `val`
+  // should be updated after every SQLFetch call.
+  SQLRETURN ret = SQLBindCol(this->stmt, 1, SQL_C_LONG, val, buf_len, ind);
+
+  ret =
+      SQLSetStmtAttr(this->stmt, SQL_ROWSET_SIZE, reinterpret_cast<SQLPOINTER>(rows), 0);
+
+  std::wstring wsql =
+      LR"(
+    SELECT 1 AS small_table
+    UNION ALL
+    SELECT 2
+    UNION ALL
+    SELECT 3;
+  )";
+  std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());
+
+  ret = SQLExecDirect(this->stmt, &sql0[0], static_cast<SQLINTEGER>(sql0.size()));
+  EXPECT_EQ(ret, SQL_SUCCESS);
+
+  // Fetch row 1-3.
+  SQLULEN row_count;
+  SQLUSMALLINT row_status[rows];
+
+  ret = SQLExtendedFetch(this->stmt, SQL_FETCH_NEXT, 0, &row_count, row_status);
+  EXPECT_EQ(ret, SQL_SUCCESS);
+  EXPECT_EQ(row_count, 3);
+
+  for (int i = 0; i < rows; i++) {
+    EXPECT_EQ(row_status[i], SQL_SUCCESS);
+  }
+
+  // Verify 1 is returned for row 1
+  EXPECT_EQ(val[0], 1);
+  // Verify 2 is returned for row 2
+  EXPECT_EQ(val[1], 2);
+  // Verify 3 is returned for row 3
+  EXPECT_EQ(val[2], 3);
+
+  // Verify result set has no more data beyond row 3
+  SQLULEN row_count2;
+  SQLUSMALLINT row_status2[rows];
+  ret = SQLExtendedFetch(this->stmt, SQL_FETCH_NEXT, 0, &row_count2, row_status2);
+  EXPECT_EQ(ret, SQL_NO_DATA);
+
+  this->disconnect();
+}
+
+TEST_F(FlightSQLODBCRemoteTestBase, TestSQLExtendedFetchQueryNullIndicator) {
+  // GH-47110: SQLExtendedFetch should return SQL_SUCCESS_WITH_INFO for 22002
+  // Limitation on mock test server prevents null from working properly, so use remote
+  // server instead. Mock server has type `DENSE_UNION` for null column data.
+  GTEST_SKIP();
+  this->connect();
+
+  SQLINTEGER val;
+
+  SQLRETURN ret = SQLBindCol(this->stmt, 1, SQL_C_LONG, &val, 0, 0);
+
+  std::wstring wsql = L"SELECT null as null_col;";
+  std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());
+
+  ret = SQLExecDirect(this->stmt, &sql0[0], static_cast<SQLINTEGER>(sql0.size()));
+  EXPECT_EQ(ret, SQL_SUCCESS);
+
+  SQLULEN row_count1;
+  SQLUSMALLINT row_status1[1];
+
+  // SQLExtendedFetch should return SQL_SUCCESS_WITH_INFO for 22002 state
+  ret = SQLExtendedFetch(this->stmt, SQL_FETCH_NEXT, 0, &row_count1, row_status1);
+  EXPECT_EQ(ret, SQL_SUCCESS_WITH_INFO);
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, error_state_22002);
+
+  this->disconnect();
+}
+
 }  // namespace arrow::flight::sql::odbc