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 SQLGetData(SQLHSTMT stmt, SQLUSMALLINT recordNumber, SQLSMALLINT cType,
                              SQLPOINTER dataPtr, SQLLEN bufferLength,
                              SQLLEN* indicatorPtr) {
@@ -150,15 +157,11 @@ SQLRETURN SQL_API SQLPrepare(SQLHSTMT stmt, SQLWCHAR* queryText, SQLINTEGER text
 
 SQLRETURN SQL_API SQLExecute(SQLHSTMT stmt) { return arrow::SQLExecute(stmt); }
 
-SQLRETURN SQL_API SQLBindCol(SQLHSTMT stmt, SQLUSMALLINT columnNumber,
-                             SQLSMALLINT targetType, SQLPOINTER targetValuePtr,
-                             SQLLEN bufferLength, SQLLEN* strLen_or_IndPtr) {
-  LOG_DEBUG(
-      "SQLBindCol called with stmt: {}, columnNumber: {}, targetType: {}, "
-      "targetValuePtr: {}, bufferLength: {}, strLen_or_IndPtr: {}",
-      stmt, columnNumber, targetType, targetValuePtr, bufferLength,
-      fmt::ptr(strLen_or_IndPtr));
-  return SQL_ERROR;
+SQLRETURN SQL_API SQLBindCol(SQLHSTMT stmt, SQLUSMALLINT recordNumber, SQLSMALLINT cType,
+                             SQLPOINTER dataPtr, SQLLEN bufferLength,
+                             SQLLEN* indicatorPtr) {
+  return arrow::SQLBindCol(stmt, recordNumber, cType, dataPtr, bufferLength,
+                           indicatorPtr);
 }
 
 SQLRETURN SQL_API SQLCancel(SQLHSTMT stmt) {

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

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

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

@@ -197,9 +197,16 @@ SQLRETURN SQLFreeStmt(SQLHSTMT handle, SQLUSMALLINT option) {
       return SQLFreeHandle(SQL_HANDLE_STMT, handle);
     }
 
-    // TODO Implement SQLBindCol
     case SQL_UNBIND: {
-      return SQL_SUCCESS;
+      using ODBC::ODBCDescriptor;
+      using ODBC::ODBCStatement;
+      return ODBCStatement::ExecuteWithDiagnostics(handle, SQL_ERROR, [=]() {
+        ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(handle);
+        ODBCDescriptor* ard = statement->GetARD();
+        // Unbind columns
+        ard->SetHeaderField(SQL_DESC_COUNT, (void*)0, 0);
+        return SQL_SUCCESS;
+      });
     }
 
     // SQLBindParameter is not supported
@@ -976,6 +983,56 @@ 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 SQLBindCol(SQLHSTMT stmt, SQLUSMALLINT recordNumber, SQLSMALLINT cType,
+                     SQLPOINTER dataPtr, SQLLEN bufferLength, SQLLEN* indicatorPtr) {
+  LOG_DEBUG(
+      "SQLBindCol called with stmt: {}, recordNumber: {}, cType: {}, "
+      "dataPtr: {}, bufferLength: {}, strLen_or_IndPtr: {}",
+      stmt, recordNumber, cType, dataPtr, bufferLength, fmt::ptr(indicatorPtr));
+  using ODBC::ODBCDescriptor;
+  using ODBC::ODBCStatement;
+  return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
+    // GH-47021: implement driver to return indicator value when data pointer is null
+    ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
+    ODBCDescriptor* ard = statement->GetARD();
+    ard->BindCol(recordNumber, cType, dataPtr, bufferLength, indicatorPtr);
+    return SQL_SUCCESS;
+  });
+}
+
 SQLRETURN SQLGetData(SQLHSTMT stmt, SQLUSMALLINT recordNumber, SQLSMALLINT cType,
                      SQLPOINTER dataPtr, SQLLEN bufferLength, SQLLEN* indicatorPtr) {
   // GH-46979: support SQL_C_GUID data type

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

@@ -68,6 +68,11 @@ 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 SQLBindCol(SQLHSTMT stmt, SQLUSMALLINT recordNumber, SQLSMALLINT cType,
+                     SQLPOINTER dataPtr, SQLLEN bufferLength, SQLLEN* indicatorPtr);
 SQLRETURN SQLGetData(SQLHSTMT stmt, SQLUSMALLINT recordNumber, SQLSMALLINT cType,
                      SQLPOINTER dataPtr, SQLLEN bufferLength, SQLLEN* indicatorPtr);
 SQLRETURN SQLMoreResults(SQLHSTMT stmt);

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* row_status_array = 0;
+  if (rowStatusArray) {
+    // For SQLExtendedFetch only
+    row_status_array = rowStatusArray;
+  } else {
+    row_status_array = m_ird->GetArrayStatusPtr();
+  }
+
+  size_t rowsFetched =
+      m_currenResult->Move(rows, m_currentArd->GetBindOffset(),
+                           m_currentArd->GetBoundStructOffset(), row_status_array);
   m_ird->SetRowsProcessed(static_cast<SQLULEN>(rowsFetched));
 
+  if (rowCountPtr) {
+    // For SQLExtendedFetch only
+    *rowCountPtr = rowsFetched;
+  }
+
   m_rowNumber += rowsFetched;
   m_hasReachedEndOfResult = rowsFetched != rows;
   return rowsFetched != 0;