FIX: Prevent connection pool size drift when close races failed open (#746)

mssql_python/pybind/connection/connection.cpp

@@ -47,7 +47,10 @@ static SqlHandlePtr getEnvHandle() {
 // transaction control, and autocommit configuration.
 //-------------------------------------------------------------------------------------------------
 Connection::Connection(const std::u16string& conn_str, bool use_pool)
-    : _connStr(conn_str), _autocommit(false), _fromPool(use_pool) {
+    : _connStr(conn_str),
+      _autocommit(false),
+      _fromPool(use_pool),
+      _lastUsed(std::chrono::steady_clock::now()) {
     PERF_TIMER("Connection::Connection");
     allocateDbcHandle();
 }
@@ -76,6 +79,10 @@ void Connection::allocateDbcHandle() {
 
 void Connection::connect(const py::dict& attrs_before) {
     PERF_TIMER("Connection::connect");
+    if (_isMock) {
+        updateLastUsed();
+        return;
+    }
     LOG("Connecting to database");
     // Apply access token before connect
     if (!attrs_before.is_none() && py::len(attrs_before) > 0) {
@@ -103,6 +110,10 @@ void Connection::connect(const py::dict& attrs_before) {
 
 void Connection::disconnect() {
     PERF_TIMER("Connection::disconnect");
+    if (_isMock) {
+        _dbcHandle.reset();
+        return;
+    }
     // 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
@@ -507,6 +518,9 @@ void Connection::applyAttrsBefore(const py::dict& attrs) {
 }
 
 bool Connection::isAlive() const {
+    if (_isMock) {
+        return true;
+    }
     if (!_dbcHandle) {
         ThrowStdException("Connection handle not allocated");
     }
@@ -517,6 +531,9 @@ bool Connection::isAlive() const {
 }
 
 bool Connection::reset() {
+    if (_isMock) {
+        return true;
+    }
     if (!_dbcHandle) {
         ThrowStdException("Connection handle not allocated");
     }
@@ -642,6 +659,9 @@ ConnectionHandle::ConnectionHandle(const std::u16string& connStr, bool usePool,
     }
     if (!_usePool) {
         _conn = std::make_shared<Connection>(_connStr, false);
+        if (ConnectionPoolManager::getInstance().mock_mode()) {
+            _conn->setMock(true);
+        }
         // 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

mssql_python/pybind/connection/connection.h

@@ -55,6 +55,17 @@ class Connection {
     bool reset();
     void updateLastUsed();
     std::chrono::steady_clock::time_point lastUsed() const;
+    void setMock(bool mock) { _isMock = mock; }
+    bool isMock() const { return _isMock; }
+    void setPoolOrigin(uint64_t pool_id, uint64_t generation) {
+        _originPoolId = pool_id;
+        _originGeneration = generation;
+    }
+    bool matchesPoolOrigin(uint64_t pool_id, uint64_t generation) const {
+        return _originPoolId == pool_id && _originGeneration == generation;
+    }
+    uint64_t originPoolId() const { return _originPoolId; }
+    uint64_t originGeneration() const { return _originGeneration; }
 
     // Materialize connect-attrs from a Python token-factory callback.
     // The factory may return either a bare attrs dict (legacy) or a
@@ -100,6 +111,9 @@ class Connection {
     std::u16string _connStr;
     bool _fromPool = false;
     bool _autocommit = true;
+    bool _isMock = false;
+    uint64_t _originPoolId = 0;
+    uint64_t _originGeneration = 0;
     SqlHandlePtr _dbcHandle;
     std::chrono::steady_clock::time_point _lastUsed;
     // POSIX-epoch expiry (seconds) of the access token this connection last
@@ -150,6 +164,9 @@ class ConnectionHandle {
     // Get information about the driver and data source
     py::object getInfo(SQLUSMALLINT infoType) const;
 
+    uint64_t originGeneration() const { return _conn ? _conn->originGeneration() : 0; }
+    uint64_t originPoolId() const { return _conn ? _conn->originPoolId() : 0; }
+
   private:
     std::shared_ptr<Connection> _conn;
     bool _usePool;