Merge branch 'main' of https://github.com/microsoft/mssql-python into release/0.7.0

Author: bewithgaurav

mssql_python/__init__.py

@@ -54,16 +54,21 @@
 threadsafety = 1
 
 from .pooling import PoolingManager
-def pooling(max_size=100, idle_timeout=600):
+def pooling(max_size=100, idle_timeout=600, enabled=True):
 #     """
 #     Enable connection pooling with the specified parameters.
-    
+#     By default:
+#         - If not explicitly called, pooling will be auto-enabled with default values.
+
 #     Args:
 #         max_size (int): Maximum number of connections in the pool.
 #         idle_timeout (int): Time in seconds before idle connections are closed.
 
 #     Returns:
 #         None
 #     """
-    PoolingManager.enable(max_size, idle_timeout) 
+    if not enabled:
+        PoolingManager.disable()
+    else:
+        PoolingManager.enable(max_size, idle_timeout)
 

mssql_python/connection.py

@@ -58,6 +58,9 @@ def __init__(self, connection_str: str = "", autocommit: bool = False, attrs_bef
             connection_str, **kwargs
         )
         self._attrs_before = attrs_before or {}
+        # Auto-enable pooling if user never called
+        if not PoolingManager.is_initialized():
+            PoolingManager.enable()
         self._pooling = PoolingManager.is_enabled()
         self._conn = ddbc_bindings.Connection(self.connection_str, self._pooling, self._attrs_before)
         self.setautocommit(autocommit)

mssql_python/pooling.py

@@ -1,9 +1,11 @@
 # mssql_python/pooling.py
+import atexit
 from mssql_python import ddbc_bindings
 import threading
 
 class PoolingManager:
     _enabled = False
+    _initialized = False 
     _lock = threading.Lock()
     _config = {
         "max_size": 100,
@@ -23,7 +25,23 @@ def enable(cls, max_size=100, idle_timeout=600):
             cls._config["max_size"] = max_size
             cls._config["idle_timeout"] = idle_timeout
             cls._enabled = True
+            cls._initialized = True
+
+    @classmethod
+    def disable(cls):
+        with cls._lock:
+            cls._enabled = False
+            cls._initialized = True
 
     @classmethod
     def is_enabled(cls):
         return cls._enabled
+
+    @classmethod
+    def is_initialized(cls):
+        return cls._initialized
+    
+@atexit.register
+def shutdown_pooling():
+    if PoolingManager.is_enabled():
+        ddbc_bindings.close_pooling()

mssql_python/pybind/connection/connection_pool.cpp

@@ -83,6 +83,25 @@ void ConnectionPool::release(std::shared_ptr<Connection> conn) {
     }
 }
 
+void ConnectionPool::close() {
+    std::vector<std::shared_ptr<Connection>> to_close;
+    {
+        std::lock_guard<std::mutex> lock(_mutex);
+        while (!_pool.empty()) {
+            to_close.push_back(_pool.front());
+            _pool.pop_front();
+        }
+        _current_size = 0;
+    }
+    for (auto& conn : to_close) {
+        try {
+            conn->disconnect();
+        } catch (const std::exception& ex) {
+            LOG("ConnectionPool::close: disconnect failed: {}", ex.what());
+        }
+    }
+}
+
 ConnectionPoolManager& ConnectionPoolManager::getInstance() {
     static ConnectionPoolManager manager;
     return manager;
@@ -110,4 +129,14 @@ void ConnectionPoolManager::configure(int max_size, int idle_timeout_secs) {
     std::lock_guard<std::mutex> lock(_manager_mutex);
     _default_max_size = max_size;
     _default_idle_secs = idle_timeout_secs;
-}
+}
+
+void ConnectionPoolManager::closePools() {
+    std::lock_guard<std::mutex> lock(_manager_mutex);
+    for (auto& [conn_str, pool] : _pools) {
+        if (pool) {
+            pool->close();
+        }
+    }
+    _pools.clear();
+}

mssql_python/pybind/connection/connection_pool.h

@@ -24,6 +24,9 @@ class ConnectionPool {
     // Returns a connection to the pool for reuse
     void release(std::shared_ptr<Connection> conn);
 
+    // Closes all connections in the pool, releasing resources
+    void close();
+
 private:
     size_t _max_size;       // Maximum number of connections allowed
     int _idle_timeout_secs; // Idle time before connections are considered stale
@@ -46,6 +49,9 @@ class ConnectionPoolManager {
     // Returns a connection to its original pool
     void returnConnection(const std::wstring& conn_str, std::shared_ptr<Connection> conn);
 
+    // Closes all pools and their connections
+    void closePools();
+
 private:
     ConnectionPoolManager() = default;    
     ~ConnectionPoolManager() = default;

mssql_python/pybind/ddbc_bindings.cpp

@@ -2109,6 +2109,7 @@ PYBIND11_MODULE(ddbc_bindings, m) {
         .def("get_autocommit", &ConnectionHandle::getAutocommit)
         .def("alloc_statement_handle", &ConnectionHandle::allocStatementHandle);
     m.def("enable_pooling", &enable_pooling, "Enable global connection pooling");
+    m.def("close_pooling", []() {ConnectionPoolManager::getInstance().closePools();});
     m.def("DDBCSQLExecDirect", &SQLExecDirect_wrap, "Execute a SQL query directly");
     m.def("DDBCSQLExecute", &SQLExecute_wrap, "Prepare and execute T-SQL statements");
     m.def("DDBCSQLRowCount", &SQLRowCount_wrap,