Reuse connections across test suite

Author: justing-bq

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

@@ -42,6 +42,11 @@ set(ARROW_FLIGHT_SQL_ODBC_TEST_SRCS
     # GH-46889: move protobuf_test_util to a more common location
     ../../../../engine/substrait/protobuf_test_util.cc)
 
+# Resolve segmentation fault error on Windows platform due to Arrow Compute Library Initialization
+if(WIN32)
+  list(APPEND ARROW_FLIGHT_SQL_ODBC_TEST_SRCS ../../../../compute/test_env.cc)
+endif()
+
 # GH-49652: TODO support static test linkage on macOS
 if(ARROW_TEST_LINKAGE STREQUAL "static")
   set(ARROW_FLIGHT_SQL_ODBC_TEST_LINK_LIBS arrow_flight_sql_odbc_static

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

@@ -87,15 +87,18 @@ TYPED_TEST(ConnectionAttributeTest, TestSQLSetConnectAttrEnlistInDtcUnsupported)
 }
 
 TYPED_TEST(ConnectionAttributeTest, TestSQLSetConnectAttrOdbcCursorsDMOnly) {
-  this->AllocEnvConnHandles();
+  SQLHENV test_env = SQL_NULL_HENV;
+  SQLHDBC test_conn = SQL_NULL_HDBC;
+  this->AllocEnvConnHandles(test_env, test_conn);
 
   // Verify DM-only attribute is settable via Driver Manager
   ASSERT_EQ(SQL_SUCCESS,
-            SQLSetConnectAttr(conn, SQL_ATTR_ODBC_CURSORS,
+            SQLSetConnectAttr(test_conn, SQL_ATTR_ODBC_CURSORS,
                               reinterpret_cast<SQLPOINTER>(SQL_CUR_USE_DRIVER), 0));
 
   std::string connect_str = this->GetConnectionString();
-  this->ConnectWithString(connect_str);
+  this->ConnectWithString(connect_str, test_conn);
+  this->Disconnect(test_env, test_conn);
 }
 
 TYPED_TEST(ConnectionAttributeTest, TestSQLSetConnectAttrQuietModeReadOnly) {

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

@@ -614,7 +614,7 @@ TYPED_TEST(ConnectionInfoTest, TestSQLGetInfoAlterTable) {
 TYPED_TEST(ConnectionInfoHandleTest, TestSQLGetInfoCatalogLocation) {
   // GH-49482 TODO: resolve inconsitent return value for SQL_CATALOG_LOCATION and change
   // test type to `ConnectionInfoTest`
-  this->ConnectWithString(this->GetConnectionString());
+  this->ConnectWithString(this->GetConnectionString(), conn);
 
   SQLUSMALLINT value;
   GetInfo(conn, SQL_CATALOG_LOCATION, &value);
@@ -750,7 +750,7 @@ TYPED_TEST(ConnectionInfoTest, TestSQLGetInfoDropDomain) {
 TYPED_TEST(ConnectionInfoHandleTest, TestSQLGetInfoDropSchema) {
   // GH-49482 TODO: resolve inconsitent return value for SQL_DROP_SCHEMA and change test
   // type to `ConnectionInfoTest`
-  this->ConnectWithString(this->GetConnectionString());
+  this->ConnectWithString(this->GetConnectionString(), conn);
 
   SQLUINTEGER value;
   GetInfo(conn, SQL_DROP_SCHEMA, &value);
@@ -764,7 +764,7 @@ TYPED_TEST(ConnectionInfoHandleTest, TestSQLGetInfoDropSchema) {
 TYPED_TEST(ConnectionInfoHandleTest, TestSQLGetInfoDropTable) {
   // GH-49482 TODO: resolve inconsitent return value for SQL_DROP_TABLE and change test
   // type to `ConnectionInfoTest`
-  this->ConnectWithString(this->GetConnectionString());
+  this->ConnectWithString(this->GetConnectionString(), conn);
 
   SQLUINTEGER value;
   GetInfo(conn, SQL_DROP_TABLE, &value);

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

@@ -51,28 +51,64 @@ class MockServerEnvironment : public ::testing::Environment {
   }
 };
 
-::testing::Environment* mock_env =
+bool RunningRemoteTests() { return !remote_test_connect_str.empty(); }
+
+class OdbcTestEnvironment : public ::testing::Environment {
+ public:
+  void SetUp() override {
+    remote_test_connect_str = ODBCTestBase::GetConnectionString();
+    if (RunningRemoteTests()) {
+      ODBCTestBase::Connect(remote_test_connect_str, remote_odbcv3_handles.env,
+                            remote_odbcv3_handles.conn, SQL_OV_ODBC3);
+      ODBCTestBase::Connect(remote_test_connect_str, remote_odbcv2_handles.env,
+                            remote_odbcv2_handles.conn, SQL_OV_ODBC2);
+    }
+
+    std::string mock_test_connect_str = ODBCMockTestBase::GetConnectionString();
+    ODBCMockTestBase::Connect(mock_test_connect_str, mock_odbcv3_handles.env,
+                              mock_odbcv3_handles.conn, SQL_OV_ODBC3);
+    ODBCMockTestBase::Connect(mock_test_connect_str, mock_odbcv2_handles.env,
+                              mock_odbcv2_handles.conn, SQL_OV_ODBC2);
+  }
+
+  void TearDown() override {
+    if (RunningRemoteTests()) {
+      ODBCTestBase::Disconnect(remote_odbcv3_handles.env, remote_odbcv3_handles.conn);
+      ODBCTestBase::Disconnect(remote_odbcv2_handles.env, remote_odbcv2_handles.conn);
+    }
+
+    ODBCTestBase::Disconnect(mock_odbcv3_handles.env, mock_odbcv3_handles.conn);
+    ODBCTestBase::Disconnect(mock_odbcv2_handles.env, mock_odbcv2_handles.conn);
+  }
+};
+
+::testing::Environment* mock_server_env =
     ::testing::AddGlobalTestEnvironment(new MockServerEnvironment);
 
-void ODBCTestBase::AllocEnvConnHandles(SQLINTEGER odbc_ver) {
+::testing::Environment* odbc_test_env =
+    ::testing::AddGlobalTestEnvironment(new OdbcTestEnvironment);
+
+void ODBCTestBase::AllocEnvConnHandles(SQLHENV& env_handle, SQLHDBC& conn_handle,
+                                       SQLINTEGER odbc_ver) {
   // Allocate an environment handle
-  ASSERT_EQ(SQL_SUCCESS, SQLAllocEnv(&env));
+  ASSERT_EQ(SQL_SUCCESS, SQLAllocEnv(&env_handle));
 
   ASSERT_EQ(
       SQL_SUCCESS,
-      SQLSetEnvAttr(env, SQL_ATTR_ODBC_VERSION,
+      SQLSetEnvAttr(env_handle, SQL_ATTR_ODBC_VERSION,
                     reinterpret_cast<SQLPOINTER>(static_cast<intptr_t>(odbc_ver)), 0));
 
   // Allocate a connection using alloc handle
-  ASSERT_EQ(SQL_SUCCESS, SQLAllocHandle(SQL_HANDLE_DBC, env, &conn));
+  ASSERT_EQ(SQL_SUCCESS, SQLAllocHandle(SQL_HANDLE_DBC, env_handle, &conn_handle));
 }
 
-void ODBCTestBase::Connect(std::string connect_str, SQLINTEGER odbc_ver) {
-  ASSERT_NO_FATAL_FAILURE(AllocEnvConnHandles(odbc_ver));
-  ASSERT_NO_FATAL_FAILURE(ConnectWithString(connect_str));
+void ODBCTestBase::Connect(std::string connect_str, SQLHENV& env_handle,
+                           SQLHDBC& conn_handle, SQLINTEGER odbc_ver) {
+  ASSERT_NO_FATAL_FAILURE(AllocEnvConnHandles(env_handle, conn_handle, odbc_ver));
+  ASSERT_NO_FATAL_FAILURE(ConnectWithString(connect_str, conn_handle));
 }
 
-void ODBCTestBase::ConnectWithString(std::string connect_str) {
+void ODBCTestBase::ConnectWithString(std::string connect_str, SQLHDBC& conn_handle) {
   // Connect string
   std::vector<SQLWCHAR> connect_str0(connect_str.begin(), connect_str.end());
 
@@ -81,31 +117,39 @@ void ODBCTestBase::ConnectWithString(std::string connect_str) {
 
   // Connecting to ODBC server.
   ASSERT_EQ(SQL_SUCCESS,
-            SQLDriverConnect(conn, NULL, &connect_str0[0],
+            SQLDriverConnect(conn_handle, NULL, &connect_str0[0],
                              static_cast<SQLSMALLINT>(connect_str0.size()), out_str,
                              kOdbcBufferSize, &out_str_len, SQL_DRIVER_NOPROMPT))
-      << GetOdbcErrorMessage(SQL_HANDLE_DBC, conn);
+      << GetOdbcErrorMessage(SQL_HANDLE_DBC, conn_handle);
 }
 
-void ODBCTestBase::Disconnect() {
+void ODBCTestBase::Disconnect(SQLHENV& env_handle, SQLHDBC& conn_handle) {
   // Disconnect from ODBC
-  EXPECT_EQ(SQL_SUCCESS, SQLDisconnect(conn))
-      << GetOdbcErrorMessage(SQL_HANDLE_DBC, conn);
+  if (conn_handle != SQL_NULL_HDBC) {
+    EXPECT_EQ(SQL_SUCCESS, SQLDisconnect(conn_handle))
+        << GetOdbcErrorMessage(SQL_HANDLE_DBC, conn_handle);
+  }
 
-  FreeEnvConnHandles();
+  FreeEnvConnHandles(env_handle, conn_handle);
 }
 
-void ODBCTestBase::FreeEnvConnHandles() {
+void ODBCTestBase::FreeEnvConnHandles(SQLHENV& env_handle, SQLHDBC& conn_handle) {
   // Free connection handle
-  EXPECT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_DBC, conn));
+  if (conn_handle != SQL_NULL_HDBC) {
+    EXPECT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_DBC, conn_handle));
+    conn_handle = SQL_NULL_HDBC;
+  }
 
   // Free environment handle
-  EXPECT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_ENV, env));
+  if (env_handle != SQL_NULL_HENV) {
+    EXPECT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_ENV, env_handle));
+    env_handle = SQL_NULL_HENV;
+  }
 }
 
 std::string ODBCTestBase::GetConnectionString() {
   std::string connect_str =
-      arrow::internal::GetEnvVar(kTestConnectStr.data()).ValueOrDie();
+      arrow::internal::GetEnvVar(kTestConnectStr.data()).ValueOr("");
   return connect_str;
 }
 
@@ -168,68 +212,58 @@ std::wstring ODBCTestBase::GetQueryAllDataTypes() {
 }
 
 void ODBCTestBase::SetUp() {
-  if (connected) {
-    ASSERT_EQ(SQL_SUCCESS, SQLAllocHandle(SQL_HANDLE_STMT, conn, &stmt));
-  }
+  ASSERT_EQ(SQL_SUCCESS, SQLAllocHandle(SQL_HANDLE_STMT, conn, &stmt));
 }
 
 void ODBCTestBase::TearDown() {
-  if (connected) {
-    ASSERT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_STMT, stmt));
-  }
-}
-
-void ODBCTestBase::TearDownTestSuite() {
-  if (connected) {
-    Disconnect();
-    connected = false;
-  }
-}
-
-void FlightSQLODBCRemoteTestBase::CheckForRemoteTest() {
-  if (arrow::internal::GetEnvVar(kTestConnectStr.data()).ValueOr("").empty()) {
-    skipping_test = true;
-    GTEST_SKIP() << "Skipping test: kTestConnectStr not set";
-  }
+  ASSERT_EQ(SQL_SUCCESS, SQLFreeHandle(SQL_HANDLE_STMT, stmt));
 }
 
 void FlightSQLODBCRemoteTestBase::SetUpTestSuite() {
-  CheckForRemoteTest();
-  if (skipping_test) {
+  if (!RunningRemoteTests()) {
+    GTEST_SKIP() << "Skipping Test Suite: Environment Variable " << kTestConnectStr.data()
+                 << " is not set";
     return;
   }
 
-  std::string connect_str = GetConnectionString();
-  Connect(connect_str, SQL_OV_ODBC3);
-  connected = true;
+  env = remote_odbcv3_handles.env;
+  conn = remote_odbcv3_handles.conn;
+  stmt = remote_odbcv3_handles.stmt;
 }
 
 void FlightSQLOdbcV2RemoteTestBase::SetUpTestSuite() {
-  CheckForRemoteTest();
-  if (skipping_test) {
+  if (!RunningRemoteTests()) {
+    GTEST_SKIP() << "Skipping Test Suite: Environment Variable " << kTestConnectStr.data()
+                 << " is not set";
     return;
   }
 
-  std::string connect_str = GetConnectionString();
-  Connect(connect_str, SQL_OV_ODBC2);
-  connected = true;
+  env = remote_odbcv2_handles.env;
+  conn = remote_odbcv2_handles.conn;
+  stmt = remote_odbcv2_handles.stmt;
 }
 
 void FlightSQLOdbcEnvConnHandleRemoteTestBase::SetUpTestSuite() {
-  CheckForRemoteTest();
-  if (skipping_test) {
+  if (!RunningRemoteTests()) {
+    GTEST_SKIP() << "Skipping Test Suite: Environment Variable " << kTestConnectStr.data()
+                 << " is not set";
     return;
   }
 
-  AllocEnvConnHandles();
+  AllocEnvConnHandles(remote_non_connection_handles.env,
+                      remote_non_connection_handles.conn);
+  env = remote_non_connection_handles.env;
+  conn = remote_non_connection_handles.conn;
+  stmt = remote_non_connection_handles.stmt;
 }
 
 void FlightSQLOdbcEnvConnHandleRemoteTestBase::TearDownTestSuite() {
-  if (skipping_test) {
+  if (!RunningRemoteTests()) {
     return;
   }
 
-  FreeEnvConnHandles();
+  FreeEnvConnHandles(remote_non_connection_handles.env,
+                     remote_non_connection_handles.conn);
 }
 
 std::string FindTokenInCallHeaders(const CallHeaders& incoming_headers) {
@@ -400,20 +434,27 @@ void ODBCMockTestBase::DropUnicodeTable() {
 }
 
 void FlightSQLODBCMockTestBase::SetUpTestSuite() {
-  std::string connect_str = GetConnectionString();
-  Connect(connect_str, SQL_OV_ODBC3);
-  connected = true;
+  env = mock_odbcv3_handles.env;
+  conn = mock_odbcv3_handles.conn;
+  stmt = mock_odbcv3_handles.stmt;
 }
 
 void FlightSQLOdbcV2MockTestBase::SetUpTestSuite() {
-  std::string connect_str = GetConnectionString();
-  Connect(connect_str, SQL_OV_ODBC2);
-  connected = true;
+  env = mock_odbcv2_handles.env;
+  conn = mock_odbcv2_handles.conn;
+  stmt = mock_odbcv2_handles.stmt;
 }
 
-void FlightSQLOdbcEnvConnHandleMockTestBase::SetUpTestSuite() { AllocEnvConnHandles(); }
+void FlightSQLOdbcEnvConnHandleMockTestBase::SetUpTestSuite() {
+  AllocEnvConnHandles(mock_non_connection_handles.env, mock_non_connection_handles.conn);
+  env = mock_non_connection_handles.env;
+  conn = mock_non_connection_handles.conn;
+  stmt = mock_non_connection_handles.stmt;
+}
 
-void FlightSQLOdbcEnvConnHandleMockTestBase::TearDownTestSuite() { FreeEnvConnHandles(); }
+void FlightSQLOdbcEnvConnHandleMockTestBase::TearDownTestSuite() {
+  FreeEnvConnHandles(mock_non_connection_handles.env, mock_non_connection_handles.conn);
+}
 
 bool CompareConnPropertyMap(Connection::ConnPropertyMap map1,
                             Connection::ConnPropertyMap map2) {