Skip to content

Commit 62afc09

Browse files
alinaliBQalinaliBQ
authored
GH-46096: [C++][FlightRPC] Environment and Connection Handle Allocation (#47759)
### Rationale for this change ODBC driver needs to handle environment handle and connection handle allocation and deallocation. ### What changes are included in this PR? - Implement SQLAllocHandle and SQLFreeHandle APIs for allocating and deallocating environment and connection handles. - Tests ### Are these changes tested? - Tested locally on Windows MSVC. - Will be tested in CI after #47689 is merged ### Are there any user-facing changes? No * GitHub Issue: #46096 * GitHub Issue: #46097 Authored-by: Alina (Xi) Li <[email protected]> Signed-off-by: David Li <[email protected]>
1 parent a625340 commit 62afc09

File tree

2 files changed

+240
-21
lines changed

2 files changed

+240
-21
lines changed

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

Lines changed: 165 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -36,26 +36,179 @@ SQLRETURN SQLAllocHandle(SQLSMALLINT type, SQLHANDLE parent, SQLHANDLE* result)
3636
ARROW_LOG(DEBUG) << "SQLAllocHandle called with type: " << type
3737
<< ", parent: " << parent
3838
<< ", result: " << static_cast<const void*>(result);
39-
// GH-46096 TODO: Implement SQLAllocEnv
40-
// GH-46097 TODO: Implement SQLAllocConnect, pre-requisite requires SQLAllocEnv
41-
// implementation
42-
43-
// GH-47706 TODO: Implement SQLAllocStmt, pre-requisite requires
39+
// GH-47706 TODO: Add tests for SQLAllocStmt, pre-requisite requires
4440
// SQLDriverConnect implementation
4541

46-
// GH-47707 TODO: Implement SQL_HANDLE_DESC for
42+
// GH-47707 TODO: Add tests for SQL_HANDLE_DESC implementation for
4743
// descriptor handle, pre-requisite requires SQLAllocStmt
48-
return SQL_INVALID_HANDLE;
44+
45+
*result = nullptr;
46+
47+
switch (type) {
48+
case SQL_HANDLE_ENV: {
49+
using ODBC::ODBCEnvironment;
50+
51+
*result = SQL_NULL_HENV;
52+
53+
try {
54+
static std::shared_ptr<FlightSqlDriver> odbc_driver =
55+
std::make_shared<FlightSqlDriver>();
56+
*result = reinterpret_cast<SQLHENV>(new ODBCEnvironment(odbc_driver));
57+
58+
return SQL_SUCCESS;
59+
} catch (const std::bad_alloc&) {
60+
// allocating environment failed so cannot log diagnostic error here
61+
return SQL_ERROR;
62+
}
63+
}
64+
65+
case SQL_HANDLE_DBC: {
66+
using ODBC::ODBCConnection;
67+
using ODBC::ODBCEnvironment;
68+
69+
*result = SQL_NULL_HDBC;
70+
71+
ODBCEnvironment* environment = reinterpret_cast<ODBCEnvironment*>(parent);
72+
73+
return ODBCEnvironment::ExecuteWithDiagnostics(environment, SQL_ERROR, [=]() {
74+
std::shared_ptr<ODBCConnection> conn = environment->CreateConnection();
75+
76+
if (conn) {
77+
// Inside `CreateConnection`, the shared_ptr `conn` is kept
78+
// in a `std::vector` of connections inside the environment handle.
79+
// As long as the parent environment handle is alive, the connection shared_ptr
80+
// will be kept alive unless the user frees the connection.
81+
*result = reinterpret_cast<SQLHDBC>(conn.get());
82+
83+
return SQL_SUCCESS;
84+
}
85+
86+
return SQL_ERROR;
87+
});
88+
}
89+
90+
case SQL_HANDLE_STMT: {
91+
using ODBC::ODBCConnection;
92+
using ODBC::ODBCStatement;
93+
94+
*result = SQL_NULL_HSTMT;
95+
96+
ODBCConnection* connection = reinterpret_cast<ODBCConnection*>(parent);
97+
98+
return ODBCConnection::ExecuteWithDiagnostics(connection, SQL_ERROR, [=]() {
99+
std::shared_ptr<ODBCStatement> statement = connection->CreateStatement();
100+
101+
if (statement) {
102+
*result = reinterpret_cast<SQLHSTMT>(statement.get());
103+
104+
return SQL_SUCCESS;
105+
}
106+
107+
return SQL_ERROR;
108+
});
109+
}
110+
111+
case SQL_HANDLE_DESC: {
112+
using ODBC::ODBCConnection;
113+
using ODBC::ODBCDescriptor;
114+
115+
*result = SQL_NULL_HDESC;
116+
117+
ODBCConnection* connection = reinterpret_cast<ODBCConnection*>(parent);
118+
119+
return ODBCConnection::ExecuteWithDiagnostics(connection, SQL_ERROR, [=]() {
120+
std::shared_ptr<ODBCDescriptor> descriptor = connection->CreateDescriptor();
121+
122+
if (descriptor) {
123+
*result = reinterpret_cast<SQLHDESC>(descriptor.get());
124+
125+
return SQL_SUCCESS;
126+
}
127+
128+
return SQL_ERROR;
129+
});
130+
}
131+
132+
default:
133+
break;
134+
}
135+
136+
return SQL_ERROR;
49137
}
50138

51139
SQLRETURN SQLFreeHandle(SQLSMALLINT type, SQLHANDLE handle) {
52140
ARROW_LOG(DEBUG) << "SQLFreeHandle called with type: " << type
53141
<< ", handle: " << handle;
54-
// GH-46096 TODO: Implement SQLFreeEnv
55-
// GH-46097 TODO: Implement SQLFreeConnect
56-
// GH-47706 TODO: Implement SQLFreeStmt
57-
// GH-47707 TODO: Implement SQL_HANDLE_DESC for descriptor handle
58-
return SQL_INVALID_HANDLE;
142+
// GH-47706 TODO: Add tests for SQLFreeStmt, pre-requisite requires
143+
// SQLAllocStmt tests
144+
145+
// GH-47707 TODO: Add tests for SQL_HANDLE_DESC implementation for
146+
// descriptor handle
147+
switch (type) {
148+
case SQL_HANDLE_ENV: {
149+
using ODBC::ODBCEnvironment;
150+
151+
ODBCEnvironment* environment = reinterpret_cast<ODBCEnvironment*>(handle);
152+
153+
if (!environment) {
154+
return SQL_INVALID_HANDLE;
155+
}
156+
157+
delete environment;
158+
159+
return SQL_SUCCESS;
160+
}
161+
162+
case SQL_HANDLE_DBC: {
163+
using ODBC::ODBCConnection;
164+
165+
ODBCConnection* conn = reinterpret_cast<ODBCConnection*>(handle);
166+
167+
if (!conn) {
168+
return SQL_INVALID_HANDLE;
169+
}
170+
171+
// `ReleaseConnection` does the equivalent of `delete`.
172+
// `ReleaseConnection` removes the connection `shared_ptr` from the `std::vector` of
173+
// connections, and the `shared_ptr` is automatically destructed afterwards.
174+
conn->ReleaseConnection();
175+
176+
return SQL_SUCCESS;
177+
}
178+
179+
case SQL_HANDLE_STMT: {
180+
using ODBC::ODBCStatement;
181+
182+
ODBCStatement* statement = reinterpret_cast<ODBCStatement*>(handle);
183+
184+
if (!statement) {
185+
return SQL_INVALID_HANDLE;
186+
}
187+
188+
statement->ReleaseStatement();
189+
190+
return SQL_SUCCESS;
191+
}
192+
193+
case SQL_HANDLE_DESC: {
194+
using ODBC::ODBCDescriptor;
195+
196+
ODBCDescriptor* descriptor = reinterpret_cast<ODBCDescriptor*>(handle);
197+
198+
if (!descriptor) {
199+
return SQL_INVALID_HANDLE;
200+
}
201+
202+
descriptor->ReleaseDescriptor();
203+
204+
return SQL_SUCCESS;
205+
}
206+
207+
default:
208+
break;
209+
}
210+
211+
return SQL_ERROR;
59212
}
60213

61214
SQLRETURN SQLFreeStmt(SQLHSTMT handle, SQLUSMALLINT option) {

‎cpp/src/arrow/flight/sql/odbc/tests/connection_test.cc‎

Lines changed: 75 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,30 +14,96 @@
1414
// KIND, either express or implied. See the License for the
1515
// specific language governing permissions and limitations
1616
// under the License.
17+
#include "arrow/flight/sql/odbc/tests/odbc_test_suite.h"
1718

18-
#ifdef _WIN32
19-
# include <windows.h>
20-
#endif
19+
#include "arrow/flight/sql/odbc/odbc_impl/platform.h"
2120

2221
#include <sql.h>
2322
#include <sqltypes.h>
2423
#include <sqlucode.h>
2524

26-
#include "gmock/gmock.h"
27-
#include "gtest/gtest.h"
25+
#include <gtest/gtest.h>
2826

2927
namespace arrow::flight::sql::odbc {
3028

31-
TEST(SQLAllocHandle, SQLAllocHandleEnv) {
29+
template <typename T>
30+
class ConnectionTest : public T {};
31+
32+
// GH-46574 TODO: add remote server test cases using `ConnectionRemoteTest`
33+
class ConnectionRemoteTest : public FlightSQLODBCRemoteTestBase {};
34+
using TestTypes = ::testing::Types<FlightSQLODBCMockTestBase, ConnectionRemoteTest>;
35+
TYPED_TEST_SUITE(ConnectionTest, TestTypes);
36+
37+
TEST(SQLAllocHandle, TestSQLAllocHandleEnv) {
38+
SQLHENV env;
39+
3240
// Allocate an environment handle
33-
SQLHENV env = nullptr;
3441
ASSERT_EQ(SQL_SUCCESS, SQLAllocHandle(SQL_HANDLE_ENV, SQL_NULL_HANDLE, &env));
3542

36-
// Check for valid handle
37-
ASSERT_NE(nullptr, env);
43+
ASSERT_NE(env, nullptr);
3844

3945
// Free an environment handle
4046
ASSERT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_ENV, env));
4147
}
4248

49+
TEST(SQLAllocEnv, TestSQLAllocEnv) {
50+
SQLHENV env;
51+
52+
// Allocate an environment handle
53+
ASSERT_EQ(SQL_SUCCESS, SQLAllocEnv(&env));
54+
55+
// Free an environment handle
56+
ASSERT_EQ(SQL_SUCCESS, SQLFreeEnv(env));
57+
}
58+
59+
TEST(SQLAllocHandle, TestSQLAllocHandleConnect) {
60+
SQLHENV env;
61+
SQLHDBC conn;
62+
63+
// Allocate an environment handle
64+
ASSERT_EQ(SQL_SUCCESS, SQLAllocEnv(&env));
65+
66+
// Allocate a connection using alloc handle
67+
ASSERT_EQ(SQL_SUCCESS, SQLAllocHandle(SQL_HANDLE_DBC, env, &conn));
68+
69+
// Free a connection handle
70+
ASSERT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_DBC, conn));
71+
72+
// Free an environment handle
73+
ASSERT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_ENV, env));
74+
}
75+
76+
TEST(SQLAllocConnect, TestSQLAllocHandleConnect) {
77+
SQLHENV env;
78+
SQLHDBC conn;
79+
80+
// Allocate an environment handle
81+
ASSERT_EQ(SQL_SUCCESS, SQLAllocEnv(&env));
82+
83+
// Allocate a connection using alloc handle
84+
ASSERT_EQ(SQL_SUCCESS, SQLAllocConnect(env, &conn));
85+
86+
// Free a connection handle
87+
ASSERT_EQ(SQL_SUCCESS, SQLFreeConnect(conn));
88+
89+
// Free an environment handle
90+
ASSERT_EQ(SQL_SUCCESS, SQLFreeEnv(env));
91+
}
92+
93+
TEST(SQLFreeHandle, TestFreeNullHandles) {
94+
SQLHENV env = NULL;
95+
SQLHDBC conn = NULL;
96+
SQLHSTMT stmt = NULL;
97+
98+
// Verifies attempt to free invalid handle does not cause segfault
99+
// Attempt to free null statement handle
100+
ASSERT_EQ(SQL_INVALID_HANDLE, SQLFreeHandle(SQL_HANDLE_STMT, stmt));
101+
102+
// Attempt to free null connection handle
103+
ASSERT_EQ(SQL_INVALID_HANDLE, SQLFreeHandle(SQL_HANDLE_DBC, conn));
104+
105+
// Attempt to free null environment handle
106+
ASSERT_EQ(SQL_INVALID_HANDLE, SQLFreeHandle(SQL_HANDLE_ENV, env));
107+
}
108+
43109
} // namespace arrow::flight::sql::odbc

0 commit comments

Comments
 (0)