[superset] branch ssh-tunnel-refactor-get-sqla-engine created (now f6490faee9)

[hu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hugh pushed a change to branch ssh-tunnel-refactor-get-sqla-engine
in repository https://gitbox.apache.org/repos/asf/superset.git


      at f6490faee9 create function for get_sqla_engine with context

This branch includes the following new commits:

     new f6490faee9 create function for get_sqla_engine with context

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.

[superset] 01/01: create function for get_sqla_engine with context

[hu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hugh pushed a commit to branch ssh-tunnel-refactor-get-sqla-engine
in repository https://gitbox.apache.org/repos/asf/superset.git

commit f6490faee98835e5cd4556cb27cef8b9dd859c72
Author: hughhhh <hu...@gmail.com>
AuthorDate: Wed Oct 12 16:28:03 2022 -0400

    create function for get_sqla_engine with context
---
 superset/models/core.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

diff --git a/superset/models/core.py b/superset/models/core.py
index 008230ef48..f566acc526 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -21,7 +21,7 @@ import json
 import logging
 import textwrap
 from ast import literal_eval
-from contextlib import closing
+from contextlib import closing, contextmanager
 from copy import deepcopy
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type
@@ -362,6 +362,60 @@ class Database(
             else None
         )
 
+    @contextmanager
+    def get_sqla_engine_with_context(
+        self,
+        schema: Optional[str] = None,
+        nullpool: bool = True,
+        source: Optional[utils.QuerySource] = None,
+    ) -> Engine:
+        extra = self.get_extra()
+        sqlalchemy_url = make_url_safe(self.sqlalchemy_uri_decrypted)
+        sqlalchemy_url = self.db_engine_spec.adjust_database_uri(sqlalchemy_url, schema)
+        effective_username = self.get_effective_user(sqlalchemy_url)
+        # If using MySQL or Presto for example, will set url.username
+        # If using Hive, will not do anything yet since that relies on a
+        # configuration parameter instead.
+        sqlalchemy_url = self.db_engine_spec.get_url_for_impersonation(
+            sqlalchemy_url, self.impersonate_user, effective_username
+        )
+
+        masked_url = self.get_password_masked_url(sqlalchemy_url)
+        logger.debug("Database.get_sqla_engine(). Masked URL: %s", str(masked_url))
+
+        params = extra.get("engine_params", {})
+        if nullpool:
+            params["poolclass"] = NullPool
+
+        connect_args = params.get("connect_args", {})
+        if self.impersonate_user:
+            self.db_engine_spec.update_impersonation_config(
+                connect_args, str(sqlalchemy_url), effective_username
+            )
+
+        if connect_args:
+            params["connect_args"] = connect_args
+
+        self.update_params_from_encrypted_extra(params)
+
+        if DB_CONNECTION_MUTATOR:
+            if not source and request and request.referrer:
+                if "/superset/dashboard/" in request.referrer:
+                    source = utils.QuerySource.DASHBOARD
+                elif "/explore/" in request.referrer:
+                    source = utils.QuerySource.CHART
+                elif "/superset/sqllab/" in request.referrer:
+                    source = utils.QuerySource.SQL_LAB
+
+            sqlalchemy_url, params = DB_CONNECTION_MUTATOR(
+                sqlalchemy_url, params, effective_username, security_manager, source
+            )
+
+        try:
+            yield create_engine(sqlalchemy_url, **params)
+        except Exception as ex:
+            raise self.db_engine_spec.get_dbapi_mapped_exception(ex)
+
     def get_sqla_engine(
         self,
         schema: Optional[str] = None,