Extract SQLExecDirect, SQLExecute, SQLPrepare implementation

alinaliBQ · alinaliBQ · commit 72e8db0aabf8 · 2025-11-03T11:57:53.000-08:00
Co-Authored-By: alinalibq &lt;alina.li@improving.com&gt;
diff --git a/cpp/src/arrow/flight/sql/odbc/odbc_api.cc b/cpp/src/arrow/flight/sql/odbc/odbc_api.cc
@@ -803,22 +803,49 @@ SQLRETURN SQLExecDirect(SQLHSTMT stmt, SQLWCHAR* query_text, SQLINTEGER text_len
   ARROW_LOG(DEBUG) << "SQLExecDirectW called with stmt: " << stmt
                    << ", query_text: " << static_cast<const void*>(query_text)
                    << ", text_length: " << text_length;
-  // GH-47711 TODO: Implement SQLExecDirect
-  return SQL_INVALID_HANDLE;
+
+  using ODBC::ODBCStatement;
+  // The driver is built to handle SELECT statements only.
+  return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
+    ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
+    std::string query = ODBC::SqlWcharToString(query_text, text_length);
+
+    statement->Prepare(query);
+    statement->ExecutePrepared();
+
+    return SQL_SUCCESS;
+  });
 }
 
 SQLRETURN SQLPrepare(SQLHSTMT stmt, SQLWCHAR* query_text, SQLINTEGER text_length) {
   ARROW_LOG(DEBUG) << "SQLPrepareW called with stmt: " << stmt
                    << ", query_text: " << static_cast<const void*>(query_text)
                    << ", text_length: " << text_length;
-  // GH-47712 TODO: Implement SQLPrepare
-  return SQL_INVALID_HANDLE;
+
+  using ODBC::ODBCStatement;
+  // The driver is built to handle SELECT statements only.
+  return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
+    ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
+    std::string query = ODBC::SqlWcharToString(query_text, text_length);
+
+    statement->Prepare(query);
+
+    return SQL_SUCCESS;
+  });
 }
 
 SQLRETURN SQLExecute(SQLHSTMT stmt) {
   ARROW_LOG(DEBUG) << "SQLExecute called with stmt: " << stmt;
-  // GH-47712 TODO: Implement SQLExecute
-  return SQL_INVALID_HANDLE;
+
+  using ODBC::ODBCStatement;
+  // The driver is built to handle SELECT statements only.
+  return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
+    ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
+
+    statement->ExecutePrepared();
+
+    return SQL_SUCCESS;
+  });
 }
 
 SQLRETURN SQLFetch(SQLHSTMT stmt) {
diff --git a/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_statement.cc b/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_statement.cc
@@ -41,9 +41,10 @@ using util::ThrowIfNotOK;
 
 namespace {
 
-void ClosePreparedStatementIfAny(std::shared_ptr<PreparedStatement>& prepared_statement) {
+void ClosePreparedStatementIfAny(std::shared_ptr<PreparedStatement>& prepared_statement,
+                                 const FlightCallOptions& options) {
   if (prepared_statement != nullptr) {
-    ThrowIfNotOK(prepared_statement->Close());
+    ThrowIfNotOK(prepared_statement->Close(options));
     prepared_statement.reset();
   }
 }
@@ -66,6 +67,10 @@ FlightSqlStatement::FlightSqlStatement(const Diagnostics& diagnostics,
   call_options_.timeout = TimeoutDuration{-1};
 }
 
+FlightSqlStatement::~FlightSqlStatement() {
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
+}
+
 bool FlightSqlStatement::SetAttribute(StatementAttributeId attribute,
                                       const Attribute& value) {
   switch (attribute) {
@@ -97,7 +102,7 @@ boost::optional<Statement::Attribute> FlightSqlStatement::GetAttribute(
 
 boost::optional<std::shared_ptr<ResultSetMetadata>> FlightSqlStatement::Prepare(
     const std::string& query) {
-  ClosePreparedStatementIfAny(prepared_statement_);
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
 
   Result<std::shared_ptr<PreparedStatement>> result =
       sql_client_.Prepare(call_options_, query);
@@ -113,7 +118,8 @@ boost::optional<std::shared_ptr<ResultSetMetadata>> FlightSqlStatement::Prepare(
 bool FlightSqlStatement::ExecutePrepared() {
   assert(prepared_statement_.get() != nullptr);
 
-  Result<std::shared_ptr<FlightInfo>> result = prepared_statement_->Execute();
+  Result<std::shared_ptr<FlightInfo>> result =
+      prepared_statement_->Execute(call_options_);
   ThrowIfNotOK(result.status());
 
   current_result_set_ = std::make_shared<FlightSqlResultSet>(
@@ -124,7 +130,7 @@ bool FlightSqlStatement::ExecutePrepared() {
 }
 
 bool FlightSqlStatement::Execute(const std::string& query) {
-  ClosePreparedStatementIfAny(prepared_statement_);
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
 
   Result<std::shared_ptr<FlightInfo>> result = sql_client_.Execute(call_options_, query);
   ThrowIfNotOK(result.status());
@@ -146,7 +152,7 @@ std::shared_ptr<ResultSet> FlightSqlStatement::GetTables(
     const std::string* catalog_name, const std::string* schema_name,
     const std::string* table_name, const std::string* table_type,
     const ColumnNames& column_names) {
-  ClosePreparedStatementIfAny(prepared_statement_);
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
 
   std::vector<std::string> table_types;
 
@@ -199,7 +205,7 @@ std::shared_ptr<ResultSet> FlightSqlStatement::GetTables_V3(
 std::shared_ptr<ResultSet> FlightSqlStatement::GetColumns_V2(
     const std::string* catalog_name, const std::string* schema_name,
     const std::string* table_name, const std::string* column_name) {
-  ClosePreparedStatementIfAny(prepared_statement_);
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
 
   Result<std::shared_ptr<FlightInfo>> result = sql_client_.GetTables(
       call_options_, catalog_name, schema_name, table_name, true, nullptr);
@@ -220,7 +226,7 @@ std::shared_ptr<ResultSet> FlightSqlStatement::GetColumns_V2(
 std::shared_ptr<ResultSet> FlightSqlStatement::GetColumns_V3(
     const std::string* catalog_name, const std::string* schema_name,
     const std::string* table_name, const std::string* column_name) {
-  ClosePreparedStatementIfAny(prepared_statement_);
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
 
   Result<std::shared_ptr<FlightInfo>> result = sql_client_.GetTables(
       call_options_, catalog_name, schema_name, table_name, true, nullptr);
@@ -239,7 +245,7 @@ std::shared_ptr<ResultSet> FlightSqlStatement::GetColumns_V3(
 }
 
 std::shared_ptr<ResultSet> FlightSqlStatement::GetTypeInfo_V2(int16_t data_type) {
-  ClosePreparedStatementIfAny(prepared_statement_);
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
 
   Result<std::shared_ptr<FlightInfo>> result = sql_client_.GetXdbcTypeInfo(call_options_);
   ThrowIfNotOK(result.status());
@@ -257,7 +263,7 @@ std::shared_ptr<ResultSet> FlightSqlStatement::GetTypeInfo_V2(int16_t data_type)
 }
 
 std::shared_ptr<ResultSet> FlightSqlStatement::GetTypeInfo_V3(int16_t data_type) {
-  ClosePreparedStatementIfAny(prepared_statement_);
+  ClosePreparedStatementIfAny(prepared_statement_, call_options_);
 
   Result<std::shared_ptr<FlightInfo>> result = sql_client_.GetXdbcTypeInfo(call_options_);
   ThrowIfNotOK(result.status());
diff --git a/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_statement.h b/cpp/src/arrow/flight/sql/odbc/odbc_impl/flight_sql_statement.h
@@ -48,6 +48,7 @@ class FlightSqlStatement : public Statement {
   FlightSqlStatement(const Diagnostics& diagnostics, FlightSqlClient& sql_client,
                      FlightCallOptions call_options,
                      const MetadataSettings& metadata_settings);
+  ~FlightSqlStatement();
 
   bool SetAttribute(StatementAttributeId attribute, const Attribute& value) override;
 
diff --git a/cpp/src/arrow/flight/sql/odbc/tests/CMakeLists.txt b/cpp/src/arrow/flight/sql/odbc/tests/CMakeLists.txt
@@ -35,6 +35,7 @@ add_arrow_test(flight_sql_odbc_test
                odbc_test_suite.cc
                odbc_test_suite.h
                connection_test.cc
+               statement_test.cc
                # Enable Protobuf cleanup after test execution
                # GH-46889: move protobuf_test_util to a more common location
                ../../../../engine/substrait/protobuf_test_util.cc
diff --git a/cpp/src/arrow/flight/sql/odbc/tests/statement_test.cc b/cpp/src/arrow/flight/sql/odbc/tests/statement_test.cc
@@ -0,0 +1,114 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+#include "arrow/flight/sql/odbc/tests/odbc_test_suite.h"
+
+#include "arrow/flight/sql/odbc/odbc_impl/platform.h"
+
+#include <sql.h>
+#include <sqltypes.h>
+#include <sqlucode.h>
+
+#include <limits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace arrow::flight::sql::odbc {
+
+template <typename T>
+class StatementTest : public T {};
+
+class StatementMockTest : public FlightSQLODBCMockTestBase {};
+class StatementRemoteTest : public FlightSQLODBCRemoteTestBase {};
+using TestTypes = ::testing::Types<StatementMockTest, StatementRemoteTest>;
+TYPED_TEST_SUITE(StatementTest, TestTypes);
+
+TYPED_TEST(StatementTest, TestSQLExecDirectSimpleQuery) {
+  std::wstring wsql = L"SELECT 1;";
+  std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());
+
+  ASSERT_EQ(SQL_SUCCESS,
+            SQLExecDirect(this->stmt, &sql0[0], static_cast<SQLINTEGER>(sql0.size())));
+
+  ASSERT_EQ(SQL_SUCCESS, SQLFetch(this->stmt));
+
+  SQLINTEGER val;
+
+  ASSERT_EQ(SQL_SUCCESS, SQLGetData(this->stmt, 1, SQL_C_LONG, &val, 0, 0));
+  // Verify 1 is returned
+  EXPECT_EQ(1, val);
+
+  ASSERT_EQ(SQL_NO_DATA, SQLFetch(this->stmt));
+
+  ASSERT_EQ(SQL_ERROR, SQLGetData(this->stmt, 1, SQL_C_LONG, &val, 0, 0));
+  // Invalid cursor state
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, kErrorState24000);
+}
+
+TYPED_TEST(StatementTest, TestSQLExecDirectInvalidQuery) {
+  std::wstring wsql = L"SELECT;";
+  std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());
+
+  ASSERT_EQ(SQL_ERROR,
+            SQLExecDirect(this->stmt, &sql0[0], static_cast<SQLINTEGER>(sql0.size())));
+  // ODBC provides generic error code HY000 to all statement errors
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, kErrorStateHY000);
+}
+
+TYPED_TEST(StatementTest, TestSQLExecuteSimpleQuery) {
+  std::wstring wsql = L"SELECT 1;";
+  std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());
+
+  ASSERT_EQ(SQL_SUCCESS,
+            SQLPrepare(this->stmt, &sql0[0], static_cast<SQLINTEGER>(sql0.size())));
+
+  ASSERT_EQ(SQL_SUCCESS, SQLExecute(this->stmt));
+
+  // GH-47713 TODO: Uncomment call to SQLFetch SQLGetData after implementation
+  /*
+  // Fetch data
+  ASSERT_EQ(SQL_SUCCESS, SQLFetch(this->stmt));
+
+  SQLINTEGER val;
+  ASSERT_EQ(SQL_SUCCESS, SQLGetData(this->stmt, 1, SQL_C_LONG, &val, 0, 0));
+
+  // Verify 1 is returned
+  EXPECT_EQ(1, val);
+
+  ASSERT_EQ(SQL_NO_DATA, SQLFetch(this->stmt));
+
+  ASSERT_EQ(SQL_ERROR, SQLGetData(this->stmt, 1, SQL_C_LONG, &val, 0, 0));
+  // Invalid cursor state
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, kErrorState24000);
+  */
+}
+
+TYPED_TEST(StatementTest, TestSQLPrepareInvalidQuery) {
+  std::wstring wsql = L"SELECT;";
+  std::vector<SQLWCHAR> sql0(wsql.begin(), wsql.end());
+
+  ASSERT_EQ(SQL_ERROR,
+            SQLPrepare(this->stmt, &sql0[0], static_cast<SQLINTEGER>(sql0.size())));
+  // ODBC provides generic error code HY000 to all statement errors
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, kErrorStateHY000);
+
+  ASSERT_EQ(SQL_ERROR, SQLExecute(this->stmt));
+  // Verify function sequence error state is returned
+  VerifyOdbcErrorState(SQL_HANDLE_STMT, this->stmt, kErrorStateHY010);
+}
+
+}  // namespace arrow::flight::sql::odbc
