Skip to content
33 changes: 30 additions & 3 deletions mssql_python/pybind/connection/connection.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,20 @@ void Connection::connect(const py::dict& attrs_before) {
void Connection::disconnect() {
if (_dbcHandle) {
LOG("Disconnecting from database");

// CRITICAL FIX: Mark all child statement handles as implicitly freed
// When we free the DBC handle below, the ODBC driver will automatically free
// all child STMT handles. We need to tell the SqlHandle objects about this
// so they don't try to free the handles again during their destruction.
LOG("Marking %zu child statement handles as implicitly freed",
_childStatementHandles.size());
for (auto& weakHandle : _childStatementHandles) {
if (auto handle = weakHandle.lock()) {
handle->markImplicitlyFreed();
}
}
_childStatementHandles.clear();

SQLRETURN ret = SQLDisconnect_ptr(_dbcHandle->get());
checkError(ret);
// triggers SQLFreeHandle via destructor, if last owner
Expand Down Expand Up @@ -173,7 +187,20 @@ SqlHandlePtr Connection::allocStatementHandle() {
SQLHANDLE stmt = nullptr;
SQLRETURN ret = SQLAllocHandle_ptr(SQL_HANDLE_STMT, _dbcHandle->get(), &stmt);
checkError(ret);
return std::make_shared<SqlHandle>(static_cast<SQLSMALLINT>(SQL_HANDLE_STMT), stmt);
auto stmtHandle = std::make_shared<SqlHandle>(static_cast<SQLSMALLINT>(SQL_HANDLE_STMT), stmt);

// Track this child handle so we can mark it as implicitly freed when connection closes
// Use weak_ptr to avoid circular references and allow normal cleanup
_childStatementHandles.push_back(stmtHandle);

// Clean up expired weak_ptrs periodically to avoid unbounded growth
// Remove entries where the weak_ptr is expired (object was already destroyed)
_childStatementHandles.erase(
std::remove_if(_childStatementHandles.begin(), _childStatementHandles.end(),
[](const std::weak_ptr<SqlHandle>& wp) { return wp.expired(); }),
_childStatementHandles.end());

return stmtHandle;
}

SQLRETURN Connection::setAttribute(SQLINTEGER attribute, py::object value) {
Expand Down Expand Up @@ -308,7 +335,7 @@ bool Connection::reset() {
disconnect();
return false;
}

// SQL_ATTR_RESET_CONNECTION does NOT reset the transaction isolation level.
// Explicitly reset it to the default (SQL_TXN_READ_COMMITTED) to prevent
// isolation level settings from leaking between pooled connection usages.
Expand All @@ -320,7 +347,7 @@ bool Connection::reset() {
disconnect();
return false;
}

updateLastUsed();
return true;
}
Expand Down
4 changes: 4 additions & 0 deletions mssql_python/pybind/connection/connection.h
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,10 @@ class Connection {
std::chrono::steady_clock::time_point _lastUsed;
std::wstring wstrStringBuffer; // wstr buffer for string attribute setting
std::string strBytesBuffer; // string buffer for byte attributes setting

// Track child statement handles to mark them as implicitly freed when connection closes
// Uses weak_ptr to avoid circular references and allow normal cleanup
std::vector<std::weak_ptr<SqlHandle>> _childStatementHandles;
};

class ConnectionHandle {
Expand Down
36 changes: 21 additions & 15 deletions mssql_python/pybind/ddbc_bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1144,6 +1144,10 @@ SQLSMALLINT SqlHandle::type() const {
return _type;
}

void SqlHandle::markImplicitlyFreed() {
_implicitly_freed = true;
}

/*
* IMPORTANT: Never log in destructors - it causes segfaults.
* During program exit, C++ destructors may run AFTER Python shuts down.
Expand All @@ -1169,16 +1173,19 @@ void SqlHandle::free() {
return;
}

// Always clean up ODBC resources, regardless of Python state
// CRITICAL FIX: Check if handle was already implicitly freed by parent handle
// When Connection::disconnect() frees the DBC handle, the ODBC driver automatically
// frees all child STMT handles. We track this state to avoid double-free attempts.
// This approach avoids calling ODBC functions on potentially-freed handles, which
// would cause use-after-free errors.
if (_implicitly_freed) {
_handle = nullptr; // Just clear the pointer, don't call ODBC functions
return;
}

// Handle is valid and not implicitly freed, proceed with normal freeing
SQLFreeHandle_ptr(_type, _handle);
_handle = nullptr;

// Only log if Python is not shutting down (to avoid segfault)
if (!pythonShuttingDown) {
// Don't log during destruction - even in normal cases it can be
// problematic If logging is needed, use explicit close() methods
// instead
}
}
}

Expand Down Expand Up @@ -2893,7 +2900,6 @@ SQLRETURN SQLGetData_wrap(SqlHandlePtr StatementHandle, SQLUSMALLINT colCount, p

// Cache decimal separator to avoid repeated system calls


for (SQLSMALLINT i = 1; i <= colCount; ++i) {
SQLWCHAR columnName[256];
SQLSMALLINT columnNameLen;
Expand Down Expand Up @@ -3615,8 +3621,6 @@ SQLRETURN FetchBatchData(SQLHSTMT hStmt, ColumnBuffers& buffers, py::list& colum
columnInfos[col].processedColumnSize + 1; // +1 for null terminator
}



// Performance: Build function pointer dispatch table (once per batch)
// This eliminates the switch statement from the hot loop - 10,000 rows × 10
// cols reduces from 100,000 switch evaluations to just 10 switch
Expand Down Expand Up @@ -4033,8 +4037,8 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch
lobColumns.push_back(i + 1); // 1-based
}
}
// Initialized to 0 for LOB path counter; overwritten by ODBC in non-LOB path;

// Initialized to 0 for LOB path counter; overwritten by ODBC in non-LOB path;
SQLULEN numRowsFetched = 0;
// If we have LOBs → fall back to row-by-row fetch + SQLGetData_wrap
if (!lobColumns.empty()) {
Expand Down Expand Up @@ -4066,7 +4070,7 @@ SQLRETURN FetchMany_wrap(SqlHandlePtr StatementHandle, py::list& rows, int fetch
LOG("FetchMany_wrap: Error when binding columns - SQLRETURN=%d", ret);
return ret;
}

SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROW_ARRAY_SIZE, (SQLPOINTER)(intptr_t)fetchSize, 0);
SQLSetStmtAttr_ptr(hStmt, SQL_ATTR_ROWS_FETCHED_PTR, &numRowsFetched, 0);

Expand Down Expand Up @@ -4360,7 +4364,9 @@ PYBIND11_MODULE(ddbc_bindings, m) {
.def_readwrite("ddbcErrorMsg", &ErrorInfo::ddbcErrorMsg);

py::class_<SqlHandle, SqlHandlePtr>(m, "SqlHandle")
.def("free", &SqlHandle::free, "Free the handle");
.def("free", &SqlHandle::free, "Free the handle")
.def("markImplicitlyFreed", &SqlHandle::markImplicitlyFreed,
"Mark handle as implicitly freed by parent handle");

py::class_<ConnectionHandle>(m, "Connection")
.def(py::init<const std::string&, bool, const py::dict&>(), py::arg("conn_str"),
Expand Down
6 changes: 6 additions & 0 deletions mssql_python/pybind/ddbc_bindings.h
Original file line number Diff line number Diff line change
Expand Up @@ -379,9 +379,15 @@ class SqlHandle {
SQLSMALLINT type() const;
void free();

// Mark this handle as implicitly freed (freed by parent handle)
// This prevents double-free attempts when the ODBC driver automatically
// frees child handles (e.g., STMT handles when DBC handle is freed)
void markImplicitlyFreed();

private:
SQLSMALLINT _type;
SQLHANDLE _handle;
bool _implicitly_freed = false; // Tracks if handle was freed by parent
};
using SqlHandlePtr = std::shared_ptr<SqlHandle>;

Expand Down
Loading
Loading