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 <[email protected]> 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)
