Rebase

Author: betodealmeida

UPDATING.md

@@ -24,6 +24,41 @@ assists people when migrating to a new version.
 
 ## Next
 
+### Engine Manager for Connection Pooling
+
+A new `EngineManager` class has been introduced to centralize SQLAlchemy engine creation and management. This enables connection pooling for analytics databases and provides a more flexible architecture for engine configuration.
+
+#### Breaking Changes
+
+1. **Removed `SSH_TUNNEL_MANAGER_CLASS` config**: SSH tunnel handling is now integrated into the EngineManager. If you have custom SSH tunnel managers, you'll need to migrate to the new architecture.
+
+2. **Removed `nullpool` parameter**: The `get_sqla_engine()` and `get_raw_connection()` methods on the `Database` model no longer accept a `nullpool` parameter. Pool configuration is now controlled through the engine manager.
+
+3. **Removed `_get_sqla_engine()` method**: The private `_get_sqla_engine()` method has been removed from the `Database` model. All engine creation now goes through the `EngineManager`.
+
+#### New Configuration Options
+
+```python
+# Engine manager mode:
+# - EngineModes.NEW: Creates a new engine for every connection (default, original behavior)
+# - EngineModes.SINGLETON: Reuses engines with connection pooling
+from superset.engines.manager import EngineModes
+ENGINE_MANAGER_MODE = EngineModes.NEW
+
+# Cleanup interval for abandoned locks (default: 5 minutes)
+from datetime import timedelta
+ENGINE_MANAGER_CLEANUP_INTERVAL = timedelta(minutes=5)
+
+# Automatically start cleanup thread for SINGLETON mode (default: True)
+ENGINE_MANAGER_AUTO_START_CLEANUP = True
+```
+
+#### Migration Guide
+
+- If you were using the `nullpool` parameter, remove it from your calls
+- If you had a custom `SSH_TUNNEL_MANAGER_CLASS`, refactor to use the new EngineManager architecture
+- If you need connection pooling, set `ENGINE_MANAGER_MODE = EngineModes.SINGLETON` and configure the pool in your database's `extra` JSON field
+
 ### WebSocket config for GAQ with Docker
 
 [35896](https://github.com/apache/superset/pull/35896) and [37624](https://github.com/apache/superset/pull/37624) updated documentation on how to run and configure Superset with Docker. Specifically for the WebSocket configuration, a new `docker/superset-websocket/config.example.json` was added to the repo, so that users could copy it to create a `docker/superset-websocket/config.json` file. The existing `docker/superset-websocket/config.json` was removed and git-ignored, so if you're using GAQ / WebSocket make sure to:

superset/engines/manager.py

@@ -320,7 +320,8 @@ def _get_engine_args(
         uri = make_url_safe(database.sqlalchemy_uri_decrypted)
 
         extra = database.get_extra(source)
-        kwargs = extra.get("engine_params", {})
+        # Make a copy to avoid mutating the original extra dict
+        kwargs = dict(extra.get("engine_params", {}))
 
         # get pool class
         if self.mode == EngineModes.NEW or "poolclass" not in kwargs:
@@ -336,7 +337,7 @@ def _get_engine_args(
             kwargs["poolclass"] = pools.get(extra["poolclass"], pool.QueuePool)
 
         # update URI for specific catalog/schema
-        connect_args = extra.setdefault("connect_args", {})
+        connect_args = dict(extra.get("connect_args", {}))
         uri, connect_args = database.db_engine_spec.adjust_engine_params(
             uri,
             connect_args,

tests/unit_tests/engines/manager_test.py

@@ -22,7 +22,6 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from sqlalchemy.pool import NullPool
 
 from superset.engines.manager import _LockManager, EngineManager, EngineModes
 
@@ -51,7 +50,7 @@ def test_cleanup_removes_unused_locks(self):
         manager = _LockManager()
 
         # Create locks
-        lock1 = manager.get_lock("key1")
+        _ = manager.get_lock("key1")  # noqa: F841
         lock2 = manager.get_lock("key2")
 
         # Cleanup with only key1 active
@@ -111,7 +110,7 @@ def mock_database(self):
         """Create a mock database."""
         database = MagicMock()
         database.sqlalchemy_uri_decrypted = "postgresql://user:pass@localhost/test"
-        database.get_extra.return_value = {"engine_params": {"poolclass": NullPool}}
+        database.get_extra.return_value = {"engine_params": {}}
         database.get_effective_user.return_value = "test_user"
         database.impersonate_user = False
         database.update_params_from_encrypted_extra = MagicMock()
@@ -231,7 +230,7 @@ def test_ssh_tunnel_creation(self, mock_tunnel_class, engine_manager):
         ssh_tunnel.server_address = "ssh.example.com"
         ssh_tunnel.server_port = 22
         ssh_tunnel.username = "ssh_user"
-        ssh_tunnel.password = "ssh_pass"
+        ssh_tunnel.password = "ssh_pass"  # noqa: S105
         ssh_tunnel.private_key = None
         ssh_tunnel.private_key_password = None
 
@@ -266,7 +265,7 @@ def test_ssh_tunnel_recreation_when_inactive(
         ssh_tunnel.server_address = "ssh.example.com"
         ssh_tunnel.server_port = 22
         ssh_tunnel.username = "ssh_user"
-        ssh_tunnel.password = "ssh_pass"
+        ssh_tunnel.password = "ssh_pass"  # noqa: S105
         ssh_tunnel.private_key = None
         ssh_tunnel.private_key_password = None
 

tests/unit_tests/sql/execution/conftest.py

@@ -202,7 +202,6 @@ def setup_mock_raw_connection(
     def _raw_connection(
         catalog: str | None = None,
         schema: str | None = None,
-        nullpool: bool = True,
         source: Any | None = None,
     ):
         yield mock_connection