Extract SQLNumResultCols implementation

Update docs

Co-Authored-By: alinalibq <[email protected]>

Author: alinaliBQ

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

@@ -1149,8 +1149,13 @@ SQLRETURN SQLNumResultCols(SQLHSTMT stmt, SQLSMALLINT* column_count_ptr) {
   ARROW_LOG(DEBUG) << "SQLNumResultCols called with stmt: " << stmt
                    << ", column_count_ptr: "
                    << static_cast<const void*>(column_count_ptr);
-  // GH-47713 TODO: Implement SQLNumResultCols
-  return SQL_INVALID_HANDLE;
+
+  using ODBC::ODBCStatement;
+  return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
+    ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
+    statement->GetColumnCount(column_count_ptr);
+    return SQL_SUCCESS;
+  });
 }
 
 SQLRETURN SQLRowCount(SQLHSTMT stmt, SQLLEN* row_count_ptr) {

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

@@ -761,6 +761,16 @@ SQLRETURN ODBCStatement::GetData(SQLSMALLINT record_number, SQLSMALLINT c_type,
                                   data_ptr, buffer_length, indicator_ptr);
 }
 
+void ODBCStatement::GetColumnCount(SQLSMALLINT* column_count_ptr) {
+  if (!column_count_ptr) {
+    // column count pointer is not valid, do nothing as ODBC spec does not mention this as
+    // an error
+    return;
+  }
+  size_t column_count = ird_->GetRecords().size();
+  *column_count_ptr = static_cast<SQLSMALLINT>(column_count);
+}
+
 void ODBCStatement::ReleaseStatement() {
   CloseCursor(true);
   connection_.DropStatement(this);

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

@@ -78,6 +78,9 @@ class ODBCStatement : public ODBCHandle<ODBCStatement> {
   SQLRETURN GetData(SQLSMALLINT record_number, SQLSMALLINT c_type, SQLPOINTER data_ptr,
                     SQLLEN buffer_length, SQLLEN* indicator_ptr);
 
+  /// \brief Return number of columns from data set
+  void GetColumnCount(SQLSMALLINT* column_count_ptr);
+
   /// \brief Closes the cursor. This does _not_ un-prepare the statement or change
   /// bindings.
   void CloseCursor(bool suppress_errors);

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

@@ -1281,6 +1281,50 @@ TYPED_TEST(StatementTest, TestSQLNativeSqlReturnsErrorOnBadInputs) {
   VerifyOdbcErrorState(SQL_HANDLE_DBC, this->conn, kErrorStateHY090);
 }
 
+TYPED_TEST(StatementTest, SQLNumResultColsReturnsColumnsOnSelect) {
+  SQLSMALLINT column_count = 0;
+  SQLSMALLINT expected_value = 3;
+  SQLWCHAR sql_query[] = L"SELECT 1 AS col1, 'One' AS col2, 3 AS col3";
+  SQLINTEGER query_length = static_cast<SQLINTEGER>(wcslen(sql_query));
+
+  ASSERT_EQ(SQL_SUCCESS, SQLExecDirect(this->stmt, sql_query, query_length));
+
+  ASSERT_EQ(SQL_SUCCESS, SQLFetch(this->stmt));
+
+  CheckIntColumn(this->stmt, 1, 1);
+  CheckStringColumnW(this->stmt, 2, L"One");
+  CheckIntColumn(this->stmt, 3, 3);
+
+  ASSERT_EQ(SQL_SUCCESS, SQLNumResultCols(this->stmt, &column_count));
+
+  EXPECT_EQ(expected_value, column_count);
+}
+
+TYPED_TEST(StatementTest, SQLNumResultColsReturnsSuccessOnNullptr) {
+  SQLWCHAR sql_query[] = L"SELECT 1 AS col1, 'One' AS col2, 3 AS col3";
+  SQLINTEGER query_length = static_cast<SQLINTEGER>(wcslen(sql_query));
+
+  ASSERT_EQ(SQL_SUCCESS, SQLExecDirect(this->stmt, sql_query, query_length));
+
+  ASSERT_EQ(SQL_SUCCESS, SQLFetch(this->stmt));
+
+  CheckIntColumn(this->stmt, 1, 1);
+  CheckStringColumnW(this->stmt, 2, L"One");
+  CheckIntColumn(this->stmt, 3, 3);
+
+  ASSERT_EQ(SQL_SUCCESS, SQLNumResultCols(this->stmt, nullptr));
+}
+
+TYPED_TEST(StatementTest, SQLNumResultColsFunctionSequenceErrorOnNoQuery) {
+  SQLSMALLINT column_count = 0;
+  SQLSMALLINT expected_value = 0;
+
+  ASSERT_EQ(SQL_ERROR, SQLNumResultCols(this->stmt, &column_count));
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, kErrorStateHY010);
+
+  EXPECT_EQ(expected_value, column_count);
+}
+
 TYPED_TEST(StatementTest, TestSQLCloseCursor) {
   std::wstring wsql = L"SELECT 1;";
   std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());