[superset] branch fix_get_sqla_engine created (now c5fe359544)

[be...@apache.org]

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

beto pushed a change to branch fix_get_sqla_engine
in repository https://gitbox.apache.org/repos/asf/superset.git


      at c5fe359544 fix: calls to _get_sqla_engine

This branch includes the following new commits:

     new c5fe359544 fix: calls to _get_sqla_engine

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: fix: calls to _get_sqla_engine

[be...@apache.org]

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

beto pushed a commit to branch fix_get_sqla_engine
in repository https://gitbox.apache.org/repos/asf/superset.git

commit c5fe3595444fe83b91f085514b33ebfa7b66e621
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Thu Aug 10 16:47:46 2023 -0700

    fix: calls to _get_sqla_engine
---
 superset/models/core.py | 50 +++++++++++++++++++++++++------------------------
 1 file changed, 26 insertions(+), 24 deletions(-)

diff --git a/superset/models/core.py b/superset/models/core.py
index e76da0dcd5..bf4651fed9 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -563,7 +563,8 @@ class Database(
         mutator: Callable[[pd.DataFrame], None] | None = None,
     ) -> pd.DataFrame:
         sqls = self.db_engine_spec.parse_sql(sql)
-        engine = self._get_sqla_engine(schema)
+        with self.get_sqla_engine_with_context(schema) as engine:
+            engine_url = engine.url
         mutate_after_split = config["MUTATE_AFTER_SPLIT"]
         sql_query_mutator = config["SQL_QUERY_MUTATOR"]
 
@@ -577,7 +578,7 @@ class Database(
         def _log_query(sql: str) -> None:
             if log_query:
                 log_query(
-                    engine.url,
+                    engine_url,
                     sql,
                     schema,
                     __name__,
@@ -624,13 +625,12 @@ class Database(
             return df
 
     def compile_sqla_query(self, qry: Select, schema: str | None = None) -> str:
-        engine = self._get_sqla_engine(schema=schema)
+        with self.get_sqla_engine_with_context(schema) as engine:
+            sql = str(qry.compile(engine, compile_kwargs={"literal_binds": True}))
 
-        sql = str(qry.compile(engine, compile_kwargs={"literal_binds": True}))
-
-        # pylint: disable=protected-access
-        if engine.dialect.identifier_preparer._double_percents:  # noqa
-            sql = sql.replace("%%", "%")
+            # pylint: disable=protected-access
+            if engine.dialect.identifier_preparer._double_percents:  # noqa
+                sql = sql.replace("%%", "%")
 
         return sql
 
@@ -645,18 +645,18 @@ class Database(
         cols: list[ResultSetColumnType] | None = None,
     ) -> str:
         """Generates a ``select *`` statement in the proper dialect"""
-        eng = self._get_sqla_engine(schema=schema, source=utils.QuerySource.SQL_LAB)
-        return self.db_engine_spec.select_star(
-            self,
-            table_name,
-            schema=schema,
-            engine=eng,
-            limit=limit,
-            show_cols=show_cols,
-            indent=indent,
-            latest_partition=latest_partition,
-            cols=cols,
-        )
+        with self.get_sqla_engine_with_context(schema) as engine:
+            return self.db_engine_spec.select_star(
+                self,
+                table_name,
+                schema=schema,
+                engine=engine,
+                limit=limit,
+                show_cols=show_cols,
+                indent=indent,
+                latest_partition=latest_partition,
+                cols=cols,
+            )
 
     def apply_limit_to_sql(
         self, sql: str, limit: int = 1000, force: bool = False
@@ -670,8 +670,8 @@ class Database(
 
     @property
     def inspector(self) -> Inspector:
-        engine = self._get_sqla_engine()
-        return sqla.inspect(engine)
+        with self.get_sqla_engine_with_context() as engine:
+            return sqla.inspect(engine)
 
     @cache_util.memoized_func(
         key="db:{self.id}:schema:{schema}:table_list",
@@ -955,8 +955,10 @@ class Database(
         return view_name in view_names
 
     def has_view(self, view_name: str, schema: str | None = None) -> bool:
-        engine = self._get_sqla_engine()
-        return engine.run_callable(self._has_view, engine.dialect, view_name, schema)
+        with self.get_sqla_engine_with_context(schema) as engine:
+            return engine.run_callable(
+                self._has_view, engine.dialect, view_name, schema
+            )
 
     def has_view_by_name(self, view_name: str, schema: str | None = None) -> bool:
         return self.has_view(view_name=view_name, schema=schema)