refining

Author: Saumya Garg

mssql_python/pybind/connection/connection.cpp

@@ -12,17 +12,24 @@
 // This class wraps low-level ODBC operations like connect/disconnect,
 // transaction control, and autocommit configuration.
 //-------------------------------------------------------------------------------------------------
-Connection::Connection(const std::wstring& conn_str, bool autocommit) : _conn_str(conn_str) , _autocommit(autocommit) {}
+Connection::Connection(const std::wstring& conn_str, bool autocommit)
+    : _conn_str(conn_str) , _autocommit(autocommit) {}
 
 Connection::~Connection() {
-    close(); // Ensure the connection is closed when the object is destroyed.
+    close();    // Ensure the connection is closed when the object is destroyed.
 }
 
 SQLRETURN Connection::connect() {
-    SQLHANDLE env = nullptr;
-    SQLHANDLE dbc = nullptr;
+    allocEnvHandle();
+    setEnvAttributes();
+    allocDbcHandle();
+    return connectToDb();
+}
 
-    LOG("Allocate SQL Handle");
+// Allocates environment handle
+void Connection::allocEnvHandle() {
+    SQLHANDLE env = nullptr;
+    LOG("Allocating Environment Handle");
     if (!SQLAllocHandle_ptr) {
         LOG("Function pointer not initialized. Loading the driver.");
         DriverLoader::getInstance().loadDriver();
@@ -34,21 +41,35 @@ SQLRETURN Connection::connect() {
     }
     _env_handle = std::make_shared<SqlHandle>(SQL_HANDLE_ENV, env);
 
-    ret =  SQLSetEnvAttr_ptr(env, SQL_ATTR_ODBC_VERSION, (void*)SQL_OV_ODBC3_80, 0);
+}
+
+// Sets environment attributes
+void Connection::setEnvAttributes() {
+    LOG("Setting environment attributes");
+    SQLRETURN ret =  SQLSetEnvAttr_ptr(_env_handle->get(), SQL_ATTR_ODBC_VERSION, (void*)SQL_OV_ODBC3_80, 0);
     if (!SQL_SUCCEEDED(ret)) {
         LOG("Failed to set environment attribute");
         throw std::runtime_error("Failed to set environment attribute");
     }
+}
 
+// Allocates DBC handle
+void Connection::allocDbcHandle() {
+    SQLHANDLE dbc = nullptr;
     LOG("Allocate SQL Connection Handle");
-    ret = SQLAllocHandle_ptr(SQL_HANDLE_DBC, env, &dbc);
+    SQLRETURN ret = SQLAllocHandle_ptr(SQL_HANDLE_DBC, _env_handle->get(), &dbc);
     if (!SQL_SUCCEEDED(ret)) {
         LOG("Failed to allocate connection handle");
         throw std::runtime_error("Failed to allocate connection handle");
     }
     _dbc_handle = std::make_shared<SqlHandle>(SQL_HANDLE_DBC, dbc);
 
-    ret = SQLDriverConnect_ptr(dbc, nullptr,
+}
+
+// Connects to the database
+SQLRETURN Connection::connectToDb() {
+    LOG("Connecting to database");
+    SQLRETURN ret = SQLDriverConnect_ptr(_dbc_handle->get(), nullptr,
                                          (SQLWCHAR*)_conn_str.c_str(), SQL_NTS,
                                          nullptr, 0, nullptr, SQL_DRIVER_NOPROMPT);
     if (!SQL_SUCCEEDED(ret)) {
@@ -70,13 +91,24 @@ SQLRETURN Connection::close() {
     return SQLDisconnect_ptr(_dbc_handle->get());
 }
 
-SQLRETURN Connection::end_transaction(SQLSMALLINT completion_type) {
-    LOG(completion_type == SQL_COMMIT ? "End SQL Transaction (Commit)" : "End SQL Transaction (Rollback)");
-    if (!SQLEndTran_ptr) {
-        LOG("Function pointer not initialized. Loading the driver.");
-        DriverLoader::getInstance().loadDriver();
+SQLRETURN Connection::commit() {
+    LOG("Commit transaction");
+    SQLRETURN ret = SQLEndTran_ptr(SQL_HANDLE_DBC, _dbc_handle->get(), SQL_COMMIT);
+    if (!SQL_SUCCEEDED(ret)) {
+        LOG("Failed to commit transaction");
+        throw std::runtime_error("Failed to commit transaction");
     }
-    return SQLEndTran_ptr(_dbc_handle->type(), _dbc_handle->get(), completion_type);
+    return ret;
+}
+
+SQLRETURN Connection::rollback() {
+    LOG("Rollback transaction");
+    SQLRETURN ret = SQLEndTran_ptr(SQL_HANDLE_DBC, _dbc_handle->get(), SQL_ROLLBACK);
+    if (!SQL_SUCCEEDED(ret)) {
+        LOG("Failed to rollback transaction");
+        throw std::runtime_error("Failed to rollback transaction");
+    }
+    return ret;
 }
 
 SQLRETURN Connection::set_autocommit(bool enable) {

mssql_python/pybind/connection/connection.h

@@ -24,8 +24,11 @@ class Connection {
     // Close the connection and free resources.
     SQLRETURN close();
 
-    // End the transaction with the specified completion type.
-    SQLRETURN end_transaction(SQLSMALLINT completion_type);
+    // Commit the current transaction.
+    SQLRETURN commit();
+
+    // Rollback the current transaction.
+    SQLRETURN rollback();
 
     // Enable or disable autocommit mode.
     SQLRETURN set_autocommit(bool value);
@@ -36,11 +39,14 @@ class Connection {
     SqlHandlePtr alloc_statement_handle(); // Will later be moved to cursor c++ class
 
 private:
+    void allocEnvHandle();
+    void setEnvAttributes();
+    void allocDbcHandle();
+    SQLRETURN connectToDb();
 
     std::wstring _conn_str;     // Connection string
     SqlHandlePtr _env_handle;   // Environment handle
     SqlHandlePtr _dbc_handle;   // Connection handle
-
     bool _autocommit = false;
 };
 #endif // CONNECTION_H