Extract SQLExecDirect, SQLExecute, SQLPrepare implementation

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

Author: alinaliBQ

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

@@ -1005,22 +1005,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) {

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

@@ -69,6 +69,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) {
@@ -119,7 +123,6 @@ bool FlightSqlStatement::ExecutePrepared() {
 
   Result<std::shared_ptr<FlightInfo>> result =
       prepared_statement_->Execute(call_options_);
-
   ThrowIfNotOK(result.status());
 
   current_result_set_ = std::make_shared<FlightSqlResultSet>(

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

@@ -49,6 +49,7 @@ class FlightSqlStatement : public Statement {
   FlightSqlStatement(const Diagnostics& diagnostics, FlightSqlClient& sql_client,
                      FlightClientOptions client_options, FlightCallOptions call_options,
                      const MetadataSettings& metadata_settings);
+  ~FlightSqlStatement();
 
   bool SetAttribute(StatementAttributeId attribute, const Attribute& value) override;
 

cpp/src/arrow/flight/sql/odbc/tests/CMakeLists.txt

@@ -37,6 +37,7 @@ add_arrow_test(flight_sql_odbc_test
                connection_attr_test.cc
                connection_test.cc
                statement_attr_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

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

@@ -0,0 +1,117 @@
+// 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())));
+
+  // GH-47713 TODO: Uncomment call to SQLFetch SQLGetData after implementation
+  /*
+  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