connection.cpp

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

#include "connection/connection.h"
#include "connection/connection_pool.h"
#include "utf_utils.h"
#include <algorithm>
#include <cstdio>
#include <memory>
#include <pybind11/pybind11.h>
#include <regex>
#include <string>
#include <utility>
#include <vector>

#define SQL_MAX_SMALL_INT 32767        // Maximum value for SQLSMALLINT

// Logging uses LOG() macro for all diagnostic output
#include "logger_bridge.hpp"
#include "performance_counter.hpp"

static bool isPythonFinalizing() {
    if (Py_IsInitialized() == 0) {
        return true;
    }
#if PY_VERSION_HEX >= 0x030D0000
    return Py_IsFinalizing() != 0;
#else
    return _Py_IsFinalizing() != 0;
#endif
}

static SqlHandlePtr getEnvHandle() {
    static SqlHandlePtr envHandle = []() -> SqlHandlePtr {
        LOG("Allocating ODBC environment handle");
        if (!SQLAllocHandle_ptr) {
            LOG("Function pointers not initialized, loading driver");
            DriverLoader::getInstance().loadDriver();
        }
        SQLHANDLE env = nullptr;
        SQLRETURN ret = SQLAllocHandle_ptr(SQL_HANDLE_ENV, SQL_NULL_HANDLE, &env);
        if (!SQL_SUCCEEDED(ret)) {
            ThrowStdException("Failed to allocate environment handle");
        }
        ret = SQLSetEnvAttr_ptr(env, SQL_ATTR_ODBC_VERSION,
                                reinterpret_cast<void*>(SQL_OV_ODBC3_80), 0);
        if (!SQL_SUCCEEDED(ret)) {
            ThrowStdException("Failed to set environment attributes");
        }
        return std::make_shared<SqlHandle>(static_cast<SQLSMALLINT>(SQL_HANDLE_ENV), env);
    }();

    return envHandle;
}

//-------------------------------------------------------------------------------------------------
// Implements the Connection class declared in connection.h.
// This class wraps low-level ODBC operations like connect/disconnect,
// transaction control, and autocommit configuration.
//-------------------------------------------------------------------------------------------------
Connection::Connection(const std::u16string& conn_str, bool use_pool)
    : _connStr(conn_str), _autocommit(false), _fromPool(use_pool) {
    PERF_TIMER("Connection::Connection");
    allocateDbcHandle();
}

Connection::~Connection() noexcept {
    disconnectNoThrow();
}

// Allocates connection handle
void Connection::allocateDbcHandle() {
    PERF_TIMER("Connection::allocateDbcHandle");
    // Fetch/initialize the shared env handle without holding the GIL (#671):
    // its first-time initialization runs under a C++ static-init guard and
    // emits log records; a thread waiting on that guard while holding the GIL
    // would deadlock the initializing thread that needs the GIL to log.
    auto envHandle = [&] {
        py::gil_scoped_release gil_release;
        return getEnvHandle();
    }();
    SQLHANDLE dbc = nullptr;
    LOG("Allocating SQL Connection Handle");
    SQLRETURN ret = SQLAllocHandle_ptr(SQL_HANDLE_DBC, envHandle->get(), &dbc);
    checkError(ret);
    _dbcHandle = std::make_shared<SqlHandle>(static_cast<SQLSMALLINT>(SQL_HANDLE_DBC), dbc);
}

void Connection::connect(const py::dict& attrs_before) {
    PERF_TIMER("Connection::connect");
    LOG("Connecting to database");
    // Apply access token before connect
    if (!attrs_before.is_none() && py::len(attrs_before) > 0) {
        LOG("Apply attributes before connect");
        applyAttrsBefore(attrs_before);
        if (_autocommit) {
            setAutocommit(_autocommit);
        }
    }
    SQLWCHAR* connStrPtr = reinterpretU16stringAsSqlWChar(_connStr);
    SQLRETURN ret;
    {
        // Release the GIL during the blocking ODBC connect call.
        // SQLDriverConnect involves DNS resolution, TCP handshake, TLS negotiation,
        // and SQL Server authentication — all pure I/O that doesn't need the GIL.
        // This allows other Python threads to run concurrently.
        py::gil_scoped_release release;
        PERF_TIMER("Connection::connect::SQLDriverConnect_call");
        ret = SQLDriverConnect_ptr(_dbcHandle->get(), nullptr, connStrPtr, SQL_NTS, nullptr,
                                   0, nullptr, SQL_DRIVER_NOPROMPT);
    }
    checkError(ret);
    updateLastUsed();
}

void Connection::disconnect(bool rollbackBeforeDisconnect) {
    PERF_TIMER("Connection::disconnect");
    clearResultMetadata();
    // Determine GIL state once, up front. disconnect() runs both from
    // pybind11-bound methods (GIL held) and from GIL-less destructor / shutdown
    // paths: Connection::~Connection() dropping the last shared_ptr, or teardown
    // running after the interpreter has been finalized. Every LOG()/LOG_ERROR()
    // below is gated on hasGil because LOG() acquires the GIL internally via
    // py::gil_scoped_acquire, which is unsafe when the GIL is not held — it can
    // hang or std::terminate during interpreter shutdown / stack unwinding.
    // Py_IsInitialized() is checked first: after Py_Finalize() the interpreter is
    // gone and PyGILState_Check() is unreliable, so treat "not initialized" as
    // "no GIL" and skip all Python calls. (#671 follow-up)
    bool hasGil = !isPythonFinalizing() && PyGILState_Check() != 0;
    if (_dbcHandle) {
        if (hasGil) {
            LOG("Disconnecting from database");
        }

        std::vector<SqlHandlePtr> childHandles;
        size_t originalSize = 0, afterCompactSize = 0, badHandleCount = 0;
        auto disconnectNative = [&]() {
            // Serialize explicit child free() calls as well as destruction.
            // This lock must be released before reacquiring the GIL or logging.
            std::lock_guard<std::mutex> cleanupLock(_cleanupState->mutex);
            {
                std::lock_guard<std::mutex> lock(_childHandlesMutex);
                originalSize = _childStatementHandles.size();
                _childStatementHandles.erase(
                    std::remove_if(_childStatementHandles.begin(), _childStatementHandles.end(),
                                   [](const std::weak_ptr<SqlHandle>& wp) { return wp.expired(); }),
                    _childStatementHandles.end());
                afterCompactSize = _childStatementHandles.size();
                childHandles.reserve(afterCompactSize);
                for (auto& weakHandle : _childStatementHandles) {
                    if (auto handle = weakHandle.lock()) {
                        if (handle->type() != SQL_HANDLE_STMT) {
                            ++badHandleCount;
                            continue;
                        }
                        childHandles.push_back(std::move(handle));
                    }
                }
            }
            if (rollbackBeforeDisconnect) {
                // Explicit SQL transactions need manual mode for SQLEndTran.
                // Never turn autocommit on here: that could commit abandoned work.
                SQLSetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_AUTOCOMMIT,
                                     reinterpret_cast<SQLPOINTER>(SQL_AUTOCOMMIT_OFF), 0);
                SQLEndTran_ptr(SQL_HANDLE_DBC, _dbcHandle->get(), SQL_ROLLBACK);
            }
            SQLRETURN result = SQLDisconnect_ptr(_dbcHandle->get());
            if (SQL_SUCCEEDED(result)) {
                // Also cover children whose weak_ptr expired as their destructor
                // began waiting for this gate: they cannot appear in the snapshot.
                _cleanupState->disconnected = true;
                std::lock_guard<std::mutex> lock(_childHandlesMutex);
                for (const auto& handle : childHandles) {
                    handle->markImplicitlyFreed();
                }
                _childStatementHandles.clear();
                _allocationsSinceCompaction = 0;
            }
            return result;
        };

        SQLRETURN ret;
        if (hasGil) {
            py::gil_scoped_release release;
            ret = disconnectNative();
        } else {
            ret = disconnectNative();
        }
        if (!SQL_SUCCEEDED(ret)) {
            if (hasGil) {
                checkError(ret);
            } else {
                std::fputs("mssql-python: native disconnect failed\n", stderr);
            }
            // Keep ownership and child-handle tracking intact for a cleanup retry.
            return;
        }
        // Log after releasing _childHandlesMutex (#671): LOG()/LOG_ERROR() acquire
        // the GIL and must not run while a native mutex is held. Also gated on
        // hasGil so the GIL-less destructor / shutdown path never tries to log.
        if (hasGil) {
            LOG("Compacted child handles: %zu -> %zu (removed %zu expired)",
                originalSize, afterCompactSize, originalSize - afterCompactSize);
            LOG("Marking %zu child statement handles as implicitly freed", afterCompactSize);
            if (badHandleCount > 0) {
                LOG_ERROR("CRITICAL: %zu non-STMT handle(s) found in _childStatementHandles. "
                          "This will cause a handle leak!", badHandleCount);
            }
        }

        // triggers SQLFreeHandle via destructor, if last owner
        _dbcHandle.reset();
    } else if (hasGil) {
        LOG("No connection handle to disconnect");
    }
}

void Connection::disconnectNoThrow() noexcept {
    try {
        if (isPythonFinalizing()) {
            abandonDuringFinalization();
            return;
        }
        if (!_dbcHandle) {
            return;
        }
        // disconnect() already supports GIL-less cleanup. Drop the GIL once so
        // neither its diagnostics nor handle destruction can enter Python.
        if (PyGILState_Check()) {
            py::gil_scoped_release release;
            disconnect(true);
        } else {
            disconnect(true);
        }
    } catch (...) {
        std::fputs("mssql-python: unexpected failure during native connection cleanup\n", stderr);
    }
}

void Connection::abandonDuringFinalization() noexcept {
    {
        std::lock_guard<std::mutex> lock(_childHandlesMutex);
        _childStatementHandles.clear();
        _allocationsSinceCompaction = 0;
    }
    // SqlHandle::free() already suppresses SQLFreeHandle during finalization.
    // Clearing the shared pointer leaves process teardown to the operating system.
    _dbcHandle.reset();
}

// TODO(microsoft): Add an exception class in C++ for error handling,
// DB spec compliant
void Connection::checkError(SQLRETURN ret) const {
    if (!SQL_SUCCEEDED(ret)) {
        // Format: "SQLSTATE:XXXXX:<odbc_error_message>" — parsed by _raise_connection_error()
        ErrorInfo err = SQLCheckError_Wrap(SQL_HANDLE_DBC, _dbcHandle, ret);
        std::string sqlState = err.sqlState;
        std::string errorMsg = err.ddbcErrorMsg;
        // Only add SQLSTATE prefix if we have a valid 5-character code
        if (sqlState.length() == 5) {
            ThrowStdException("SQLSTATE:" + sqlState + ":" + errorMsg);
        } else {
            // No valid SQLSTATE (e.g., SQL_INVALID_HANDLE) — throw clean error message
            ThrowStdException(errorMsg);
        }
    }
}

void Connection::clearResultMetadata() {
    std::vector<SqlHandlePtr> handles;
    {
        std::lock_guard<std::mutex> lock(_childHandlesMutex);
        handles.reserve(_childStatementHandles.size());
        for (const auto& weakHandle : _childStatementHandles) {
            if (auto handle = weakHandle.lock()) {
                handles.push_back(std::move(handle));
            }
        }
    }
    // Releasing the last handle can acquire the connection cleanup gate.
    // Keep that destruction outside the child-list lock.
    for (const auto& handle : handles) {
        handle->resultMetadata.clear();
    }
}

void Connection::commit() {
    PERF_TIMER("Connection::commit");
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }
    updateLastUsed();
    LOG("Committing transaction");
    clearResultMetadata();
    SQLRETURN ret;
    {
        // Release the GIL during the blocking SQLEndTran network round-trip.
        py::gil_scoped_release release;
        ret = SQLEndTran_ptr(SQL_HANDLE_DBC, _dbcHandle->get(), SQL_COMMIT);
    }
    checkError(ret);
}

void Connection::rollback() {
    PERF_TIMER("Connection::rollback");
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }
    updateLastUsed();
    LOG("Rolling back transaction");
    clearResultMetadata();
    SQLRETURN ret;
    {
        // Release the GIL during the blocking SQLEndTran network round-trip.
        py::gil_scoped_release release;
        ret = SQLEndTran_ptr(SQL_HANDLE_DBC, _dbcHandle->get(), SQL_ROLLBACK);
    }
    checkError(ret);
}

void Connection::setAutocommit(bool enable) {
    PERF_TIMER("Connection::setAutocommit");
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }
    clearResultMetadata();
    SQLINTEGER value = enable ? SQL_AUTOCOMMIT_ON : SQL_AUTOCOMMIT_OFF;
    LOG("Setting autocommit=%d", enable);
    SQLRETURN ret;
    {
        // Release the GIL during the blocking ODBC call. Holding the GIL
        // here can deadlock when the network path goes through another
        // Python thread (e.g. an in-process SSH tunnel via paramiko +
        // sshtunnel), since that thread also needs the GIL to run.
        py::gil_scoped_release release;
        ret = SQLSetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_AUTOCOMMIT,
                                    reinterpret_cast<SQLPOINTER>(static_cast<SQLULEN>(value)), 0);
    }
    checkError(ret);
    if (value == SQL_AUTOCOMMIT_ON) {
        LOG("Autocommit enabled");
    } else {
        LOG("Autocommit disabled");
    }
    _autocommit = enable;
}

bool Connection::getAutocommit() const {
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }
    LOG("Getting autocommit attribute");
    SQLINTEGER value;
    SQLINTEGER string_length;
    SQLRETURN ret = SQLGetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_AUTOCOMMIT, &value,
                                          sizeof(value), &string_length);
    checkError(ret);
    return value == SQL_AUTOCOMMIT_ON;
}

SqlHandlePtr Connection::allocStatementHandle() {
    PERF_TIMER("Connection::allocStatementHandle");
    LOG("Allocating statement handle");
    // Keep the wrapper outside the lock scope: unwinding a failed registration
    // frees the statement through the same cleanup gate.
    SqlHandlePtr stmtHandle;
    bool compacted = false;
    size_t compactBefore = 0, compactAfter = 0;
    {
        py::gil_scoped_release release;
        std::lock_guard<std::mutex> cleanupLock(_cleanupState->mutex);
        if (_cleanupState->disconnected || !_dbcHandle) {
            ThrowStdException("Connection handle not allocated");
        }
        updateLastUsed();
        SQLHANDLE stmt = nullptr;
        SQLRETURN ret = SQLAllocHandle_ptr(SQL_HANDLE_STMT, _dbcHandle->get(), &stmt);
        if (!SQL_SUCCEEDED(ret)) {
            // Snapshot diagnostics before disconnect can overwrite/free the DBC.
            ErrorInfo err = SQLReadError(SQL_HANDLE_DBC, _dbcHandle->get(), ret);
            ThrowStdException(err.sqlState.length() == 5
                                  ? "SQLSTATE:" + err.sqlState + ":" + err.ddbcErrorMsg
                                  : err.ddbcErrorMsg);
        }
        stmtHandle = std::make_shared<SqlHandle>(static_cast<SQLSMALLINT>(SQL_HANDLE_STMT),
                                                stmt, _cleanupState);
        std::lock_guard<std::mutex> lock(_childHandlesMutex);
        
        // 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);
        _allocationsSinceCompaction++;

        // Compact expired weak_ptrs only periodically to avoid O(n²) overhead
        // This keeps allocation fast (O(1) amortized) while preventing unbounded growth
        // disconnect() also compacts, so this is just for long-lived connections with many cursors
        if (_allocationsSinceCompaction >= COMPACTION_INTERVAL) {
            compactBefore = _childStatementHandles.size();
            _childStatementHandles.erase(
                std::remove_if(_childStatementHandles.begin(), _childStatementHandles.end(),
                               [](const std::weak_ptr<SqlHandle>& wp) { return wp.expired(); }),
                _childStatementHandles.end());
            compactAfter = _childStatementHandles.size();
            _allocationsSinceCompaction = 0;
            compacted = true;
        }
    }  // Release lock

    // Log after releasing _childHandlesMutex (#671): LOG() acquires the GIL and
    // must not run while a native mutex is held.
    if (compacted) {
        LOG("Periodic compaction: %zu -> %zu handles (removed %zu expired)",
            compactBefore, compactAfter, compactBefore - compactAfter);
    }

    return stmtHandle;
}

SQLRETURN Connection::setAttribute(SQLINTEGER attribute, py::object value) {
    clearResultMetadata();
    LOG("Setting SQL attribute=%d", attribute);
    // SQLPOINTER ptr = nullptr;
    // SQLINTEGER length = 0;

    // Fail closed on a non-binary access token. SQL_COPT_SS_ACCESS_TOKEN (1256)
    // MUST be the raw [DWORD byte-length][UTF-16LE token] struct passed as
    // bytes/bytearray. If a caller supplies it as a py::str, the str->UTF-16
    // cast in the string branch below would mangle that struct; worse, the
    // Python identity-aware pool-key logic only hashes bytes/bytearray tokens,
    // so a str token slips through with the bare connection-string pool key and
    // two callers passing different str tokens against the same server could
    // share a pooled, authenticated connection. Reject any non-binary token at
    // this native boundary so the cross-identity invariant ("a token is present
    // => the pool key is never the bare connStr") holds regardless of how the
    // Connection was constructed.
    if (attribute == SQL_COPT_SS_ACCESS_TOKEN && !py::isinstance<py::bytes>(value) &&
        !py::isinstance<py::bytearray>(value)) {
        LOG("Rejecting non-binary SQL_COPT_SS_ACCESS_TOKEN (attribute=%d): access token "
            "must be bytes/bytearray",
            attribute);
        return SQL_ERROR;
    }

    if (py::isinstance<py::int_>(value)) {
        // Get the integer value
        int64_t longValue = value.cast<int64_t>();

        SQLRETURN ret;
        {
            // Release the GIL around the ODBC call for consistency with the
            // other connection-attribute paths; some attributes can block.
            py::gil_scoped_release release;
            ret = SQLSetConnectAttr_ptr(
                _dbcHandle->get(), attribute,
                reinterpret_cast<SQLPOINTER>(static_cast<SQLULEN>(longValue)), SQL_IS_INTEGER);
        }

        if (!SQL_SUCCEEDED(ret)) {
            LOG("Failed to set integer attribute=%d, ret=%d", attribute, ret);
        } else {
            LOG("Set integer attribute=%d successfully", attribute);
        }
        return ret;
    } else if (py::isinstance<py::str>(value)) {
        try {
            // Store the value in a Connection-owned, per-attribute member
            // buffer so the memory remains valid for the lifetime of the
            // Connection object. Some ODBC connect attributes (notably
            // SQL_COPT_SS_ACCESS_TOKEN, 1256) are "deferred": the MS driver
            // stores the caller's pointer at SQLSetConnectAttr time and
            // dereferences it later during SQLDriverConnect to build the
            // FedAuth login packet. A stack-local buffer freed when this
            // function returns would cause a use-after-free during connect
            // (issue #594). Keying by attribute id also prevents a second
            // deferred attribute from invalidating the pointer stored for
            // the first.
            //
            // Lifetime: the buffer MUST outlive every potential dereference
            // of the deferred-attribute pointer by the driver, which
            // includes paths beyond the initial connect (Idle Connection
            // Resiliency re-auth on a dropped socket, transparent pool
            // checkout re-handshake). SQL_ATTR_RESET_CONNECTION (see
            // Connection::reset()) only wipes per-session state and does
            // NOT tear down the driver-side authentication context, so the
            // per-attribute buffers are NOT cleared on reset()/checkin;
            // they are released only when the Connection object itself is
            // destroyed.
            //
            // Note: attrs_before is applied once, sequentially, during
            // connect(); the Connection's attribute setters are not designed
            // for concurrent mutation from multiple threads.
            auto& buf = this->_attrStringBuffers[attribute];
            buf = value.cast<std::u16string>();

            SQLPOINTER ptr = reinterpretU16stringAsSqlWChar(buf);
            SQLINTEGER length =
                static_cast<SQLINTEGER>(buf.length() * sizeof(SQLWCHAR));

            SQLRETURN ret;
            {
                py::gil_scoped_release release;
                ret = SQLSetConnectAttr_ptr(_dbcHandle->get(), attribute, ptr, length);
            }
            if (!SQL_SUCCEEDED(ret)) {
                LOG("Failed to set string attribute=%d, ret=%d", attribute, ret);
            } else {
                LOG("Set string attribute=%d successfully", attribute);
            }
            return ret;
        } catch (const std::exception& e) {
            LOG("Exception during string attribute=%d setting: %s", attribute, e.what());
            return SQL_ERROR;
        }
    } else if (py::isinstance<py::bytes>(value) || py::isinstance<py::bytearray>(value)) {
        try {
            // Store the value in a Connection-owned, per-attribute member
            // buffer so the memory remains valid for the lifetime of the
            // Connection object. SQL_COPT_SS_ACCESS_TOKEN (1256) is a
            // deferred attribute: the driver stores this pointer at
            // SQLSetConnectAttr time and dereferences it later during
            // SQLDriverConnect. A stack-local buffer freed when this
            // function returns would cause a use-after-free during connect
            // (issue #594, symptoms: SIGBUS on macOS, "Authentication
            // token is missing in the federated authentication message"
            // on Windows, TCP reset 0x2746 against Azure SQL). Keying by
            // attribute id also prevents a second deferred attribute from
            // invalidating the pointer stored for the first.
            //
            // Lifetime: the buffer MUST outlive every potential dereference
            // of the deferred-attribute pointer by the driver, which
            // includes paths beyond the initial connect:
            //   * Idle Connection Resiliency (ICR): if the underlying TCP
            //     connection drops while the connection sits idle in the
            //     pool, the driver transparently re-establishes it on the
            //     next use and re-runs the Login7 / FedAuth handshake,
            //     dereferencing the same stashed token pointer.
            //   * SQL_ATTR_RESET_CONNECTION pool checkin (see
            //     Connection::reset()) only wipes per-session state; the
            //     driver-side authentication context and the stashed
            //     deferred-attribute pointer are intentionally retained.
            // For these reasons the per-attribute buffers are NOT cleared
            // on reset()/checkin; they are released only when the
            // Connection object itself is destroyed.
            //
            // Note: attrs_before is applied once, sequentially, during
            // connect(); concurrent setAttribute() on the same Connection
            // from different threads is not a supported pattern.
            auto& buf = this->_attrBytesBuffers[attribute];
            buf = value.cast<std::string>();
            SQLPOINTER ptr = const_cast<char*>(buf.data());
            SQLINTEGER length = static_cast<SQLINTEGER>(buf.size());

            SQLRETURN ret;
            {
                py::gil_scoped_release release;
                ret = SQLSetConnectAttr_ptr(_dbcHandle->get(), attribute, ptr, length);
            }
            if (!SQL_SUCCEEDED(ret)) {
                LOG("Failed to set binary attribute=%d, ret=%d", attribute, ret);
            } else {
                LOG("Set binary attribute=%d successfully (length=%d)", attribute, length);
            }
            return ret;
        } catch (const std::exception& e) {
            LOG("Exception during binary attribute=%d setting: %s", attribute, e.what());
            return SQL_ERROR;
        }
    } else {
        LOG("Unsupported attribute value type for attribute=%d", attribute);
        return SQL_ERROR;
    }
}

void Connection::applyAttrsBefore(const py::dict& attrs) {
    for (const auto& item : attrs) {
        int key;
        try {
            key = py::cast<int>(item.first);
        } catch (...) {
            continue;
        }

        // Apply all supported attributes
        SQLRETURN ret = setAttribute(key, py::reinterpret_borrow<py::object>(item.second));
        if (!SQL_SUCCEEDED(ret)) {
            std::string attrName = std::to_string(key);
            std::string errorMsg = "Failed to set attribute " + attrName + " before connect";
            ThrowStdException(errorMsg);
        }
    }
}

bool Connection::isAlive() const {
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }
    SQLUINTEGER status;
    SQLRETURN ret =
        SQLGetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_CONNECTION_DEAD, &status, 0, nullptr);
    return SQL_SUCCEEDED(ret) && status == SQL_CD_FALSE;
}

bool Connection::reset() {
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }
    clearResultMetadata();
    LOG("Resetting connection via SQL_ATTR_RESET_CONNECTION");
    // NOTE: SQL_ATTR_RESET_CONNECTION is a pool-checkin reset: it asks the
    // driver to wipe per-session state (temp tables, open cursors, SET
    // options, etc.) on the next use. It does NOT tear down the underlying
    // TCP/TLS connection nor the driver-side authentication context, and
    // it does NOT discard the deferred connect attributes the driver has
    // stashed (e.g., the SQL_COPT_SS_ACCESS_TOKEN pointer used to build
    // the FedAuth Login7 packet). The driver may still dereference those
    // pointers after this reset on Idle Connection Resiliency re-auth or
    // a transparent reconnect, so the per-attribute buffers owned by this
    // Connection (_attrStringBuffers / _attrBytesBuffers) are intentionally
    // retained here. Clearing them would reintroduce issue #594 in a new
    // form (UAF during silent reconnect).
    SQLRETURN ret;
    {
        // Release the GIL around the ODBC call for consistency with the
        // other connection-attribute paths; some attributes can block.
        py::gil_scoped_release release;
        ret = SQLSetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_RESET_CONNECTION,
                                    (SQLPOINTER)SQL_RESET_CONNECTION_YES, SQL_IS_INTEGER);
    }
    if (!SQL_SUCCEEDED(ret)) {
        LOG("Failed to reset connection (ret=%d). Marking as dead.", ret);
        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.
    LOG("Resetting transaction isolation level to READ COMMITTED");
    {
        py::gil_scoped_release release;
        ret = SQLSetConnectAttr_ptr(_dbcHandle->get(), SQL_ATTR_TXN_ISOLATION,
                                    (SQLPOINTER)SQL_TXN_READ_COMMITTED, SQL_IS_INTEGER);
    }
    if (!SQL_SUCCEEDED(ret)) {
        LOG("Failed to reset transaction isolation level (ret=%d). Marking as dead.", ret);
        return false;
    }

    updateLastUsed();
    return true;
}

void Connection::prepareForPool(bool transactionAlreadyRolledBack) {
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }

    // Explicit BEGIN TRANSACTION is valid while ODBC autocommit is on, but
    // SQLEndTran does not end that transaction until the connection enters
    // manual-commit mode.
    if (getAutocommit()) {
        setAutocommit(false);
    }
    if (!transactionAlreadyRolledBack) {
        rollback();
    }
    // The SQL Server ODBC driver can leave an empty transaction visible after
    // SQLEndTran while manual-commit mode remains enabled, so always park the
    // physical connection in autocommit mode.
    setAutocommit(true);
}

void Connection::updateLastUsed() {
    _lastUsed = std::chrono::steady_clock::now();
}

std::chrono::steady_clock::time_point Connection::lastUsed() const {
    return _lastUsed;
}

py::dict Connection::invokeTokenFactory(const py::object& tokenFactory,
                                        long long& outExpiryEpoch) {
    outExpiryEpoch = 0;
    py::object result = tokenFactory();
    // New contract: factory returns (attrs, expires_on). Remain
    // backward compatible with the legacy contract where it returned a
    // bare attrs dict.
    if (py::isinstance<py::tuple>(result)) {
        py::tuple parts = result.cast<py::tuple>();
        // Defensive: a well-formed factory always returns at least (attrs,).
        // Guard the index so a misbehaving/empty tuple falls through to the
        // cast below (which raises a clear tuple->dict error) instead of an
        // out-of-range access on parts[0].
        if (parts.size() >= 1) {
            py::dict attrs = parts[0].cast<py::dict>();
            if (parts.size() > 1 && !parts[1].is_none()) {
                outExpiryEpoch = parts[1].cast<long long>();
            }
            return attrs;
        }
    }
    return result.cast<py::dict>();
}

void Connection::setTokenExpiry(long long epochSeconds) {
    _tokenExpiryEpoch = epochSeconds;
}

bool Connection::isTokenNearExpiry(int thresholdSecs) const {
    if (_tokenExpiryEpoch == 0) {
        // Unknown expiry. Fail closed when we actually hold a token whose
        // validity we cannot prove: reusing it risks handing back a token that
        // expires mid-query, so force a refresh check instead (matching
        // tokenExpirySafelyBeyond()'s fail-closed treatment of an unknown
        // expiry). With no token present (an empty access token, e.g. a factory
        // that supplies non-token attrs for SQL auth) there is nothing to
        // expire, so the connection stays reusable. Real credentials always
        // report expires_on, so the fail-closed arm is a safety net.
        return !currentAccessToken().empty();
    }
    const long long now = static_cast<long long>(
        std::chrono::duration_cast<std::chrono::seconds>(
            std::chrono::system_clock::now().time_since_epoch())
            .count());
    return (now + static_cast<long long>(thresholdSecs)) >= _tokenExpiryEpoch;
}

std::string Connection::currentAccessToken() const {
    auto it = _attrBytesBuffers.find(SQL_COPT_SS_ACCESS_TOKEN);
    return it != _attrBytesBuffers.end() ? it->second : std::string();
}

ConnectionHandle::ConnectionHandle(const std::u16string& connStr, bool usePool,
                                   const py::dict& attrsBefore, const std::u16string& poolKey,
                                   const py::object& tokenFactory)
    : _usePool(usePool), _connStr(connStr), _poolKey(poolKey.empty() ? connStr : poolKey) {
    PERF_TIMER("ConnectionHandle::ConnectionHandle");
    if (_usePool) {
        _conn = ConnectionPoolManager::getInstance().acquireConnection(_connStr, attrsBefore,
                                                                       _poolKey, tokenFactory,
                                                                       &_originPool);
        // acquireConnection returns nullptr when pooling was disabled out from
        // under us (a disable_pooling() won the race). Fall back to a non-pooled
        // connection and flip _usePool so close() disconnects it directly rather
        // than trying to return it to a pool that no longer exists.
        if (!_conn) {
            _usePool = false;
        }
    }
    if (!_usePool) {
        _conn = std::make_shared<Connection>(_connStr, false);
        // Non-pooled connect still honors the lazy token factory: a
        // token is materialized only when a physical connection is opened. The
        // factory may also carry the token expiry, but a non-pooled
        // connection is never reused, so expiry-aware checkout does not
        // apply and the expiry is intentionally not recorded here.
        if (tokenFactory && !tokenFactory.is_none()) {
            long long expiry = 0;
            py::dict connect_attrs = Connection::invokeTokenFactory(tokenFactory, expiry);
            _conn->connect(connect_attrs);
        } else {
            _conn->connect(attrsBefore);
        }
    }
}

ConnectionHandle::~ConnectionHandle() {
    if (_conn) {
        if (isPythonFinalizing()) {
            _conn->abandonDuringFinalization();
            _conn = nullptr;
            return;
        }
        try {
            // Discard ends abandoned work without returning this connection to
            // the pool or entering Python from a native destructor.
            ConnectionPoolManager::getInstance().discardConnection(_originPool, _conn);
        } catch (...) {
            std::fputs("mssql-python: failed to release native connection pool capacity\n", stderr);
            _conn->disconnectNoThrow();
        }
    }
}

void ConnectionHandle::close(bool transactionAlreadyRolledBack) {
    PERF_TIMER("ConnectionHandle::close");
    if (!_conn) {
        ThrowStdException("Connection object is not initialized");
    }
    if (_usePool) {
        try {
            _conn->prepareForPool(transactionAlreadyRolledBack);
        } catch (...) {
            // Never retain a connection whose transaction state could not be
            // sanitized. Discarding also releases this connection's reserved
            // pool capacity. Preserve the original check-in error.
            try {
                ConnectionPoolManager::getInstance().discardConnection(_originPool, _conn);
            } catch (...) {
            }
            _conn = nullptr;
            throw;
        }
        ConnectionPoolManager::getInstance().returnConnection(_poolKey, _originPool, _conn);
    } else {
        _conn->disconnect();
    }
    _conn = nullptr;
}

void ConnectionHandle::commit() {
    PERF_TIMER("ConnectionHandle::commit");
    if (!_conn) {
        ThrowStdException("Connection object is not initialized");
    }
    _conn->commit();
}

void ConnectionHandle::rollback() {
    PERF_TIMER("ConnectionHandle::rollback");
    if (!_conn) {
        ThrowStdException("Connection object is not initialized");
    }
    _conn->rollback();
}

void ConnectionHandle::setAutocommit(bool enabled) {
    PERF_TIMER("ConnectionHandle::setAutocommit");
    if (!_conn) {
        ThrowStdException("Connection object is not initialized");
    }
    _conn->setAutocommit(enabled);
}

bool ConnectionHandle::getAutocommit() const {
    if (!_conn) {
        ThrowStdException("Connection object is not initialized");
    }
    return _conn->getAutocommit();
}

SqlHandlePtr ConnectionHandle::allocStatementHandle() {
    PERF_TIMER("ConnectionHandle::allocStatementHandle");
    // close() can detach _conn while allocation waits without the GIL.
    auto conn = _conn;
    if (!conn) {
        ThrowStdException("Connection object is not initialized");
    }
    return conn->allocStatementHandle();
}

py::object Connection::getInfo(SQLUSMALLINT infoType) const {
    if (!_dbcHandle) {
        ThrowStdException("Connection handle not allocated");
    }

    // First call with NULL buffer to get required length
    SQLSMALLINT requiredLen = 0;
    SQLRETURN ret = SQLGetInfo_ptr(_dbcHandle->get(), infoType, NULL, 0, &requiredLen);

    if (!SQL_SUCCEEDED(ret)) {
        checkError(ret);
        return py::none();
    }

    // For zero-length results
    if (requiredLen == 0) {
        py::dict result;
        result["data"] = py::bytes("", 0);
        result["length"] = 0;
        result["info_type"] = infoType;
        return result;
    }

    // Cap buffer allocation to SQL_MAX_SMALL_INT to prevent excessive
    // memory usage
    SQLSMALLINT allocSize = requiredLen + 10;
    if (allocSize > SQL_MAX_SMALL_INT) {
        allocSize = SQL_MAX_SMALL_INT;
    }
    std::vector<char> buffer(allocSize, 0);  // Extra padding for safety

    // Get the actual data - avoid using std::min
    SQLSMALLINT bufferSize = requiredLen + 10;
    if (bufferSize > SQL_MAX_SMALL_INT) {
        bufferSize = SQL_MAX_SMALL_INT;
    }

    SQLSMALLINT returnedLen = 0;
    ret = SQLGetInfo_ptr(_dbcHandle->get(), infoType, buffer.data(), bufferSize, &returnedLen);

    if (!SQL_SUCCEEDED(ret)) {
        checkError(ret);
        return py::none();
    }

    // Create a dictionary with the raw data
    py::dict result;

    // IMPORTANT: Pass exactly what SQLGetInfo returned
    // No null-terminator manipulation, just pass the raw data
    result["data"] = py::bytes(buffer.data(), returnedLen);
    result["length"] = returnedLen;
    result["info_type"] = infoType;

    return result;
}

py::object ConnectionHandle::getInfo(SQLUSMALLINT infoType) const {
    PERF_TIMER("ConnectionHandle::getInfo");
    if (!_conn) {
        ThrowStdException("Connection object is not initialized");
    }
    return _conn->getInfo(infoType);
}

void ConnectionHandle::setAttr(int attribute, py::object value) {
    PERF_TIMER("ConnectionHandle::setAttr");
    if (!_conn) {
        ThrowStdException("Connection not established");
    }

    // Use existing setAttribute with better error handling
    SQLRETURN ret = _conn->setAttribute(static_cast<SQLINTEGER>(attribute), value);
    if (!SQL_SUCCEEDED(ret)) {
        // Get detailed error information from ODBC
        try {
            ErrorInfo errorInfo = SQLCheckError_Wrap(SQL_HANDLE_DBC, _conn->getDbcHandle(), ret);

            std::string errorMsg =
                "Failed to set connection attribute " + std::to_string(attribute);
            if (!errorInfo.ddbcErrorMsg.empty()) {
                errorMsg += ": " + errorInfo.ddbcErrorMsg;
            }

            LOG("Connection setAttribute failed: %s", errorMsg.c_str());
            ThrowStdException(errorMsg);
        } catch (...) {
            // Fallback to generic error if detailed error retrieval fails
            std::string errorMsg =
                "Failed to set connection attribute " + std::to_string(attribute);
            LOG("Connection setAttribute failed: %s", errorMsg.c_str());
            ThrowStdException(errorMsg);
        }
    }
}