betodealmeida commented on code in PR #34826:
URL: https://github.com/apache/superset/pull/34826#discussion_r2765573230
##########
superset/models/core.py:
##########
@@ -424,130 +419,31 @@ def get_effective_user(self, object_url: URL) -> str |
None:
)
@contextmanager
- def get_sqla_engine( # pylint: disable=too-many-arguments
+ def get_sqla_engine(
self,
catalog: str | None = None,
schema: str | None = None,
- nullpool: bool = True,
source: utils.QuerySource | None = None,
- ) -> Engine:
+ ) -> Iterator[Engine]:
"""
Context manager for a SQLAlchemy engine.
- This method will return a context manager for a SQLAlchemy engine.
Using the
- context manager (as opposed to the engine directly) is important
because we need
- to potentially establish SSH tunnels before the connection is created,
and clean
- them up once the engine is no longer used.
+ This method will return a context manager for a SQLAlchemy engine. The
engine
+ manager handles connection pooling, SSH tunnels, and other connection
details
+ based on the configured mode (NEW or SINGLETON).
"""
+ # Import here to avoid circular imports
+ from superset.extensions import engine_manager_extension
- sqlalchemy_uri = self.sqlalchemy_uri_decrypted
-
- ssh_context_manager = (
- ssh_manager_factory.instance.create_tunnel(
- ssh_tunnel=self.ssh_tunnel,
- sqlalchemy_database_uri=sqlalchemy_uri,
- )
- if self.ssh_tunnel
- else nullcontext()
- )
-
- with ssh_context_manager as ssh_context:
- if ssh_context:
- logger.info(
- "[SSH] Successfully created tunnel w/ %s tunnel_timeout +
%s "
- "ssh_timeout at %s",
- sshtunnel.TUNNEL_TIMEOUT,
- sshtunnel.SSH_TIMEOUT,
- ssh_context.local_bind_address,
- )
- sqlalchemy_uri = ssh_manager_factory.instance.build_sqla_url(
- sqlalchemy_uri,
- ssh_context,
- )
-
- engine_context_manager = app.config["ENGINE_CONTEXT_MANAGER"]
- with engine_context_manager(self, catalog, schema):
- with check_for_oauth2(self):
- yield self._get_sqla_engine(
- catalog=catalog,
- schema=schema,
- nullpool=nullpool,
- source=source,
- sqlalchemy_uri=sqlalchemy_uri,
- )
-
- def _get_sqla_engine( # pylint: disable=too-many-locals # noqa: C901
- self,
- catalog: str | None = None,
- schema: str | None = None,
- nullpool: bool = True,
- source: utils.QuerySource | None = None,
- sqlalchemy_uri: str | None = None,
- ) -> Engine:
- sqlalchemy_url = make_url_safe(
- sqlalchemy_uri if sqlalchemy_uri else self.sqlalchemy_uri_decrypted
- )
- self.db_engine_spec.validate_database_uri(sqlalchemy_url)
-
- extra = self.get_extra(source)
- engine_kwargs = extra.get("engine_params", {})
- if nullpool:
- engine_kwargs["poolclass"] = NullPool
- connect_args = engine_kwargs.setdefault("connect_args", {})
-
- # modify URL/args for a specific catalog/schema
- sqlalchemy_url, connect_args =
self.db_engine_spec.adjust_engine_params(
- uri=sqlalchemy_url,
- connect_args=connect_args,
+ # Use the engine manager to get the engine
+ engine_manager = engine_manager_extension.manager
Review Comment:
`EngineManager.get_engine` does this:
```python
@contextmanager
def get_engine(...) -> Iterator[Engine]:
with self.engine_context_manager(database, catalog, schema):
with check_for_oauth2(database):
yield self._get_engine(database, catalog, schema, source)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
