Extract SQLRowCount implementation

Update doc

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

Author: alinaliBQ

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

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

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

@@ -761,6 +761,17 @@ SQLRETURN ODBCStatement::GetData(SQLSMALLINT record_number, SQLSMALLINT c_type,
                                   data_ptr, buffer_length, indicator_ptr);
 }
 
+void ODBCStatement::GetRowCount(SQLLEN* row_count_ptr) {
+  if (!row_count_ptr) {
+    // row count pointer is not valid, do nothing as ODBC spec does not mention this as an
+    // error
+    return;
+  }
+  // Will always be -1 (meaning number of rows unknown) since only SELECT is supported by
+  // driver
+  *row_count_ptr = -1;
+}
+
 void ODBCStatement::ReleaseStatement() {
   CloseCursor(true);
   connection_.DropStatement(this);

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

@@ -78,6 +78,11 @@ 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 rows affected by an UPDATE, INSERT, or DELETE statement\
+  ///
+  ///  -1 is returned as driver only supports SELECT statement
+  void GetRowCount(SQLLEN* row_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, SQLRowCountReturnsNegativeOneOnSelect) {
+  SQLLEN row_count = 0;
+  SQLLEN expected_value = -1;
+  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, SQLRowCount(this->stmt, &row_count));
+
+  EXPECT_EQ(expected_value, row_count);
+}
+
+TYPED_TEST(StatementTest, SQLRowCountReturnsSuccessOnNullptr) {
+  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, SQLRowCount(this->stmt, nullptr));
+}
+
+TYPED_TEST(StatementTest, SQLRowCountFunctionSequenceErrorOnNoQuery) {
+  SQLLEN row_count = 0;
+  SQLLEN expected_value = 0;
+
+  ASSERT_EQ(SQL_ERROR, SQLRowCount(this->stmt, &row_count));
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, kErrorStateHY010);
+
+  EXPECT_EQ(expected_value, row_count);
+}
+
 TYPED_TEST(StatementTest, TestSQLCloseCursor) {
   std::wstring wsql = L"SELECT 1;";
   std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());