Skip to content

Commit 26304de

Browse files
committed
Extract SQLFetch & SQLGetData
Co-Authored-By: alinalibq <[email protected]>
1 parent 42f27ab commit 26304de

File tree

8 files changed

+1363
-47
lines changed

8 files changed

+1363
-47
lines changed

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

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -823,8 +823,24 @@ SQLRETURN SQLExecute(SQLHSTMT stmt) {
823823

824824
SQLRETURN SQLFetch(SQLHSTMT stmt) {
825825
ARROW_LOG(DEBUG) << "SQLFetch called with stmt: " << stmt;
826-
// GH-47713 TODO: Implement SQLFetch
827-
return SQL_INVALID_HANDLE;
826+
827+
using ODBC::ODBCDescriptor;
828+
using ODBC::ODBCStatement;
829+
return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
830+
ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
831+
832+
// The SQL_ATTR_ROW_ARRAY_SIZE statement attribute specifies the number of rows in the
833+
// rowset.
834+
ODBCDescriptor* ard = statement->GetARD();
835+
size_t rows = static_cast<size_t>(ard->GetArraySize());
836+
837+
if (statement->Fetch(rows)) {
838+
return SQL_SUCCESS;
839+
} else {
840+
// Reached the end of rowset
841+
return SQL_NO_DATA;
842+
}
843+
});
828844
}
829845

830846
SQLRETURN SQLExtendedFetch(SQLHSTMT stmt, SQLUSMALLINT fetch_orientation,
@@ -876,8 +892,14 @@ SQLRETURN SQLGetData(SQLHSTMT stmt, SQLUSMALLINT record_number, SQLSMALLINT c_ty
876892
<< ", record_number: " << record_number << ", c_type: " << c_type
877893
<< ", data_ptr: " << data_ptr << ", buffer_length: " << buffer_length
878894
<< ", indicator_ptr: " << static_cast<const void*>(indicator_ptr);
879-
// GH-47713 TODO: Implement SQLGetData
880-
return SQL_INVALID_HANDLE;
895+
896+
using ODBC::ODBCStatement;
897+
898+
return ODBCStatement::ExecuteWithDiagnostics(stmt, SQL_ERROR, [=]() {
899+
ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(stmt);
900+
return statement->GetData(record_number, c_type, data_ptr, buffer_length,
901+
indicator_ptr);
902+
});
881903
}
882904

883905
SQLRETURN SQLMoreResults(SQLHSTMT stmt) {

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
#include "arrow/flight/sql/odbc/odbc_impl/flight_sql_result_set.h"
1919

20+
#include <sql.h>
21+
2022
#include <utility>
2123
#include "arrow/flight/types.h"
2224
#include "arrow/scalar.h"
@@ -212,14 +214,14 @@ void FlightSqlResultSet::Cancel() {
212214
current_chunk_.data = nullptr;
213215
}
214216

215-
bool FlightSqlResultSet::GetData(int column_n, int16_t target_type, int precision,
216-
int scale, void* buffer, size_t buffer_length,
217-
ssize_t* str_len_buffer) {
217+
SQLRETURN FlightSqlResultSet::GetData(int column_n, int16_t target_type, int precision,
218+
int scale, void* buffer, size_t buffer_length,
219+
ssize_t* str_len_buffer) {
218220
reset_get_data_ = true;
219221
// Check if the offset is already at the end.
220222
int64_t& value_offset = get_data_offsets_[column_n - 1];
221223
if (value_offset == -1) {
222-
return false;
224+
return SQL_NO_DATA;
223225
}
224226

225227
ColumnBinding binding(util::ConvertCDataTypeFromV2ToV3(target_type), precision, scale,
@@ -235,7 +237,11 @@ bool FlightSqlResultSet::GetData(int column_n, int16_t target_type, int precisio
235237
diagnostics_, nullptr);
236238

237239
// If there was truncation, the converter would have reported it to the diagnostics.
238-
return diagnostics_.HasWarning();
240+
if (diagnostics_.HasWarning()) {
241+
return SQL_SUCCESS_WITH_INFO;
242+
} else {
243+
return SQL_SUCCESS;
244+
}
239245
}
240246

241247
std::shared_ptr<ResultSetMetadata> FlightSqlResultSet::GetMetadata() { return metadata_; }

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,8 @@ class FlightSqlResultSet : public ResultSet {
6060

6161
void Cancel() override;
6262

63-
bool GetData(int column_n, int16_t target_type, int precision, int scale, void* buffer,
64-
size_t buffer_length, ssize_t* str_len_buffer) override;
63+
SQLRETURN GetData(int column_n, int16_t target_type, int precision, int scale,
64+
void* buffer, size_t buffer_length, ssize_t* str_len_buffer) override;
6565

6666
size_t Move(size_t rows, size_t bind_offset, size_t bind_type,
6767
uint16_t* row_status_array) override;

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

Lines changed: 40 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,9 @@ SQLSMALLINT getc_typeForSQLType(const DescriptorRecord& record) {
129129
case SQL_WLONGVARCHAR:
130130
return SQL_C_WCHAR;
131131

132+
case SQL_BIT:
133+
return SQL_C_BIT;
134+
132135
case SQL_BINARY:
133136
case SQL_VARBINARY:
134137
case SQL_LONGVARBINARY:
@@ -146,13 +149,20 @@ SQLSMALLINT getc_typeForSQLType(const DescriptorRecord& record) {
146149
case SQL_BIGINT:
147150
return record.is_unsigned ? SQL_C_UBIGINT : SQL_C_SBIGINT;
148151

152+
case SQL_NUMERIC:
153+
case SQL_DECIMAL:
154+
return SQL_C_NUMERIC;
155+
156+
case SQL_FLOAT:
149157
case SQL_REAL:
150158
return SQL_C_FLOAT;
151159

152-
case SQL_FLOAT:
153160
case SQL_DOUBLE:
154161
return SQL_C_DOUBLE;
155162

163+
case SQL_GUID:
164+
return SQL_C_GUID;
165+
156166
case SQL_DATE:
157167
case SQL_TYPE_DATE:
158168
return SQL_C_TYPE_DATE;
@@ -165,32 +175,32 @@ SQLSMALLINT getc_typeForSQLType(const DescriptorRecord& record) {
165175
case SQL_TYPE_TIMESTAMP:
166176
return SQL_C_TYPE_TIMESTAMP;
167177

168-
case SQL_C_INTERVAL_DAY:
169-
return SQL_INTERVAL_DAY;
170-
case SQL_C_INTERVAL_DAY_TO_HOUR:
171-
return SQL_INTERVAL_DAY_TO_HOUR;
172-
case SQL_C_INTERVAL_DAY_TO_MINUTE:
173-
return SQL_INTERVAL_DAY_TO_MINUTE;
174-
case SQL_C_INTERVAL_DAY_TO_SECOND:
175-
return SQL_INTERVAL_DAY_TO_SECOND;
176-
case SQL_C_INTERVAL_HOUR:
177-
return SQL_INTERVAL_HOUR;
178-
case SQL_C_INTERVAL_HOUR_TO_MINUTE:
179-
return SQL_INTERVAL_HOUR_TO_MINUTE;
180-
case SQL_C_INTERVAL_HOUR_TO_SECOND:
181-
return SQL_INTERVAL_HOUR_TO_SECOND;
182-
case SQL_C_INTERVAL_MINUTE:
183-
return SQL_INTERVAL_MINUTE;
184-
case SQL_C_INTERVAL_MINUTE_TO_SECOND:
185-
return SQL_INTERVAL_MINUTE_TO_SECOND;
186-
case SQL_C_INTERVAL_SECOND:
187-
return SQL_INTERVAL_SECOND;
188-
case SQL_C_INTERVAL_YEAR:
189-
return SQL_INTERVAL_YEAR;
190-
case SQL_C_INTERVAL_YEAR_TO_MONTH:
191-
return SQL_INTERVAL_YEAR_TO_MONTH;
192-
case SQL_C_INTERVAL_MONTH:
193-
return SQL_INTERVAL_MONTH;
178+
case SQL_INTERVAL_DAY:
179+
return SQL_C_INTERVAL_DAY;
180+
case SQL_INTERVAL_DAY_TO_HOUR:
181+
return SQL_C_INTERVAL_DAY_TO_HOUR;
182+
case SQL_INTERVAL_DAY_TO_MINUTE:
183+
return SQL_C_INTERVAL_DAY_TO_MINUTE;
184+
case SQL_INTERVAL_DAY_TO_SECOND:
185+
return SQL_C_INTERVAL_DAY_TO_SECOND;
186+
case SQL_INTERVAL_HOUR:
187+
return SQL_C_INTERVAL_HOUR;
188+
case SQL_INTERVAL_HOUR_TO_MINUTE:
189+
return SQL_C_INTERVAL_HOUR_TO_MINUTE;
190+
case SQL_INTERVAL_HOUR_TO_SECOND:
191+
return SQL_C_INTERVAL_HOUR_TO_SECOND;
192+
case SQL_INTERVAL_MINUTE:
193+
return SQL_C_INTERVAL_MINUTE;
194+
case SQL_INTERVAL_MINUTE_TO_SECOND:
195+
return SQL_C_INTERVAL_MINUTE_TO_SECOND;
196+
case SQL_INTERVAL_SECOND:
197+
return SQL_C_INTERVAL_SECOND;
198+
case SQL_INTERVAL_YEAR:
199+
return SQL_C_INTERVAL_YEAR;
200+
case SQL_INTERVAL_YEAR_TO_MONTH:
201+
return SQL_C_INTERVAL_YEAR_TO_MONTH;
202+
case SQL_INTERVAL_MONTH:
203+
return SQL_C_INTERVAL_MONTH;
194204

195205
default:
196206
throw DriverException("Unknown SQL type: " + std::to_string(record.concise_type),
@@ -691,9 +701,9 @@ void ODBCStatement::CloseCursor(bool suppress_errors) {
691701
has_reached_end_of_result_ = false;
692702
}
693703

694-
bool ODBCStatement::GetData(SQLSMALLINT record_number, SQLSMALLINT c_type,
695-
SQLPOINTER data_ptr, SQLLEN buffer_length,
696-
SQLLEN* indicator_ptr) {
704+
SQLRETURN ODBCStatement::GetData(SQLSMALLINT record_number, SQLSMALLINT c_type,
705+
SQLPOINTER data_ptr, SQLLEN buffer_length,
706+
SQLLEN* indicator_ptr) {
697707
if (record_number == 0) {
698708
throw DriverException("Bookmarks are not supported", "07009");
699709
} else if (record_number > ird_->GetRecords().size()) {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,8 @@ class ODBCStatement : public ODBCHandle<ODBCStatement> {
7777

7878
inline SQLULEN GetRowsetSize() { return rowset_size_; }
7979

80-
bool GetData(SQLSMALLINT record_number, SQLSMALLINT c_type, SQLPOINTER data_ptr,
81-
SQLLEN buffer_length, SQLLEN* indicator_ptr);
80+
SQLRETURN GetData(SQLSMALLINT record_number, SQLSMALLINT c_type, SQLPOINTER data_ptr,
81+
SQLLEN buffer_length, SQLLEN* indicator_ptr);
8282

8383
/**
8484
* @brief Closes the cursor. This does _not_ un-prepare the statement or change

cpp/src/arrow/flight/sql/odbc/odbc_impl/spi/result_set.h

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424

2525
#include "arrow/flight/sql/odbc/odbc_impl/types.h"
2626

27+
#include <sqltypes.h>
28+
2729
namespace arrow::flight::sql::odbc {
2830

2931
class ResultSetMetadata;
@@ -87,10 +89,10 @@ class ResultSet {
8789
/// \param buffer Target buffer to be populated.
8890
/// \param buffer_length Target buffer length.
8991
/// \param strlen_buffer Buffer that holds the length of value being fetched.
90-
/// \returns true if there is more data to fetch from the current cell;
91-
/// false if the whole value was already fetched.
92-
virtual bool GetData(int column, int16_t target_type, int precision, int scale,
93-
void* buffer, size_t buffer_length, ssize_t* strlen_buffer) = 0;
92+
/// \returns SQLRETURN for SQLGetData.
93+
virtual SQLRETURN GetData(int column, int16_t target_type, int precision, int scale,
94+
void* buffer, size_t buffer_length,
95+
ssize_t* strlen_buffer) = 0;
9496
};
9597

9698
} // namespace arrow::flight::sql::odbc

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ add_arrow_test(flight_sql_odbc_test
3535
odbc_test_suite.cc
3636
odbc_test_suite.h
3737
connection_test.cc
38+
statement_test.cc
3839
# Enable Protobuf cleanup after test execution
3940
# GH-46889: move protobuf_test_util to a more common location
4041
../../../../engine/substrait/protobuf_test_util.cc

0 commit comments

Comments
 (0)